PaddlePaddle
diff --git a/‎paddle/fluid/eager/api/manual/eager_manual/dygraph_forward_api.h
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions b/‎paddle/fluid/eager/api/manual/eager_manual/dygraph_forward_api.h
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions
diff --git a/‎paddle/fluid/eager/api/manual/eager_manual/forwards/CMakeLists.txt
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/fluid/eager/api/manual/eager_manual/forwards/CMakeLists.txt
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/fluid/eager/api/manual/eager_manual/forwards/reshard_fwd_func.cc
Copy file name to clipboard
+88Lines changed: 88 additions & 0 deletions b/‎paddle/fluid/eager/api/manual/eager_manual/forwards/reshard_fwd_func.cc
Copy file name to clipboard
+88Lines changed: 88 additions & 0 deletions
diff --git a/‎paddle/fluid/eager/api/manual/eager_manual/nodes/CMakeLists.txt
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/fluid/eager/api/manual/eager_manual/nodes/CMakeLists.txt
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/fluid/eager/api/manual/eager_manual/nodes/nodes.h
Copy file name to clipboard
+47Lines changed: 47 additions & 0 deletions b/‎paddle/fluid/eager/api/manual/eager_manual/nodes/nodes.h
Copy file name to clipboard
+47Lines changed: 47 additions & 0 deletions
diff --git a/‎paddle/fluid/eager/api/manual/eager_manual/nodes/reshard_node.cc
Copy file name to clipboard
+106Lines changed: 106 additions & 0 deletions b/‎paddle/fluid/eager/api/manual/eager_manual/nodes/reshard_node.cc
Copy file name to clipboard
+106Lines changed: 106 additions & 0 deletions
diff --git a/‎paddle/fluid/pybind/auto_parallel_py.cc
Copy file name to clipboard
+2-27Lines changed: 2 additions & 27 deletions b/‎paddle/fluid/pybind/auto_parallel_py.cc
Copy file name to clipboard
+2-27Lines changed: 2 additions & 27 deletions
diff --git a/‎paddle/phi/api/include/tensor_utils.h
Copy file name to clipboard
+22Lines changed: 22 additions & 0 deletions b/‎paddle/phi/api/include/tensor_utils.h
Copy file name to clipboard
+22Lines changed: 22 additions & 0 deletions
@@ -15,6 +15,7 @@
 #pragma once
 
 #include "paddle/phi/api/include/tensor.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
 
 paddle::Tensor add_n_ad_func(const std::vector<paddle::Tensor>& x);
 
@@ -49,6 +50,11 @@ sync_batch_norm__ad_func(const paddle::Tensor& x,
                          std::string data_layout,
                          bool use_global_stats,
                          bool trainable_statistics);
