PaddlePaddle
diff --git a/‎paddle/fluid/eager/api/manual/eager_manual/dygraph_forward_api.h
Copy file name to clipboard
+4-1Lines changed: 4 additions & 1 deletion b/‎paddle/fluid/eager/api/manual/eager_manual/dygraph_forward_api.h
Copy file name to clipboard
+4-1Lines changed: 4 additions & 1 deletion
diff --git a/‎paddle/fluid/eager/api/manual/eager_manual/forwards/dtensor_to_local_fwd_func.cc
Copy file name to clipboard
+9-1Lines changed: 9 additions & 1 deletion b/‎paddle/fluid/eager/api/manual/eager_manual/forwards/dtensor_to_local_fwd_func.cc
Copy file name to clipboard
+9-1Lines changed: 9 additions & 1 deletion
diff --git a/‎paddle/fluid/eager/api/manual/eager_manual/nodes/dtensor_to_local_node.cc
Copy file name to clipboard
+2-4Lines changed: 2 additions & 4 deletions b/‎paddle/fluid/eager/api/manual/eager_manual/nodes/dtensor_to_local_node.cc
Copy file name to clipboard
+2-4Lines changed: 2 additions & 4 deletions
diff --git a/‎paddle/fluid/eager/api/manual/eager_manual/nodes/nodes.h
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions b/‎paddle/fluid/eager/api/manual/eager_manual/nodes/nodes.h
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_api.cc
Copy file name to clipboard
+17-3Lines changed: 17 additions & 3 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_api.cc
Copy file name to clipboard
+17-3Lines changed: 17 additions & 3 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_api.h
Copy file name to clipboard
+7-1Lines changed: 7 additions & 1 deletion b/‎paddle/fluid/pir/dialect/distributed/ir/dist_api.h
Copy file name to clipboard
+7-1Lines changed: 7 additions & 1 deletion
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_op.cc
Copy file name to clipboard
+17-4Lines changed: 17 additions & 4 deletions b/‎paddle/fluid/pir/dialect/distributed/ir/dist_op.cc
Copy file name to clipboard
+17-4Lines changed: 17 additions & 4 deletions
diff --git a/‎paddle/fluid/pir/dialect/distributed/ir/dist_op.h
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion b/‎paddle/fluid/pir/dialect/distributed/ir/dist_op.h
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion
diff --git a/‎paddle/fluid/pir/dialect/operator/utils/utils.cc
Copy file name to clipboard
+47-43Lines changed: 47 additions & 43 deletions b/‎paddle/fluid/pir/dialect/operator/utils/utils.cc
Copy file name to clipboard
+47-43Lines changed: 47 additions & 43 deletions
diff --git a/‎paddle/fluid/pybind/auto_parallel_py.cc
Copy file name to clipboard
+6-2Lines changed: 6 additions & 2 deletions b/‎paddle/fluid/pybind/auto_parallel_py.cc
Copy file name to clipboard
+6-2Lines changed: 6 additions & 2 deletions
@@ -57,7 +57,10 @@ paddle::Tensor reshard_ad_function(
     const paddle::Tensor& tensor,
     const phi::distributed::TensorDistAttr dist_attr);
 
