[Auto Parallel] Add spmd rule No.6 for unique ops. (#72824)

ooooo-create · web-flow · commit 44de935c0241 · 2025-05-29T20:44:49.000+08:00
* add spmd_rule for unique

* update

* refine code

* apply review

* try
diff --git a/paddle/fluid/pybind/auto_parallel_py.cc b/paddle/fluid/pybind/auto_parallel_py.cc
@@ -944,6 +944,11 @@ static void parse_single_pyobject(PyObject *obj,
                                   phi::distributed::InferSpmdContext *ctx,
                                   const size_t arg_pos) {
   if (PyList_Check(obj)) {  // list inputs, spmd not allow tuple inputs
+    Py_ssize_t list_size = PyList_Size(obj);
+    if (list_size == 0) {
+      ctx->EmplaceBackAttr(std::vector<int64_t>());
+      return;
+    }
     PyObject *first_item = PyList_GetItem(obj, 0);
     if (PyObject_TypeCheck(first_item, g_dist_tensor_spec_pytype)) {
       parse_tensors(obj, ctx, arg_pos);
diff --git a/paddle/phi/core/distributed/auto_parallel/inferspmd_utils.cc b/paddle/phi/core/distributed/auto_parallel/inferspmd_utils.cc
@@ -92,6 +92,9 @@ std::vector<int> InferSpmdContext::AttrAt(size_t idx) const {
     if (attr.type() == typeid(std::vector<bool>)) {
       std::vector<bool> val = PADDLE_GET_CONST(std::vector<bool>, attr);
       return std::vector<int>(val.begin(), val.end());
+    } else if (attr.type() == typeid(std::vector<int64_t>) &&
+               paddle::get<std::vector<int64_t>>(attr).empty()) {
+      return std::vector<int>();
     } else {
       return paddle::get<std::vector<int>>(attr);
     }
diff --git a/paddle/phi/infermeta/spmd_rules/rules.cc b/paddle/phi/infermeta/spmd_rules/rules.cc
@@ -708,6 +708,9 @@ PD_REGISTER_SPMD_RULE(cumsum,
                       PD_INFER_SPMD(phi::distributed::CumSumInferSpmd),
                       PD_INFER_SPMD(phi::distributed::CumSumInferSpmdReverse));
 
+// unique
+PD_REGISTER_SPMD_RULE(unique, PD_INFER_SPMD(phi::distributed::UniqueInferSpmd));
+
 // argmin
 PD_REGISTER_SPMD_RULE(
     argmin,
diff --git a/paddle/phi/infermeta/spmd_rules/rules.h b/paddle/phi/infermeta/spmd_rules/rules.h
@@ -75,5 +75,6 @@ limitations under the License. */
 #include "paddle/phi/infermeta/spmd_rules/transpose.h"
 #include "paddle/phi/infermeta/spmd_rules/triu.h"
 #include "paddle/phi/infermeta/spmd_rules/unbind.h"
+#include "paddle/phi/infermeta/spmd_rules/unique.h"
 #include "paddle/phi/infermeta/spmd_rules/unsqueeze.h"
 #include "paddle/phi/infermeta/spmd_rules/where.h"
diff --git a/paddle/phi/infermeta/spmd_rules/unique.cc b/paddle/phi/infermeta/spmd_rules/unique.cc
@@ -0,0 +1,86 @@
+/* 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
+
+    https://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/unique.h"
+#include "glog/logging.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/infermeta/spmd_rules/spmd_rule_macro_define.h"
+#include "paddle/phi/infermeta/spmd_rules/utils.h"
+
+namespace phi {
+namespace distributed {
+
+SpmdInfo UniqueInferSpmd(const DistMetaTensor& x,
+                         bool return_index,
+                         bool return_inverse,
+                         bool return_counts,
+                         const std::vector<int>& axis,
+                         DataType dtype) {
+  // Verify input args
+  EXTRACT_SHAPE_AND_DIST_ATTR(x);
+  std::vector<int64_t> x_dims_mapping_dst(x_ndim, -1);
+  std::vector<int64_t> out_dims_mapping_dst(x_dims_mapping_dst);
+  TensorDistAttr x_dist_attr_dst = CopyTensorDistAttrForOutput(x_dist_attr_src);
+  x_dist_attr_dst.set_dims_mapping(x_dims_mapping_dst);
+
+  if (axis.empty()) {
+    out_dims_mapping_dst = {-1};
+  }
+  TensorDistAttr out_dist_attr_dst =
+      CopyTensorDistAttrForOutput(x_dist_attr_src);
+  out_dist_attr_dst.set_dims_mapping(out_dims_mapping_dst);
+
+  TensorDistAttr indices_dist_attr_dst = TensorDistAttr();
+  if (return_index) {
+    indices_dist_attr_dst = CopyTensorDistAttrForOutput(x_dist_attr_src);
+    indices_dist_attr_dst.set_dims_mapping({-1});
+  }
+
+  TensorDistAttr inverse_dist_attr_dst = TensorDistAttr();
+  if (return_inverse) {
+    inverse_dist_attr_dst = CopyTensorDistAttrForOutput(x_dist_attr_src);
+    inverse_dist_attr_dst.set_dims_mapping({-1});
+    // TODO(dev): https://github.com/PaddlePaddle/Paddle/issues/72822
+    // if (axis.empty()) {
+    //   inverse_dist_attr_dst.set_dims_mapping(x_dims_mapping_dst);
+    // }
+  }
+
+  TensorDistAttr counts_dist_attr_dst = TensorDistAttr();
+  if (return_counts) {
+    counts_dist_attr_dst = CopyTensorDistAttrForOutput(x_dist_attr_src);
+    counts_dist_attr_dst.set_dims_mapping({-1});
+  }
+
+  VLOG(4) << "UniqueInferSpmd: All input and output TensorDistAttr are set to "
+             "fully replicated status.";
+  return {{x_dist_attr_dst},
+          {out_dist_attr_dst,
+           indices_dist_attr_dst,
+           inverse_dist_attr_dst,
+           counts_dist_attr_dst}};
+}
+
+SpmdInfo UniqueInferSpmdStatic(const DistMetaTensor& x,
+                               bool return_index,
+                               bool return_inverse,
+                               bool return_counts,
+                               const std::vector<int>& axis,
+                               DataType dtype,
+                               bool is_sorted) {
+  return UniqueInferSpmd(
+      x, return_index, return_inverse, return_counts, axis, dtype);
+}
+}  // namespace distributed
+}  // namespace phi
diff --git a/paddle/phi/infermeta/spmd_rules/unique.h b/paddle/phi/infermeta/spmd_rules/unique.h
@@ -0,0 +1,38 @@
+/* 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
+
+    https://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"
+
+namespace phi {
+namespace distributed {
+
+SpmdInfo UniqueInferSpmd(const DistMetaTensor& x,
+                         bool return_index,
+                         bool return_inverse,
+                         bool return_counts,
+                         const std::vector<int>& axis,
+                         DataType dtype);
+
+SpmdInfo UniqueInferSpmdStatic(const DistMetaTensor& x,
+                               bool return_index,
+                               bool return_inverse,
+                               bool return_counts,
+                               const std::vector<int>& axis,
+                               DataType dtype,
+                               bool is_sorted);
+}  // namespace distributed
+}  // namespace phi
diff --git a/paddle/phi/ops/yaml/inconsistent/dygraph_ops.yaml b/paddle/phi/ops/yaml/inconsistent/dygraph_ops.yaml
@@ -381,6 +381,7 @@
   output : Tensor(out), Tensor(indices), Tensor(inverse), Tensor(counts)
   infer_meta :
     func : UniqueInferMeta
+    spmd_rule : UniqueInferSpmd
   kernel :
     func : unique
     data_type : x
diff --git a/paddle/phi/ops/yaml/inconsistent/static_ops.yaml b/paddle/phi/ops/yaml/inconsistent/static_ops.yaml
@@ -850,6 +850,7 @@
   optional : indices, counts
   infer_meta :
     func : UniqueRawInferMeta
+    spmd_rule : UniqueInferSpmdStatic
   kernel :
     func : unique
     data_type : x
diff --git a/test/auto_parallel/spmd_rules/CMakeLists.txt b/test/auto_parallel/spmd_rules/CMakeLists.txt
@@ -44,6 +44,7 @@ if(WITH_DISTRIBUTE)
   py_test_modules(test_logsumexp_rule MODULES test_logsumexp_rule)
   py_test_modules(test_nonzero_rule MODULES test_nonzero_rule)
   if(NOT WITH_ROCM)
+    py_test_modules(test_unique_rule MODULES test_unique_rule)
     py_test_modules(test_topk_rule MODULES test_topk_rule)
     py_test_modules(test_add_n_rule MODULES test_add_n_rule)
     py_test_modules(test_mean_all_rule MODULES test_mean_all_rule)
diff --git a/test/auto_parallel/spmd_rules/test_unique_rule.py b/test/auto_parallel/spmd_rules/test_unique_rule.py
@@ -0,0 +1,97 @@
+# 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
+#
+#     https://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 collections import OrderedDict
+
+from paddle.distributed.auto_parallel.static.dist_attribute import (
+    DistTensorSpec,
+    TensorDistAttr,
+)
+from paddle.distributed.fleet import auto
+from paddle.framework import convert_np_dtype_to_dtype_, core
+
+
+class TestUniqueSPMDRule(unittest.TestCase):
+    def setUp(self):
+        self.rule = core.get_phi_spmd_rule("unique")
+        x_shape = [4, 8]
+        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)
+        self.attrs = OrderedDict()
+        self.attrs["return_index"] = True
+        self.attrs["return_inverse"] = True
+        self.attrs["return_counts"] = True
+        self.attrs["axis"] = []
+        self.attrs['dtype'] = convert_np_dtype_to_dtype_("int32")
+
+    def test_infer_forward(self):
+        # return_index=True, return_inverse=True, return_counts=True, axis={}
+        # [0, -1] --> [-1,-1], [-1], [-1], [-1], [-1]
+        self.x_dist_tensor_spec.set_dims_mapping([0, -1])
+        result_dist_attrs = self.rule.infer_forward(
+            self.x_dist_tensor_spec,
+            self.attrs["return_index"],
+            self.attrs["return_inverse"],
+            self.attrs["return_counts"],
+            self.attrs["axis"],
+            self.attrs['dtype'],
+        )
+
+        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), 4)
+
+        self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [-1, -1])
+        self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [-1])
+        self.assertEqual(inferred_output_dist_attrs[1].dims_mapping, [-1])
+        self.assertEqual(inferred_output_dist_attrs[2].dims_mapping, [-1])
+        self.assertEqual(inferred_output_dist_attrs[3].dims_mapping, [-1])
+
+        # return_index=True, return_inverse=True, return_counts=True, axis={0}
+        # [0, -1] --> [-1,-1], [-1，-1], [-1], [-1], [-1]
+        self.x_dist_tensor_spec.set_dims_mapping([0, -1])
+        self.attrs["axis"] = [0]
+        result_dist_attrs = self.rule.infer_forward(
+            self.x_dist_tensor_spec,
+            self.attrs["return_index"],
+            self.attrs["return_inverse"],
+            self.attrs["return_counts"],
+            self.attrs["axis"],
+            self.attrs['dtype'],
+        )
+
+        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), 4)
+
+        self.assertEqual(inferred_input_dist_attrs[0].dims_mapping, [-1, -1])
+        self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [-1, -1])
+        self.assertEqual(inferred_output_dist_attrs[1].dims_mapping, [-1])
+        self.assertEqual(inferred_output_dist_attrs[2].dims_mapping, [-1])
+        self.assertEqual(inferred_output_dist_attrs[3].dims_mapping, [-1])
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/test/cpp/auto_parallel/spmd_rule_test.cc b/test/cpp/auto_parallel/spmd_rule_test.cc
@@ -2940,6 +2940,62 @@ TEST(ArgSortInferSpmd, Ctor) {
           << std::endl;
 }
 
+TEST(Unique, 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);
+  bool return_index = true;
+  bool return_inverse = true;
+  bool return_counts = true;
+  std::vector<int> axis = {};
+
+  // test forward
+  // return_index=True, return_inverse=True, return_counts=True, axis={}
+  // [0, -1] --> [-1,-1], [-1], [-1], [-1], [-1]
+  auto x_dist_attr = TensorDistAttr();
+  x_dist_attr.set_process_mesh(process_mesh);
+  x_dist_attr.set_dims_mapping({0, -1});
+  x_dist_attr.set_dynamic_dims({false, false});
+  phi::distributed::DistMetaTensor x =
+      phi::distributed::DistMetaTensor(common::make_ddim({4, 8}), x_dist_attr);
+
+  phi::distributed::SpmdInfo forward_info =
+      phi::distributed::UniqueInferSpmd(x,
+                                        return_index,
+                                        return_inverse,
+                                        return_counts,
+                                        axis,
+                                        phi::DataType::FLOAT32);
+
+  EXPECT_EQ(forward_info.first.size(), 1UL);
+  EXPECT_EQ(forward_info.second.size(), 4UL);
+
+  check_dim_mapping(forward_info.first[0], {-1, -1});
+  check_dim_mapping(forward_info.second[0], {-1});
+  check_dim_mapping(forward_info.second[1], {-1});
+  check_dim_mapping(forward_info.second[2], {-1});
+  check_dim_mapping(forward_info.second[3], {-1});
+
+  // return_index=True, return_inverse=True, return_counts=True, axis={0}
+  // [0, -1] --> [-1, -1], [-1, -1], [-1], [-1], [-1]
+  axis = {0};
+  forward_info = phi::distributed::UniqueInferSpmd(x,
+                                                   return_index,
+                                                   return_inverse,
+                                                   return_counts,
+                                                   axis,
+                                                   phi::DataType::FLOAT32);
+
+  EXPECT_EQ(forward_info.first.size(), 1UL);
+  EXPECT_EQ(forward_info.second.size(), 4UL);
+
+  check_dim_mapping(forward_info.first[0], {-1, -1});
+  check_dim_mapping(forward_info.second[0], {-1, -1});
+  check_dim_mapping(forward_info.second[1], {-1});
+  check_dim_mapping(forward_info.second[2], {-1});
+  check_dim_mapping(forward_info.second[3], {-1});
+}
 }  // namespace auto_parallel
 }  // namespace distributed
 }  // namespace paddle
