PaddlePaddle
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/attribute_storage.h
Copy file name to clipboard
+9-4Lines changed: 9 additions & 4 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/attribute_storage.h
Copy file name to clipboard
+9-4Lines changed: 9 additions & 4 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_attribute.cc
Copy file name to clipboard
+5-2Lines changed: 5 additions & 2 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_attribute.cc
Copy file name to clipboard
+5-2Lines changed: 5 additions & 2 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_attribute.h
Copy file name to clipboard
+8-3Lines changed: 8 additions & 3 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_attribute.h
Copy file name to clipboard
+8-3Lines changed: 8 additions & 3 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_dialect.cc
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_dialect.cc
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/transforms/dist_to_dense_pass.cc
Copy file name to clipboard
+5Lines changed: 5 additions & 0 deletions b/‎paddle/fluid/pir/dialect/distributed/transforms/dist_to_dense_pass.cc
Copy file name to clipboard
+5Lines changed: 5 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/serialize_deserialize/include/deserialize_utils.h
Copy file name to clipboard
+4-1Lines changed: 4 additions & 1 deletion b/‎paddle/fluid/pir/serialize_deserialize/include/deserialize_utils.h
Copy file name to clipboard
+4-1Lines changed: 4 additions & 1 deletion
diff --git a/‎paddle/fluid/pir/serialize_deserialize/include/serialize_utils.h
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/fluid/pir/serialize_deserialize/include/serialize_utils.h
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/fluid/pybind/dist_api.cc
Copy file name to clipboard
+20-2Lines changed: 20 additions & 2 deletions b/‎paddle/fluid/pybind/dist_api.cc
Copy file name to clipboard
+20-2Lines changed: 20 additions & 2 deletions
diff --git a/‎python/paddle/autograd/ir_backward.py
Copy file name to clipboard
+26Lines changed: 26 additions & 0 deletions b/‎python/paddle/autograd/ir_backward.py
Copy file name to clipboard
+26Lines changed: 26 additions & 0 deletions
diff --git a/‎python/paddle/base/framework.py
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion b/‎python/paddle/base/framework.py
Copy file name to clipboard
+1-1Lines changed: 1 addition & 1 deletion
@@ -122,11 +122,14 @@ class OperationDistAttrStorage : public pir::AttributeStorage {
   ///
   using ParamKey = std::tuple<ProcessMeshAttribute,
                               std::vector<pir::Attribute>,
-                              std::vector<pir::Attribute>>;
-  OperationDistAttrStorage(ParamKey&& param)  // NOLINT
+                              std::vector<pir::Attribute>,
+                              int64_t>;
+
+  explicit OperationDistAttrStorage(ParamKey&& param)
       : mesh_attr(std::get<0>(param)),
         operands(std::get<1>(param)),
-        results(std::get<2>(param)) {}
+        results(std::get<2>(param)),
+        chunk_id(std::get<3>(param)) {}
 
   ///
   /// \brief Each derived TypeStorage must define a Construct method, which
@@ -149,6 +152,7 @@ class OperationDistAttrStorage : public pir::AttributeStorage {
       auto tmp_value = std::hash<pir::Attribute>()(iter);
       hash_value = pir::detail::hash_combine(hash_value, tmp_value);
     }
+    hash_value = pir::detail::hash_combine(hash_value, std::get<3>(key));
     return hash_value;
   }
 
@@ -157,12 +161,13 @@ class OperationDistAttrStorage : public pir::AttributeStorage {
   ///
   bool operator==(const ParamKey& key) const {
     return mesh_attr == std::get<0>(key) && operands == std::get<1>(key) &&
-           results == std::get<2>(key);
+           results == std::get<2>(key) && chunk_id == std::get<3>(key);
   }
 
   ProcessMeshAttribute mesh_attr;
   std::vector<pir::Attribute> operands;
   std::vector<pir::Attribute> results;
+  int64_t chunk_id;
 };
 
 }  // namespace dialect
 
@@ -126,11 +126,14 @@ uint32_t OperationDistAttribute::num_results() const {
   return results().size();
 }
 
