【PIR Dist Op Reg No.21】 reg c_allreduce_prod (#60790)

enkilee · web-flow · commit 1d988829db85 · 2024-01-16T19:34:19.000+08:00
* fix

* fix
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -114,6 +114,8 @@
     'c_allgather',
     'c_allreduce_max',
     'c_allreduce_sum',
+    'c_allreduce_prod',
+    'c_allreduce_prod_',
     'c_embedding',
     'c_identity',
     'c_reduce_sum',
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
@@ -164,6 +164,16 @@
     func : c_allreduce_max
   inplace : (x -> out)
 
+- op : c_allreduce_prod
+  args : (Tensor x, int ring_id, bool use_calc_stream, bool use_model_parallel)
+  output : Tensor(out)
+  infer_meta :
+    func : AllReduceInferMeta
+    param : [x]
+  kernel :
+    func : c_allreduce_prod
+  inplace : (x -> out)
+
 - op : c_allreduce_sum
   args : (Tensor x, int ring_id, bool use_calc_stream, bool use_model_parallel)
   output : Tensor(out)
diff --git a/paddle/fluid/pir/dialect/operator/utils/utils.cc b/paddle/fluid/pir/dialect/operator/utils/utils.cc
@@ -49,6 +49,7 @@ const std::unordered_set<std::string> LegacyOpList = {
     DpsgdOp::name(),
     SendV2Op::name(),
     RecvV2Op::name(),
+    CAllreduceProd_Op::name(),
     CAllreduceSumOp::name(),
     CAllreduceSum_Op::name(),
     CReduceSumOp::name(),
diff --git a/paddle/phi/api/yaml/legacy_ops.yaml b/paddle/phi/api/yaml/legacy_ops.yaml
@@ -144,6 +144,16 @@
     func : c_allreduce_max
   inplace : (x -> out)
 
+- op : c_allreduce_prod
+  args : (Tensor x, int ring_id, bool use_calc_stream, bool use_model_parallel)
+  output : Tensor(out)
+  infer_meta :
+    func : AllReduceInferMeta
+    param : [x]
+  kernel :
+    func : c_allreduce_prod
+  inplace : (x -> out)
+
 - op : c_allreduce_sum
   args : (Tensor x, int ring_id, bool use_calc_stream, bool use_model_parallel)
   output : Tensor(out)
diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml
@@ -3397,6 +3397,12 @@
   outputs :
     out: Out
 
+- op: c_allreduce_prod
+  inputs :
+    x : X
+  outputs :
+    out: Out
+
 - op: c_allreduce_sum
   inputs :
     x : X
diff --git a/test/ir/pir/translator/CMakeLists.txt b/test/ir/pir/translator/CMakeLists.txt
@@ -5,6 +5,7 @@ file(
 string(REPLACE ".py" "" TEST_INTERP_CASES "${TEST_INTERP_CASES}")
 
 set(DISTRIBUTED_OP_TRANSLATOR_TEST test_c_reduce_min_translator)
+list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_c_allreduce_prod_translator)
 
 if(NOT WITH_DISTRIBUTE)
   list(REMOVE_ITEM TEST_INTERP_CASES ${DISTRIBUTED_OP_TRANSLATOR_TEST})
diff --git a/test/ir/pir/translator/test_c_allreduce_prod_translator.py b/test/ir/pir/translator/test_c_allreduce_prod_translator.py
@@ -0,0 +1,46 @@
+# Copyright (c) 2024 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
+
+import test_op_translator
+
+import paddle
+from paddle.base.layer_helper import LayerHelper
+
+
+class TestCAllReduceProdOpTranslator(test_op_translator.TestOpTranslator):
+    def append_op(self):
+        self.op_type = "c_allreduce_prod"
+        x = paddle.ones(shape=(100, 2, 3), dtype='float32')
+        y = paddle.ones(shape=(100, 2, 3), dtype='float32')
+        attrs = {
+            'ring_id': 0,
+            'use_calc_stream': False,
+            'use_model_parallel': False,
+        }
+        helper = LayerHelper(self.op_type)
+        helper.append_op(
+            type=self.op_type,
+            inputs={"X": x},
+            outputs={"Out": y},
+            attrs=attrs,
+        )
+
+    def test_translator(self):
+        self.check()
+
+
+if __name__ == "__main__":
+    unittest.main()
