【Semi-Auto】Adapt split spmd rule to phi (PaddlePaddle#57467)

* adapt split rule to phi * fix bugs and modify apis in unit test * fix codestyle * bug fix
Frida-a · Oct 14, 2023 · ed4812f · ed4812f
1 parent 9138b69
commit ed4812f
Show file tree

Hide file tree

Showing 7 changed files with 338 additions and 312 deletions.
diff --git a/paddle/fluid/distributed/auto_parallel/spmd_rules/rules.h b/paddle/fluid/distributed/auto_parallel/spmd_rules/rules.h
@@ -18,7 +18,6 @@
 #include "paddle/fluid/distributed/auto_parallel/spmd_rules/cross_entropy_with_softmax_spmd_rule.h"
 #include "paddle/fluid/distributed/auto_parallel/spmd_rules/replicated_spmd_rule.h"
 #include "paddle/fluid/distributed/auto_parallel/spmd_rules/softmax_spmd_rule.h"
-#include "paddle/fluid/distributed/auto_parallel/spmd_rules/split_spmd_rule.h"
 #include "paddle/fluid/distributed/auto_parallel/spmd_rules/transpose_spmd_rule.h"
 
 // TODO(ljz) Automatic this process in cmake file.
@@ -37,10 +36,6 @@ REGISTER_SPMD_RULE(log_softmax, SoftmaxSPMDRule);
 REGISTER_SPMD_RULE(cross_entropy_with_softmax, CrossEntropyWithSoftmaxSPMDRule);
 REGISTER_SPMD_RULE(softmax_with_cross_entropy, CrossEntropyWithSoftmaxSPMDRule);
 
-// split rule
-REGISTER_SPMD_RULE(split, SplitSPMDRule);
-REGISTER_SPMD_RULE(split_with_num, SplitSPMDRule);
-
 // transpose rule
 REGISTER_SPMD_RULE(transpose, TransposeSPMDRule);
 

diff --git a/paddle/fluid/distributed/auto_parallel/spmd_rules/split_spmd_rule.cc b/paddle/fluid/distributed/auto_parallel/spmd_rules/split_spmd_rule.cc
diff --git a/paddle/fluid/distributed/auto_parallel/spmd_rules/split_spmd_rule.h b/paddle/fluid/distributed/auto_parallel/spmd_rules/split_spmd_rule.h
diff --git a/paddle/phi/infermeta/spmd_rules/rules.h b/paddle/phi/infermeta/spmd_rules/rules.h
@@ -24,6 +24,7 @@ limitations under the License. */
 #include "paddle/phi/infermeta/spmd_rules/reduction.h"
 #include "paddle/phi/infermeta/spmd_rules/replicated.h"
 #include "paddle/phi/infermeta/spmd_rules/reshape.h"
+#include "paddle/phi/infermeta/spmd_rules/split.h"
 
 /**
  * Design Notes:
@@ -485,5 +486,14 @@ PD_REGISTER_SPMD_RULE(
     PD_INFER_SPMD(phi::distributed::EmbeddingInferSpmd),
     PD_INFER_SPMD(phi::distributed::EmbeddingInferSpmdReverse));
 
+// split rule
+PD_REGISTER_SPMD_RULE(split,
+                      PD_INFER_SPMD(phi::distributed::SplitInferSpmd),
+                      PD_INFER_SPMD(phi::distributed::SplitInferSpmdReverse));
+PD_REGISTER_SPMD_RULE(
+    split_with_num,
+    PD_INFER_SPMD(phi::distributed::SplitWithNumInferSpmd),
+    PD_INFER_SPMD(phi::distributed::SplitWithNumInferSpmdReverse));
+
 }  // namespace distributed
 }  // namespace phi