[AutoParallel] Add reduce_any spmd rule (#73175)

ooooo-create · web-flow · commit ad43bc040c4b · 2025-06-10T14:20:46.000+08:00
* [AutoParallel] Add reduce_any spmd rule

* refine

* refine
diff --git a/paddle/phi/core/distributed/gloo_utils.h b/paddle/phi/core/distributed/gloo_utils.h
@@ -149,6 +149,13 @@ void SetReduceFunc(P* opts, int reduce_type) {
           static_cast<void (*)(void*, const void*, const void*, size_t)>(
               &gloo::min<T>));
       break;
+    case ReduceType::kRedAny:
+      // NOTE(ooooo): There is no reduce_any math function for gloo, just use
+      // max to replace
+      opts->setReduceFunction(
+          static_cast<void (*)(void*, const void*, const void*, size_t)>(
+              &gloo::max<T>));
+      break;
     default:
       PADDLE_THROW(
           errors::InvalidArgument("Unsupported reduce type: %d.", reduce_type));
diff --git a/paddle/phi/infermeta/spmd_rules/reduction.cc b/paddle/phi/infermeta/spmd_rules/reduction.cc
@@ -171,6 +171,12 @@ SpmdInfo ReductionAllInferSpmdDynamic(const DistMetaTensor& x,
   return ReductionInferSpmdBase(
       x, axis.GetData(), keep_dim, static_cast<int>(ReduceType::kRedAll));
 }
+SpmdInfo ReductionAnyInferSpmdDynamic(const DistMetaTensor& x,
+                                      const IntArray& axis,
+                                      bool keep_dim) {
+  return ReductionInferSpmdBase(
+      x, axis.GetData(), keep_dim, static_cast<int>(ReduceType::kRedAny));
+}
 
 SpmdInfo ReductionInferSpmdReverse(const DistMetaTensor& x,
                                    const DistMetaTensor& out,
diff --git a/paddle/phi/infermeta/spmd_rules/reduction.h b/paddle/phi/infermeta/spmd_rules/reduction.h
@@ -55,6 +55,9 @@ SpmdInfo ReductionMinInferSpmdDynamic(const DistMetaTensor& x,
 SpmdInfo ReductionAllInferSpmdDynamic(const DistMetaTensor& x,
                                       const IntArray& axis,
                                       bool keep_dim);
+SpmdInfo ReductionAnyInferSpmdDynamic(const DistMetaTensor& x,
+                                      const IntArray& axis,
+                                      bool keep_dim);
 
 SpmdInfo ReductionInferSpmdReverse(const DistMetaTensor& x,
                                    const DistMetaTensor& out,
diff --git a/paddle/phi/kernels/gpu/all_reduce_kernel.cu b/paddle/phi/kernels/gpu/all_reduce_kernel.cu
@@ -67,6 +67,11 @@ void AllReduceKernel(const Context& dev_ctx,
       // min to replace
       red_type = ncclMin;
       break;
+    case ReduceType::kRedAny:
+      // NOTE(ooooo): There is no reduce_any type of ncclRedOp_t, just use
+      // max to replace
+      red_type = ncclMax;
+      break;
   }
   comm_ctx->AllReduce(out, x, red_type, stream);
 #else
diff --git a/paddle/phi/ops/yaml/ops.yaml b/paddle/phi/ops/yaml/ops.yaml
@@ -273,6 +273,7 @@
   output : Tensor(out)
   infer_meta :
     func : ReduceInferMeta
+    spmd_rule : ReductionAnyInferSpmdDynamic
   kernel :
     func : any
   traits : paddle::dialect::ForwardOnlyTrait
diff --git a/test/auto_parallel/semi_auto_parallel_for_reduction.py b/test/auto_parallel/semi_auto_parallel_for_reduction.py
@@ -37,10 +37,10 @@ def test_body(
     ):
         paddle.seed(self._seed)
         np.random.seed(self._seed)
-        is_op_func_all = op_func == paddle.all
+        is_op_func_all_or_any = op_func == paddle.all or op_func == paddle.any
 
         x = paddle.randn(x_shape, self._dtype)
-        if is_op_func_all:
+        if is_op_func_all_or_any:
             x = x > 0
         x.stop_gradient = False
 
@@ -52,7 +52,7 @@ def test_body(
         self.check_tensor_eq(out, dist_out)
         np.testing.assert_equal(dist_out.shape, out_shape, verbose=True)
 
-        if not is_op_func_all:
+        if not is_op_func_all_or_any:
             dist_out.backward()
             out.backward()
             self.check_tensor_eq(x.grad, dist_x.grad)
@@ -137,6 +137,26 @@ def test_all_x_shard_on_axis(self):
             op_func=paddle.all,
         )
 
+    def test_any_x_shard(self):
+        self.test_body(
+            x_shape=[4, 8, 6],
+            out_shape=[4, 6],
+            x_placements=[dist.Shard(0)],
+            axis=1,
+            keepdim=False,
+            op_func=paddle.any,
+        )
+
+    def test_any_x_shard_on_axis(self):
+        self.test_body(
+            x_shape=[4, 8, 6],
+            out_shape=[4, 6],
+            x_placements=[dist.Shard(1)],
+            axis=1,
+            keepdim=False,
+            op_func=paddle.any,
+        )
+
     def run_test_case(self):
         if self._backend == "cpu":
             paddle.set_device("cpu")
@@ -153,6 +173,8 @@ def run_test_case(self):
         self.test_max_x_shard_on_axis()
         self.test_all_x_shard()
         self.test_all_x_shard_on_axis()
+        self.test_any_x_shard()
+        self.test_any_x_shard_on_axis()
 
 
 if __name__ == '__main__':
diff --git a/test/cpp/auto_parallel/spmd_rule_test.cc b/test/cpp/auto_parallel/spmd_rule_test.cc
@@ -1022,6 +1022,27 @@ TEST(ReduceAllRule, Ctor) {
   check_partial_dims(forward_info.second[0], {0});
 }
 
+TEST(ReduceAnyRule, Ctor) {
+  std::vector<int64_t> mesh_shape = {2};
+  std::vector<int64_t> process_ids = {0, 1};
+  std::vector<std::string> dim_names = {"x"};
+  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
+
+  // test forward
+  auto t_dist_attr = TensorDistAttr();
+  t_dist_attr.set_process_mesh(process_mesh);
+  t_dist_attr.set_dims_mapping({-1, 0, -1});
+  t_dist_attr.set_dynamic_dims({false, false, false});
+  phi::distributed::DistMetaTensor x =
+      phi::distributed::DistMetaTensor(phi::make_ddim({4, 6, 8}), t_dist_attr);
+  phi::IntArray axis = {1};
+  bool keep_dim = false;
+  phi::distributed::SpmdInfo forward_info =
+      phi::distributed::ReductionAnyInferSpmdDynamic(x, axis, keep_dim);
+  check_dim_mapping(forward_info.second[0], {-1, -1});
+  check_partial_dims(forward_info.second[0], {0});
+}
+
 TEST(Numel, Ctor) {
   std::vector<int64_t> mesh_shape = {2, 2};
   std::vector<int64_t> process_ids = {0, 1, 2, 3};
