[Auto Parallel] Add spmd rule No.12 for label_smooth and label_smooth ops.

paddle/phi/infermeta/spmd_rules/label_smooth.cc

@@ -0,0 +1,85 @@
+/* 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/label_smooth.h"
+#include "glog/logging.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#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/elementwise.h"
+#include "paddle/phi/infermeta/spmd_rules/spmd_rule_macro_define.h"
+#include "paddle/phi/infermeta/spmd_rules/utils.h"
+
+namespace phi {
+namespace distributed {
+
+std::vector<int64_t> ReplicateWithExcludeAxes(
+    const std::vector<int64_t>& x, const std::vector<int64_t>& exclude_axes) {
+  std::vector<int64_t> res(x);
+  for (size_t i = 0; i < x.size(); ++i) {
+    if (res[i] == -1) {
+      continue;
+    }
+
+    if (std::find(exclude_axes.begin(), exclude_axes.end(), res[i]) !=
+        exclude_axes.end()) {
+      res[i] = -1;
+    }
+  }
+  return res;
+}
+
+SpmdInfo LabelSmoothInferSpmd(const DistMetaTensor& label,
+                              const DistMetaTensor& prior_dist,
+                              float epsilon) {
+  if (prior_dist.initialized()) {
+    const auto& label_dims_attr = label.dist_attr();
+    const auto& prior_dims_attr = prior_dist.dist_attr();
+    const auto& label_dims_mapping = label_dims_attr.dims_mapping();
+    const auto& prior_dims_mapping = prior_dims_attr.dims_mapping();
+    size_t label_rank = label.dims().size();
+    size_t prior_rank = prior_dist.dims().size();
+
+    if (label_rank > prior_rank) {
+      size_t extra_dims_count = label_rank - prior_rank;
+      std::vector<std::int64_t> extra_dims_mapping_values;
+      extra_dims_mapping_values.reserve(extra_dims_count);
+      for (size_t i = 0; i < extra_dims_count; ++i) {
+        extra_dims_mapping_values.push_back(label_dims_mapping[i]);
+      }
+      const auto& new_prior_dims_mapping = ReplicateWithExcludeAxes(
+          prior_dims_mapping, extra_dims_mapping_values);
+      TensorDistAttr new_prior_dist_attr = prior_dims_attr;
+      new_prior_dist_attr.set_dims_mapping(new_prior_dims_mapping);
+      DistMetaTensor modified_prior_dist(prior_dist.dims(),
+                                         new_prior_dist_attr);
+      VLOG(4) << "LabelSmoothInferSpmd call ElementwiseBinaryInferSpmd:";
+      return ElementwiseBinaryInferSpmd(label, modified_prior_dist);
+    }
+    VLOG(4) << "LabelSmoothInferSpmd call ElementwiseBinaryInferSpmd:";
+    return ElementwiseBinaryInferSpmd(label, prior_dist);
+  }
+  VLOG(4) << "LabelSmoothInferSpmd call ElementwiseUnaryInferSpmd:";
+  SpmdInfo unary_spmd_info = ElementwiseUnaryInferSpmd(label);
+  unary_spmd_info.first.push_back(TensorDistAttr());
+  return unary_spmd_info;
+}
+
+SpmdInfo LabelSmoothGradInferSpmd(const DistMetaTensor& out_grad,
+                                  float epsilon) {
+  VLOG(4) << "LabelSmoothGradInferSpmd call ElementwiseUnaryGradInferSpmd:";
+  return ElementwiseUnaryGradInferSpmd(out_grad);
+}
+}  // namespace distributed
+}  // namespace phi

paddle/phi/infermeta/spmd_rules/label_smooth.h

@@ -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. */
+
+#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 LabelSmoothInferSpmd(const DistMetaTensor& label,
+                              const DistMetaTensor& prior_dist,
+                              float epsilon);
+
+SpmdInfo LabelSmoothGradInferSpmd(const DistMetaTensor& out_grad,
+                                  float epsilon);
+
+}  // namespace distributed
+}  // namespace phi

paddle/phi/infermeta/spmd_rules/rules.cc

