PaddlePaddle
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_attribute.cc
Copy file name to clipboard
+11-5Lines changed: 11 additions & 5 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_attribute.cc
Copy file name to clipboard
+11-5Lines changed: 11 additions & 5 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_dialect.cc
Copy file name to clipboard
+12Lines changed: 12 additions & 0 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_dialect.cc
Copy file name to clipboard
+12Lines changed: 12 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_op.cc
Copy file name to clipboard
+38-8Lines changed: 38 additions & 8 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_op.cc
Copy file name to clipboard
+38-8Lines changed: 38 additions & 8 deletions
diff --git a/‎paddle/fluid/pybind/pir.cc
Copy file name to clipboard
+38Lines changed: 38 additions & 0 deletions b/‎paddle/fluid/pybind/pir.cc
Copy file name to clipboard
+38Lines changed: 38 additions & 0 deletions
diff --git a/‎python/paddle/distributed/auto_parallel/static/engine.py
Copy file name to clipboard
+13-4Lines changed: 13 additions & 4 deletions b/‎python/paddle/distributed/auto_parallel/static/engine.py
Copy file name to clipboard
+13-4Lines changed: 13 additions & 4 deletions
diff --git a/‎python/paddle/distributed/auto_parallel/static/pir_pass.py
Copy file name to clipboard
+37-2Lines changed: 37 additions & 2 deletions b/‎python/paddle/distributed/auto_parallel/static/pir_pass.py
Copy file name to clipboard
+37-2Lines changed: 37 additions & 2 deletions
diff --git a/‎python/paddle/distributed/auto_parallel/static/reshard_funcs/p_to_r_reshard_func.py
Copy file name to clipboard
+6-3Lines changed: 6 additions & 3 deletions b/‎python/paddle/distributed/auto_parallel/static/reshard_funcs/p_to_r_reshard_func.py
Copy file name to clipboard
+6-3Lines changed: 6 additions & 3 deletions
diff --git a/‎python/paddle/distributed/auto_parallel/static/reshard_funcs/reshard_func_register.py
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎python/paddle/distributed/auto_parallel/static/reshard_funcs/reshard_func_register.py
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
@@ -121,12 +121,18 @@ OperationDistAttribute OperationDistAttribute::get(
     const std::vector<pir::Attribute>& result_attrs) {
   auto check_dist_attr = [=](pir::Attribute attr) {
     auto dist_attr = attr.dyn_cast<TensorDistAttribute>();
-    PADDLE_ENFORCE_EQ(mesh,
-                      dist_attr.process_mesh_attr(),
-                      common::errors::PreconditionNotMet(
-                          "operand_dist_attrs element's mesh(%s) not equal "
-                          "to input mesh(%s)"));
+    auto ids = mesh.process_ids();
+    for (const auto& id : dist_attr.process_mesh_attr().process_ids()) {
+      PADDLE_ENFORCE_EQ(std::find(ids.begin(), ids.end(), id) != ids.end(),
+                        true,
+                        common::errors::PreconditionNotMet(
+                            "operand_dist_attrs element's mesh(%s) not belong "
+                            "to input mesh(%s)",
+                            dist_attr.process_mesh_attr(),
+                            mesh));
+    }
   };
+
   for (auto attr : operand_attrs) {
     // NOTE: The operand dist attr maybe empty while the corresponding input is
     // optional.
 
@@ -74,6 +74,10 @@ void DistDialect::PrintAttribute(pir::Attribute attr, std::ostream &os) const {
               phi::distributed::auto_parallel::str_join(
                   tensor_dist_attr.process_mesh_attr().shape()) +
               "]";
+    os << ",process_ids:[" +
+              phi::distributed::auto_parallel::str_join(
+                  tensor_dist_attr.process_mesh_attr().process_ids()) +
+              "]";
     os << ",dims_mappings:[" +
               phi::distributed::auto_parallel::str_join(
                   tensor_dist_attr.dims_mapping()) +
@@ -111,6 +115,10 @@ void DistDialect::PrintAttribute(pir::Attribute attr, std::ostream &os) const {
                   phi::distributed::auto_parallel::str_join(
                       dist_attr.process_mesh_attr().shape()) +
                   "],";
+        os << "process_ids:[" +
+                  phi::distributed::auto_parallel::str_join(
+                      dist_attr.process_mesh_attr().process_ids()) +
+                  "],";
       }
       os << "dims_maping:[" +
                 phi::distributed::auto_parallel::str_join(
@@ -145,6 +153,10 @@ void DistDialect::PrintAttribute(pir::Attribute attr, std::ostream &os) const {
                   phi::distributed::auto_parallel::str_join(
                       dist_attr.process_mesh_attr().shape()) +
                   "],";
+        os << "process_ids:[" +
+                  phi::distributed::auto_parallel::str_join(
+                      dist_attr.process_mesh_attr().process_ids()) +
+                  "],";
       }
       os << "dims_maping:[" +
                 phi::distributed::auto_parallel::str_join(
 
@@ -214,10 +214,11 @@ void ReshardOp::VerifySig() {
         1u,
         common::errors::PreconditionNotMet(
             "The size %d of inputs must be equal to 1.", input_size));
-    PADDLE_ENFORCE_EQ((*this)
-                          ->operand_source(0)
-                          .type()
-                          .isa<paddle::dialect::DistDenseTensorType>(),
+    PADDLE_ENFORCE_EQ(!(*this)->operand_source(0) ||
+                          (*this)  // reshard allow NULL TYPE as input
+                              ->operand_source(0)
+                              .type()
+                              .isa<paddle::dialect::DistDenseTensorType>(),
                       true,
                       common::errors::PreconditionNotMet(
                           "Type validation failed for the 0th input."));
@@ -241,7 +242,13 @@ void ReshardOp::VerifySig() {
         common::errors::PreconditionNotMet(
             "The size %d of outputs must be equal to 1.", output_size));
     PADDLE_ENFORCE_EQ(
-        (*this)->result(0).type().isa<paddle::dialect::DistDenseTensorType>(),
+        !(*this)->result(0) ||
+            (*this)
+                ->result(0)
+                .type()
+                .isa<paddle::dialect::DistDenseTensorType>(),  // reshard allow
+                                                               // NULL TYPE as
+                                                               // output
         true,
         common::errors::PreconditionNotMet(
             "Type validation failed for the 0th output."));
@@ -267,12 +274,34 @@ void ReshardOp::VerifySig() {
   VLOG(4) << "End Verifying for: ShardTensorOp.";
 }
 
+ProcessMeshAttribute MergeMeshes(const ProcessMeshAttribute& mesh1,
+                                 const ProcessMeshAttribute& mesh2) {
+  if (mesh1 == mesh2) return mesh1;
+  // Combine the two ids
+  std::vector<int64_t> merged_ids;
+  std::vector<int64_t> ids1 = mesh1.process_ids();
+  std::vector<int64_t> ids2 = mesh2.process_ids();
+
+  merged_ids.reserve(ids1.size() + ids2.size());
+  merged_ids.insert(merged_ids.end(), ids1.begin(), ids1.end());
+  merged_ids.insert(merged_ids.end(), ids2.begin(), ids2.end());
+
+  // Remove duplicates
+  std::sort(merged_ids.begin(), merged_ids.end());
+  auto last = std::unique(merged_ids.begin(), merged_ids.end());
+  merged_ids.erase(last, merged_ids.end());
+
+  return ProcessMeshAttribute::get(
+      pir::IrContext::Instance(),
+      {static_cast<int64_t>(merged_ids.size())},  // flatten mesh shape
+      merged_ids,
+      {"merged"});
+}
+
 void ReshardOp::Build(pir::Builder& builder,
                       pir::OperationArgument& argument,
                       pir::Value input,
                       TensorDistAttribute tensor_dist_attr) {
-  VLOG(4) << "Start build ReshardOp";
-
   paddle::dialect::DistDenseTensorType input_tensor_type;
   if (input.type().isa<paddle::dialect::DistDenseTensorType>()) {
     input_tensor_type =
@@ -288,7 +317,8 @@ void ReshardOp::Build(pir::Builder& builder,
   VLOG(4) << "Builder construction attributes";
   pir::Attribute op_dist_attr = OperationDistAttribute::get(
       pir::IrContext::Instance(),
-      input_tensor_type.tensor_dist_attr().process_mesh_attr(),
+      MergeMeshes(input_tensor_type.tensor_dist_attr().process_mesh_attr(),
+                  tensor_dist_attr.process_mesh_attr()),
       std::vector<pir::Attribute>{input_tensor_type.tensor_dist_attr()},
       std::vector<pir::Attribute>{tensor_dist_attr});
   argument.AddAttribute("op_dist_attr", op_dist_attr);
 
@@ -16,6 +16,7 @@
 
 #include <Python.h>
 #include <algorithm>
+#include <iterator>
 #include <memory>
 #include <sstream>
 #include <string>
@@ -1103,6 +1104,43 @@ struct PyInsertionPoint {
 void BindInsertionPoint(pybind11::module *m) {
   py::class_<PyInsertionPoint> ir_insertion_point(*m, "InsertionPoint", R"DOC(
     InsertionPoint class represents the insertion point in the Builder.)DOC");
+
+  ir_insertion_point
+      .def(
+          "next",
+          [](PyInsertionPoint &self) -> Operation & {
+            if (self.value.second == self.value.first->end()) {
+              PADDLE_THROW(common::errors::InvalidArgument(
+                  "The insertion point is already at the end and can't call "
+                  "next()."));
+            }
+            return *(self.value.second++);
+          },
+          return_value_policy::reference)
+      .def(
+          "prev",
+          [](PyInsertionPoint &self) -> Operation & {
+            if (self.value.second == self.value.first->begin()) {
+              PADDLE_THROW(common::errors::InvalidArgument(
+                  "The insertion point is already at the begin and can't call "
+                  "prev()."));
+            }
+            return *(self.value.second--);
+          },
+          return_value_policy::reference)
+      .def(
+          "get_operation",
+          [](PyInsertionPoint &self) -> Operation & {
+            if (self.value.second == self.value.first->begin()) {
+              PADDLE_THROW(common::errors::InvalidArgument(
+                  "The insertion point is already at the begin."));
+            } else if (self.value.second == self.value.first->end()) {
+              PADDLE_THROW(common::errors::InvalidArgument(
+                  "The insertion point is already at the end."));
+            }
+            return *(self.value.second);
+          },
+          return_value_policy::reference);
 }
 
 std::list<Operation *>::const_iterator list_offset(const Block *block,
 
@@ -1782,9 +1782,15 @@ def run(self, data=None, feed=None, fetch_list=None, mode=None):
         use_cache = self._strategy.use_cache
         if self._in_pir_mode:
             use_cache = False
-            fetch_names = [
-                self.main_program.get_output_value_by_name(self._loss_names[0])
-            ]
+            no_fetch = False  # not last rank should not fetch loss in pipeline parallel
+            loss_value = self.main_program.get_output_value_by_name(
+                self._loss_names[0]
+            )
+            if paddle.pir.is_fake_value(loss_value):
+                no_fetch = True
+                fetch_names = []
+            else:
+                fetch_names = [loss_value]
 
         outs = self._executor.run(
             self.main_program,
@@ -1795,7 +1801,10 @@ def run(self, data=None, feed=None, fetch_list=None, mode=None):
         )
 
         if self._in_pir_mode:
-            logs = {"outputs": outs[0], "loss": outs[0]}
+            if no_fetch:
+                logs = {"outputs": None, "loss": None}
+            else:
+                logs = {"outputs": outs[0], "loss": outs[0]}
             return logs
 
         logs = self._prepare_logger(
 
@@ -60,6 +60,7 @@ def apply_partition_pass(program):
                         prev_var, operand_dist_attr
                     )
                     var.set_source(reshard_var)
+
             for var, result_dist_attr in zip(
                 op.results(), op.dist_attr.result_dist_attrs()
             ):
@@ -71,6 +72,40 @@ def apply_partition_pass(program):
                     reshard_var = paddle._C_ops.reshard_v2(var, old_dist_attr)
                     var.replace_all_uses_with(reshard_var)
                     reshard_var.get_defining_op().operand(0).set_source(var)
+
+        # pruning op and value not belong to cur rank
+        cur_rank = paddle.distributed.get_rank()
+        for op in new_program.global_block().ops[::-1]:
+            if cur_rank not in op.dist_attr.process_mesh.process_ids:
+                new_program.global_block().remove_op(op)
+            else:
+                # set the operand as null when it is not belong to cur rank
+                if (
+                    op.name() == 'dist_op.reshard'
+                    and cur_rank
+                    not in op.operand(0)
+                    .source()
+                    .dist_attr()
+                    .process_mesh.process_ids
+                ):
+                    op.operand(0).set_source(paddle.pir.fake_value())
+
+        # merge pd.data ops for
+        lr_ops = []
+        for op in new_program.global_block().ops[::-1]:
+            if (
+                op.name() == 'pd_op.data'
+                and "learning_rate" in op.attrs()["name"]
+            ):
+                lr_ops.append(op)
+
+        if len(lr_ops) > 1:
+            lr_value = lr_ops[0].result(0)
+            for op in lr_ops[1:]:
+                lr = op.result(0)
+                lr.replace_all_uses_with(lr_value)
+                new_program.global_block().remove_op(op)
+
     return new_program
 
 
@@ -84,14 +119,14 @@ def apply_reshard_pass(program):
                 src_dist_attr = op_dist_attr.operand_dist_attr(0)
                 dst_dist_attr = op_dist_attr.result_dist_attr(0)
                 assert (
-                    var.source().dist_attr() == src_dist_attr
+                    paddle.pir.is_fake_value(var.source())
+                    or var.source().dist_attr() == src_dist_attr
                 ), f"The dist_attr of reshard op's input and operand should be equal, but got {var.source().dist_attr()} and {src_dist_attr}"
 
                 reshard_func = choose_reshard_func(src_dist_attr, dst_dist_attr)
                 assert (
                     reshard_func is not None
                 ), f'There is no reshard function that matches src_dist_attr: {src_dist_attr} and dst_dist_attr: {dst_dist_attr}'
-
                 reshard_func.reshard(
                     new_program, op, src_dist_attr, dst_dist_attr
                 )
 
@@ -103,16 +103,19 @@ def reshard(self, program, op, src_dist_attr, dst_dist_attr):
             src_dist_attr.dims_mapping,
             src_dist_attr.partial_status,
         )
-        out, out_dist_attr = same_status_func.reshard(
+        pre_op, out_dist_attr = same_status_func.reshard(
             program, op, src_dist_attr, tmp_dist_attr
         )
 
+        if pre_op is None:
+            return None, out_dist_attr
+
         curr_global_rank = paddle.distributed.get_rank()
         if curr_global_rank in dst_dist_attr.process_mesh.process_ids:
             p_to_r_func = PToRReshardFunction()
             assert p_to_r_func.is_suitable(
                 out_dist_attr, dst_dist_attr
-            ), f"Invoke the p to r reshard function is not valid from {out.dist_attr()} to {dst_dist_attr}"
+            ), f"Invoke the p to r reshard function is not valid from {pre_op.dist_attr()} to {dst_dist_attr}"
             p_to_r_func.reshard(
-                program, out, out_dist_attr, dst_dist_attr, False
+                program, pre_op, out_dist_attr, dst_dist_attr, False
             )
@@ -17,11 +17,13 @@
     PToRReshardFunction,
     PToRReshardFunctionCrossMesh,
 )
+from .same_status_reshard_func import SameStatusReshardFunction
 
 
 def register_reshard_funcs():
     register_reshard_func(PToRReshardFunction())
     register_reshard_func(PToRReshardFunctionCrossMesh())
+    register_reshard_func(SameStatusReshardFunction())
 
 
 register_reshard_funcs()