【PIR Dist Op Reg No.7、23】 reg c_reduce_prod c_reduce_max (#62270)

enkilee · web-flow · commit 476403b570fd · 2024-03-04T11:17:09.000+08:00
* fix

* fix

* add reduce_max
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -158,8 +158,12 @@
     'soft_relu',
     'uniform_random_batch_size_like',
     'match_matrix_tensor',
+    'c_reduce_max',
+    'c_reduce_max_',
     'c_reduce_min',
     'c_reduce_min_',
+    'c_reduce_prod',
+    'c_reduce_prod_',
     'push_sparse_v2',
     'push_sparse_v2_',
     'partial_send',
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
@@ -218,6 +218,16 @@
     func : c_identity
   inplace : (x -> out)
 
+- op : c_reduce_max
+  args : (Tensor x, int ring_id, int root_id, bool use_calc_stream)
+  output : Tensor(out)
+  infer_meta :
+    func : DistReduceInferMeta
+    param : [x]
+  kernel :
+    func : c_reduce_max
+  inplace : (x -> out)
+
 - op : c_reduce_min
   args : (Tensor x, int ring_id, int root_id, bool use_calc_stream)
   output : Tensor(out)
@@ -228,6 +238,16 @@
     func : c_reduce_min
   inplace : (x -> out)
 
+- op : c_reduce_prod
+  args : (Tensor x, int ring_id, int root_id, bool use_calc_stream)
+  output : Tensor(out)
+  infer_meta :
+    func : DistReduceInferMeta
+    param : [x]
+  kernel :
+    func : c_reduce_prod
+  inplace : (x -> out)
+
 - op : c_reduce_sum
   args : (Tensor x, int ring_id, int root_id, bool use_calc_stream)
   output : Tensor(out)
diff --git a/paddle/fluid/pir/dialect/operator/utils/utils.cc b/paddle/fluid/pir/dialect/operator/utils/utils.cc
@@ -86,7 +86,9 @@ const std::unordered_set<std::string> LegacyOpList = {
     paddle::onednn::dialect::MultiGruOp::name(),
     paddle::onednn::dialect::FusionLstmOp::name(),
 #endif
+    CReduceMaxOp::name(),
     CReduceMinOp::name(),
+    CReduceProdOp::name(),
     PushSparseV2Op::name(),
     PartialSendOp::name()};
 
diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml
@@ -3549,12 +3549,24 @@
   outputs :
     out: Out
 
+- op: c_reduce_max
+  inputs :
+    x : X
+  outputs :
+    out: Out
+
 - op: c_reduce_min
   inputs :
     x : X
   outputs :
     out: Out
 
+- op: c_reduce_prod
+  inputs :
+    x : X
+  outputs :
+    out: Out
+
 - op: c_reduce_sum
   inputs :
     x : X
diff --git a/test/ir/pir/translator/CMakeLists.txt b/test/ir/pir/translator/CMakeLists.txt
@@ -10,6 +10,8 @@ list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_c_allreduce_prod_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST
      test_distributed_lookup_table_translate)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_partial_send_translator)
+list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_c_reduce_max_translator)
+list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_c_reduce_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_reduce_max_translator.py b/test/ir/pir/translator/test_c_reduce_max_translator.py
@@ -0,0 +1,42 @@
+# 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 TestCReduceMaxOpTranslator(test_op_translator.TestOpTranslator):
+    def append_op(self):
+        self.op_type = "c_reduce_max"
+        x = paddle.ones(shape=(100, 2, 3), dtype='float32')
+        y = paddle.ones(shape=(100, 2, 3), dtype='float32')
+        attrs = {'ring_id': 0, 'root_id': 0, 'use_calc_stream': 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()
diff --git a/test/ir/pir/translator/test_c_reduce_prod_translator.py b/test/ir/pir/translator/test_c_reduce_prod_translator.py
@@ -0,0 +1,42 @@
+# 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 TestCReduceProdOpTranslator(test_op_translator.TestOpTranslator):
+    def append_op(self):
+        self.op_type = "c_reduce_prod"
+        x = paddle.ones(shape=(100, 2, 3), dtype='float32')
+        y = paddle.ones(shape=(100, 2, 3), dtype='float32')
+        attrs = {'ring_id': 0, 'root_id': 0, 'use_calc_stream': 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()
