Elementwise like rule (#58563)

liuzhenhai93 · web-flow · commit 7ddd29d430b3 · 2023-11-28T10:17:50.000+08:00
* polish

* polish
diff --git a/paddle/phi/api/yaml/backward.yaml b/paddle/phi/api/yaml/backward.yaml
@@ -1719,6 +1719,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param: [x]
+    spmd_rule: PowGradInferSpmd
   kernel :
     func : pow_grad
     data_type : out_grad
diff --git a/paddle/phi/api/yaml/legacy_ops.yaml b/paddle/phi/api/yaml/legacy_ops.yaml
@@ -507,6 +507,7 @@
   infer_meta :
     func : CreateLikeInferMeta
     param : [x, dtype]
+    spmd_rule : FullLikeInferSpmd
   kernel :
     func : full_like
     param : [x, value, dtype]
diff --git a/paddle/phi/api/yaml/ops.yaml b/paddle/phi/api/yaml/ops.yaml
@@ -1999,6 +1999,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param: [x]
+    spmd_rule: PowInferSpmd
   kernel :
     func : pow
     data_type : x
@@ -2209,6 +2210,7 @@
   infer_meta :
     func : UnchangedInferMeta
     param : [x]
+    spmd_rule : ScaleInferSpmd
   kernel :
     func : scale {dense -> dense},
            scale_sr {selected_rows -> selected_rows}
diff --git a/paddle/phi/infermeta/spmd_rules/full_like.cc b/paddle/phi/infermeta/spmd_rules/full_like.cc
@@ -0,0 +1,26 @@
+/* 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. */
+
+#include "paddle/phi/infermeta/spmd_rules/full_like.h"
+#include "paddle/phi/infermeta/spmd_rules/elementwise.h"
+
+namespace phi {
+namespace distributed {
+SpmdInfo FullLikeInferSpmd(const DistMetaTensor& x,
+                           const Scalar& y,
+                           phi::DataType dtype) {
+  return ElementwiseUnaryInferSpmd(x);
+}
+}  // namespace distributed
+}  // namespace phi
diff --git a/paddle/phi/infermeta/spmd_rules/full_like.h b/paddle/phi/infermeta/spmd_rules/full_like.h
@@ -0,0 +1,29 @@
+/* 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. */
+
+#pragma once
+
+#include <vector>
+
+#include "paddle/phi/common/scalar.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h"
+#include "paddle/phi/core/distributed/type_defs.h"
+
+namespace phi {
+namespace distributed {
+SpmdInfo FullLikeInferSpmd(const DistMetaTensor& x,
+                           const Scalar& y,
+                           phi::DataType dtype);
+}
+}  // namespace phi
diff --git a/paddle/phi/infermeta/spmd_rules/pow.cc b/paddle/phi/infermeta/spmd_rules/pow.cc
@@ -0,0 +1,27 @@
+/* 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. */
+
+#include "paddle/phi/infermeta/spmd_rules/pow.h"
+#include "paddle/phi/infermeta/spmd_rules/elementwise.h"
+
+namespace phi {
+namespace distributed {
+SpmdInfo PowInferSpmd(const DistMetaTensor& x, const Scalar& y) {
+  return ElementwiseUnaryInferSpmd(x);
+}
+SpmdInfo PowGradInferSpmd(const DistMetaTensor& x,
+                          const DistMetaTensor& out_grad,
+                          const Scalar y) {
+  return ElementwiseUnaryGradInferSpmd(x, out_grad);
+}
+}  // namespace distributed
+}  // namespace phi
diff --git a/paddle/phi/infermeta/spmd_rules/pow.h b/paddle/phi/infermeta/spmd_rules/pow.h
@@ -0,0 +1,27 @@
+/* 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. */
+
+#pragma once
+
+#include <vector>
+
+#include "paddle/phi/common/scalar.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h"
+#include "paddle/phi/core/distributed/type_defs.h"
+
+namespace phi {
+namespace distributed {
+SpmdInfo PowInferSpmd(const DistMetaTensor& x, const Scalar& y);
+SpmdInfo PowGradInferSpmd(const DistMetaTensor& x,
+                          const DistMetaTensor& out_grad,
+                          const Scalar y);
+}  // namespace distributed
+}  // namespace phi
diff --git a/paddle/phi/infermeta/spmd_rules/rules.h b/paddle/phi/infermeta/spmd_rules/rules.h
@@ -13,23 +13,24 @@ See the License for the specific language governing permissions and
 limitations under the License. */
 
 #pragma once
-
 #include "paddle/phi/core/distributed/auto_parallel/inferspmd_utils.h"
-
 #include "paddle/phi/infermeta/spmd_rules/cast.h"
 #include "paddle/phi/infermeta/spmd_rules/concat.h"
 #include "paddle/phi/infermeta/spmd_rules/default_data_parallel.h"
 #include "paddle/phi/infermeta/spmd_rules/elementwise.h"
 #include "paddle/phi/infermeta/spmd_rules/embedding.h"
 #include "paddle/phi/infermeta/spmd_rules/flash_attention.h"
 #include "paddle/phi/infermeta/spmd_rules/flatten.h"
+#include "paddle/phi/infermeta/spmd_rules/full_like.h"
 #include "paddle/phi/infermeta/spmd_rules/layer_norm.h"
 #include "paddle/phi/infermeta/spmd_rules/matmul.h"
 #include "paddle/phi/infermeta/spmd_rules/numel.h"
 #include "paddle/phi/infermeta/spmd_rules/optimizer.h"
+#include "paddle/phi/infermeta/spmd_rules/pow.h"
 #include "paddle/phi/infermeta/spmd_rules/reduction.h"
 #include "paddle/phi/infermeta/spmd_rules/replicated.h"
 #include "paddle/phi/infermeta/spmd_rules/reshape.h"
+#include "paddle/phi/infermeta/spmd_rules/scale.h"
 #include "paddle/phi/infermeta/spmd_rules/slice.h"
 #include "paddle/phi/infermeta/spmd_rules/softmax.h"
 #include "paddle/phi/infermeta/spmd_rules/split.h"
diff --git a/paddle/phi/infermeta/spmd_rules/scale.cc b/paddle/phi/infermeta/spmd_rules/scale.cc
@@ -0,0 +1,24 @@
+/* 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. */
+
+#include "paddle/phi/infermeta/spmd_rules/scale.h"
+#include "paddle/phi/infermeta/spmd_rules/elementwise.h"
+
+namespace phi {
+namespace distributed {
+SpmdInfo ScaleInferSpmd(const DistMetaTensor& x,
+                        const Scalar& scale,
+                        float bias,
+                        bool bias_after_scale) {
+  return ElementwiseUnaryInferSpmd(x);
+}
+}  // namespace distributed
+}  // namespace phi
diff --git a/paddle/phi/infermeta/spmd_rules/scale.h b/paddle/phi/infermeta/spmd_rules/scale.h
@@ -0,0 +1,30 @@
+// 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.
+
+#pragma once
+
+#include <vector>
+
+#include "paddle/phi/common/scalar.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_meta_tensor.h"
+#include "paddle/phi/core/distributed/type_defs.h"
+
+namespace phi {
+namespace distributed {
+SpmdInfo ScaleInferSpmd(const DistMetaTensor& x,
+                        const Scalar& scale,
+                        float bias,
+                        bool bias_after_scale);
+}
+}  // namespace phi
diff --git a/test/auto_parallel/semi_auto_parallel_for_unary_elementwise_like.py b/test/auto_parallel/semi_auto_parallel_for_unary_elementwise_like.py
@@ -0,0 +1,90 @@
+# 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.
+
+
+from semi_auto_parallel_util import SemiAutoParallelTestBase
+
+import paddle
+import paddle.distributed as dist
+
+
+class TestElementwiseLikeApiForSemiAutoParallel(SemiAutoParallelTestBase):
+    def __init__(self):
+        super().__init__()
+
+    def check_specs_unchanged(self, input, output):
+        input_dist_attr = input.dist_attr
+        output_dist_attr = output.dist_attr
+        assert input_dist_attr.dims_mapping == output_dist_attr.dims_mapping
+
+    def test_pow_shard(self):
+        x_shape = [16, 32]
+        x_specs = ['x', None]
+        inputs, outputs = self.runfunc_and_check(
+            x_shape, x_specs, op_func=paddle.pow, with_backward=True, y=2
+        )
+        self.check_specs_unchanged(inputs, outputs)
+
+    def test_cast_shard(self):
+        x_shape = [16, 32]
+        x_specs = ['x', None]
+        inputs, outputs = self.runfunc_and_check(
+            x_shape,
+            x_specs,
+            op_func=paddle.cast,
+            with_backward=True,
+            dtype="float64",
+        )
+        self.check_specs_unchanged(inputs, outputs)
+
+    def test_full_like_shard(self):
+        x_shape = [16, 32]
+        x_specs = ['x', None]
+        inputs, outputs = self.runfunc_and_check(
+            x_shape,
+            x_specs,
+            op_func=paddle.full_like,
+            with_backward=False,
+            fill_value=1.0,
+        )
+        self.check_specs_unchanged(inputs, outputs)
+
+    def test_scale_shard(self):
+        x_shape = [16, 32]
+        x_specs = ['x', None]
+        inputs, outputs = self.runfunc_and_check(
+            x_shape,
+            x_specs,
+            op_func=paddle.scale,
+            with_backward=True,
+            scale=2.0,
+        )
+        self.check_specs_unchanged(inputs, outputs)
+
+    def run_test_case(self):
+        if self._backend == "cpu":
+            paddle.set_device("cpu")
+        elif self._backend == "gpu":
+            paddle.set_device("gpu:" + str(dist.get_rank()))
+        else:
+            raise ValueError("Only support cpu or gpu backend.")
+
+        self.test_pow_shard()
+        self.test_full_like_shard()
+        self.test_cast_shard()
+        self.test_scale_shard()
+
+
+if __name__ == '__main__':
+    TestElementwiseLikeApiForSemiAutoParallel().run_test_case()
diff --git a/test/auto_parallel/test_semi_auto_parallel_basic.py b/test/auto_parallel/test_semi_auto_parallel_basic.py
@@ -168,6 +168,16 @@ def test_transpose_api(self):
                 user_defined_envs=envs,
             )
 