+int64_t OperationDistAttribute::chunk_id() const { return storage()->chunk_id; }
+
 OperationDistAttribute OperationDistAttribute::get(
     pir::IrContext* ctx,
     ProcessMeshAttribute mesh,
     const std::vector<pir::Attribute>& operands,
-    const std::vector<pir::Attribute>& results) {
+    const std::vector<pir::Attribute>& results,
+    const int64_t& chunk_id) {
   auto check_dist_attr = [=](pir::Attribute attr) {
     auto dist_attr = attr.dyn_cast<TensorDistAttribute>();
     auto ids = mesh.process_ids();
@@ -156,7 +159,7 @@ OperationDistAttribute OperationDistAttribute::get(
       check_dist_attr(attr);
     }
   }
-  return Base::get(ctx, mesh, operands, results);
+  return Base::get(ctx, mesh, operands, results, chunk_id);
 }
 
 }  // namespace dialect
 
@@ -121,16 +121,21 @@ class OperationDistAttribute : public pir::AttrBase<OperationDistAttribute,
 
   uint32_t num_results() const;
 
+  int64_t chunk_id() const;
+
   static OperationDistAttribute get(pir::IrContext* ctx,
                                     ProcessMeshAttribute mesh,
                                     const std::vector<Attribute>& operands,
-                                    const std::vector<Attribute>& results);
+                                    const std::vector<Attribute>& results,
+                                    const int64_t& chunk_id = -1);
 
   static OperationDistAttribute get(pir::IrContext* ctx,
                                     const phi::distributed::ProcessMesh& mesh,
                                     const std::vector<Attribute>& operands,
-                                    const std::vector<Attribute>& results) {
-    return get(ctx, ProcessMeshAttribute::get(ctx, mesh), operands, results);
+                                    const std::vector<Attribute>& results,
+                                    const int64_t& chunk_id = -1) {
+    return get(
+        ctx, ProcessMeshAttribute::get(ctx, mesh), operands, results, chunk_id);
   }
 
   static std::string name() { return "a_op_dist"; }
 
@@ -113,6 +113,7 @@ void DistDialect::PrintAttribute(pir::Attribute attr, std::ostream &os) const {
       os << ",result(" + std::to_string(i) + "):{" << op_dist_attr.result(i)
          << "}";
     }