-paddle::Tensor dtensor_to_local_ad_function(const paddle::Tensor& input);
+paddle::Tensor dtensor_to_local_ad_function(
+    const paddle::Tensor& input,
+    const phi::distributed::ProcessMesh& processmesh,
+    const phi::distributed::Placements& placements);
 
 paddle::Tensor dtensor_from_local_ad_function(
     const paddle::Tensor& input,
 
@@ -18,7 +18,10 @@
 #include "paddle/fluid/eager/api/utils/global_utils.h"
 #include "paddle/phi/core/platform/profiler/event_tracing.h"
 
-paddle::Tensor dtensor_to_local_ad_function(const paddle::Tensor& input) {
+paddle::Tensor dtensor_to_local_ad_function(
+    const paddle::Tensor& input,
+    const phi::distributed::ProcessMesh& process_mesh,
+    const phi::distributed::Placements& placements) {
 #ifdef PADDLE_WITH_DISTRIBUTE
   VLOG(3) << "Running AD API: "
           << "dtensor_to_local dygraph";
@@ -49,6 +52,11 @@ paddle::Tensor dtensor_to_local_ad_function(const paddle::Tensor& input) {
 
     // Set TensorWrappers for Forward Inputs if needed
     grad_node->SetTensorWrapperNoNeedBuffer_Input(input);
+
+    phi::distributed::TensorDistAttr grad_dist_attr =
+        ToTensorDistAttr(process_mesh, placements, input.dims());
+
+    grad_node->SetGradDistAttr(grad_dist_attr);
   }
 
   // Forward API Call
 
@@ -47,9 +47,7 @@ DtensorToLocalGradNode::operator()(
 
   // Collect GradIn Tensors, Attrs and Recovered TensorWrappers
   auto input = egr::EagerUtils::RecoverTensorWrapper(&this->input_);
-  const auto& dist_attr =
-      std::static_pointer_cast<phi::distributed::DistTensor>(input.impl())
-          ->dist_attr();
+
   auto& grad_out = hooked_grad[0][0];
   // Prepare Grad function call
 
@@ -82,7 +80,7 @@ DtensorToLocalGradNode::operator()(
 
   // Backward call dtensor_to_local_func function
   auto dist_grad_ptr = std::make_shared<phi::distributed::DistTensor>(
-      grad_out.dims(), dist_attr);
+      grad_out.dims(), grad_dist_attr_);
 
   *(dist_grad_ptr->unsafe_mutable_value()) =
       *(static_cast<phi::DenseTensor*>(grad_out.impl().get()));
 
@@ -489,9 +489,15 @@ class DtensorToLocalGradNode : public egr::GradNodeBase {
     input_ = egr::TensorWrapper(input, true);
   }
 
+  void SetGradDistAttr(const phi::distributed::TensorDistAttr& dist_attr) {
+    grad_dist_attr_ = dist_attr;
+  }
+
  private:
   // TensorWrappers
   egr::TensorWrapper input_;
+
+  phi::distributed::TensorDistAttr grad_dist_attr_;
 };
 
 class DtensorFromLocalGradNode : public egr::GradNodeBase {
 
@@ -85,9 +85,23 @@ pir::Value dtensor_from_local(const pir::Value& x,
       .result(0);
 }
 
-pir::Value dtensor_to_local(const pir::Value& x) {
-  return ApiBuilder::Instance().GetBuilder()->Build<DtensorToLocalOp>(x).result(
-      0);
+pir::Value dtensor_to_local(
+    const pir::Value& x,
+    const phi::distributed::ProcessMesh& process_mesh,
+    const std::vector<int64_t>& dims_mapping,
+    const flat_hash_map<int64_t, phi::ReduceType>& partial_status) {
+  pir::IrContext* ctx = pir::IrContext::Instance();
+  TensorDistAttribute grad_dist_attr =
+      TensorDistAttribute::get(ctx, process_mesh, dims_mapping, partial_status);
+  return dtensor_to_local(x, grad_dist_attr);
+}
+
+pir::Value dtensor_to_local(const pir::Value& x,
+                            const TensorDistAttribute& grad_dist_attr) {
+  return ApiBuilder::Instance()
+      .GetBuilder()
+      ->Build<DtensorToLocalOp>(x, grad_dist_attr)
+      .result(0);
 }
 
 std::vector<pir::Value> moe_sub_mesh_tensors(
 
@@ -50,7 +50,13 @@ pir::Value dtensor_from_local(
 pir::Value dtensor_from_local(const pir::Value& x,
                               const TensorDistAttribute& tensor_dist_attr);
 
-pir::Value dtensor_to_local(const pir::Value& x);
+pir::Value dtensor_to_local(
+    const pir::Value& x,
+    const phi::distributed::ProcessMesh& process_mesh,
+    const std::vector<int64_t>& dims_mapping,
+    const flat_hash_map<int64_t, phi::ReduceType>& partial_status = {});
+pir::Value dtensor_to_local(const pir::Value& x,
+                            const TensorDistAttribute& grad_dist_attr);
 
 std::vector<pir::Value> moe_sub_mesh_tensors(
     const pir::Value& input,
 
@@ -413,7 +413,8 @@ std::vector<std::vector<pir::Value>> DtensorFromLocalOp::Vjp(
                    ->result(0);
   }
 
-  auto grad_op = builder.Build<DtensorToLocalOp>(out_grad);
+  auto grad_op = builder.Build<DtensorToLocalOp>(
+      out_grad, dist_type.tensor_dist_attr() /*unused*/);
 
   VLOG(6) << "End call vjp for dtensor_from_local op.";
 
@@ -422,13 +423,23 @@ std::vector<std::vector<pir::Value>> DtensorFromLocalOp::Vjp(
 
 void DtensorToLocalOp::Build(pir::Builder& builder,
                              pir::OperationArgument& argument,
-                             pir::Value input) {
+                             pir::Value input,
+                             TensorDistAttribute grad_dist_attr) {
   VLOG(4) << "Start build DtensorToLocalOp";
+  paddle::dialect::DistDenseTensorType input_tensor_type;
+  if (input.type().isa<paddle::dialect::DistDenseTensorType>()) {
+    input_tensor_type =
+        input.type().dyn_cast<paddle::dialect::DistDenseTensorType>();
+  } else {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "Only support paddle::dialect::DistDenseTensorType"));
+  }
 
   VLOG(4) << "Builder construction inputs";
   argument.AddInput(input);
 
   VLOG(4) << "Builder construction attributes";
+  argument.AddAttribute("grad_dist_attr", grad_dist_attr);
 
   VLOG(4) << "Builder construction outputs";
 
@@ -494,9 +505,11 @@ std::vector<std::vector<pir::Value>> DtensorToLocalOp::Vjp(
           "dtensor_from_local op's outputs grad[0] size should be 1"));
 
   auto& builder = *ApiBuilder::Instance().GetBuilder();
+  const auto& grad_dist_attr =
+      op->attribute<paddle::dialect::TensorDistAttribute>("grad_dist_attr");
 
-  auto grad_op = builder.Build<DtensorFromLocalOp>(
-      out_grads[0][0], dist_type.tensor_dist_attr());
+  auto grad_op =
+      builder.Build<DtensorFromLocalOp>(out_grads[0][0], grad_dist_attr);
 
   VLOG(6) << "End call vjp for dtensor_from_local op.";
 
 
@@ -94,7 +94,8 @@ class DtensorToLocalOp
   static constexpr uint32_t attributes_num = 0;
   TEST_API static void Build(pir::Builder& builder,             // NOLINT
                              pir::OperationArgument& argument,  // NOLINT
-                             pir::Value input);
+                             pir::Value input,
+                             TensorDistAttribute grad_dist_attr);
 
   static OpInfoTuple GetOpInfo();
   static std::vector<std::vector<pir::Value>> Vjp(
 
@@ -18,6 +18,7 @@
 
 #include "paddle/common/errors.h"
 #include "paddle/fluid/framework/phi_utils.h"
+#include "paddle/fluid/pir/dialect/distributed/ir/dist_attribute.h"
 #include "paddle/fluid/pir/dialect/kernel/ir/kernel_type.h"
 #include "paddle/fluid/pir/dialect/operator/ir/manual_op.h"
 #include "paddle/fluid/pir/dialect/operator/ir/op_attribute.h"
@@ -57,43 +58,43 @@ enum class AttrType {
   DOUBLE,
 
   ARRAY,
+  STRING,
+  TENSOR_NAME,
+  DATA_TYPE,
   INT_ARRAY,
+  PLACE,
+  TensorDist,
 
   SCALAR,
-  DATA_TYPE,
   DATA_LAYOUT,
-  PLACE,
-
-  STRING,
-
-  TENSOR_NAME,
-
   NUM_ATTR_TYPES,
 };
 
 static inline AttrType GetAttributeType(const pir::Attribute& attr) {
   if (attr.isa<pir::BoolAttribute>()) {
     return AttrType::BOOL;
-  } else if (attr.isa<pir::FloatAttribute>()) {
-    return AttrType::FLOAT;
-  } else if (attr.isa<pir::DoubleAttribute>()) {
-    return AttrType::DOUBLE;
   } else if (attr.isa<pir::Int32Attribute>()) {
     return AttrType::INT32;
   } else if (attr.isa<pir::Int64Attribute>()) {
     return AttrType::INT64;
+  } else if (attr.isa<pir::FloatAttribute>()) {
+    return AttrType::FLOAT;
+  } else if (attr.isa<pir::DoubleAttribute>()) {
+    return AttrType::DOUBLE;
   } else if (attr.isa<pir::ArrayAttribute>()) {
     return AttrType::ARRAY;
   } else if (attr.isa<pir::StrAttribute>()) {
     return AttrType::STRING;
-  } else if (attr.isa<paddle::dialect::IntArrayAttribute>()) {
-    return AttrType::INT_ARRAY;
+  } else if (attr.isa<pir::TensorNameAttribute>()) {
+    return AttrType::TENSOR_NAME;
   } else if (attr.isa<paddle::dialect::DataTypeAttribute>()) {
     return AttrType::DATA_TYPE;
+  } else if (attr.isa<paddle::dialect::IntArrayAttribute>()) {
+    return AttrType::INT_ARRAY;
   } else if (attr.isa<paddle::dialect::PlaceAttribute>()) {
     return AttrType::PLACE;
-  } else if (attr.isa<pir::TensorNameAttribute>()) {
-    return AttrType::TENSOR_NAME;
+  } else if (attr.isa<paddle::dialect::TensorDistAttribute>()) {
+    return AttrType::TensorDist;
   } else {
     PADDLE_THROW(common::errors::Unimplemented(
         "Unsupported ir Attribute type when casting it into "
@@ -110,14 +111,6 @@ static std::function<T(const pir::Attribute& attr)> GetAttrCast(
            [](const pir::Attribute& attr) {
              return T{attr.dyn_cast<pir::BoolAttribute>().data()};
            }},
-          {AttrType::FLOAT,
-           [](const pir::Attribute& attr) {
-             return T{attr.dyn_cast<pir::FloatAttribute>().data()};
-           }},
-          {AttrType::DOUBLE,
-           [](const pir::Attribute& attr) {
-             return T{attr.dyn_cast<pir::DoubleAttribute>().data()};
-           }},
           {AttrType::INT32,
            [](const pir::Attribute& attr) {
              return T{attr.dyn_cast<pir::Int32Attribute>().data()};
@@ -126,28 +119,13 @@ static std::function<T(const pir::Attribute& attr)> GetAttrCast(
            [](const pir::Attribute& attr) {
              return T{attr.dyn_cast<pir::Int64Attribute>().data()};
            }},
-          {AttrType::INT_ARRAY,
-           [](const pir::Attribute& attr) {
-             return T{attr.dyn_cast<paddle::dialect::IntArrayAttribute>()
-                          .data()
-                          .GetData()};
-           }},
-          {AttrType::STRING,
-           [](const pir::Attribute& attr) {
-             return T{attr.dyn_cast<pir::StrAttribute>().AsString()};
-           }},
-          {AttrType::DATA_TYPE,
-           [](const pir::Attribute& attr) {
-             return T{
-                 attr.dyn_cast<paddle::dialect::DataTypeAttribute>().data()};
-           }},
-          {AttrType::PLACE,
+          {AttrType::FLOAT,
            [](const pir::Attribute& attr) {
-             return T{attr.dyn_cast<paddle::dialect::PlaceAttribute>().data()};
+             return T{attr.dyn_cast<pir::FloatAttribute>().data()};
            }},
-          {AttrType::TENSOR_NAME,
+          {AttrType::DOUBLE,
            [](const pir::Attribute& attr) {
-             return T{attr.dyn_cast<pir::TensorNameAttribute>().data()};
+             return T{attr.dyn_cast<pir::DoubleAttribute>().data()};
            }},
           {AttrType::ARRAY,
            [](const pir::Attribute& attr) {
@@ -211,7 +189,33 @@ static std::function<T(const pir::Attribute& attr)> GetAttrCast(
                    "vector."));
              }
            }},
-      };
+          {AttrType::STRING,
+           [](const pir::Attribute& attr) {
+             return T{attr.dyn_cast<pir::StrAttribute>().AsString()};
+           }},
+
+          {AttrType::TENSOR_NAME,
+           [](const pir::Attribute& attr) {
+             return T{attr.dyn_cast<pir::TensorNameAttribute>().data()};
+           }},
+          {AttrType::DATA_TYPE,
+           [](const pir::Attribute& attr) {
+             return T{
+                 attr.dyn_cast<paddle::dialect::DataTypeAttribute>().data()};
+           }},
+          {AttrType::INT_ARRAY,
+           [](const pir::Attribute& attr) {
+             return T{attr.dyn_cast<paddle::dialect::IntArrayAttribute>()
+                          .data()
+                          .GetData()};
+           }},
+          {AttrType::PLACE,
+           [](const pir::Attribute& attr) {
+             return T{attr.dyn_cast<paddle::dialect::PlaceAttribute>().data()};
+           }},
+          {AttrType::TensorDist, [](const pir::Attribute& attr) {
+             return T{attr.dyn_cast<paddle::dialect::TensorDistAttribute>()};
+           }}};
   return kAttrCastMap[attr_type];
 }
 
 
@@ -803,9 +803,13 @@ void BindAutoParallel(py::module *m) {
 
   m->def(
       "dtensor_to_local",
-      [](py::handle py_tensor) {
+      [](py::handle py_tensor,
+         py::handle py_process_mesh,
+         py::handle py_placements) {
         auto tensor = CastPyArg2Tensor(py_tensor.ptr(), 0);
-        return dtensor_to_local_ad_function(tensor);
+        auto process_mesh = CastPyArg2ProcessMesh(py_process_mesh.ptr(), 1);
+        auto placements = CastPyArg2VectorOfPlacement(py_placements.ptr(), 2);
+        return dtensor_to_local_ad_function(tensor, process_mesh, placements);
       },
       py::return_value_policy::reference);