PaddlePaddle
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_attribute.cc
Copy file name to clipboard
+32Lines changed: 32 additions & 0 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_attribute.cc
Copy file name to clipboard
+32Lines changed: 32 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_attribute.h
Copy file name to clipboard
+3-1Lines changed: 3 additions & 1 deletion b/‎paddle/fluid/pir/dialect/distributed/ir/dist_attribute.h
Copy file name to clipboard
+3-1Lines changed: 3 additions & 1 deletion
diff --git a/‎paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎paddle/fluid/pir/dialect/op_generator/ops_api_gen.py
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
diff --git a/‎paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎paddle/fluid/pir/transforms/pd_op_to_kernel_pass.cc
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
diff --git a/‎paddle/fluid/pybind/dist_api.cc
Copy file name to clipboard
+12-2Lines changed: 12 additions & 2 deletions b/‎paddle/fluid/pybind/dist_api.cc
Copy file name to clipboard
+12-2Lines changed: 12 additions & 2 deletions
diff --git a/‎paddle/fluid/pybind/pir.cc
Copy file name to clipboard
+29-21Lines changed: 29 additions & 21 deletions b/‎paddle/fluid/pybind/pir.cc
Copy file name to clipboard
+29-21Lines changed: 29 additions & 21 deletions
diff --git a/‎python/paddle/distributed/__init__.py
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎python/paddle/distributed/__init__.py
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
@@ -60,6 +60,38 @@ TensorDistAttribute::partial_status() const {
   return storage()->partial_status;
 }
 
