PaddlePaddle
diff --git a/‎paddle/fluid/pir/transforms/general/inplace_pass.cc
Copy file name to clipboard
+5-2Lines changed: 5 additions & 2 deletions b/‎paddle/fluid/pir/transforms/general/inplace_pass.cc
Copy file name to clipboard
+5-2Lines changed: 5 additions & 2 deletions
diff --git a/‎paddle/fluid/pybind/pir.cc
Copy file name to clipboard
+15Lines changed: 15 additions & 0 deletions b/‎paddle/fluid/pybind/pir.cc
Copy file name to clipboard
+15Lines changed: 15 additions & 0 deletions
diff --git a/‎paddle/pir/src/core/program.cc
Copy file name to clipboard
+11Lines changed: 11 additions & 0 deletions b/‎paddle/pir/src/core/program.cc
Copy file name to clipboard
+11Lines changed: 11 additions & 0 deletions
diff --git a/‎python/paddle/base/executor.py
Copy file name to clipboard
+38-6Lines changed: 38 additions & 6 deletions b/‎python/paddle/base/executor.py
Copy file name to clipboard
+38-6Lines changed: 38 additions & 6 deletions
diff --git a/‎python/paddle/distributed/auto_parallel/static/engine.py
Copy file name to clipboard
+72-3Lines changed: 72 additions & 3 deletions b/‎python/paddle/distributed/auto_parallel/static/engine.py
Copy file name to clipboard
+72-3Lines changed: 72 additions & 3 deletions
@@ -211,6 +211,10 @@ bool IsNoNeedBuffer(pir::Operation* op, pir::Value value) {
 std::unordered_set<pir::Value> GetSkipDeletionValues(const pir::Block& block) {
   std::unordered_set<pir::Value> skip_dels;
   for (auto& op : block) {
+    if (op.name() == "builtin.shadow_output") {
+      skip_dels.insert(op.operand_source(0));
+      continue;
+    }
     if (op.dialect()->name() != paddle::dialect::KernelDialect::name()) {
       continue;
     }
@@ -228,8 +232,7 @@ std::unordered_set<pir::Value> GetSkipDeletionValues(const pir::Block& block) {
       continue;
     }
     // TODO(chenxi67) add logic for shadow_feed_tensors op
-    if (upper_op_name == "pd_op.fetch" ||
-        upper_op_name == "builtin.shadow_output") {
+    if (upper_op_name == "pd_op.fetch") {
       skip_dels.insert(op.operand_source(0));
       continue;
     }
 
@@ -1049,6 +1049,21 @@ void BindOperation(py::module *m) {
           },
           [](Operation &self, OperationDistAttribute op_dist_attr) {
             self.set_attribute(kAttrOpDistAttr, op_dist_attr);
+          })
+      .def_property(
+          "op_role",
+          [](Operation &self) -> py::object {
+            auto int_attr = self.attribute<Int32Attribute>("op_role");
+            if (int_attr) {
+              return py::cast(int_attr.data());
+            } else {
+              return py::cast<py::none>(Py_None);
+            }
+          },
+          [](Operation &self, const int &op_role) {
+            self.set_attribute(
+                "op_role",
+                Int32Attribute::get(pir::IrContext::Instance(), op_role));
           });
   py::class_<Operation::BlockContainer> block_container(
       *m, "Operation_BlockContainer", R"DOC(
 
@@ -67,6 +67,17 @@ std::shared_ptr<Program> Program::Clone(IrMapping& ir_mapping) const {
   pir::IrContext* ctx = pir::IrContext::Instance();
   auto new_program = std::make_shared<Program>(ctx);
   auto clone_options = CloneOptions::All();
+
+  // deal kwargs
+  for (auto [key, value] : block()->kwargs()) {
+    auto new_arg = new_program->block()->AddKwarg(key, value.type());
+    auto tmp_block_arg = value.dyn_cast<BlockArgument>();
+    for (auto [name, attr_value] : tmp_block_arg.attributes()) {
+      new_arg.set_attribute(name, attr_value);
+    }
+    ir_mapping.Add(value, new_arg);
+  }
+
   for (const auto& op : *block()) {
     auto* new_op = op.Clone(ir_mapping, clone_options);
     new_program->block()->push_back(new_op);
 
@@ -533,6 +533,20 @@ def _add_pir_fetch_ops(program, fetch_list, fetch_var_name):
                 out.persistable = True
 
 
+def _add_single_pir_fetch_op(program, fetch_value, fetch_name, fetch_col):
+    import paddle
+
+    global_block = program.global_block()
+    fetch_op = "pd_op.fetch"
+    need_fetch_info = has_fetch_operations(
+        global_block, [fetch_value], fetch_name, fetch_op
+    )
+    if need_fetch_info:
+        with paddle.static.program_guard(program):
+            out = paddle._pir_ops.fetch(fetch_value, fetch_name, fetch_col)
+            out.persistable = True
+
+
 def _merge_tensors(tensor, micro_batch_num):
     if micro_batch_num <= 1:
         return tensor
@@ -856,6 +870,7 @@ def __init__(
             fetch_var_name,
             place,
             scope,
+            plan=None,
         ):
             self.program = program
             self.feed = feed
@@ -864,6 +879,7 @@ def __init__(
             self.fetch_var_name = fetch_var_name
             self.place = place
             self.scope = scope
+            self.plan = plan
 
             # NOTE(Ruibiao): Not all changeable item is considered for key at present,
             # ONLY: program, feed, and fetch_list
@@ -1114,10 +1130,17 @@ def get_pir_program_and_executor(
         fetch_var_name,
         place,
         scope,
+        plan,
     ):
-        _add_pir_fetch_ops(
-            program, fetch_list=fetch_list, fetch_var_name=fetch_var_name
-        )
+        if plan is None:
+            _add_pir_fetch_ops(
+                program, fetch_list=fetch_list, fetch_var_name=fetch_var_name
+            )
+        else:
+            for i, value in enumerate(fetch_list):
+                _add_single_pir_fetch_op(
+                    value.block.program, value, fetch_var_name + str(i), i
+                )
         return self._get_cached_program_and_executor_pir_mode(
             self._CachedData(
                 program,
@@ -1127,6 +1150,7 @@ def get_pir_program_and_executor(
                 fetch_var_name,
                 place,
                 scope,
+                plan,
             )
         )
 
@@ -1139,9 +1163,12 @@ def _get_pir_program_and_executor(self, cached_data):
         place = cached_data.place
         scope = cached_data.scope
 
-        default_job = core.Job("default")
-        type_to_program = {"default": program}
-        plan = core.Plan([default_job], type_to_program)
+        if cached_data.plan is None:
+            default_job = core.Job("default")
+            type_to_program = {"default": program}
+            plan = core.Plan([default_job], type_to_program)
+        else:
+            plan = cached_data.plan
 
         new_exe = _StandaloneExecutor(place, plan, scope)
 
@@ -1252,6 +1279,7 @@ def __init__(self, place=None):
         self._closed = False
         self.pruned_program_scope_caches = {}
         self._prepare_to_run_called = False
+        self.plan = None
 
         self._auto_checkpoint_name = unique_name.generate(
             "__auto_checkpoint_executor__"
@@ -1280,6 +1308,9 @@ def __del__(self):
         self.close()
         self._executor_cache.clear()
 
+    def _set_plan(self, plan):
+        self.plan = plan
+
     def _get_scope_cache(self, program_cache_key):
         return self.scope_caches.get(program_cache_key, None)
 
@@ -2094,6 +2125,7 @@ def _run_pir_impl(
             fetch_var_name,
             self.place,
             scope,
+            self.plan,
         )
         self._pir_feed_data(program, feed, scope, data_op_infos)
 
 
@@ -58,6 +58,8 @@
 from .pir_pass import (
     apply_partition_pass,
     apply_reshard_pass,
+    complete_op_role,
+    pipeline_pass,
     remove_other_rank_input_output_pass,
     remove_other_rank_op_pass,
     remove_unuseful_comm_op_pass,
@@ -260,6 +262,7 @@ def __init__(
         self._dygraph_mode = False
         self._tuning = self._strategy.tuning
         self._acc_steps = 1
+        self._pipeline_plan = None
         self._in_pir_mode = paddle.base.framework.get_flags(
             "FLAGS_enable_pir_api"
         )["FLAGS_enable_pir_api"]
@@ -657,13 +660,17 @@ def _parallel_pir(self, mode):
         mix_fw_program = self._fwd_main_progs[mode]
         startup_program = self._startup_progs[mode]
 
+        forward_op_start_idx = 0
+        backward_op_start_idx = -1
+        opt_op_start_idx = -1
         # Part 1: Complete program
         # Step 1.1: Mix2Dist Pass
         # TODO(JZ-LIANG) regulization pass with pass management.
         dist_program = mix_fw_program.clone()
         apply_mix2dist_pass(dist_program)
 
         # Step 1.2: pir backward
+        last_forward_op = dist_program.global_block().ops[-1]
         if mode == "train" and self._loss and self._optimizer:
             loss = dist_program.get_output_value_by_name(self._loss_names[0])
             if loss.initialized():
@@ -701,16 +708,58 @@ def _parallel_pir(self, mode):
                             and self._strategy.amp.dtype != 'bfloat16',
                         )
                         scaled = scaler.scale(loss)
-                        scaler.minimize(self._optimizer, scaled)
+                        last_forward_op = dist_program.global_block().ops[-1]
+                        optimizer_ops, params_grads = scaler.minimize(
+                            self._optimizer, scaled
+                        )
+                        first_opt_op = optimizer_ops[0]
+                        backward_op_start_idx = (
+                            dist_program.global_block().ops.index(
+                                last_forward_op
+                            )
+                            + 1
+                        )
+                        opt_op_start_idx = (
+                            dist_program.global_block().ops.index(first_opt_op)
+                        )
                         # print('after minimize', dist_program, flush=1)
-
+                        complete_op_role(
+                            dist_program,
+                            [
+                                [forward_op_start_idx, backward_op_start_idx],
+                                [backward_op_start_idx, opt_op_start_idx],
+                                [opt_op_start_idx, dist_program.num_ops()],
+                            ],
+                        )
                     else:
                         params_grads = (
                             paddle.autograd.ir_backward.append_backward(loss)
                         )
+                        last_backward_op = dist_program.global_block().ops[-1]
                         self._optimizer._apply_optimize(
                             loss, startup_program, params_grads=params_grads
                         )
+
+                        backward_op_start_idx = (
+                            dist_program.global_block().ops.index(
+                                last_forward_op
+                            )
+                            + 1
+                        )
+                        opt_op_start_idx = (
+                            dist_program.global_block().ops.index(
+                                last_backward_op
+                            )
+                            + 1
+                        )
+                        complete_op_role(
+                            dist_program,
+                            [
+                                [forward_op_start_idx, backward_op_start_idx],
+                                [backward_op_start_idx, opt_op_start_idx],
+                                [opt_op_start_idx, dist_program.num_ops()],
+                            ],
+                        )
                         # self._optimizer.minimize(loss, startup_program=startup_program)
 
             else:
@@ -786,6 +835,11 @@ def _parallel_pir(self, mode):
         paddle.base.libpaddle.pir.apply_dist2dense_pass(dense_program)
         remove_unuseful_comm_op_pass(dense_program)
 
+        if self._strategy.pipeline.enable:
+            self._pipeline_plan = pipeline_pass(
+                [dense_program], [dense_program], self._strategy.pipeline
+            )
+
         self._pir_dense_main_progs[mode] = dense_program
         self._pir_dist_main_progs[mode] = dist_program
 
@@ -1206,6 +1260,11 @@ def _initialize(self, mode, init_parameters=True):
                 for del_op in del_ops:
                     del_op.erase()
                 self._executor.run(startup_prog)
+                if self._pipeline_plan is not None:
+                    # pipeline scheduling should be enabled after running
+                    # startup program, otherwise the startup program cannot
+                    # run correctly.
+                    self._executor._set_plan(self._pipeline_plan)
             return
 
         if self._strategy.seed:
@@ -1947,7 +2006,17 @@ def run(self, data=None, feed=None, fetch_list=None, mode=None):
         if self._in_pir_mode:
             use_cache = False
             no_fetch = False  # not last rank should not fetch loss in pipeline parallel
-            loss_value = self.main_program.get_output_value_by_name(
+            if self._pipeline_plan is None:
+                program_for_executor = self.main_program
+            else:
+                # NOTE: If pipeline scheduling is enabled, The program_for_executor
+                # is used to tell the executor where to feed data and add fetch op,
+                # not the program to be executed. The ``plan`` object is already
+                # constructed, and the programs to be executed are  stored in the
+                # ``plan`` object.
+                program_for_executor = self._pipeline_plan.ir_program("forward")
+
+            loss_value = program_for_executor.get_output_value_by_name(
                 self._loss_names[0]
             )
             if paddle.pir.is_fake_value(loss_value):