+
+paddle::Tensor reshard_ad_function(
+    const paddle::Tensor& tensor,
+    const phi::distributed::TensorDistAttr dist_attr);
+
 namespace sparse {
 std::tuple<paddle::Tensor,
            paddle::Tensor&,
 
@@ -3,4 +3,5 @@ set(eager_manual_functions
     ${PADDLE_SOURCE_DIR}/paddle/fluid/eager/api/manual/eager_manual/forwards/conv2d_fwd_function.cc
     ${PADDLE_SOURCE_DIR}/paddle/fluid/eager/api/manual/eager_manual/forwards/sync_batch_norm_fwd_func.cc
     ${PADDLE_SOURCE_DIR}/paddle/fluid/eager/api/manual/eager_manual/forwards/multiply_fwd_func.cc
+    ${PADDLE_SOURCE_DIR}/paddle/fluid/eager/api/manual/eager_manual/forwards/reshard_fwd_func.cc
     PARENT_SCOPE)
@@ -0,0 +1,88 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "paddle/fluid/eager/accumulation/accumulation_node.h"
+#include "paddle/fluid/eager/api/manual/eager_manual/dygraph_forward_api.h"
+#include "paddle/fluid/eager/api/manual/eager_manual/nodes/nodes.h"
+#include "paddle/fluid/eager/api/utils/global_utils.h"
+#include "paddle/fluid/platform/profiler/event_tracing.h"
+
+paddle::Tensor reshard_ad_function(
+    const paddle::Tensor& input,
+    const phi::distributed::TensorDistAttr dist_attr) {
+#ifdef PADDLE_WITH_DISTRIBUTE
+  VLOG(3) << "Running AD API: "
+          << "reshard dygraph";
+  // Dygraph Record Event
+  paddle::platform::RecordEvent dygraph_entrance_record_event(
+      "reshard dygraph", paddle::platform::TracerEventType::Communication, 1);
+
+  // Get Input AutoGradMeta
+  egr::AutogradMeta* input_autograd_meta =
+      egr::EagerUtils::nullable_autograd_meta(input);
+  bool trace_backward = egr::Controller::Instance().HasGrad();
+  bool require_any_grad =
+      egr::EagerUtils::ComputeRequireGrad(trace_backward, input_autograd_meta);
+
+  // Node Declaration
+  std::shared_ptr<ReshardGradNode> grad_node;
+
+  // Set grad_node before API Call
+  if (require_any_grad) {
+    paddle::platform::RecordEvent node_creation_record_event(
+        "reshard node_creation",
+        paddle::platform::TracerEventType::Communication,
+        1);
+
+    // Node Construction
+    grad_node =
+        std::shared_ptr<ReshardGradNode>(new ReshardGradNode(1, 1));  // NOLINT
+
+    // Set TensorWrappers for Forward Inputs if needed
+    grad_node->SetTensorWrapperNoNeedBufferInput(input);
+  }
+
+  // Forward API Call
+  // reshard_func(input, api_result, dist_attr);
+  auto dist_out_ptr = paddle::reshard(input, dist_attr);
+  auto api_result = paddle::Tensor(dist_out_ptr);
+
+  // Get Outputs
+  auto& out = api_result;
+
+  // Get Output AutoGradMeta
+  egr::AutogradMeta* out_autograd_meta = egr::EagerUtils::autograd_meta(&out);
+
+  // Set grad_node after API call
+  if (require_any_grad) {
+    egr::EagerUtils::PassStopGradient(false, out_autograd_meta);
+
+    // SetGradOutMeta & SetEdges
+    grad_node->SetGradOutMeta(input, 0);
+    // SetOutRank & SetHistory & SetGradInMeta
+    if (out_autograd_meta) {
+      egr::EagerUtils::SetOutRankWithSlot(out_autograd_meta, 0);
+      egr::EagerUtils::SetHistory(out_autograd_meta, grad_node);
+    }
+    grad_node->SetGradInMeta(out, 0);
+  }
+
+  return out;
+#else
+  PADDLE_THROW(phi::errors::Unavailable(
+      "Reshard is not supported in this version of Paddle. Try to recompile it "
+      "with WITH_DISTRIBTUE=ON and reinstall this package."));
+  return paddle::Tensor();
+#endif
+}
@@ -3,4 +3,5 @@ set(eager_manual_nodes
     ${PADDLE_SOURCE_DIR}/paddle/fluid/eager/api/manual/eager_manual/nodes/add_n_node.cc
     ${PADDLE_SOURCE_DIR}/paddle/fluid/eager/api/manual/eager_manual/nodes/sync_batch_norm_node.cc
     ${PADDLE_SOURCE_DIR}/paddle/fluid/eager/api/manual/eager_manual/nodes/multiply_node.cc
+    ${PADDLE_SOURCE_DIR}/paddle/fluid/eager/api/manual/eager_manual/nodes/reshard_node.cc
     PARENT_SCOPE)
@@ -396,6 +396,53 @@ class SyncBatchNormGradNode : public egr::GradNodeBase {
   bool trainable_statistics_;
 };
 
+class ReshardGradNode : public egr::GradNodeBase {
+ public:
+  ReshardGradNode() : egr::GradNodeBase() {
+    VLOG(3) << " Construct ReshardGrad Node.";
+  }
+
+  ReshardGradNode(size_t bwd_in_slot_num, size_t bwd_out_slot_num)
+      : egr::GradNodeBase(bwd_in_slot_num, bwd_out_slot_num) {
+    VLOG(3) << " Construct ReshardGrad Node, bwd_in_slot_num: "
+            << bwd_in_slot_num << ", bwd_out_slot_num: " << bwd_out_slot_num;
+  }
+
+  ~ReshardGradNode() override { VLOG(3) << " Destruct ReshardGrad Node."; }
+
+  virtual paddle::small_vector<std::vector<paddle::Tensor>,
+                               egr::kSlotSmallVectorSize>
+  operator()(paddle::small_vector<std::vector<paddle::Tensor>,
+                                  egr::kSlotSmallVectorSize>& grads,  // NOLINT
+             bool create_graph = false,
+             bool is_new_grad = false) override;
+
+  void ClearTensorWrappers() override {
+    input_.clear();
+    SetIsTensorWrappersCleared(true);
+  }
+
+  std::string name() override { return "ReshardGradNode"; }
+
+  std::shared_ptr<GradNodeBase> Copy() const override {
+    {
+      auto copied_node =
+          std::shared_ptr<ReshardGradNode>(new ReshardGradNode(*this));
+      return copied_node;
+    }
+  }
+
+  // SetTensorWrapperX
+  // Only input's meta is needed.
+  void SetTensorWrapperNoNeedBufferInput(const paddle::Tensor& input) {
+    input_ = egr::TensorWrapper(input, true);
+  }
+
+ private:
+  // TensorWrappers
+  egr::TensorWrapper input_;
+};
+
 namespace sparse {
 class SyncBatchNormGradNode : public egr::GradNodeBase {
  public:
 
@@ -0,0 +1,106 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include "glog/logging.h"
+#include "paddle/fluid/eager/api/manual/eager_manual/dygraph_forward_api.h"
+#include "paddle/fluid/eager/api/manual/eager_manual/nodes/nodes.h"
+#include "paddle/fluid/eager/api/utils/global_utils.h"
+#include "paddle/fluid/eager/utils.h"
+#include "paddle/fluid/imperative/tracer.h"
+
+paddle::small_vector<std::vector<paddle::Tensor>,
+                     egr::kSlotSmallVectorSize>  // NOLINT
+ReshardGradNode::operator()(
+    paddle::small_vector<std::vector<paddle::Tensor>,
+                         egr::kSlotSmallVectorSize>& grads,
+    bool create_graph,
+    bool is_new_grad) {
+#ifdef PADDLE_WITH_DISTRIBUTE
+  VLOG(3) << "Running AD API GRAD: "
+          << "reshard_grad";
+
+  // Apply Gradient Hooks
+  auto hooked_grad = ApplyGradientHooks(grads);
+
+  // 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
+
+  const auto& out_metas = OutputMeta();
+  paddle::small_vector<std::vector<paddle::Tensor>, egr::kSlotSmallVectorSize>
+      returns(1);
+
+  out_metas[0].size() == 0 ? returns[0].resize(1)
+                           : returns[0].resize(out_metas[0].size());
+
+  auto& grad_input = returns[0][0];
+
+  VLOG(5) << "Running C++ API: "
+          << "reshard_func";
+
+  if (VLOG_IS_ON(3)) {
+    const char* INPUT_PRINT_TEMPLATE = "{ Input: [%s]} ";
+
+    std::string input_str = "";
+    const char* TENSOR_OUT_GRAD_TEMPLATE = " \n( out_grad , [%s]), ";
+    std::string input_out_grad_str = paddle::string::Sprintf(
+        TENSOR_OUT_GRAD_TEMPLATE, egr::EagerUtils::TensorStr(grad_out));
+    input_str += input_out_grad_str;
+    const char* TENSOR_X_TEMPLATE = " \n( x , [%s]), ";
+    std::string input_x_str = paddle::string::Sprintf(
+        TENSOR_X_TEMPLATE, egr::EagerUtils::TensorStr(input));
+    input_str += input_x_str;
+    VLOG(3) << paddle::string::Sprintf(INPUT_PRINT_TEMPLATE, input_str);
+  }
+
+  // Backward call reshard_func function
+  auto dist_out_ptr = paddle::reshard(grad_out, dist_attr);
+  grad_input.set_impl(dist_out_ptr);
+
+  VLOG(5) << "Finish C++ API: reshard_func";
+  VLOG(6) << "gradnode_ptr = " << this;
+
+  if (VLOG_IS_ON(4)) {
+    const char* INPUT_PRINT_TEMPLATE = "{ Input: [%s],  \n Output: [%s] } ";
+    std::string input_str = "";
+    std::string output_str = "";
+    const char* TENSOR_OUT_GRAD_TEMPLATE = " \n( out_grad , [%s]), ";
+    std::string input_out_grad_str = paddle::string::Sprintf(
+        TENSOR_OUT_GRAD_TEMPLATE, egr::EagerUtils::TensorStr(grad_out));
+    input_str += input_out_grad_str;
+    const char* TENSOR_X_TEMPLATE = " \n( x , [%s]), ";
+    std::string input_x_str = paddle::string::Sprintf(
+        TENSOR_X_TEMPLATE, egr::EagerUtils::TensorStr(input));
+    input_str += input_x_str;
+    const char* TENSOR_X_GRAD_TEMPLATE = " \n ( input_grad , [%s]), ";
+    std::string output_x_grad_str = paddle::string::Sprintf(
+        TENSOR_X_GRAD_TEMPLATE, egr::EagerUtils::TensorStr(grad_input));
+    output_str += output_x_grad_str;
+    VLOG(4) << paddle::string::Sprintf(
+        INPUT_PRINT_TEMPLATE, input_str, output_str);
+  }
+
+  return returns;
+#else
+  PADDLE_THROW(phi::errors::Unavailable(
+      "ReshardGrad is not supported in this version of Paddle. Try to "
+      "recompile it with WITH_DISTRIBTUE=ON and reinstall this package."));
+  return paddle::small_vector<std::vector<paddle::Tensor>,
+                              egr::kSlotSmallVectorSize>(1);
+#endif
+}
@@ -35,6 +35,7 @@
 
 #include "paddle/fluid/distributed/auto_parallel/spmd_rules/common.h"
 #include "paddle/fluid/distributed/auto_parallel/spmd_rules/dist_tensor_spec.h"
+#include "paddle/fluid/eager/api/manual/eager_manual/dygraph_forward_api.h"
 #include "paddle/phi/api/lib/data_transform.h"
 #include "paddle/phi/backends/context_pool.h"
 #include "paddle/phi/core/dense_tensor.h"
@@ -571,33 +572,7 @@ void BindAutoParallel(py::module *m) {
       "reshard",
       [](py::handle py_tensor, const TensorDistAttr &dist_attr) {
         auto tensor = CastPyArg2Tensor(py_tensor.ptr(), 0);
-        auto dev_ctx = phi::DeviceContextPool::Instance().Get(tensor.place());
-        std::shared_ptr<phi::distributed::DistTensor> dist_out_ptr = nullptr;
-        if (phi::distributed::DistTensor::classof(tensor.impl().get())) {
-          auto tensor_in = tensor.impl();
-          if (tensor_in) {
-            phi::distributed::DistTensor *dist_tensor =
-                static_cast<phi::distributed::DistTensor *>(tensor_in.get());
-            if (dist_tensor->dist_attr() != dist_attr) {
-              VLOG(6) << "reshard func, reshard tensor from "
-                      << dist_tensor->dist_attr() << " to " << dist_attr;
-              auto *func = phi::distributed::ChooseProperReshardFunction(
-                  *dist_tensor, dist_attr);
-              dist_out_ptr = func->Eval(dev_ctx, *dist_tensor, dist_attr);
-            } else {
-              dist_out_ptr =
-                  std::static_pointer_cast<phi::distributed::DistTensor>(
-                      tensor_in);
-            }
-          }
-          return paddle::Tensor(dist_out_ptr);
-        } else {
-          PADDLE_THROW(phi::errors::InvalidArgument(
-              "The input tensor of shard function should be "
-              "``phi::distributed::DistTensor``. "
-              "However it's %s",
-              typeid(tensor.impl().get()).name()));
-        }
+        return reshard_ad_function(tensor, dist_attr);
       },
       py::return_value_policy::reference);
 
 
@@ -17,6 +17,10 @@ limitations under the License. */
 #include <functional>
 
 #include "paddle/phi/api/include/tensor.h"
+#ifdef PADDLE_WITH_DISTRIBUTE
+#include "paddle/phi/core/distributed/auto_parallel/reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard_utils.h"
+#endif
 
 namespace paddle {
 
@@ -50,4 +54,22 @@ PADDLE_API Tensor from_blob(void* data,
                             const phi::Place& place = phi::Place(),
                             const Deleter& deleter = nullptr);
 
+#ifdef PADDLE_WITH_DISTRIBUTE
+/**
+ * @brief Reshard a DistTensor by given DistAttr.
+ *
+ * @note Input of `Reshard` should be a `paddle::Tensor` whose impl is
+ * shared_ptr of DistTensor. According to the given DistAttr, input will be
+ * reshard to wanted distributed state. And it will return shared_ptr of a new
+ * DistTensor as outptut.
+ *
+ * @param input The input tensor to be resharded.
+ * @param dist_attr The dist_attr to be resharded.
+ * @return Shared_ptr of a new DistTensor
+ */
+// TODO(GhostScreaming): All APIs should call this unified function later.
+PADDLE_API std::shared_ptr<phi::distributed::DistTensor> reshard(
+    const paddle::Tensor& input,
+    const phi::distributed::TensorDistAttr& dist_attr);
+#endif
 }  // namespace paddle