PaddlePaddle
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_tools.cc
Copy file name to clipboard
+6-6Lines changed: 6 additions & 6 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_tools.cc
Copy file name to clipboard
+6-6Lines changed: 6 additions & 6 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_tools.h
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/fluid/pir/dialect/distributed/ir/dist_tools.h
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/fluid/pybind/pir.cc
Copy file name to clipboard
+16-1Lines changed: 16 additions & 1 deletion b/‎paddle/fluid/pybind/pir.cc
Copy file name to clipboard
+16-1Lines changed: 16 additions & 1 deletion
diff --git a/‎python/paddle/distributed/auto_parallel/static/engine.py
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎python/paddle/distributed/auto_parallel/static/engine.py
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
diff --git a/‎python/paddle/distributed/auto_parallel/static/pir_pass.py
Copy file name to clipboard
+34-38Lines changed: 34 additions & 38 deletions b/‎python/paddle/distributed/auto_parallel/static/pir_pass.py
Copy file name to clipboard
+34-38Lines changed: 34 additions & 38 deletions
@@ -159,10 +159,10 @@ pir::Attribute CreateReplicatedDistAttr(pir::Type prim_type,
   }
   return nullptr;
 }
-pir::Type CvtToPirDistType(pir::Type prim_type, pir::Attribute dist_attr) {
-  if (!prim_type) return nullptr;
+pir::Type CvtToPirDistType(pir::Type global_type, pir::Attribute dist_attr) {
+  if (!global_type) return nullptr;
   auto ctx = pir::IrContext::Instance();
-  if (auto dense_tensor_type = prim_type.dyn_cast<pir::DenseTensorType>()) {
+  if (auto dense_tensor_type = global_type.dyn_cast<pir::DenseTensorType>()) {
     auto tensor_dist_attr = dist_attr.dyn_cast<TensorDistAttribute>();
     if (!tensor_dist_attr) {
       VLOG(0) << "Convert dense tensor type to dist type with attribute {"
@@ -172,7 +172,7 @@ pir::Type CvtToPirDistType(pir::Type prim_type, pir::Attribute dist_attr) {
           "with non-empty TensorDistAttr"));
     }
     return DistDenseTensorType::get(ctx, dense_tensor_type, tensor_dist_attr);
-  } else if (auto vec_type = prim_type.dyn_cast<pir::VectorType>()) {
+  } else if (auto vec_type = global_type.dyn_cast<pir::VectorType>()) {
     auto array_attr = dist_attr.dyn_cast<pir::ArrayAttribute>();
     if (!array_attr) {
       VLOG(0) << "Convert vector type to dist type with attribute {"
@@ -192,8 +192,8 @@ pir::Type CvtToPirDistType(pir::Type prim_type, pir::Attribute dist_attr) {
     }
     return pir::VectorType::get(ctx, dist_vec_type);
   } else {
-    VLOG(0) << "Convert type{" << prim_type << "} to dist type with attribute {"
-            << dist_attr << "}.";
+    VLOG(0) << "Convert type{" << global_type
+            << "} to dist type with attribute {" << dist_attr << "}.";
     PADDLE_THROW(common::errors::InvalidArgument(
         "Currently only support convert dense_tensor_type r vector type to "
         "dist."));
 
@@ -37,7 +37,7 @@ pir::Attribute CvtToPirAttr(const phi::distributed::ArgDistAttr& dist_attr);
 pir::Attribute CreateReplicatedDistAttr(pir::Type prim_type,
                                         ProcessMeshAttribute mesh);
 
-pir::Type CvtToPirDistType(pir::Type prim_type, pir::Attribute dist_attr);
+pir::Type CvtToPirDistType(pir::Type global_type, pir::Attribute dist_attr);
 
 ///
 /// When the following conditions are met:
 
@@ -116,6 +116,7 @@ using pir::Program;
 using pir::StrAttribute;
 using pir::Type;
 using pir::Value;
+using pir::VectorType;
 using pybind11::return_value_policy;
 
 COMMON_DECLARE_bool(print_ir);
@@ -1330,6 +1331,13 @@ void BindType(py::module *m) {
             PADDLE_THROW(phi::errors::InvalidArgument(
                 "can't set _local_shape when building static graph"));
           })
+      .def("as_vec_type",
+           [](Type self) -> py::object {
+             if (auto vec_type = self.dyn_cast<VectorType>()) {
+               return py::cast(vec_type);
+             }
+             return py::cast<py::none>(Py_None);
+           })
       .def("__str__", [](Type &self) {
         std::ostringstream print_stream;
         print_stream << self;
@@ -1364,7 +1372,13 @@ void BindType(py::module *m) {
            }
          });
 }
-
+void BindVectorType(py::module *m) {
+  py::class_<VectorType, Type> vec_type(*m, "VectorType");
+  vec_type.def("as_list", &VectorType::data);
+  m->def("create_vec_type", [](std::vector<Type> &types) {
+    return VectorType::get(pir::IrContext::Instance(), types);
+  });
+}
 void BindAttribute(py::module *m) {
   py::class_<Attribute> ir_attr(*m, "Attribute", py::module_local());
   ir_attr.def("__eq__", &Attribute::operator==)
@@ -2496,6 +2510,7 @@ void BindPir(pybind11::module *module) {
   BindOperation(&ir_module);
   BindOpOperand(&ir_module);
   BindType(&ir_module);
+  BindVectorType(&ir_module);
   BindAttribute(&ir_module);
   BindInsertionPoint(&ir_module);
   BindUtils(&ir_module);
 
@@ -58,6 +58,7 @@
 from .pir_pass import (
     apply_partition_pass,
     apply_reshard_pass,
+    remove_other_rank_op_pass,
     remove_unuseful_comm_op_pass,
 )
 from .planner_v2 import Planner
@@ -696,6 +697,8 @@ def _parallel_pir(self, mode):
         #   collect the communicator created during resolution.
         apply_reshard_pass(dist_program)
 
+        remove_other_rank_op_pass(dist_program)
+
         # Part 4: Optimization Pass
         # NOTE Only those Optimization Pass that related to Parallelism (need dist attr) should be placed here and all the Pass should be Optional.
 
 
@@ -89,44 +89,6 @@ def apply_partition_pass(program):
                 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 program.global_block().ops[::-1]:
-        if op.name() in partition_skip_op_list:
-            can_delete = True
-            for val in op.results():
-                if not val.use_empty():
-                    can_delete = False
-            if can_delete:
-                op.erase()
-            continue
-        if cur_rank not in op.dist_attr.process_mesh.process_ids:
-            op.erase()
-        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(None)
-
-    # merge pd.data ops for
-    lr_ops = []
-    for op in 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)
-            op.erase()
-
 
 def apply_reshard_pass(program):
     for op in program.global_block().ops:
@@ -160,6 +122,40 @@ def apply_reshard_pass(program):
                 op.erase()
 
 
+# pruning op and value not belong to cur rank
+def remove_other_rank_op_pass(dist_program):
+    cur_rank = paddle.distributed.get_rank()
+    for op in dist_program.global_block().ops[::-1]:
+        if op.name() in partition_skip_op_list:
+            can_delete = True
+            for val in op.results():
+                if not val.use_empty():
+                    can_delete = False
+            if can_delete:
+                op.erase()
+            continue
+        if cur_rank not in op.dist_attr.process_mesh.process_ids:
+            op.erase()
+        elif op.name() == "dist_op.reshard":
+            assert op.result(
+                0
+            ).use_empty(), f'There should not have useful dist.reshard op in remove_other_rank_op_pass. but find : {op}'
+            op.erase()
+
+    # merge pd.data ops for
+    lr_ops = []
+    for op in dist_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)
+            op.erase()
+
+
 # Note: this is the pass in the dense program
 comm_ops = ["pd_op.c_allreduce_sum_", "pd_op.c_allgather"]