fix elementwise inferspmd (#59707)

LiYuRio · web-flow · commit 53b59191646e · 2023-12-07T13:05:24.000+08:00
diff --git a/paddle/phi/infermeta/spmd_rules/elementwise.cc b/paddle/phi/infermeta/spmd_rules/elementwise.cc
@@ -321,6 +321,37 @@ SpmdInfo ElementwiseUnaryGradInferSpmd(const DistMetaTensor& x,
           {out_grad.dist_attr()}};
 }
 
+bool DimsNotEqualOrHasBroadcastDim(const DistMetaTensor& x,
+                                   const DistMetaTensor& out) {
+  if (x.dims() != out.dims()) {
+    return true;
+  }
+
+  // Now the dims of x must equal to out.
+  const auto& out_dims_mapping = out.dist_attr().dims_mapping();
+  for (int64_t i = x.dims().size(); i >= 0; --i) {
+    if ((x.dims()[i] == 1) && (out_dims_mapping[i] != -1)) {
+      return true;
+    }
+  }
+  return false;
+}
+
+std::vector<int64_t> GetExplicitReduceDim(const DistMetaTensor& x,
+                                          const DistMetaTensor& out) {
+  std::vector<int64_t> reduce_dims;
+  const auto& out_dims_mapping = out.dist_attr().dims_mapping();
+  int64_t diff = out.dims().size() - x.dims().size();
+
+  for (int64_t i = x.dims().size(); i >= 0; --i) {
+    if ((x.dims()[i] == 1) && (out_dims_mapping[i + diff] != -1)) {
+      reduce_dims.emplace_back(i);
+    }
+  }
+
+  return reduce_dims;
+}
+
 SpmdInfo ElementwiseBinaryGradInferSpmd(const DistMetaTensor& x,
                                         const DistMetaTensor& y,
                                         const DistMetaTensor& out_grad,
@@ -349,43 +380,80 @@ SpmdInfo ElementwiseBinaryGradInferSpmd(const DistMetaTensor& x,
                                    "the rank of input as [%d].",
                                    out_grad.dims().size(),
                                    y.dims().size()));
-
   // The backward rule of elementwise follows the princple: the dist_attr
   // of input should equal to out_grad.
   // Caution the special case when the inputs calculate together with different
   // shape it means one of the input is broadcast to same shape with the other
   // first. When doing backward the input_grad with broadcast input is in
   // partial status, which need to do communicate and get the right result.
