[Auto Parallel] Add spmd_rule for 11 ops. (#72312)

Glencsa · web-flow · commit c8552a651dd7 · 2025-05-19T15:21:14.000+08:00
* add unary ops which have spmd_rule but not add in yaml file.

* add spmd_rule for 11 ops.

* add CI test for 11 ops of spmd rule and add spmd rule for nonzero ops.

* add spmd rule for 12 backward ops and add some CPP CI test.

* remove py CI

* fix bug
diff --git a/paddle/phi/infermeta/spmd_rules/elementwise.cc b/paddle/phi/infermeta/spmd_rules/elementwise.cc
@@ -534,4 +534,117 @@ SpmdInfo ElementwiseBinaryGradInferSpmd(const DistMetaTensor& x,
   info.first.emplace(info.first.begin() + 2, out_grad.dist_attr());
   return info;
 }
+// round
+SpmdInfo RoundInfoSpmd(const DistMetaTensor& x, const int decimals) {
+  return ElementwiseUnaryInferSpmd(x);
+}
+
+// mish
+SpmdInfo MishInfoSpmd(const DistMetaTensor& x, const float lambda) {
+  return ElementwiseUnaryInferSpmd(x);
+}
+SpmdInfo MishGradInfoSpmd(const DistMetaTensor& x,
+                          const DistMetaTensor& out_grad,
+                          const float lambda) {
+  return ElementwiseUnaryGradInferSpmd(x, out_grad);
+}
+
+// elu
+SpmdInfo EluInfoSpmd(const DistMetaTensor& x, const float alpha) {
+  return ElementwiseUnaryInferSpmd(x);
+}
+
+SpmdInfo EluGradInfoSpmd(const DistMetaTensor& x,
+                         const DistMetaTensor& out,
+                         const DistMetaTensor& out_grad,
+                         const float alpha) {
+  return ElementwiseUnaryGradInferSpmd(x, out, out_grad);
+}
+// selu
+SpmdInfo SeluInfoSpmd(const DistMetaTensor& x,
+                      const float alpha,
+                      const float scale) {
+  return ElementwiseUnaryInferSpmd(x);
+}
+
+SpmdInfo SeluGradInfoSpmd(const DistMetaTensor& x,
+                          const DistMetaTensor& out_grad,
+                          const float scale,
+                          const float alpha) {
+  return ElementwiseUnaryGradInferSpmd(x, out_grad);
+}
+// celu
+SpmdInfo CeluInfoSpmd(const DistMetaTensor& x, const float alpha) {
+  return ElementwiseUnaryInferSpmd(x);
+}
+
+SpmdInfo CeluGradInfoSpmd(const DistMetaTensor& x,
+                          const DistMetaTensor& out_grad,
+                          const float alpha) {
+  return ElementwiseUnaryGradInferSpmd(x, out_grad);
+}
+// stanh
+SpmdInfo StanhInfoSpmd(const DistMetaTensor& x,
+                       const float scale_a,
+                       const float scale_b) {
+  return ElementwiseUnaryInferSpmd(x);
+}
+
+SpmdInfo StanhGradInfoSpmd(const DistMetaTensor& x,
+                           const DistMetaTensor& out_grad,
+                           const float scale_a,
+                           const float scale_b) {
+  return ElementwiseUnaryGradInferSpmd(x, out_grad);
+}
+
+// softplus
+SpmdInfo SoftplusInfoSpmd(const DistMetaTensor& x,
+                          const float beta,
+                          const float threshold) {
+  return ElementwiseUnaryInferSpmd(x);
+}
+
+SpmdInfo SoftplusGradInfoSpmd(const DistMetaTensor& x,
+                              const DistMetaTensor& out_grad,
+                              const float beta,
+                              const float threshold) {
+  return ElementwiseUnaryGradInferSpmd(x, out_grad);
+}
+
+// softshrink
+SpmdInfo SoftshrinkInfoSpmd(const DistMetaTensor& x, const float threshold) {
+  return ElementwiseUnaryInferSpmd(x);
+}
+
+SpmdInfo SoftshrinkGradInfoSpmd(const DistMetaTensor& x,
+                                const DistMetaTensor& out_grad,
+                                const float threshold) {
+  return ElementwiseUnaryGradInferSpmd(x, out_grad);
+}
+
+// thresholded_relu
+SpmdInfo ThresholdedReluInfoSpmd(const DistMetaTensor& x,
+                                 const float threshold,
+                                 const float value) {
+  return ElementwiseUnaryInferSpmd(x);
+}
+
+SpmdInfo ThresholdedReluGradInfoSpmd(const DistMetaTensor& x,
+                                     const DistMetaTensor& out_grad,
+                                     const float threshold,
+                                     const float value) {
+  return ElementwiseUnaryGradInferSpmd(x, out_grad);
+}
+
+// logit
+SpmdInfo LogitInfoSpmd(const DistMetaTensor& x, const float eps) {
+  return ElementwiseUnaryInferSpmd(x);
+}
+
+SpmdInfo LogitGradInfoSpmd(const DistMetaTensor& x,
+                           const DistMetaTensor& out_grad,
+                           const float eps) {
+  return ElementwiseUnaryGradInferSpmd(x, out_grad);
+}
+
 }  // namespace phi::distributed