+phi::distributed::Placements TensorDistAttribute::placements() const {
+  auto process_mesh = process_mesh_attr();
+  phi::distributed::Placements placements;
+  placements.resize(process_mesh.ndim(),
+                    std::make_shared<phi::distributed::Replicate>());
+
+  for (const auto& pair : partial_status()) {
+    placements[pair.first] =
+        std::make_shared<phi::distributed::Partial>(pair.second);
+  }
+
+  auto& dim_mapping = dims_mapping();
+  for (size_t i = 0; i < dim_mapping.size(); ++i) {
+    auto& mesh_id = dim_mapping[i];
+    if (mesh_id >= 0) {
+      auto& p = placements[mesh_id];
+      if (p->is_shard()) {
+        PADDLE_THROW(phi::errors::PreconditionNotMet(
+            "ProcessMesh dimension cann't be mapped to two  dimension of the "
+            "same tensor: {%d} and {%d}",
+            i,
+            dynamic_cast<phi::distributed::Shard&>(*p).get_dim()));
+      } else if (p->is_partial()) {
+        PADDLE_THROW(phi::errors::PreconditionNotMet(
+            "ProcessMesh dimension {%d} cannot be both shard and partial!",
+            mesh_id));
+      }
+      placements[mesh_id] = std::make_shared<phi::distributed::Shard>(i);
+    }
+  }
+  return placements;
+}
 TensorDistAttribute TensorDistAttribute::get(
     pir::IrContext* ctx,
     ProcessMeshAttribute mesh,
 
@@ -15,13 +15,13 @@
 #pragma once
 
 #include "paddle/phi/common/reduce_type.h"
+#include "paddle/phi/core/distributed/auto_parallel/placement_types.h"
 #include "paddle/phi/core/distributed/auto_parallel/process_mesh.h"
 #include "paddle/pir/include/core/attribute.h"
 #include "paddle/pir/include/core/builtin_attribute.h"
 #include "paddle/pir/include/core/builtin_attribute_storage.h"
 #include "paddle/pir/include/core/utils.h"
 #include "paddle/utils/flat_hash_map.h"
-
 namespace paddle {
 namespace dialect {
 class ProcessMeshAttrStorage;
@@ -76,6 +76,8 @@ class TensorDistAttribute : public pir::AttrBase<TensorDistAttribute,
 
   const flat_hash_map<int64_t, phi::ReduceType>& partial_status() const;
 
+  phi::distributed::Placements placements() const;
+
   // construct a new attribute with new mesh attribute.
   TensorDistAttribute CopyWithNewMesh(ProcessMeshAttribute mesh) const {
     return get(ir_context(), mesh, dims_mapping(), partial_status());
 
@@ -93,6 +93,7 @@
     'c_allreduce_avg_',
     'c_reduce_avg',
     'c_reduce_avg_',
+    'c_reducescatter',
     'c_allreduce_min_',
     'c_allreduce_prod_',
     'distributed_fused_lamb_init',
@@ -163,7 +164,6 @@
     'c_embedding',
     'c_identity',
     'c_reduce_sum',
-    'c_reducescatter',
     'c_softmax_with_cross_entropy',
     'c_split',
     'decayed_adagrad',
 
@@ -3078,6 +3078,9 @@ void ProcessBlock(
   auto inputs_by_data_op = GetInputsByDataOp(block);
   for (auto& [keyword, arg] : block->kwargs()) {
     auto new_arg = new_block->AddKwarg(keyword, arg.type());
+    for (auto& [name, attr] : arg.dyn_cast<pir::BlockArgument>().attributes()) {
+      new_arg.set_attribute(name, attr);
+    }
     (*map_value_pair)[arg] = new_arg;
     if (auto dense_tensor_type = arg.type().dyn_cast<DenseTensorType>()) {
       new_arg.set_type(
 
@@ -38,6 +38,7 @@ struct type_caster<paddle::flat_hash_map<Key, Value, Hash, Equal, Alloc>>
                  Value> {};
 }  // namespace pybind11::detail
 
+using paddle::dialect::DistTypeInterface;
 using paddle::dialect::OperationDistAttribute;
 using paddle::dialect::ProcessMeshAttribute;
 using paddle::dialect::TensorDistAttribute;
@@ -86,11 +87,19 @@ void BindTensorDistAttribute(py::module *m) {
       .def_property_readonly(
           "partial_status",
           [](TensorDistAttribute &self) { return self.partial_status(); })
-      .def_property_readonly("partial_dims", [](TensorDistAttribute &self) {
-        return self.partial_dims();
+      .def_property_readonly(
+          "partial_dims",
+          [](TensorDistAttribute &self) { return self.partial_dims(); })
+      .def_property_readonly("placements", [](TensorDistAttribute &self) {
+        return self.placements();
       });
 }
 
+void BindDistType(py::module *m) {
+  py::class_<DistTypeInterface, pir::Type> dist_type(*m, "DistType");
+  dist_type.def("dist_attr", &DistTypeInterface::tensor_dist_attr);
+}
+
 void BindDistOpsAPI(pybind11::module *module) {
   {
     if (PyModule_AddFunctions(module->ptr(), DistOpsAPI) < 0) {
@@ -140,6 +149,7 @@ void BindDistApi(pybind11::module *module) {
   auto ir_module = module->def_submodule("pir");
   BindOperationDistAttribute(&ir_module);
   BindTensorDistAttribute(&ir_module);
+  BindDistType(&ir_module);
   BindDistUtils(&ir_module);
   BindDistPassAPI(&ir_module);
   auto ops_modules = ir_module.def_submodule("ops");
 
@@ -40,6 +40,7 @@
 #include "paddle/fluid/pir/dialect/operator/ir/api_builder.h"
 #include "paddle/fluid/pir/dialect/operator/ir/control_flow_op.h"
 #include "paddle/fluid/pir/dialect/operator/ir/manual_pylayer_op.h"
+#include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_dialect.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_type.h"
 #include "paddle/fluid/pir/dialect/operator/ir/pd_api.h"
@@ -96,9 +97,9 @@ using paddle::dialect::SparseCsrTensorType;
 using paddle::dialect::WhileOp;
 using pir::TuplePopOp;
 
+using paddle::dialect::IntArrayAttribute;
 using paddle::dialect::OperationDistAttribute;
 using paddle::dialect::TensorDistAttribute;
-
 using pir::ArrayAttribute;
 using pir::Attribute;
 using pir::Block;
@@ -775,6 +776,14 @@ void BindOperation(py::module *m) {
                  attr_name,
                  pir::BoolAttribute::get(pir::IrContext::Instance(), flag));
            })
+      .def("set_int_array_attr",
+           [](Operation &self,
+              std::string &attr_name,
+              const std::vector<int64_t> &val) {
+             auto attr = IntArrayAttribute::get(pir::IrContext::Instance(),
+                                                phi::IntArray(val));
+             self.set_attribute(attr_name, attr);
+           })
       .def("attrs",
            [](Operation &self) -> py::dict {
              py::dict attrs_dict;
@@ -1297,20 +1306,20 @@ void BindValue(py::module *m) {
            [](Value self) { return self.type().isa<pir::VectorType>(); })
       .def("is_dist",
            [](Value self) { return self.type().isa<DistTypeInterface>(); })
-      // TODO(winter-wang): Move to python end: return self.type().dist_attr().
-      .def("dist_attr",
-           [](Value self) -> py::object {
-             auto type = self.type();
-             if (auto dist_type = type.dyn_cast<DistTypeInterface>()) {
-               return py::cast(dist_type.tensor_dist_attr());
-             } else {
-               return py::cast<py::none>(Py_None);
-             }
-           })
       // The function will calculate the new local shape based on the global
       // shape and the dist_attr argument.
-      .def("update_dist_attr", [](Value &self, TensorDistAttribute dist_attr) {
-        self.set_type(dialect::CvtToPirDistType(self.type(), dist_attr));
+      .def("update_dist_attr",
+           [](Value &self, TensorDistAttribute dist_attr) {
+             self.set_type(dialect::CvtToPirDistType(self.type(), dist_attr));
+           })
+      .def_property_readonly("process_mesh", [](Value &self) -> py::object {
+        auto type = self.type();
+        if (auto dist_type = type.dyn_cast<DistTypeInterface>()) {
+          return py::cast(
+              dist_type.tensor_dist_attr().process_mesh_attr().process_mesh());
+        } else {
+          return py::cast<py::none>(Py_None);
+        }
       });
 }
 
@@ -1356,14 +1365,6 @@ void BindType(py::module *m) {
             PADDLE_THROW(phi::errors::InvalidArgument(
                 "can't set dtype when building static graph"));
           })
-      .def("dist_attr",
-           [](Type self) -> py::object {
-             if (auto dist_type = self.dyn_cast<DistTypeInterface>()) {
-               return py::cast(dist_type.tensor_dist_attr());
-             } else {
-               return py::cast<py::none>(Py_None);
-             }
-           })
       .def_property(
           "_local_shape",
           [](Type self) {
@@ -1385,6 +1386,13 @@ void BindType(py::module *m) {
              }
              return py::cast<py::none>(Py_None);
            })
+      .def("as_dist_type",
+           [](Type &self) -> py::object {
+             if (auto dist_type = self.dyn_cast<DistTypeInterface>()) {
+               return py::cast(dist_type);
+             }
+             return py::cast<py::none>(Py_None);
+           })
       .def("__str__", [](Type &self) {
         std::ostringstream print_stream;
         print_stream << self;
 
@@ -14,6 +14,9 @@
 
 import atexit  # noqa: F401
 
+from .value_patch import monkey_patch_value_in_dist
+
+monkey_patch_value_in_dist()
 from paddle.base.core import Placement, ReduceType
 from paddle.distributed.fleet.base.topology import ParallelMode
 from paddle.distributed.fleet.dataset import InMemoryDataset, QueueDataset