【PIR Dist Op Reg No.17】 reg barrier

paddle/fluid/pir/dialect/op_generator/ops_api_gen.py

@@ -122,6 +122,7 @@
     'add_n_',
     'all_reduce',
     'all_reduce_',
+    'barrier',
     'c_allgather',
     'c_allreduce_avg',
     'c_allreduce_avg_',

paddle/fluid/pir/dialect/operator/ir/ops.yaml

@@ -133,6 +133,12 @@
     data_type : dtype
     backend : place > output
 
+- op : barrier
+  args : (Tensor x, int ring_id=0)
+  output : Tensor(out)
+  kernel :
+    func : barrier
+
 - op : batch_norm
   args : (Tensor x, Tensor mean, Tensor variance, Tensor scale, Tensor bias, bool is_test, float momentum, float epsilon, str data_format, bool use_global_stats, bool trainable_statistics)
   output : Tensor(out), Tensor(mean_out), Tensor(variance_out), Tensor(saved_mean), Tensor(saved_variance), Tensor(reserve_space)

paddle/phi/api/yaml/op_compat.yaml

@@ -329,6 +329,12 @@
   outputs :
     {auc : AUC, stat_pos_out : StatPosOut, stat_neg_out : StatNegOut}
 
+- op : barrier
+  inputs :
+    {x : X}
+  outputs :
+    out : Out
+
 - op : batch_norm
   backward : batch_norm_grad, batch_norm_double_grad(batch_norm_grad_grad)
   inputs:

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_all_reduce_translator)
+list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_barrier_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_c_reduce_min_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_c_allreduce_min_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_c_allreduce_prod_translator)

test/ir/pir/translator/test_barrier_translator.py

@@ -0,0 +1,44 @@
+# 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
+#
+#     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
+
+import test_op_translator
+
+import paddle
+from paddle.base.layer_helper import LayerHelper
+
+
+class TestBarrierOpTranslator(test_op_translator.TestOpTranslator):
+    def append_op(self):
+        self.op_type = "barrier"
+        x = paddle.ones(shape=(100, 2, 3), dtype='float32')
+        y = paddle.ones(shape=(100, 2, 3), dtype='float32')
+        attrs = {
+            'ring_id': 0,
+        }
+        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()