feat(pir): reg send_and_recv (#64203)

xiaoyewww · web-flow · commit 32a4fe767b0c · 2024-05-11T16:15:50.000+08:00
diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
@@ -215,6 +215,8 @@
     'pull_box_sparse_',
     'push_box_sparse',
     'push_box_sparse_',
+    'send_and_recv',
+    'send_and_recv_',
 ]
 
 
diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
@@ -1563,6 +1563,17 @@
   kernel:
     func: seed
 
+- op : send_and_recv
+  args : (Tensor[] x, str message_name, str[] send_var_name, str[] recv_var_name, int trainer_id = 0, str mode="forward", str[] endpoints={"127.0.0.1:6164"}, str[] next_endpoints={"127.0.0.1:6164"}, str[] previous_endpoints={"127.0.0.1:6164"})
+  output : Tensor[](out){x.size()}
+  infer_meta :
+    func : UnchangedVectorInferMeta
+    param : [x]
+  kernel :
+    func : send_and_recv
+    param : [x, message_name, mode, send_var_name, recv_var_name, trainer_id, endpoints, next_endpoints, previous_endpoints]
+  inplace: (x -> out)
+
 - op : send_v2
   args : (Tensor x, int ring_id = 0, int peer = 0, bool use_calc_stream = false, bool dynamic_shape = false)
   output :
diff --git a/paddle/fluid/pir/dialect/operator/utils/utils.cc b/paddle/fluid/pir/dialect/operator/utils/utils.cc
@@ -105,7 +105,8 @@ const std::unordered_set<std::string> LegacyOpList = {
     PushBoxSparseOp::name(),
     PushSparseV2Op::name(),
     PartialSendOp::name(),
-    PartialRecvOp::name()};
+    PartialRecvOp::name(),
+    SendAndRecvOp::name()};
 
 enum class AttrType {
   UNDEFINED = 0,
diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml
@@ -3053,6 +3053,12 @@
   outputs :
     out : Out
 
+- op : send_and_recv
+  inputs :
+    x: X
+  outputs :
+    out : Out
+
 - op : send_u_recv(graph_send_recv)
   backward : send_u_recv_grad(graph_send_recv_grad)
   inputs :
diff --git a/test/deprecated/ir/pir/translator/CMakeLists.txt b/test/deprecated/ir/pir/translator/CMakeLists.txt
@@ -31,6 +31,7 @@ list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_pull_gpups_sparse_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_pull_sparse_v2_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_random_routing_translator)
+list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_send_and_recv_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_limit_by_capacity_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_global_scatter_translator)
 list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_global_gather_translator)
diff --git a/test/deprecated/ir/pir/translator/test_send_and_recv_translator.py b/test/deprecated/ir/pir/translator/test_send_and_recv_translator.py
@@ -0,0 +1,48 @@
+# 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 TestCReduceMinOpTranslator(test_op_translator.TestOpTranslator):
+    def append_op(self):
+        self.op_type = "send_and_recv"
+        x = paddle.ones(shape=(100, 2, 3), dtype='float32')
+        y = paddle.ones(shape=(100, 2, 3), dtype='float32')
+
+        attrs = {
+            'trainer_id': 0,
+            'message_name': "x",
+            'send_var_name': ["x"],
+            'recv_var_name': ["x"],
+        }
+        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()

Original file line number	Diff line number	Diff line change
`@@ -215,6 +215,8 @@`
`215`	`215`	`'pull_box_sparse_',`
`216`	`216`	`'push_box_sparse',`
`217`	`217`	`'push_box_sparse_',`
	`218`	`+ 'send_and_recv',`
	`219`	`+ 'send_and_recv_',`
`218`	`220`	`]`
`219`	`221`
`220`	`222`