Fix

Author: co63oc

paddle/fluid/framework/new_executor/interpreter/execution_config.cc

@@ -27,8 +27,8 @@
 // FLAGS_force_sync_ops is used to finer control the op-sync in executor.
 // The format is: "micro_batch_id, job_name, op_id, op_name | micro_batch_id,
 // job_name, op_id, op_name | ...". Keep spaces to syncs all name/id. Example:
-// 1. sync the recv_v2 op in the second backward-job of 1F1B scheduling:
-// FLAGS_force_sync_ops="1, backward, , recv_v2"
+// 1. sync the p_recv op in the second backward-job of 1F1B scheduling:
+// FLAGS_force_sync_ops="1, backward, , p_recv"
 // 2. sync the full op with op_id=5: FLAGS_force_sync_ops=" , , 5, full"
 // 3. sync all ops in the first default-job: FLAGS_force_sync_ops="0,default,,
 // 4. sync all ops in the forward-job and backward-job: FLAGS_force_sync_ops=" ,

paddle/fluid/framework/new_executor/interpreter/interpreter_util.cc

@@ -142,6 +142,7 @@ bool IsCommunicationOp(const OperatorBase* op) {
       "recv",
       "send_v2",
       "recv_v2",
+      "p_recv",
   };
   const std::string communication_op_prefix = "c_";
   if (op_name.find(communication_op_prefix) != std::string::npos ||
@@ -170,6 +171,7 @@ bool IsCommunicationOp(const ::pir::Operation* op) {
   const std::set<std::string> special_comm_op_set = {
       paddle::dialect::SendV2Op::name(),
       paddle::dialect::RecvV2Op::name(),
+      paddle::dialect::PRecvOp::name(),
   };
   const std::string communication_op_prefix = "c_";
   if (op_name.find(communication_op_prefix) != std::string::npos ||

paddle/fluid/framework/new_executor/pir_interpreter.cc

@@ -551,6 +551,7 @@ void PirInterpreter::UpdateNcclOpNum() {
       "pd_op.partial_recv",
       "pd_op.partial_allgather",
       "pd_op.recv_v2",
+      "pd_op.p_recv",
       "pd_op.send_v2",
       "pd_op.mp_allreduce_sum",
       "pd_op.barrier",
@@ -586,7 +587,7 @@ void PirInterpreter::UpdateNcclOpNum() {
       "pd_op.partial_send_grad",
       "pd_op.partial_recv_grad",
       "pd_op.partial_allgather_grad",
-      "pd_op.recv_v2_grad",
+      "pd_op.p_recv_grad",
       "pd_op.send_v2_grad",
       "pd_op.mp_allreduce_sum_grad",
       "pd_op.barrier_grad",
@@ -625,7 +626,7 @@ void PirInterpreter::UpdateNcclOpNum() {
       "pd_op.partial_send_",
       "pd_op.partial_recv_",
       "pd_op.partial_allgather_",
-      "pd_op.recv_v2_",
+      "pd_op.p_recv_",
       "pd_op.send_v2_",
       "pd_op.mp_allreduce_sum_",
       "pd_op.barrier_",
@@ -661,7 +662,7 @@ void PirInterpreter::UpdateNcclOpNum() {
       "pd_op.partial_send_grad_",
       "pd_op.partial_recv_grad_",
       "pd_op.partial_allgather_grad_",
-      "pd_op.recv_v2_grad_",
+      "pd_op.p_recv_grad_",
       "pd_op.send_v2_grad_",
       "pd_op.mp_allreduce_sum_grad_",
       "pd_op.barrier_grad_",

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

@@ -140,6 +140,7 @@
     'coalesce_tensor_',
     'send_v2',
     'recv_v2',
+    'p_recv',
     'sequence_expand',
     'sequence_softmax',
     'qkv_unpack_mha',

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.cc

@@ -565,6 +565,58 @@ bool RecvV2OpInferSymbolicShape(pir::Operation *op,
   return true;
 }
 
+bool PRecvOpInferSymbolicShape(pir::Operation *op,
+                               pir::InferSymbolicShapeContext *infer_context) {
+  const int ring_id = op->attribute<pir::Int32Attribute>("ring_id").data();
+  const bool dynamic_shape =
+      op->attribute<pir::BoolAttribute>("dynamic_shape").data();
+  const int peer = op->attribute<pir::Int32Attribute>("peer").data();
+
+  PADDLE_ENFORCE_GE(
+      peer,
+      0,
+      common::errors::InvalidArgument(
+          "The peer (%d) for p_recv op must be non-negative.", peer));
+
+  PADDLE_ENFORCE_GE(
+      ring_id,
+      0,
+      common::errors::InvalidArgument(
+          "The ring_id (%d) for p_recv op must be non-negative.", ring_id));
+
+  const std::vector<int> out_shape =
+      paddle::dialect::details::GetVectorAttr<int>(op, "out_shape");
+  if (!dynamic_shape) {
+    PADDLE_ENFORCE_GE(out_shape.size(),
+                      1,
+                      common::errors::InvalidArgument(
+                          "The size of the output shape must be greater than 0 "
+                          "but the value given is %d.",
+                          out_shape.size()));
+
+    std::vector<symbol::DimExpr> output_shape;
+    for (size_t i = 0; i < out_shape.size(); ++i) {
+      PADDLE_ENFORCE_GE(out_shape[i],
+                        1,
+                        common::errors::InvalidArgument(
+                            "The shape attribute for p_recv must be set "
+                            "explicitly, but the %dth element is %d which "
+                            "is less than 1. Or dynamic_shape should be set to "
+                            "True for both send_v2 and p_recv.",
+                            i,
+                            out_shape[i]));
+      output_shape.push_back(symbol::DimExpr(out_shape[i]));
+    }
+
+    infer_context->SetShapeOrDataForValue(
+        op->result(0),
+        symbol::ShapeOrDataDimExprs{
+            symbol::TensorShapeOrDataDimExprs(output_shape)});
+  }
+
+  return true;
+}
+
 bool SeedOpInferSymbolicShape(pir::Operation *op,
                               pir::InferSymbolicShapeContext *infer_context) {
   std::vector<symbol::DimExpr> dims = {symbol::DimExpr(1)};

paddle/fluid/pir/dialect/operator/interface/infer_symbolic_shape/nullary_infer_sym.h

@@ -34,6 +34,7 @@ OP_DECLARE_INFER_SYMBOLIC_SHAPE(Randperm)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(ReadFile)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(Seed)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(RecvV2)
+OP_DECLARE_INFER_SYMBOLIC_SHAPE(PRecv)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TrilIndices)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TriuIndices)
 OP_DECLARE_INFER_SYMBOLIC_SHAPE(TruncatedGaussianRandom)

paddle/fluid/pybind/eager_generator.cc

@@ -3395,6 +3395,7 @@ std::map<std::string, std::set<std::string>> op_passing_outs_map = {
     {"accuracy", {"Correct", "Total"}},
     {"fill_constant", {"Out"}},
     {"recv_v2", {"Out"}},
+    {"p_recv", {"Out"}},
     {"partial_recv", {"Out"}},
     {"matmul", {"Out"}},
     {"c_broadcast", {"Out"}},

paddle/phi/infermeta/nullary.cc

@@ -274,7 +274,7 @@ void PRecvArrayInferMeta(int peer,
         errors::InvalidArgument("The shape attribute for recv must be set "
                                 "explicitly, but the %dth element is %d which "
                                 "is less than 1. Or dynamic_shape should be "
-                                "set to True for both send_v2 and recv_v2.",
+                                "set to True for both send_v2 and p_recv.",
                                 i,
                                 out_shape[i]));
   }
@@ -291,13 +291,13 @@ void RecvV2InferMeta(const int ring_id,
       peer,
       0,
       errors::InvalidArgument(
-          "The peer (%d) for recv_v2 op must be non-negative.", peer));
+          "The peer (%d) for p_recv op must be non-negative.", peer));
 
   PADDLE_ENFORCE_GE(
       ring_id,
       0,
       errors::InvalidArgument(
-          "The ring_id (%d) for recv_v2 op must be non-negative.", ring_id));
+          "The ring_id (%d) for p_recv op must be non-negative.", ring_id));
 
   if (!dynamic_shape) {
     PADDLE_ENFORCE_GE(out_shape.size(),
@@ -310,10 +310,10 @@ void RecvV2InferMeta(const int ring_id,
       PADDLE_ENFORCE_GE(out_shape[i],
                         1,
                         errors::InvalidArgument(
-                            "The shape attribute for recv_v2 must be set "
+                            "The shape attribute for p_recv must be set "
                             "explicitly, but the %dth element is %d which "
                             "is less than 1. Or dynamic_shape should be "
-                            "set to True for both send_v2 and recv_v2.",
+                            "set to True for both send_v2 and p_recv.",
                             i,
                             out_shape[i]));
     }

paddle/phi/ops/yaml/inconsistent/static_ops.yaml

@@ -597,6 +597,19 @@
   interfaces : paddle::dialect::InferSymbolicShapeInterface
   traits : paddle::dialect::ForwardOnlyTrait
 
+- op : p_recv
+  args : (int ring_id = 0, int peer = 0, DataType dtype = DataType::FLOAT32, int[] out_shape = {}, bool dynamic_shape = false)
+  output : Tensor(out)
+  infer_meta :
+    func : PRecvInferMeta
+    param : [peer, dtype, out_shape, dynamic_shape]
+  kernel :
+    func : p_recv
+    param : [peer, dtype, out_shape, dynamic_shape]
+    data_type : dtype
+  interfaces : paddle::dialect::InferSymbolicShapeInterface
+  traits : paddle::dialect::ForwardOnlyTrait
+
 - op : partial_recv
   args : (int ring_id = 0, int peer = 0, DataType dtype=DataType::FLOAT32, int[] out_shape= {}, int num = 1, int id = 0)
   output : Tensor(out)

paddle/phi/ops/yaml/op_compat.yaml

@@ -4469,6 +4469,10 @@
   outputs :
     out : Out
 
+- op: p_recv
+  outputs :
+    out : Out
+
 - op: partial_send
   inputs :
     x : X

python/paddle/distributed/auto_parallel/static/cost/base_cost.py

@@ -27,6 +27,7 @@
 COMM_OP_TYPE = [
     "send_v2",
     "recv_v2",
+    "p_recv",
     "broadcast",
     "all_gather",
     "c_allreduce_sum",

python/paddle/distributed/auto_parallel/static/cost/comm_op_cost.py

@@ -181,7 +181,7 @@ def calc_time(self):
 
 @register_op_cost
 class RecvOpCost(CommOpCost):
-    OP_TYPE = "recv_v2"
+    OP_TYPE = "p_recv"
 
     def __init__(self, op=None, op_desc=None, comm_context=None):
         super().__init__(op=op, op_desc=op_desc, comm_context=comm_context)

python/paddle/distributed/auto_parallel/static/dist_saver.py

@@ -174,7 +174,7 @@ def save_inference_model(self, path, feed_vars, fetch_vars, exe, **kwargs):
         for idx, op in enumerate(ops):
             if op.attr(op_role_key) != op_role_forward:
                 continue
-            if op.type == "read" or op.type == "feed" or op.type == 'recv_v2':
+            if op.type == "read" or op.type == "feed" or op.type == 'p_recv':
                 feed_vars_names += op.output("Out")
             if op.type == "send_v2":
                 fetch_vars_names += op.input("X")

python/paddle/distributed/auto_parallel/static/mapper.py

@@ -50,7 +50,7 @@ def is_collective_comm_op(op):
 
 
 def is_p2p_comm_op(op):
-    comm_list = ["send_v2", "recv_v2"]
+    comm_list = ["send_v2", "p_recv"]
     if op.type in comm_list:
         return True
     else:
@@ -87,7 +87,7 @@ def get_comm_volume(comm_op, src_rank, tgt_rank):
     if src_rank == tgt_rank:
         return comm_volume
     comm_op_type = comm_op.type
-    if comm_op_type != "recv_v2":
+    if comm_op_type != "p_recv":
         tensor_name = comm_op.input_arg_names[0]
     else:
         tensor_name = comm_op.output_arg_names[0]
@@ -128,7 +128,7 @@ def get_comm_volume(comm_op, src_rank, tgt_rank):
             comm_volume = tensor_bytes
         else:
             comm_volume = None
-    elif "recv_v2" in comm_op_type:
+    elif "p_recv" in comm_op_type:
         comm_volume = None
     else:
         raise ValueError("Unrecognized communication operator.")

python/paddle/distributed/auto_parallel/static/parallelizer.py

@@ -490,7 +490,7 @@ def parallelize(
             if self._dist_strategy.auto_search:
                 is_pipeline = False
                 for op in dist_main_prog.global_block().ops:
-                    if op.type == "send_v2" or op.type == "recv_v2":
+                    if op.type == "send_v2" or op.type == "p_recv":
                         is_pipeline = True
                         break
                 if is_pipeline:

python/paddle/distributed/auto_parallel/static/reshard.py

@@ -387,7 +387,7 @@ def insert_send_op(block, idx, tensor, src, dst, op_role, sync=True):
     @staticmethod
     def insert_recv_op(block, idx, tensor, src, dst, op_role, sync=True):
         """Insert recv op into block at the given index."""
-        op_type = 'recv_v2'
+        op_type = 'p_recv'
         insert_operation = (
             block._insert_op if sync else block._insert_op_without_sync
         )

python/paddle/distributed/communication/stream/recv.py

@@ -48,7 +48,7 @@ def _recv_in_dygraph(
 def _recv_in_static_mode(
     tensor, src_rank_in_group, group, sync_op, use_calc_stream
 ):
-    op_type = 'recv_v2'
+    op_type = 'p_recv'
     data_feeder.check_variable_and_dtype(
         tensor,
         'tensor',

python/paddle/distributed/fleet/meta_optimizers/sharding_optimizer.py

@@ -1025,7 +1025,7 @@ def _prune_main_program(self, block, shard, rings):
                 "c_gen_xccl_id",
                 "c_comm_init",
                 'send_v2',
-                'recv_v2',
+                'p_recv',
             ]:
                 pass
             elif op.type == "conditional_block":

python/paddle/distributed/fleet/utils/hybrid_parallel_inference.py

@@ -637,7 +637,7 @@ def _insert_send_recv(cur_id, prev_id):
 
                     block._insert_op_without_sync(
                         index=index + extra_index_info['index'],
-                        type='recv_v2',
+                        type='p_recv',
                         outputs={'Out': [var]},
                         attrs={
                             'out_shape': var_shape,
@@ -737,7 +737,7 @@ def _insert_sendrecv_ops_in_while_block(
                         )
                     block._insert_op_without_sync(
                         index=index,
-                        type='recv_v2',
+                        type='p_recv',
                         outputs={'Out': [var]},
                         attrs={
                             'out_shape': var_shape,

python/paddle/distributed/passes/auto_parallel_grad_clip.py

@@ -265,7 +265,7 @@ def _is_pure_data_parallel(self):
                 and not is_data_parallel_reduce_op(op)
             ):
                 return False
-            if op.type in ["send_v2", "recv_v2"]:
+            if op.type in ["send_v2", "p_recv"]:
                 return False
 
         return True

python/paddle/distributed/passes/pass_utils.py

@@ -447,7 +447,7 @@ def _pir_overlap_send_recv(program):
     The finally target of this function is as follows:
         1. no need to insert the 'c_sync_calc' and 'c_sync_calc' operators
         2. 'send_v2' operator uses 'dist_attr.execution_stream' to set stream of its own.
-        3. 'recv_v2' operator uses 'dist_attr.execution_stream' to set stream of its own.
+        3. 'p_recv' operator uses 'dist_attr.execution_stream' to set stream of its own.
     """
     for block in program.blocks:
         for op in block.ops:
@@ -457,7 +457,7 @@ def _pir_overlap_send_recv(program):
                 ring_id = op.attrs()["ring_id"]
                 op.set_execution_stream(f"send_stream_{ring_id}")
                 op.set_scheduling_priority(0)
-            elif op.name() == "pd_op.recv_v2":
+            elif op.name() == "pd_op.p_recv":
                 op.set_bool_attr("dynamic_shape", False)
                 op.set_bool_attr("use_calc_stream", True)
                 op.set_execution_stream("recv_stream")
@@ -468,7 +468,7 @@ def _insert_sync_for_fthenb_1f1b(program, dist_context=None):
     """
     This implementation refers to lots of Paddle/python/paddle/base/optimizer.py.
     The difference between this function with 'PipelineOptimizer' is that
-    'send_v2' op and 'recv_v2' op have been inserted in program by 'reshard'.
+    'send_v2' op and 'p_recv' op have been inserted in program by 'reshard'.
     """
 
     for block in program.blocks:
@@ -482,7 +482,7 @@ def _insert_sync_for_fthenb_1f1b(program, dist_context=None):
         # insert sync ops
         for index, op in enumerate(list(block.ops)):
             # NOTE: pipeline might hang when dynamic_shape is True
-            if op.type in ['send_v2', 'recv_v2']:
+            if op.type in ['send_v2', 'p_recv']:
                 op._set_attr("dynamic_shape", False)
             # set send op on comm stream
             if op.type == 'send_v2':
@@ -556,7 +556,7 @@ def _insert_sync_for_fthenb_1f1b(program, dist_context=None):
         offset = 0
         backward_recv_index = None
         for index, op in enumerate(block.ops):
-            if op.type == "recv_v2" and is_backward_op(op):
+            if op.type == "p_recv" and is_backward_op(op):
                 backward_recv_index = index
                 break
         if backward_recv_index is None:
@@ -917,7 +917,7 @@ def _add_event_dependency(recorder_op, waiter_op):
     '''
     Add the extra event dependency of the two operators.
     This function mainly aims for the cross-programs in pipeline parallelism,
-    especial for the 'send_v2' 'recv_v2' etc.
+    especial for the 'send_v2' 'p_recv' etc.
     '''
     if not recorder_op.dist_attr.force_record_event:
         recorder_op.dist_attr.force_record_event = True