add partitial_send to prim gen blacklist

PaddlePaddle · Feb 22, 2024 · 2117ed7 · 2117ed7
1 parent 71a73f5
commit 2117ed7
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git 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
@@ -52,6 +52,7 @@
     "embedding_sparse_grad",
     "embedding_grad",
     "full",
+    "partial_send",
 ]
 
 # prim op with one input and one output, with no attribute