【PIR Dist Op Reg No.13】 reg partial_send (#60484)

* reg partial_send

* fix cmake

* fix import

* fix cmake

* try to remove out

* add partitial_send to prim gen blacklist

* fix typo

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

@@ -156,6 +156,7 @@
     'c_reduce_min_',
     'push_sparse_v2',
     'push_sparse_v2_',
+    'partial_send',
 ]
 
 

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

@@ -1654,6 +1654,16 @@
   kernel:
     func: onednn_to_paddle_layout
 
+- op: partial_send
+  args: (Tensor x, int ring_id = 0, int peer = 0, bool use_calc_stream = false, int num = 1, int id = 0)
+  output :
+  infer_meta:
+    func: PartialSendInferMeta
+    param: [x, ring_id, peer, use_calc_stream, num, id]
+  kernel:
+    func: partial_send
+    param: [x, ring_id, peer, use_calc_stream, num, id]
+
 - op: sparse_momentum
   args: (Tensor param, Tensor grad, Tensor velocity, Tensor index, Tensor learning_rate, Tensor master_param,float mu, Scalar axis=0, bool use_nesterov=false,str regularization_method="", float regularization_coeff=0.0f, bool multi_precision=false, float rescale_grad=1.0f)
   output: Tensor(param_out), Tensor(velocity_out), Tensor(master_param_out)

paddle/fluid/pir/dialect/operator/utils/utils.cc

@@ -80,7 +80,8 @@ const std::unordered_set<std::string> LegacyOpList = {
     paddle::onednn::dialect::MultiGruOp::name(),
 #endif
     CReduceMinOp::name(),
-    PushSparseV2Op::name()};
+    PushSparseV2Op::name(),
+    PartialSendOp::name()};
 
 enum class AttrType {
   UNDEFINED = 0,

paddle/fluid/primitive/codegen/gen.py

@@ -52,6 +52,7 @@
     "embedding_sparse_grad",
     "embedding_grad",
     "full",
+    "partial_send",
 ]
 
 # prim op with one input and one output, with no attribute

paddle/phi/api/yaml/op_compat.yaml

@@ -3653,6 +3653,10 @@
   outputs :
     out : Out
 
+- op: partial_send
+  inputs :
+    x : X
+
 - op: read_from_array
   inputs:
     array : X

paddle/phi/infermeta/unary.cc

@@ -2915,6 +2915,34 @@ void Pad3dInferMeta(const MetaTensor& x,
   out->share_lod(x);
 }
 
+void PartialSendInferMeta(const MetaTensor& x,
+                          int ring_id,
+                          int peer,
+                          bool use_calc_stream,
+                          int num,
+                          int id) {
+  PADDLE_ENFORCE_GE(
+      peer,
+      0,
+      phi::errors::InvalidArgument(
+          "The peer (%d) for partial_send op must be non-negative.", peer));
+  PADDLE_ENFORCE_GE(
+      ring_id,
+      0,
+      phi::errors::InvalidArgument(
+          "The ring_id (%d) for partial_send op must be non-negative.",
+          ring_id));
+  PADDLE_ENFORCE_GE(num,
+                    1,
+                    phi::errors::InvalidArgument(
+                        "The num (%d) for partial_send op must >=1", num));
+  PADDLE_ENFORCE_EQ(
+      (id >= 0 && id < num),
+      true,
+      phi::errors::InvalidArgument(
+          "The id (%d) for partial_send op must >=0 and <num (%d)", id, num));
+}
+
 void PixelShuffleInferMeta(const MetaTensor& x,
                            int upscale_factor,
                            const std::string& data_format,

paddle/phi/infermeta/unary.h

@@ -434,6 +434,13 @@ void Pad3dInferMeta(const MetaTensor& x,
                     MetaTensor* out,
                     MetaConfig config = MetaConfig());
 
+void PartialSendInferMeta(const MetaTensor& x,
+                          int ring_id,
+                          int peer,
+                          bool use_calc_stream,
+                          int num,
+                          int id);
+
 void PixelShuffleInferMeta(const MetaTensor& x,
                            int upscale_factor,
                            const std::string& data_format,

test/ir/pir/translator/CMakeLists.txt

@@ -9,6 +9,7 @@ list(APPEND DISTRIBUTED_OP_TRANSLATOR_TEST test_c_allreduce_min_translator)
 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)
 
 if(NOT WITH_DISTRIBUTE)
   list(REMOVE_ITEM TEST_INTERP_CASES ${DISTRIBUTED_OP_TRANSLATOR_TEST})

test/ir/pir/translator/test_partial_send_translator.py

@@ -0,0 +1,40 @@
+# 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
+#
+#     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 TestPartialSendTranslator(test_op_translator.TestOpTranslator):
+    def append_op(self):
+        self.op_type = "partial_send"
+        x = 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},
+            attrs=attrs,
+        )
+
+    def test_translator(self):
+        self.check()
+
+
+if __name__ == "__main__":
+    unittest.main()