+    os << ",chunk_id:" << op_dist_attr.chunk_id();
     os << "}";
   } else {
     os << "error_attribute_type";
 
@@ -142,9 +142,14 @@ void ProcessDistBlock(pir::Block* block) {
     //     common::errors::PreconditionNotMet("The op [%s] has not
     //     op_dist_attr.",
     //                                        op_item->name()));
+
+    int64_t chunk_id = -1;
     if (op_item->HasAttribute(kAttrOpDistAttr)) {
+      chunk_id = op_item->attribute<OperationDistAttribute>(kAttrOpDistAttr)
+                     .chunk_id();
       op_item->erase_attribute(kAttrOpDistAttr);
     }
+    op_item->set_attribute("chunk_id", pir::Int64Attribute::get(ctx, chunk_id));
 
     // TODO(2024-Q2) Handle other special dist op in future.
   }
 
@@ -339,8 +339,11 @@ paddle::dialect::OperationDistAttribute deserializeOperationDistAttr(
   for (auto& item : results_json) {
     results.push_back(parseAttr(&item));
   }
+
+  Json chunk_id_json = data_json.at(3);
+  int64_t chunk_id = chunk_id_json.get<int64_t>();
   return paddle::dialect::OperationDistAttribute::get(
-      ctx, mesh, operands, results);
+      ctx, mesh, operands, results, chunk_id);
 }
 
 pir::Attribute AttrTypeReader::ReadBuiltInAttr(const std::string attr_name,
 
@@ -344,6 +344,7 @@ Json serializeAttrToJson<paddle::dialect::OperationDistAttribute>(
     results_json.push_back(writeAttr(attr.results().at(i)));
   }
   content.push_back(results_json);
+  content.push_back(attr.chunk_id());
 
   json_obj[DATA] = content;
   return json_obj;
 
@@ -23,6 +23,7 @@
 #include "paddle/fluid/pybind/dist_static_op_function.h"
 #include "paddle/phi/core/distributed/auto_parallel/reshard/reshard_utils.h"
 #include "paddle/phi/core/enforce.h"
+#include "paddle/pir/include/core/builtin_attribute.h"
 
 namespace py = pybind11;
 
@@ -60,6 +61,9 @@ void BindOperationDistAttribute(py::module *m) {
                              [](OperationDistAttribute &self) {
                                return self.process_mesh_attr().process_mesh();
                              })
+      .def_property_readonly(
+          "chunk_id",
+          [](OperationDistAttribute &self) { return self.chunk_id(); })
       .def("num_operands", &OperationDistAttribute::num_operands)
       .def("operands", &OperationDistAttribute::operands)
       .def("operand", &OperationDistAttribute::operand)
@@ -123,9 +127,10 @@ TensorDistAttribute CreateTensorDistAttribute(
 OperationDistAttribute CreateOperationDistAttribute(
     const phi::distributed::ProcessMesh &mesh,
     const std::vector<pir::Attribute> &operands,
-    const std::vector<pir::Attribute> &results) {
+    const std::vector<pir::Attribute> &results,
+    const int64_t &chunk_id) {
   return OperationDistAttribute::get(
-      pir::IrContext::Instance(), mesh, operands, results);
+      pir::IrContext::Instance(), mesh, operands, results, chunk_id);
 }
 
 ArrayAttribute CreateArrayAttribute(
@@ -135,7 +140,20 @@ ArrayAttribute CreateArrayAttribute(
 
 void BindDistUtils(pybind11::module *m) {
   m->def("create_tensor_dist_attribute", CreateTensorDistAttribute);
+  m->def("create_array_dist_attribute", CreateArrayAttribute);
   m->def("create_op_dist_attribute", CreateOperationDistAttribute);
+  m->def("create_op_dist_attribute",
+         &CreateOperationDistAttribute,
+         py::arg("mesh"),
+         py::arg("operands"),
+         py::arg("results"),
+         py::arg("chunk_id") = -1);
+  m->def("create_process_mesh",
+         [](const std::vector<int64_t> &shape,
+            const std::vector<int64_t> &process_ids,
+            const std::vector<std::string> &dim_names) {
+           return phi::distributed::ProcessMesh(shape, process_ids, dim_names);
+         });
   m->def("create_array_attribute", CreateArrayAttribute);
   m->def("get_sub_meshes", phi::distributed::GetSubMeshes);
   m->def("cvt_to_dist_type",
 
@@ -1050,11 +1050,36 @@ def calc_gradient_helper(
         state,
         ValueDict(),
     )
+
     # now value_to_valuegrad should be value <-> value (add sum op for the same values's grad value)
     outputs_set, inputs_set, no_gradvar_set = create_backward_prune_set(
         outputs_fwd_set, inputs_fwd_set, no_grad_set, state
     )
 
+    # set struct name for grad ops
+    for op in block.ops:
+        if op in state.op_to_opgrad:
+            if op.dist_attr is None:
+                continue
+
+            op_chunk_id = op.dist_attr.chunk_id
+            if op_chunk_id == -1 and op.name() == "dist_op.reshard":
+                op_chunk_id = (
+                    op.operand_source(0).get_defining_op().dist_attr.chunk_id
+                )
+
+            for bwd_op in state.op_to_opgrad[op]:
+                if bwd_op.dist_attr is None:
+                    continue
+                bwd_op.dist_attr = (
+                    paddle.base.libpaddle.pir.create_op_dist_attribute(
+                        bwd_op.dist_attr.process_mesh,
+                        bwd_op.dist_attr.operands(),
+                        bwd_op.dist_attr.results(),
+                        op_chunk_id,
+                    )
+                )
+
     remove_ops = []
     if not is_inplace_net(backward_ops) and inputs:
         _, remove_ops = prune_ops(
@@ -1077,6 +1102,7 @@ def calc_gradient_helper(
         if bwd_op.result(0).use_empty():
             remove_op(block, bwd_op, state)
     state.turn_map()
+
     input_grad_map = state.value_to_valuegrad
 
     return input_grad_map
 
@@ -1298,7 +1298,7 @@ def name_struct(prefix=None):
                 for idx in reversed(range(op_num_before, op_num)):
                     op = all_ops[idx]
                     if op.has_attr("struct_name"):
-                        break
+                        continue
                     op.set_str_attr("struct_name", _full_name_struct())
 
             _name_struct = _name_struct.parent()
Original file line number	Diff line number	Diff line change
`@@ -126,11 +126,14 @@ uint32_t OperationDistAttribute::num_results() const {`
`126`	`126`	`return results().size();`
`127`	`127`	`}`
`128`	`128`
	`129`	`+int64_t OperationDistAttribute::chunk_id() const { return storage()->chunk_id; }`
	`130`	`+`
`129`	`131`	`OperationDistAttribute OperationDistAttribute::get(`
`130`	`132`	`pir::IrContext* ctx,`
`131`	`133`	`ProcessMeshAttribute mesh,`
`132`	`134`	`const std::vector<pir::Attribute>& operands,`
`133`		`- const std::vector<pir::Attribute>& results) {`
	`135`	`+ const std::vector<pir::Attribute>& results,`
	`136`	`+ const int64_t& chunk_id) {`
`134`	`137`	`auto check_dist_attr = [=](pir::Attribute attr) {`
`135`	`138`	`auto dist_attr = attr.dyn_cast<TensorDistAttribute>();`
`136`	`139`	`auto ids = mesh.process_ids();`
`@@ -156,7 +159,7 @@ OperationDistAttribute OperationDistAttribute::get(`
`156`	`159`	`check_dist_attr(attr);`
`157`	`160`	`}`
`158`	`161`	`}`
`159`		`- return Base::get(ctx, mesh, operands, results);`
	`162`	`+ return Base::get(ctx, mesh, operands, results, chunk_id);`
`160`	`163`	`}`
`161`	`164`
`162`	`165`	`} // namespace dialect`
Original file line number	Diff line number	Diff line change
`@@ -113,6 +113,7 @@ void DistDialect::PrintAttribute(pir::Attribute attr, std::ostream &os) const {`
`113`	`113`	`os << ",result(" + std::to_string(i) + "):{" << op_dist_attr.result(i)`
`114`	`114`	`<< "}";`
`115`	`115`	`}`
	`116`	`+ os << ",chunk_id:" << op_dist_attr.chunk_id();`
`116`	`117`	`os << "}";`
`117`	`118`	`} else {`
`118`	`119`	`os << "error_attribute_type";`
Original file line number	Diff line number	Diff line change
`@@ -142,9 +142,14 @@ void ProcessDistBlock(pir::Block* block) {`
`142`	`142`	`// common::errors::PreconditionNotMet("The op [%s] has not`
`143`	`143`	`// op_dist_attr.",`
`144`	`144`	`// op_item->name()));`
	`145`	`+`
	`146`	`+ int64_t chunk_id = -1;`
`145`	`147`	`if (op_item->HasAttribute(kAttrOpDistAttr)) {`
	`148`	`+ chunk_id = op_item->attribute<OperationDistAttribute>(kAttrOpDistAttr)`
	`149`	`+ .chunk_id();`
`146`	`150`	`op_item->erase_attribute(kAttrOpDistAttr);`
`147`	`151`	`}`
	`152`	`+ op_item->set_attribute("chunk_id", pir::Int64Attribute::get(ctx, chunk_id));`
`148`	`153`
`149`	`154`	`// TODO(2024-Q2) Handle other special dist op in future.`
`150`	`155`	`}`
Original file line number	Diff line number	Diff line change
`@@ -339,8 +339,11 @@ paddle::dialect::OperationDistAttribute deserializeOperationDistAttr(`
`339`	`339`	`for (auto& item : results_json) {`
`340`	`340`	`results.push_back(parseAttr(&item));`
`341`	`341`	`}`
	`342`	`+`
	`343`	`+ Json chunk_id_json = data_json.at(3);`
	`344`	`+ int64_t chunk_id = chunk_id_json.get<int64_t>();`
`342`	`345`	`return paddle::dialect::OperationDistAttribute::get(`
`343`		`- ctx, mesh, operands, results);`
	`346`	`+ ctx, mesh, operands, results, chunk_id);`
`344`	`347`	`}`
`345`	`348`
`346`	`349`	`pir::Attribute AttrTypeReader::ReadBuiltInAttr(const std::string attr_name,`
Original file line number	Diff line number	Diff line change
`@@ -344,6 +344,7 @@ Json serializeAttrToJson<paddle::dialect::OperationDistAttribute>(`
`344`	`344`	`results_json.push_back(writeAttr(attr.results().at(i)));`
`345`	`345`	`}`
`346`	`346`	`content.push_back(results_json);`
	`347`	`+ content.push_back(attr.chunk_id());`
`347`	`348`
`348`	`349`	`json_obj[DATA] = content;`
`349`	`350`	`return json_obj;`