@@ -218,8 +218,8 @@ PD_REGISTER_SPMD_RULE(
     PD_INFER_SPMD(phi::distributed::ElementwiseUnaryInferSpmdReverse));
 PD_REGISTER_SPMD_RULE(
     label_smooth,
-    PD_INFER_SPMD(phi::distributed::ElementwiseUnaryInferSpmd),
-    PD_INFER_SPMD(phi::distributed::ElementwiseUnaryInferSpmdReverse));
+    PD_INFER_SPMD(phi::distributed::LabelSmoothInferSpmd),
+    PD_INFER_SPMD(phi::distributed::LabelSmoothGradInferSpmd));
 PD_REGISTER_SPMD_RULE(
     leaky_relu,
     PD_INFER_SPMD(phi::distributed::ElementwiseUnaryInferSpmd),

paddle/phi/infermeta/spmd_rules/rules.h

@@ -46,6 +46,7 @@ limitations under the License. */
 #include "paddle/phi/infermeta/spmd_rules/gather.h"
 #include "paddle/phi/infermeta/spmd_rules/gather_nd.h"
 #include "paddle/phi/infermeta/spmd_rules/index_select.h"
+#include "paddle/phi/infermeta/spmd_rules/label_smooth.h"
 #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"

paddle/phi/ops/yaml/backward.yaml

@@ -1859,6 +1859,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [out_grad]
+    spmd_rule : LabelSmoothGradInferSpmd
   kernel :
     func : label_smooth_grad
 

paddle/phi/ops/yaml/ops.yaml

@@ -2995,6 +2995,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [label]
+    spmd_rule : LabelSmoothInferSpmd
   kernel :
     func : label_smooth
     data_type : label

test/auto_parallel/spmd_rules/CMakeLists.txt

@@ -54,6 +54,7 @@ if(WITH_DISTRIBUTE)
     py_test_modules(test_cummax_rule MODULES test_cummax_rule)
     py_test_modules(test_cummin_rule MODULES test_cummin_rule)
     py_test_modules(test_argsort_rule MODULES test_argsort_rule)
+    py_test_modules(test_label_smooth_rule MODULES test_label_smooth_rule)
     py_test_modules(test_index_select_rule MODULES test_index_select_rule)
     py_test_modules(test_put_along_axis_rule MODULES test_put_along_axis_rule)
   endif()

test/auto_parallel/spmd_rules/test_label_smooth_rule.py

@@ -0,0 +1,216 @@
+# 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.
+
+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 core
+
+
+class TestLabelSmoothSPMDRule(unittest.TestCase):
+    def setUp(self):
+        label_shape = [16, 16, 16]
+        out_shape = [16, 16, 16]
+        prior_dist_shape = [1, 16]
+        process_mesh = auto.ProcessMesh(mesh=[[0, 1], [2, 3]])
+
+        label_tensor_dist_attr = TensorDistAttr()
+        label_tensor_dist_attr.dims_mapping = [-1, -1, -1]
+        label_tensor_dist_attr.process_mesh = process_mesh
+
+        self.label_dist_tensor_spec = DistTensorSpec(
+            label_shape, label_tensor_dist_attr
+        )
+        prior_dist_tensor_dist_attr = TensorDistAttr()
+        prior_dist_tensor_dist_attr.dims_mapping = [-1, -1]
+        prior_dist_tensor_dist_attr.process_mesh = process_mesh
+
+        self.prior_dist_dist_tensor_spec = DistTensorSpec(
+            prior_dist_shape, prior_dist_tensor_dist_attr
+        )
+
+        out_tensor_dist_attr = TensorDistAttr()
+        out_tensor_dist_attr.dims_mapping = [-1, -1, -1]
+        out_tensor_dist_attr.process_mesh = process_mesh
+        self.out_dist_tensor_spec = DistTensorSpec(
+            out_shape, out_tensor_dist_attr
+        )
+
+        self.rule = core.get_phi_spmd_rule("label_smooth")
+        self.attrs = OrderedDict()
+        self.attrs['epsilon'] = 0.1
+
+    def test_label_smooth_forward(self):
+        # [0, 1, -1], [Fake] --> [0, 1, -1], [Fake], [0, 1, -1]
+        self.label_dist_tensor_spec.set_dims_mapping([0, 1, -1])
+
+        result_dist_attrs = self.rule.infer_forward(
+            self.label_dist_tensor_spec,
+            DistTensorSpec(),
+            self.attrs['epsilon'],
+        )
+
+        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, [0, 1, -1])
+        self.assertEqual(inferred_input_dist_attrs[1], TensorDistAttr())
+        self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, 1, -1])
+
+        # shape: [16, 16, 16], [1, 16]. [0, 1, -1], [-1, -1] --> [0, 1, -1], [-1, -1], [0, 1, -1]
+        self.prior_dist_dist_tensor_spec.set_dims_mapping([-1, -1])
+        result_dist_attrs = self.rule.infer_forward(
+            self.label_dist_tensor_spec,
+            self.prior_dist_dist_tensor_spec,
+            self.attrs['epsilon'],
+        )
+
+        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, [0, 1, -1])
+        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, -1])
+        self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, 1, -1])
+
+        # shape: [16, 16, 16], [1, 16]. [0, 1, -1], [-1, 1] --> [0, 1, -1], [-1, -1], [0, 1, -1]
+        self.prior_dist_dist_tensor_spec.set_dims_mapping([-1, 1])
+        result_dist_attrs = self.rule.infer_forward(
+            self.label_dist_tensor_spec,
+            self.prior_dist_dist_tensor_spec,
+            self.attrs['epsilon'],
+        )
+        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, [0, 1, -1])
+        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, -1])
+        self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, 1, -1])
+
+        # shape: [16, 16, 16], [1, 16]. [0, 1, -1], [-1, 0] --> [0, 1, -1], [-1, -1], [0, 1, -1]
+        self.prior_dist_dist_tensor_spec.set_dims_mapping([-1, 0])
+        result_dist_attrs = self.rule.infer_forward(
+            self.label_dist_tensor_spec,
+            self.prior_dist_dist_tensor_spec,
+            self.attrs['epsilon'],
+        )
+        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, [0, 1, -1])
+        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, -1])
+        self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, 1, -1])
+
+        # shape: [16, 16, 16], [1, 16]. [0, -1, 1], [-1, 1] --> [0, -1, 1], [-1, 1], [0, -1, 1]
+        self.label_dist_tensor_spec.set_dims_mapping([0, -1, 1])
+        self.prior_dist_dist_tensor_spec.set_dims_mapping([-1, 1])
+        result_dist_attrs = self.rule.infer_forward(
+            self.label_dist_tensor_spec,
+            self.prior_dist_dist_tensor_spec,
+            self.attrs['epsilon'],
+        )
+        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, [0, -1, 1])
+        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, 1])
+        self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, -1, 1])
+
+        # shape: [16, 16, 16], [1, 16]. [0, -1, 1], [-1, 0] --> [0, -1, 1], [-1, 1], [0, -1, 1]
+        self.prior_dist_dist_tensor_spec.set_dims_mapping([-1, 0])
+        result_dist_attrs = self.rule.infer_forward(
+            self.label_dist_tensor_spec,
+            self.prior_dist_dist_tensor_spec,
+            self.attrs['epsilon'],
+        )
+        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, [0, -1, 1])
+        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, 1])
+        self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, -1, 1])
+
+        # shape: [16, 16, 16], [1, 16]. [0, -1, 1], [-1, -1] --> [0, -1, 1], [-1, 1], [0, -1, 1]
+        self.prior_dist_dist_tensor_spec.set_dims_mapping([-1, -1])
+        result_dist_attrs = self.rule.infer_forward(
+            self.label_dist_tensor_spec,
+            self.prior_dist_dist_tensor_spec,
+            self.attrs['epsilon'],
+        )
+        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, [0, -1, 1])
+        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [-1, 1])
+        self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, -1, 1])
+
+        # shape: [16, 16, 16], [16, 16, 16]. [0, 1, -1], [0, -1, 1] --> [0, 1, -1], [0, 1, -1], [0, 1, -1]
+        self.label_dist_tensor_spec.set_dims_mapping([0, 1, -1])
+        self.prior_dist_dist_tensor_spec.shape = [16, 16, 16]
+        self.prior_dist_dist_tensor_spec.set_dims_mapping([0, -1, 1])
+        result_dist_attrs = self.rule.infer_forward(
+            self.label_dist_tensor_spec,
+            self.prior_dist_dist_tensor_spec,
+            self.attrs['epsilon'],
+        )
+        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, [0, 1, -1])
+        self.assertEqual(inferred_input_dist_attrs[1].dims_mapping, [0, 1, -1])
+        self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, 1, -1])
+
+    def test_label_smooth_backward(self):
+        # shape: [16, 16 ,16]. [0, -1, 1] --> [0, -1, 1], [0, -1, 1]
+        self.out_dist_tensor_spec.set_dims_mapping([0, -1, 1])
+        result_dist_attrs = self.rule.infer_backward(
+            self.out_dist_tensor_spec,
+            self.attrs['epsilon'],
+        )
+
+        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, 1])
+        self.assertEqual(inferred_output_dist_attrs[0].dims_mapping, [0, -1, 1])
+
+
+if __name__ == "__main__":
+    unittest.main()