PaddlePaddle
diff --git a/‎paddle/fluid/distributed/auto_parallel/spmd_rules/rules.h
Copy file name to clipboard
-5Lines changed: 0 additions & 5 deletions b/‎paddle/fluid/distributed/auto_parallel/spmd_rules/rules.h
Copy file name to clipboard
-5Lines changed: 0 additions & 5 deletions
diff --git a/‎paddle/fluid/distributed/auto_parallel/spmd_rules/split_spmd_rule.cc
Copy file name to clipboard
-218Lines changed: 0 additions & 218 deletions b/‎paddle/fluid/distributed/auto_parallel/spmd_rules/split_spmd_rule.cc
Copy file name to clipboard
-218Lines changed: 0 additions & 218 deletions
diff --git a/‎paddle/fluid/distributed/auto_parallel/spmd_rules/split_spmd_rule.h
Copy file name to clipboard
-41Lines changed: 0 additions & 41 deletions b/‎paddle/fluid/distributed/auto_parallel/spmd_rules/split_spmd_rule.h
Copy file name to clipboard
-41Lines changed: 0 additions & 41 deletions
diff --git a/‎paddle/phi/infermeta/spmd_rules/rules.h
Copy file name to clipboard
+10Lines changed: 10 additions & 0 deletions b/‎paddle/phi/infermeta/spmd_rules/rules.h
Copy file name to clipboard
+10Lines changed: 10 additions & 0 deletions
@@ -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);
 
 
@@ -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