diff --git a/paddle/phi/infermeta/spmd_rules/elementwise.h b/paddle/phi/infermeta/spmd_rules/elementwise.h
@@ -65,5 +65,63 @@ SpmdInfo SwiGLUGradInferSpmd(const DistMetaTensor& x,
                              const DistMetaTensor& y,
                              const DistMetaTensor& out_grad);
 
+SpmdInfo RoundInfoSpmd(const DistMetaTensor& x, const int decimals);
+
+SpmdInfo MishInfoSpmd(const DistMetaTensor& x, const float lambda);
+SpmdInfo MishGradInfoSpmd(const DistMetaTensor& x,
+                          const DistMetaTensor& out_grad,
+                          const float lambda);
+
+SpmdInfo EluInfoSpmd(const DistMetaTensor& x, const float alpha);
+SpmdInfo EluGradInfoSpmd(const DistMetaTensor& x,
+                         const DistMetaTensor& out,
+                         const DistMetaTensor& out_grad,
+                         const float alpha);
+
+SpmdInfo SeluInfoSpmd(const DistMetaTensor& x,
+                      const float alpha,
+                      const float scale);
+SpmdInfo SeluGradInfoSpmd(const DistMetaTensor& x,
+                          const DistMetaTensor& out_grad,
+                          const float scale,
+                          const float alpha);
+SpmdInfo CeluInfoSpmd(const DistMetaTensor& x, const float alpha);
+SpmdInfo CeluGradInfoSpmd(const DistMetaTensor& x,
+                          const DistMetaTensor& out_grad,
+                          const float alpha);
+
+SpmdInfo StanhInfoSpmd(const DistMetaTensor& x,
+                       const float scale_a,
+                       const float scale_b);
+SpmdInfo StanhGradInfoSpmd(const DistMetaTensor& x,
+                           const DistMetaTensor& out_grad,
+                           const float scale_a,
+                           const float scale_b);
+
+SpmdInfo SoftplusInfoSpmd(const DistMetaTensor& x,
+                          const float beta,
+                          const float threshold);
+SpmdInfo SoftplusGradInfoSpmd(const DistMetaTensor& x,
+                              const DistMetaTensor& out_grad,
+                              const float beta,
+                              const float threshold);
+
+SpmdInfo SoftshrinkInfoSpmd(const DistMetaTensor& x, const float threshold);
+SpmdInfo SoftshrinkGradInfoSpmd(const DistMetaTensor& x,
+                                const DistMetaTensor& out_grad,
+                                const float threshold);
+
+SpmdInfo ThresholdedReluInfoSpmd(const DistMetaTensor& x,
+                                 const float threshold,
+                                 const float value);
+SpmdInfo ThresholdedReluGradInfoSpmd(const DistMetaTensor& x,
+                                     const DistMetaTensor& out_grad,
+                                     const float threshold,
+                                     const float value);
+
+SpmdInfo LogitInfoSpmd(const DistMetaTensor& x, const float eps);
+SpmdInfo LogitGradInfoSpmd(const DistMetaTensor& x,
+                           const DistMetaTensor& out_grad,
+                           const float eps);
 }  // namespace distributed
 }  // namespace phi
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
@@ -398,6 +398,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param: [x]
+    spmd_rule : CeluGradInfoSpmd
   kernel :
     func : celu_grad
   backward : celu_double_grad
