From 2117ed774a486079f4ce122e3e7a273606ee6982 Mon Sep 17 00:00:00 2001
From: Difers <707065510@qq.com>
Date: Thu, 22 Feb 2024 19:15:10 +0800
Subject: [PATCH] add partitial_send to prim gen blacklist

---
 paddle/fluid/pir/dialect/operator/ir/ops.yaml | 3 ++-
 paddle/fluid/primitive/codegen/gen.py         | 1 +
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
index 866fb0690b3bef..4fcd90c99fe0ab 100644
--- a/paddle/fluid/pir/dialect/operator/ir/ops.yaml
+++ b/paddle/fluid/pir/dialect/operator/ir/ops.yaml
@@ -1656,12 +1656,13 @@
 
 - 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: Tensor(out)
+  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)
diff --git a/paddle/fluid/primitive/codegen/gen.py b/paddle/fluid/primitive/codegen/gen.py
index 3c6791a344a8ba..fb1579968423ac 100644
--- a/paddle/fluid/primitive/codegen/gen.py
+++ b/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