diff --git a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py b/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py index 13c656207f1b8..8328e406ae0e6 100644 --- a/paddle/fluid/pir/dialect/op_generator/ops_api_gen.py +++ b/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', ] diff --git a/paddle/fluid/pir/dialect/operator/ir/ops.yaml b/paddle/fluid/pir/dialect/operator/ir/ops.yaml index d32c1d8b7a6bd..4fcd90c99fe0a 100644 --- a/paddle/fluid/pir/dialect/operator/ir/ops.yaml +++ b/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) diff --git a/paddle/fluid/pir/dialect/operator/utils/utils.cc b/paddle/fluid/pir/dialect/operator/utils/utils.cc index ea8002c1c842f..b4bad427567b7 100644 --- a/paddle/fluid/pir/dialect/operator/utils/utils.cc +++ b/paddle/fluid/pir/dialect/operator/utils/utils.cc @@ -80,7 +80,8 @@ const std::unordered_set LegacyOpList = { paddle::onednn::dialect::MultiGruOp::name(), #endif CReduceMinOp::name(), - PushSparseV2Op::name()}; + PushSparseV2Op::name(), + PartialSendOp::name()}; enum class AttrType { UNDEFINED = 0, diff --git a/paddle/fluid/primitive/codegen/gen.py b/paddle/fluid/primitive/codegen/gen.py index 3c6791a344a8b..fb1579968423a 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 diff --git a/paddle/phi/api/yaml/op_compat.yaml b/paddle/phi/api/yaml/op_compat.yaml index 237724dabe69f..53e0cea953b87 100755 --- a/paddle/phi/api/yaml/op_compat.yaml +++ b/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 diff --git a/paddle/phi/infermeta/unary.cc b/paddle/phi/infermeta/unary.cc index 611b5239dccdf..a742ad2a83d6e 100644 --- a/paddle/phi/infermeta/unary.cc +++ b/paddle/phi/infermeta/unary.cc @@ -2917,6 +2917,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