@@ -921,6 +922,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [x]
+    spmd_rule : EluGradInfoSpmd
   kernel :
     func : elu_grad
   backward : elu_double_grad
@@ -2000,6 +2002,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [x]
+    spmd_rule : LogitGradInfoSpmd
   kernel :
     func : logit_grad
 
@@ -2227,6 +2230,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [x]
+    spmd_rule : MishGradInfoSpmd
   kernel :
     func : mish_grad
   inplace : (out_grad -> x_grad)
@@ -2888,6 +2892,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [out]
+    spmd_rule : SeluGradInfoSpmd
   kernel :
     func : selu_grad
     data_type : out
@@ -3144,6 +3149,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [x]
+    spmd_rule : SoftplusGradInfoSpmd
   kernel :
     func : softplus_grad
   backward : softplus_double_grad
@@ -3156,6 +3162,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [x]
+    spmd_rule : SoftshrinkGradInfoSpmd
   kernel :
     func : softshrink_grad
   inplace : (out_grad -> x_grad)
@@ -3315,6 +3322,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [x]
+    spmd_rule : StanhGradInfoSpmd
   kernel :
     func : stanh_grad
 
@@ -3509,6 +3517,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [x]
+    spmd_rule : ThresholdedReluGradInfoSpmd
   kernel :
     func : thresholded_relu_grad
   inplace : (out_grad -> x_grad)
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -928,6 +928,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param: [x]
+    spmd_rule : CeluInfoSpmd
   kernel :
     func : celu
   backward : celu_grad
@@ -1656,6 +1657,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [x]
+    spmd_rule : EluInfoSpmd
   kernel :
     func : elu
   inplace : (x -> out)
@@ -3192,6 +3194,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [x]
+    spmd_rule : LogitInfoSpmd
   kernel :
     func : logit
   inplace: (x -> out)
@@ -3563,6 +3566,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [x]
+    spmd_rule : MishInfoSpmd
   kernel :
     func : mish
   backward : mish_grad
@@ -3735,6 +3739,7 @@
   output : Tensor(out)
   infer_meta :
     func : NonZeroInferMeta
+    spmd_rule : NonZeroInferSpmd
   kernel :
     func : nonzero
     data_type: condition
@@ -4376,6 +4381,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [x]
+    spmd_rule : RoundInfoSpmd
   kernel :
     func : round
   inplace : (x -> out)
@@ -4490,6 +4496,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [x]
+    spmd_rule : SeluInfoSpmd
   kernel :
     func : selu
   backward : selu_grad
@@ -4776,6 +4783,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [x]
+    spmd_rule : SoftplusInfoSpmd
   kernel :
     func : softplus
   backward : softplus_grad
@@ -4788,6 +4796,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [x]
+    spmd_rule : SoftshrinkInfoSpmd
   kernel :
     func : softshrink
   backward : softshrink_grad
@@ -4945,6 +4954,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [x]
+    spmd_rule : StanhInfoSpmd
   kernel :
     func : stanh
   backward : stanh_grad
@@ -5143,6 +5153,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [x]
+    spmd_rule : ThresholdedReluInfoSpmd
   kernel :
     func : thresholded_relu
   inplace: (x -> out)
diff --git a/test/cpp/auto_parallel/spmd_rule_test.cc b/test/cpp/auto_parallel/spmd_rule_test.cc
