PaddlePaddle
diff --git a/‎paddle/fluid/pybind/eager_properties.cc
Copy file name to clipboard
+25Lines changed: 25 additions & 0 deletions b/‎paddle/fluid/pybind/eager_properties.cc
Copy file name to clipboard
+25Lines changed: 25 additions & 0 deletions
diff --git a/‎paddle/phi/core/distributed/auto_parallel/CMakeLists.txt
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions b/‎paddle/phi/core/distributed/auto_parallel/CMakeLists.txt
Copy file name to clipboard
+1Lines changed: 1 addition & 0 deletions
diff --git a/‎paddle/phi/core/distributed/auto_parallel/dist_tensor.cc
Copy file name to clipboard
+17-7Lines changed: 17 additions & 7 deletions b/‎paddle/phi/core/distributed/auto_parallel/dist_tensor.cc
Copy file name to clipboard
+17-7Lines changed: 17 additions & 7 deletions
diff --git a/‎paddle/phi/core/distributed/auto_parallel/dist_tensor.h
Copy file name to clipboard
+8-9Lines changed: 8 additions & 9 deletions b/‎paddle/phi/core/distributed/auto_parallel/dist_tensor.h
Copy file name to clipboard
+8-9Lines changed: 8 additions & 9 deletions
diff --git a/‎paddle/phi/core/distributed/auto_parallel/r_to_s_reshard_function.cc
Copy file name to clipboard
+8-14Lines changed: 8 additions & 14 deletions b/‎paddle/phi/core/distributed/auto_parallel/r_to_s_reshard_function.cc
Copy file name to clipboard
+8-14Lines changed: 8 additions & 14 deletions
diff --git a/‎paddle/phi/core/distributed/auto_parallel/r_to_s_reshard_function.h
Copy file name to clipboard
+4-4Lines changed: 4 additions & 4 deletions b/‎paddle/phi/core/distributed/auto_parallel/r_to_s_reshard_function.h
Copy file name to clipboard
+4-4Lines changed: 4 additions & 4 deletions
diff --git a/‎paddle/phi/core/distributed/auto_parallel/reshard_concat_functor.cc
Copy file name to clipboard
+55Lines changed: 55 additions & 0 deletions b/‎paddle/phi/core/distributed/auto_parallel/reshard_concat_functor.cc
Copy file name to clipboard
+55Lines changed: 55 additions & 0 deletions
diff --git a/‎paddle/phi/core/distributed/auto_parallel/reshard_concat_functor.h
Copy file name to clipboard
+30Lines changed: 30 additions & 0 deletions b/‎paddle/phi/core/distributed/auto_parallel/reshard_concat_functor.h
Copy file name to clipboard
+30Lines changed: 30 additions & 0 deletions
diff --git a/‎paddle/phi/core/distributed/auto_parallel/reshard_function.cc
Copy file name to clipboard
+45-2Lines changed: 45 additions & 2 deletions b/‎paddle/phi/core/distributed/auto_parallel/reshard_function.cc
Copy file name to clipboard
+45-2Lines changed: 45 additions & 2 deletions
diff --git a/‎paddle/phi/core/distributed/auto_parallel/reshard_function.h
Copy file name to clipboard
+31-4Lines changed: 31 additions & 4 deletions b/‎paddle/phi/core/distributed/auto_parallel/reshard_function.h
Copy file name to clipboard
+31-4Lines changed: 31 additions & 4 deletions
@@ -388,6 +388,26 @@ PyObject* tensor_properties_get_dist_attr(TensorObject* self, void* closure) {
   EAGER_CATCH_AND_THROW_RETURN_NULL
 }
 
