[Auto Parallel] Add spmd rule No.6 for mean_all and mean_all_grad ops. (PaddlePaddle#72479)

* add spmd rule for pd_op.mean_all

* Fix test config for shape and dims_mapping

* refine and add cpp tests

* reuse ReductionInfer to MeanAll

Author: ooooo-create

paddle/phi/infermeta/spmd_rules/mean_all.cc

@@ -0,0 +1,31 @@
+/* Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+    http://www.apache.org/licenses/LICENSE-2.0
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License. */
+
+#include "paddle/phi/infermeta/spmd_rules/mean_all.h"
+#include "glog/logging.h"
+
+namespace phi {
+namespace distributed {
+
+SpmdInfo MeanAllInferSpmd(const DistMetaTensor& x) {
+  VLOG(4) << "MeanAllInferSpmd Call ReductionInferSpmdBase";
+  return ReductionInferSpmdBase(
+      x, {}, false, static_cast<int>(ReduceType::kRedAvg));
+}
+
+SpmdInfo MeanAllGradInferSpmd(const DistMetaTensor& x,
+                              const DistMetaTensor& out_grad) {
+  VLOG(4) << "MeanAllGradInferSpmd Call ReductionGradInferSpmd";
+  return ReductionGradInferSpmd(x, out_grad, {}, false, true);
+}
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/infermeta/spmd_rules/mean_all.h

@@ -0,0 +1,29 @@
+// Copyright (c) 2025 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include "paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h"
+#include "paddle/phi/core/distributed/type_defs.h"
+#include "paddle/phi/infermeta/spmd_rules/reduction.h"
+
+namespace phi {
+namespace distributed {
+SpmdInfo MeanAllInferSpmd(const DistMetaTensor& x);
+
+SpmdInfo MeanAllGradInferSpmd(const DistMetaTensor& x,
+                              const DistMetaTensor& out_grad);
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/infermeta/spmd_rules/rules.cc

@@ -521,6 +521,11 @@ PD_REGISTER_SPMD_RULE(
     PD_INFER_SPMD(phi::distributed::ReductionInferSpmd),
     PD_INFER_SPMD(phi::distributed::ReductionInferSpmdReverse));
 
+// mean_all
+PD_REGISTER_SPMD_RULE(mean_all,
+                      PD_INFER_SPMD(phi::distributed::MeanAllInferSpmd),
+                      PD_INFER_SPMD(phi::distributed::MeanAllGradInferSpmd));
+
 // layer_norm
 PD_REGISTER_SPMD_RULE(
     layer_norm,

paddle/phi/infermeta/spmd_rules/rules.h

@@ -45,6 +45,7 @@ limitations under the License. */
 #include "paddle/phi/infermeta/spmd_rules/layer_norm.h"
 #include "paddle/phi/infermeta/spmd_rules/logsumexp.h"
 #include "paddle/phi/infermeta/spmd_rules/matmul.h"
+#include "paddle/phi/infermeta/spmd_rules/mean_all.h"
 #include "paddle/phi/infermeta/spmd_rules/moe_combine.h"
 #include "paddle/phi/infermeta/spmd_rules/moe_gate_dispatch.h"
 #include "paddle/phi/infermeta/spmd_rules/nonzero.h"

paddle/phi/ops/yaml/backward.yaml

@@ -2173,6 +2173,7 @@
   infer_meta :
     func : UnchangedExceptLayoutInferMeta
     param: [x]
+    spmd_rule : MeanAllGradInferSpmd
   kernel :
     func : mean_all_grad
     data_type: out_grad

paddle/phi/ops/yaml/ops.yaml

@@ -3451,6 +3451,7 @@
   output : Tensor
   infer_meta :
     func : MeanAllInferMeta
+    spmd_rule : MeanAllInferSpmd
   kernel :
     func : mean_all
   backward : mean_all_grad

test/auto_parallel/spmd_rules/CMakeLists.txt

@@ -45,6 +45,7 @@ if(WITH_DISTRIBUTE)
   py_test_modules(test_nonzero_rule MODULES test_nonzero_rule)
   if(NOT WITH_ROCM)
     py_test_modules(test_add_n_rule MODULES test_add_n_rule)
+    py_test_modules(test_mean_all_rule MODULES test_mean_all_rule)
     py_test_modules(test_argmin_rule MODULES test_argmin_rule)
   endif()
   # End of unittests WITH single card WITHOUT timeout

test/auto_parallel/spmd_rules/test_mean_all_rule.py

@@ -0,0 +1,83 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+
+from paddle.distributed.auto_parallel.static.dist_attribute import (
+    DistTensorSpec,
+    TensorDistAttr,
+)
+from paddle.distributed.fleet import auto
+from paddle.framework import core
+
+
+class TestMeanAllSPMDRule(unittest.TestCase):
+    def setUp(self):
+        self.rule = core.get_phi_spmd_rule("mean_all")
+        x_shape = [4, 8]
+        out_shape = []
+        process_mesh = auto.ProcessMesh(mesh=[[0, 1], [2, 3]])
+
+        x_tensor_dist_attr = TensorDistAttr()
+        x_tensor_dist_attr.dims_mapping = [1, 0]
+        x_tensor_dist_attr.process_mesh = process_mesh
+        self.x_dist_tensor_spec = DistTensorSpec(x_shape, x_tensor_dist_attr)
+
+        out_tensor_dist_attr = TensorDistAttr()
+        out_tensor_dist_attr.dims_mapping = []
+        out_tensor_dist_attr.process_mesh = process_mesh
+        self.out_dist_tensor_spec = DistTensorSpec(
+            out_shape, out_tensor_dist_attr
+        )
+
+    def test_infer_forward(self):
+        # [0, -1] --> [], partial_on_dim:[0]
+        self.x_dist_tensor_spec.set_dims_mapping([0, -1])
+        result_dist_attrs = self.rule.infer_forward(self.x_dist_tensor_spec)
+
+        self.assertEqual(len(result_dist_attrs), 2)
+        inferred_input_dist_attrs = result_dist_attrs[0]
+        inferred_output_dist_attrs = result_dist_attrs[1]
+
+        self.assertEqual(len(inferred_input_dist_attrs), 1)
+        self.assertEqual(len(inferred_output_dist_attrs), 1)
+
+        self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [0, -1])
+        self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [])
+        self.assertEqual(inferred_output_dist_attrs[0]._is_partial(), True)
+        self.assertEqual(inferred_output_dist_attrs[0]._partial_dims(), {0})
+
+    def test_infer_backward(self):
+        # [] --> [-1, -1], []
+        self.out_dist_tensor_spec.shape = []
+        self.out_dist_tensor_spec.set_dims_mapping([])
+        result_dist_attrs = self.rule.infer_backward(
+            self.x_dist_tensor_spec,
+            self.out_dist_tensor_spec,
+        )
+
+        self.assertEqual(len(result_dist_attrs), 2)
+        inferred_input_dist_attrs = result_dist_attrs[0]
+        inferred_output_dist_attrs = result_dist_attrs[1]
+
+        self.assertEqual(len(inferred_input_dist_attrs), 2)
+        self.assertEqual(len(inferred_output_dist_attrs), 1)
+
+        self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [-1, -1])
+        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [])
+        self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [-1, -1])
+
+
+if __name__ == "__main__":
+    unittest.main()

test/cpp/auto_parallel/spmd_rule_test.cc

@@ -2426,6 +2426,49 @@ TEST(Dropout, Ctor) {
   check_dim_mapping(backward_info.second[0], {0, -1, -1});
 }
 
+TEST(MeanAll, Ctor) {
+  std::vector<int64_t> mesh_shape = {2, 2};
+  std::vector<int64_t> process_ids = {0, 1, 2, 3};
+  std::vector<std::string> dim_names = {"x", "y"};
+  ProcessMesh process_mesh(mesh_shape, process_ids, dim_names);
+
+  // test forward
+  // [0, -1] --> [], partial_on_dim:[0]
+  auto t_dist_attr = TensorDistAttr();
+  t_dist_attr.set_process_mesh(process_mesh);
+  t_dist_attr.set_dims_mapping({0, -1});
+  t_dist_attr.set_dynamic_dims({false, false});
+  phi::distributed::DistMetaTensor x =
+      phi::distributed::DistMetaTensor(common::make_ddim({4, 8}), t_dist_attr);
+  phi::distributed::SpmdInfo forward_info =
+      phi::distributed::MeanAllInferSpmd(x);
+
+  EXPECT_EQ(forward_info.first.size(), 1UL);
+  EXPECT_EQ(forward_info.second.size(), 1UL);
+
+  check_dim_mapping(forward_info.first[0], {0, -1});
+  check_dim_mapping(forward_info.second[0], {});
+  check_partial_dims(forward_info.second[0], {0});
+
+  // test backward
+  // [] --> [-1, -1], []
+  auto out_grad_dist_attr = TensorDistAttr();
+  out_grad_dist_attr.set_process_mesh(process_mesh);
+  out_grad_dist_attr.set_dims_mapping({});
+  out_grad_dist_attr.set_dynamic_dims({});
+  phi::distributed::DistMetaTensor out_grad = phi::distributed::DistMetaTensor(
+      common::make_ddim({}), out_grad_dist_attr);
+  phi::distributed::SpmdInfo backward_info =
+      phi::distributed::MeanAllGradInferSpmd(x, out_grad);
+
+  EXPECT_EQ(backward_info.first.size(), 2UL);
+  EXPECT_EQ(backward_info.second.size(), 1UL);
+
+  check_dim_mapping(backward_info.first[0], {-1, -1});
+  check_dim_mapping(backward_info.first[1], {});
+  check_dim_mapping(backward_info.second[0], {-1, -1});
+}
+
 }  // namespace auto_parallel
 }  // namespace distributed
 }  // namespace paddle