+    def test_unary_elementwise_like_api(self):
+        envs_list = test_base.gen_product_envs_list(
+            self._default_envs, self._changeable_envs
+        )
+        for envs in envs_list:
+            self.run_test_case(
+                "semi_auto_parallel_for_unary_elementwise_like.py",
+                user_defined_envs=envs,
+            )
+
 
 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
@@ -1277,6 +1277,56 @@ TEST(Reshape, Ctor) {
   check_dim_mapping(spmd_grad.second[0], {0, 1, -1});
 }
 
+TEST(ElementwiseUnaryLike, 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);
+
+  std::vector<int64_t> shape = {16, 16, 16};
+  std::vector<int64_t> dims_mapping = {0, -1, 1};
+
+  auto t_dist_attr = TensorDistAttr();
+  t_dist_attr.set_process_mesh(process_mesh);
+  t_dist_attr.set_dims_mapping(dims_mapping);
+  t_dist_attr.set_dynamic_dims({false, false, false});
+
+  auto check_element_unary_like = [&dims_mapping](auto& spmd_info) {
+    EXPECT_EQ(spmd_info.first.size(), static_cast<size_t>(1));
+    EXPECT_EQ(spmd_info.second.size(), static_cast<size_t>(1));
+    check_dim_mapping(spmd_info.first[0], dims_mapping);
+    check_dim_mapping(spmd_info.second[0], dims_mapping);
+    check_partial_dims(spmd_info.second[0], {});
+  };
+
+  // cast
+  auto input =
+      phi::distributed::DistMetaTensor(phi::make_ddim(shape), t_dist_attr);
+  auto infered_dist_attrs =
+      phi::distributed::CastInferSpmd(input, phi::DataType::FLOAT32);
+
+  check_element_unary_like(infered_dist_attrs);
+  // full like
+  input = phi::distributed::DistMetaTensor(phi::make_ddim(shape), t_dist_attr);
+  infered_dist_attrs =
+      phi::distributed::FullLikeInferSpmd(input, 1.0, phi::DataType::FLOAT32);
+  check_element_unary_like(infered_dist_attrs);
+
+  // pow
+  input = phi::distributed::DistMetaTensor(phi::make_ddim(shape), t_dist_attr);
+  infered_dist_attrs = phi::distributed::PowInferSpmd(input, 2);
+  check_element_unary_like(infered_dist_attrs);
+
+  // pow backward
+  input = phi::distributed::DistMetaTensor(phi::make_ddim(shape), t_dist_attr);
+  infered_dist_attrs = phi::distributed::PowGradInferSpmd(input, input, 2);
+
+  // scale
+  input = phi::distributed::DistMetaTensor(phi::make_ddim(shape), t_dist_attr);
+  infered_dist_attrs = phi::distributed::ScaleInferSpmd(input, 1.0, 1.0, false);
+  check_element_unary_like(infered_dist_attrs);
+}
+
 }  // namespace auto_parallel
 }  // namespace distributed
 }  // namespace paddle