+PyObject* tensor_properties_get_local_shape(TensorObject* self, void* closure) {
+  EAGER_TRY
+  if (self->tensor.is_dist_tensor()) {
+#ifdef PADDLE_WITH_DISTRIBUTE
+    phi::distributed::DistTensor* dist_tensor =
+        static_cast<phi::distributed::DistTensor*>(self->tensor.impl().get());
+    return ToPyObject(phi::vectorize<int64_t>(dist_tensor->local_dims()));
+#else
+    PADDLE_THROW(platform::errors::Unavailable(
+        "The `_local_shape` property of (Dist)Tensor is not supported "
+        "in the current PaddlePaddle, please recompile and install "
+        "PaddlePaddle "
+        "with the option of `WITH_DISTRIBUTE=ON`."));
+#endif
+  } else {
+    RETURN_PY_NONE
+  }
+  EAGER_CATCH_AND_THROW_RETURN_NULL
+}
+
 PyDoc_STRVAR(tensor_shape__doc__,
              R"DOC(shape
 
@@ -716,6 +736,11 @@ struct PyGetSetDef variable_properties[] = {  // NOLINT
      (setter)tensor_properties_set_persistable,
      tensor_persistable__doc__,
      nullptr},
+    {"_local_shape",
+     (getter)tensor_properties_get_local_shape,
+     nullptr,
+     nullptr,
+     nullptr},
     {"shape",
      (getter)tensor_properties_get_shape,
      nullptr,
 
@@ -9,6 +9,7 @@ if(WITH_DISTRIBUTE)
     dist_tensor.cc
     reshard_function.cc
     reshard_split_functor.cc
+    reshard_concat_functor.cc
     reshard_all_gather_functor.cc
     r_to_s_reshard_function.cc
     s_to_r_reshard_function.cc)
 
@@ -14,6 +14,10 @@
 
 #include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
 
+#include "paddle/phi/backends/context_pool.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard_function.h"
+#include "paddle/phi/core/distributed/auto_parallel/reshard_utils.h"
+
 namespace phi {
 namespace distributed {
 
@@ -27,15 +31,21 @@ inline void check_defined(const DistTensor& dist_tensor,
           method_hint));
 }
 
-// TODO(chenweihang): Reshard the input global value into local value
 DistTensor::DistTensor(const phi::DenseTensor& global_value,
                        const TensorDistAttr& dist_attr)
-    : dims_(global_value.dims()), dist_attr_(dist_attr), value_(global_value) {}
-
-DistTensor::DistTensor(const phi::DenseTensor& value,
-                       const DDim& dims,
-                       const TensorDistAttr& dist_attr)
-    : dims_(dims), dist_attr_(dist_attr), value_(value) {}
+    : dims_(global_value.dims()), dist_attr_(dist_attr), value_(global_value) {
+  if (!IsDimsMappingReplicated(dist_attr_.dims_mapping())) {
+    // 1. create replicated global tensor
+    int64_t dims_size = global_value.dims().size();
+    std::vector<int64_t> dims_mapping(dims_size, -1);
+    dist_attr_.set_dims_mapping(dims_mapping);
+
+    // 2. reshard from replicated to other state
+    auto* func = ChooseProperReshardFunction(*this, dist_attr);
+    auto* dev_ctx = DeviceContextPool::Instance().Get(global_value.place());
+    func->Eval(dev_ctx, *this, dist_attr, this);
+  }
+}
 
 DistTensor::DistTensor(const DDim& dims, const TensorDistAttr& dist_attr)
     : dims_(dims), dist_attr_(dist_attr) {}
 
@@ -21,26 +21,23 @@
 
 namespace phi {
 namespace distributed {
+class ReshardFunction;
 
 class DistTensor final
     : public phi::TensorBase,
       public phi::TypeInfoTraits<phi::TensorBase, DistTensor> {
  public:
+  /// \brief Careful to create dist tensor using default constructor.
+  /// this should only used in reshard for now, and the dist properties
+  /// will be set by reshard later.
+  DistTensor() = default;
+
   /// \brief Construct a dist tensor based dense tensor.
   /// \param global_value The global dense tensor of the current tensor.
   /// \param dist_attr The distributed attributes of the current tensor.
   DistTensor(const phi::DenseTensor& global_value,
              const TensorDistAttr& dist_attr);
 
-  // TODO(chenweihang): Remove this constructor after added reshard impl
-  /// \brief Construct a dist tensor based dense tensor.
-  /// \param value The local dense tensor of the current tensor.
-  /// \param dims The global dimension of the currnet tensor.
-  /// \param dist_attr The distributed attributes of the current tensor.
-  DistTensor(const phi::DenseTensor& value,
-             const DDim& dims,
-             const TensorDistAttr& dist_attr);
-
   /// \brief Construct a empty dist tensor (for infer spmd)
   /// \param dims The global dimension of the currnet Tensor.
   /// \param dist_attr The distributed attributes of the current tensor.
@@ -109,6 +106,8 @@ class DistTensor final
                      bool fake_alloc = false) override;
 
  private:
+  friend class ReshardFunction;
+
   // The global dimensions(shape)
   DDim dims_;
   // The distributed attributes
 
@@ -46,10 +46,10 @@ bool RToSReshardFunction::IsSuitable(const DistTensor& in,
   return flag;
 }
 
-std::shared_ptr<DistTensor> RToSReshardFunction::Eval(
-    phi::DeviceContext* dev_ctx,
-    const DistTensor& in,
-    const TensorDistAttr& out_dist_attr) {
+void RToSReshardFunction::Eval(phi::DeviceContext* dev_ctx,
+                               const DistTensor& in,
+                               const TensorDistAttr& out_dist_attr,
+                               DistTensor* out) {
   const auto& out_dims_mapping = out_dist_attr.dims_mapping();
   const auto& out_process_mesh = out_dist_attr.process_mesh();
   const DenseTensor& in_physical_tensor_cur_rank = in.value();
@@ -63,14 +63,6 @@ std::shared_ptr<DistTensor> RToSReshardFunction::Eval(
   int64_t split_axis = split_axis_to_mesh_axis.begin()->first;
   int64_t mesh_axis = split_axis_to_mesh_axis.begin()->second;
 
-  PADDLE_ENFORCE_LT(
-      mesh_axis,
-      out_process_mesh.ndim(),
-      phi::errors::OutOfRange(
-          "The mesh axis %lld exceed the size of process mesh %lld.",
-          mesh_axis,
-          out_process_mesh.ndim()));
-
   int64_t num_of_process = out_process_mesh.shape()[mesh_axis];
   VLOG(3) << "RToSReshard: Tensor will be split on axis " << split_axis
           << ". Split will use axis " << mesh_axis << " of process_mesh."
@@ -86,13 +78,15 @@ std::shared_ptr<DistTensor> RToSReshardFunction::Eval(
 
   VLOG(3) << "The current process will remain the idx "
           << coord_in_mesh[mesh_axis] << " piece of tensor";
+
   out_physical_tensor_cur_rank = split_out_vec[coord_in_mesh[mesh_axis]];
   VLOG(3) << "The shape of physical tensor after split is "
           << out_physical_tensor_cur_rank.dims();
 
-  return std::make_shared<DistTensor>(
-      out_physical_tensor_cur_rank, in.dims(), out_dist_attr);
+  set_dist_props(out, out_physical_tensor_cur_rank, in.dims(), out_dist_attr);
 }
 
+REGISTER_RESHARD_FUNC(RToSReshardFunction);
+
 }  // namespace distributed
 }  // namespace phi
@@ -27,10 +27,10 @@ class RToSReshardFunction final : public ReshardFunction {
   bool IsSuitable(const DistTensor& in,
                   const TensorDistAttr& out_dist_attr) override;
 
-  std::shared_ptr<DistTensor> Eval(
-      DeviceContext* dev_ctx,
-      const DistTensor& in,
-      const TensorDistAttr& out_dist_attr) override;
+  void Eval(DeviceContext* dev_ctx,
+            const DistTensor& in,
+            const TensorDistAttr& out_dist_attr,
+            DistTensor* out) override;
 };
 
 }  // namespace distributed
 
@@ -0,0 +1,55 @@
+// 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/phi/core/distributed/auto_parallel/reshard_concat_functor.h"
+
+#include "paddle/phi/backends/all_context.h"
+#include "paddle/phi/core/dense_tensor.h"
+#include "paddle/phi/core/device_context.h"
+#include "paddle/phi/core/visit_type.h"
+#include "paddle/phi/kernels/concat_kernel.h"
+
+namespace phi {
+namespace distributed {
+
+DenseTensor ReshardConcatFunctor(const DeviceContext& dev_ctx,
+                                 const std::vector<const DenseTensor*>& input,
+                                 int64_t axis) {
+  DenseTensor result;
+  auto dtype = (*input.begin())->dtype();
+
+  if (phi::CPUContext::classof(&dev_ctx)) {
+    PD_VISIT_ALL_TYPES(
+        dtype, "Concat", ([&] {
+          Concat<data_t>(
+              static_cast<const CPUContext&>(dev_ctx), input, axis, &result);
+        }));
+    return result;
+  }
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+  if (phi::GPUContext::classof(&dev_ctx)) {
+    PD_VISIT_ALL_TYPES(
+        dtype, "Concat", ([&] {
+          Concat<data_t>(
+              static_cast<const GPUContext&>(dev_ctx), input, axis, &result);
+        }));
+    return result;
+  }
+#endif
+  PADDLE_THROW(phi::errors::Unimplemented(
+      "The concat in reshard only supported on CPU and GPU for now."));
+}
+
+}  // namespace distributed
+}  // namespace phi
@@ -0,0 +1,30 @@
+// 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.
+
+#pragma once
+
+#include <cstdint>
+#include <vector>
+
+namespace phi {
+class DeviceContext;
+class DenseTensor;
+namespace distributed {
+
+DenseTensor ReshardConcatFunctor(const DeviceContext& dev_ctx,
+                                 const std::vector<const DenseTensor*>& input,
+                                 int64_t axis);
+
+}  // namespace distributed
+}  // namespace phi
@@ -13,10 +13,53 @@
 // limitations under the License.
 
 #include "paddle/phi/core/distributed/auto_parallel/reshard_function.h"
-#include "paddle/phi/core/device_context.h"
+
 #include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
 #include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
 
 namespace phi {
-namespace distributed {}  // namespace distributed
+namespace distributed {
+
+std::shared_ptr<DistTensor> ReshardFunction::Eval(
+    DeviceContext* dev_ctx,
+    const DistTensor& in,
+    const TensorDistAttr& out_dist_attr) {
+  std::shared_ptr<DistTensor> out = std::make_shared<DistTensor>();
+  Eval(dev_ctx, in, out_dist_attr, out.get());
+  return out;
+}
+
+void ReshardFunction::set_dist_props(DistTensor* tensor,
+                                     const DenseTensor& value,
+                                     const DDim& dims,
+                                     const TensorDistAttr& dist_attr) {
+  PADDLE_ENFORCE_EQ(dist_attr.verify(vectorize(dims)),
+                    true,
+                    phi::errors::InvalidArgument(
+                        "The input dist_attr and dims are improper."));
+
+  tensor->value_ = value;
+  tensor->dims_ = dims;
+  tensor->dist_attr_ = dist_attr;
+}
+
+ReshardFunction* ChooseProperReshardFunction(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr) {
+  for (const auto& func : GetReshardFunctionList()) {
+    if (func->IsSuitable(in, out_dist_attr)) {
+      return func.get();
+    }
+  }
+  PADDLE_THROW(phi::errors::Unimplemented(
+      "Can not reshard from in_dist_attr=%s to out_dist_attr=%s.",
+      in.dist_attr().to_string(),
+      out_dist_attr.to_string()));
+}
+
+std::vector<std::unique_ptr<ReshardFunction>>& GetReshardFunctionList() {
+  static std::vector<std::unique_ptr<ReshardFunction>> func_list;
+  return func_list;
+}
+
+}  // namespace distributed
 }  // namespace phi
@@ -14,6 +14,9 @@
 
 #pragma once
 #include <memory>
+#include <vector>
+
+#include "paddle/phi/core/dense_tensor.h"
 
 namespace phi {
 class DeviceContext;
@@ -31,11 +34,35 @@ class ReshardFunction {
   virtual bool IsSuitable(const DistTensor& in,
                           const TensorDistAttr& out_dist_attr) = 0;
 
-  virtual std::shared_ptr<DistTensor> Eval(
-      DeviceContext* dev_ctx,
-      const DistTensor& in,
-      const TensorDistAttr& out_dist_attr) = 0;
+  std::shared_ptr<DistTensor> Eval(DeviceContext* dev_ctx,
+                                   const DistTensor& in,
+                                   const TensorDistAttr& out_dist_attr);
+
+  virtual void Eval(DeviceContext* dev_ctx,
+                    const DistTensor& in,
+                    const TensorDistAttr& out_dist_attr,
+                    DistTensor* out) = 0;
+
+ protected:
+  void set_dist_props(DistTensor* tensor,
+                      const DenseTensor& value,
+                      const DDim& dims,
+                      const TensorDistAttr& dist_attr);
 };
 
+std::vector<std::unique_ptr<ReshardFunction>>& GetReshardFunctionList();
+
+#define REGISTER_RESHARD_FUNC(func_type)                                    \
+  class __RegisterReshard_##func_type {                                     \
+   public:                                                                  \
+    __RegisterReshard_##func_type() {                                       \
+      GetReshardFunctionList().emplace_back(std::make_unique<func_type>()); \
+    }                                                                       \
+  };                                                                        \
+  static __RegisterReshard_##func_type local_reshard_func_##func_type
+
+ReshardFunction* ChooseProperReshardFunction(
+    const DistTensor& in, const TensorDistAttr& out_dist_attr);
+
 }  // namespace distributed
 }  // namespace phi