-  if (x.dims() != out_grad.dims()) {
+  if (DimsNotEqualOrHasBroadcastDim(x, out_grad)) {
+    VLOG(3) << "We need to do some special operations with the dist attr of "
+               "input x. "
+            << "The global dim of input x is " << x.dims()
+            << ". The global dim of out_grad is " << out_grad.dims();
+    // Step 1: remove the useless dimensions which is not appear in input x.
     int64_t diff = out_grad.dims().size() - x.dims().size();
     auto dims_mapping = x_dist_attr.dims_mapping();
     dims_mapping.erase(dims_mapping.begin(), dims_mapping.begin() + diff);
+    // Step 2: get the explicit reduce dimensions
+    std::vector<int64_t> explicit_reduce_dims =
+        GetExplicitReduceDim(x, out_grad);
+    VLOG(4) << "The explicit reduce dims has " << explicit_reduce_dims.size()
+            << " elements.";
+    for (const auto& dim : explicit_reduce_dims) {
+      VLOG(4) << "Explicit reduce dims is " << dim;
+      dims_mapping[dim] = -1;
+    }
     x_dist_attr.set_dims_mapping(dims_mapping);
     x_dist_attr.set_default_dynamic_dims(dims_mapping);
     x_grad_dist_attr.set_dims_mapping(dims_mapping);
     x_grad_dist_attr.set_default_dynamic_dims(dims_mapping);
+    // Step 3: set partial dimension
     for (int64_t i = 0; i < diff; ++i) {
       if (out_grad.dist_attr().dims_mapping()[i] != -1) {
         x_grad_dist_attr.set_partial_status(
             std::vector<int64_t>{out_grad.dist_attr().dims_mapping()[i]});
       }
     }
+    for (const auto& dim : explicit_reduce_dims) {
+      x_grad_dist_attr.set_partial_status(std::vector<int64_t>{
+          out_grad.dist_attr().dims_mapping()[diff + dim]});
+    }
   }
 
-  if (y.dims() != out_grad.dims()) {
+  if (DimsNotEqualOrHasBroadcastDim(y, out_grad)) {
+    VLOG(3) << "We need to do some special operations with the dist attr of "
+               "input y. "
+            << "The global dim of input y is " << y.dims()
+            << ". The global dim of out_grad is " << out_grad.dims();
+    // Step 1: remove the useless dimensions which is not appear in input y.
     int64_t diff = out_grad.dims().size() - y.dims().size();
     auto dims_mapping = y_dist_attr.dims_mapping();
     dims_mapping.erase(dims_mapping.begin(), dims_mapping.begin() + diff);
+    // Step 2: get the explicit reduce dimensions
+    std::vector<int64_t> explicit_reduce_dims =
+        GetExplicitReduceDim(y, out_grad);
+    VLOG(4) << "The explicit reduce dims has " << explicit_reduce_dims.size()
+            << " elements.";
+    for (const auto& dim : explicit_reduce_dims) {
+      VLOG(4) << "Explicit reduce dims is " << dim;
+      dims_mapping[dim] = -1;
+    }
     y_dist_attr.set_dims_mapping(dims_mapping);
     y_dist_attr.set_default_dynamic_dims(dims_mapping);
     y_grad_dist_attr.set_dims_mapping(dims_mapping);
     y_grad_dist_attr.set_default_dynamic_dims(dims_mapping);
+    // Step 3: set partial dimension
     for (int64_t i = 0; i < diff; ++i) {
       if (out_grad.dist_attr().dims_mapping()[i] != -1) {
         y_grad_dist_attr.set_partial_status(
             std::vector<int64_t>{out_grad.dist_attr().dims_mapping()[i]});
       }
     }
+    for (const auto& dim : explicit_reduce_dims) {
+      y_grad_dist_attr.set_partial_status(std::vector<int64_t>{
+          out_grad.dist_attr().dims_mapping()[diff + dim]});
+    }
   }
 
   return {{x_dist_attr, y_dist_attr, out_grad.dist_attr()},
diff --git a/test/auto_parallel/semi_auto_parallel_for_elementwise.py b/test/auto_parallel/semi_auto_parallel_for_elementwise.py
@@ -28,7 +28,7 @@ def __init__(self):
         self._seed = eval(os.getenv("seed"))
         self._mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
         self._rtol = 1e-6
-        self._atol = 0.0
+        self._atol = 1e-6
         paddle.seed(self._seed)
         np.random.seed(self._seed)
 
@@ -138,6 +138,19 @@ def test_add_x_y_shard_broadcast(self):
             binary_func=paddle.add,
         )
 
+    def test_add_broadcast_with_shard(self):
+        if self._backend == "cpu":
+            return
+
+        self.test_binary_body(
+            x_shape=[16, 4, 32],
+            y_shape=[16, 1, 32],
+            out_shape=[16, 4, 32],
+            x_placements=[dist.Shard(1)],
+            y_placements=[dist.Replicate()],
+            binary_func=paddle.add,
+        )
+
     def test_sub_x_y_shard_broadcast(self):
         if self._backend == "cpu":
             return
@@ -472,6 +485,7 @@ def run_test_case(self):
         self.test_add_x_shard_broadcast()
         self.test_add_x_y_shard()
         self.test_add_x_y_shard_broadcast()
+        self.test_add_broadcast_with_shard()
         self.test_sub_x_shard()
         self.test_sub_x_y_shard_broadcast()
         self.test_square_x_shard()
