[Auto Parallel] Add spmd_rule for five backward ops(round_grad, trunc_grad, ceil_grad, floor_grad, poisson_grad). (#72677)

Glencsa · web-flow · commit a15f21136247 · 2025-05-20T11:25:26.000+08:00
* add unary ops which have spmd_rule but not add in yaml file.

* add spmd_rule for backward ops which has one arg.

* fix bug.

* fix bug.

* re-run CI
diff --git a/paddle/phi/infermeta/spmd_rules/elementwise.cc b/paddle/phi/infermeta/spmd_rules/elementwise.cc
@@ -365,7 +365,11 @@ SpmdInfo ElementwiseBinaryInferSpmdReverse(const DistMetaTensor& x,
 
   return {{x_dist_attr_dst, y_dist_attr_dst}, {out_dist_attr_dst}};
 }
-
+SpmdInfo ElementwiseUnaryGradInferSpmd(const DistMetaTensor& out_grad) {
+  auto dist_attr = CopyTensorDistAttrForOutput(out_grad.dist_attr());
+  dist_attr.set_dims_mapping(out_grad.dist_attr().dims_mapping());
+  return {{dist_attr}, {dist_attr}};
+}
 SpmdInfo ElementwiseUnaryGradInferSpmd(const DistMetaTensor& x,
                                        const DistMetaTensor& out_grad) {
   auto dist_attr = CopyTensorDistAttrForOutput(out_grad.dist_attr());
diff --git a/paddle/phi/infermeta/spmd_rules/elementwise.h b/paddle/phi/infermeta/spmd_rules/elementwise.h
@@ -29,6 +29,8 @@ SpmdInfo ElementwiseUnaryWithPartialInferSpmd(const DistMetaTensor& x);
 SpmdInfo ElementwiseUnaryInferSpmdReverse(const DistMetaTensor& x,
                                           const DistMetaTensor& out);
 
+SpmdInfo ElementwiseUnaryGradInferSpmd(const DistMetaTensor& out_grad);
+
 SpmdInfo ElementwiseUnaryGradInferSpmd(const DistMetaTensor& x,
                                        const DistMetaTensor& out_grad);
 
diff --git a/paddle/phi/ops/yaml/backward.yaml b/paddle/phi/ops/yaml/backward.yaml
@@ -376,6 +376,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param: [out_grad]
+    spmd_rule : ElementwiseUnaryGradInferSpmd
   kernel :
     func : ceil_grad
   inplace : (out_grad -> x_grad)
@@ -1224,6 +1225,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param: [out_grad]
+    spmd_rule : ElementwiseUnaryGradInferSpmd
   kernel :
     func : floor_grad
   composite : floor_grad(out_grad, x_grad)
@@ -2441,6 +2443,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [out_grad]
+    spmd_rule : ElementwiseUnaryGradInferSpmd
   kernel :
     func : poisson_grad
 
@@ -2802,6 +2805,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param: [out_grad]
+    spmd_rule : ElementwiseUnaryGradInferSpmd
   kernel :
     func : round_grad
   inplace : (out_grad -> x_grad)
@@ -3628,6 +3632,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [out_grad]
+    spmd_rule : ElementwiseUnaryGradInferSpmd
   kernel :
     func : trunc_grad
 
diff --git a/test/cpp/auto_parallel/spmd_rule_test.cc b/test/cpp/auto_parallel/spmd_rule_test.cc
@@ -1478,6 +1478,7 @@ TEST(ElementwiseUnaryLike, Ctor) {
     check_dim_mapping(spmd_info.second[0], dims_mapping);
     check_partial_dims(spmd_info.second[0], {});
   };
+
   auto check_element_unary_like_backward = [&dims_mapping](auto& spmd_info) {
     EXPECT_GT(spmd_info.first.size(), static_cast<size_t>(1));
     EXPECT_EQ(spmd_info.second.size(), static_cast<size_t>(1));
@@ -1494,6 +1495,11 @@ TEST(ElementwiseUnaryLike, Ctor) {
   auto inferred_dist_attrs =
       phi::distributed::CastInferSpmd(input, phi::DataType::FLOAT32);
 
+  check_element_unary_like(inferred_dist_attrs);
+  // ElementwiseUnaryGradInferSpmd
+  input =
+      phi::distributed::DistMetaTensor(common::make_ddim(shape), t_dist_attr);
+  inferred_dist_attrs = phi::distributed::ElementwiseUnaryGradInferSpmd(input);
   check_element_unary_like(inferred_dist_attrs);
   // full like
   input =
