PaddlePaddle
diff --git a/‎paddle/fluid/eager/grad_node_info.cc
Copy file name to clipboard
+34-3Lines changed: 34 additions & 3 deletions b/‎paddle/fluid/eager/grad_node_info.cc
Copy file name to clipboard
+34-3Lines changed: 34 additions & 3 deletions
diff --git a/‎paddle/fluid/eager/grad_tensor_holder.cc
Copy file name to clipboard
+25Lines changed: 25 additions & 0 deletions b/‎paddle/fluid/eager/grad_tensor_holder.cc
Copy file name to clipboard
+25Lines changed: 25 additions & 0 deletions
diff --git a/‎paddle/fluid/pybind/eager.cc
Copy file name to clipboard
+4-2Lines changed: 4 additions & 2 deletions b/‎paddle/fluid/pybind/eager.cc
Copy file name to clipboard
+4-2Lines changed: 4 additions & 2 deletions
diff --git a/‎paddle/fluid/pybind/eager_method.cc
Copy file name to clipboard
+37Lines changed: 37 additions & 0 deletions b/‎paddle/fluid/pybind/eager_method.cc
Copy file name to clipboard
+37Lines changed: 37 additions & 0 deletions
diff --git a/‎paddle/phi/api/lib/CMakeLists.txt
Copy file name to clipboard
+36-12Lines changed: 36 additions & 12 deletions b/‎paddle/phi/api/lib/CMakeLists.txt
Copy file name to clipboard
+36-12Lines changed: 36 additions & 12 deletions
diff --git a/‎paddle/phi/api/lib/api_gen_utils.cc
Copy file name to clipboard
+28Lines changed: 28 additions & 0 deletions b/‎paddle/phi/api/lib/api_gen_utils.cc
Copy file name to clipboard
+28Lines changed: 28 additions & 0 deletions
diff --git a/‎paddle/phi/api/lib/api_gen_utils.h
Copy file name to clipboard
+12Lines changed: 12 additions & 0 deletions b/‎paddle/phi/api/lib/api_gen_utils.h
Copy file name to clipboard
+12Lines changed: 12 additions & 0 deletions
diff --git a/‎paddle/phi/api/lib/data_transform.cc
Copy file name to clipboard
+45Lines changed: 45 additions & 0 deletions b/‎paddle/phi/api/lib/data_transform.cc
Copy file name to clipboard
+45Lines changed: 45 additions & 0 deletions
@@ -26,8 +26,12 @@
 #include "paddle/fluid/platform/errors.h"
 #include "paddle/phi/common/data_type.h"
 #include "paddle/phi/core/dense_tensor.h"
+
 #include "paddle/phi/core/sparse_coo_tensor.h"
 #include "paddle/phi/core/sparse_csr_tensor.h"
+#ifdef PADDLE_WITH_DISTRIBUTE
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+#endif
 
 /**
  * Implementation of GradNodeBase, Edge and GradTensorHolder.
@@ -121,6 +125,14 @@ void GradNodeBase::SetGradInMeta(const paddle::Tensor& fwd_out,
     phi::SparseCsrTensor* csr_tensor =
         static_cast<phi::SparseCsrTensor*>(fwd_out.impl().get());
     dense_tensor = csr_tensor->mutable_non_zero_elements();
+#ifdef PADDLE_WITH_DISTRIBUTE
+  } else if (phi::distributed::DistTensor::classof(fwd_out.impl().get())) {
+    // TODO(chenweihang): DistTensor contains global and local meta, here
+    // only set the local meta now, we should set global meta later
+    dense_tensor =
+        static_cast<phi::distributed::DistTensor*>(fwd_out.impl().get())
+            ->mutable_value();
+#endif
   } else {
     VLOG(7) << "Unable to initialize the DenseTensorMeta of GradSlotMeta with "
                "non-DenseTensor argument.";
@@ -256,10 +268,28 @@ void GradNodeBase::SetGradOutMeta(const paddle::Tensor& fwd_in,
                                           "which is illegal."));
       meta.SetTensorMeta(dense_tensor->meta());
       meta.SetPlace(fwd_in.place());
+#ifdef PADDLE_WITH_DISTRIBUTE
+    } else if (phi::distributed::DistTensor::classof(fwd_in.impl().get())) {
+      phi::DenseTensor* dense_tensor =
+          static_cast<phi::distributed::DistTensor*>(fwd_in.impl().get())
+              ->mutable_value();
+      PADDLE_ENFORCE_NE(
+          dense_tensor->meta().dtype,
+          phi::DataType::UNDEFINED,
+          paddle::platform::errors::Fatal("Attempting to copy DenseTensorMeta "
+                                          "with phi::DataType::UNDEFINED,"
+                                          "which is illegal."));
+      meta.SetTensorMeta(dense_tensor->meta());
+      meta.SetPlace(fwd_in.place());
+#endif
+    } else {
+      VLOG(7)
+          << "Unable to initialize the DenseTensorMeta of GradSlotMeta with "
+             "non-DenseTensor argument.";
     }
   } else {
-    VLOG(7) << "Unable to initialize the DenseTensorMeta of GradSlotMeta with "
-               "non-DenseTensor argument.";
+    VLOG(7) << "Unable to initialize the DenseTensorMeta because the Tensor "
+               "is not initialized.";
   }
 }
 
@@ -367,7 +397,8 @@ void GradNodeBase::SetGradOutMeta(const std::vector<paddle::Tensor>& fwd_in,
     // Record TensorMeta
     if (fwd_in_tensor.impl() && fwd_in_tensor.impl().get()) {
       if (phi::DenseTensor::classof(fwd_in_tensor.impl().get())) {
-        // Only Copy Meta
+        // TODO(chenweihang): DistTensor contains global and local meta, here
+        // only set the local meta now, we should set global meta later
         phi::DenseTensor* dense_tensor =
             static_cast<phi::DenseTensor*>(fwd_in_tensor.impl().get());
         PADDLE_ENFORCE_NE(dense_tensor->dtype(),
 
@@ -20,6 +20,10 @@
 #include "paddle/fluid/imperative/gradient_accumulator.h"
 #include "paddle/phi/core/sparse_coo_tensor.h"
 #include "paddle/phi/kernels/funcs/math_function.h"
+#ifdef PADDLE_WITH_DISTRIBUTE
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+#endif
 
 namespace egr {
 
@@ -83,6 +87,23 @@ void GradTensorHolder::CopyValueFromTensor(size_t slot_id,
       } else if (t.is_sparse_csr_tensor() || t.is_sparse_coo_tensor()) {
         buffer_[slot_id][rank] =
             paddle::experimental::sparse::full_like(t, 1, t.dtype());
+#ifdef PADDLE_WITH_DISTRIBUTE
+      } else if (t.is_dist_tensor()) {
+        VLOG(6) << "Create a new dist tensor.";
+        // TODO(chenweihang): we need a shard_tensor API in C++
+        // TODO(chenweihang): replace by valid dist_attr later
+        auto temp =
+            paddle::experimental::full(t.shape(), 1, t.dtype(), t.place());
+        auto dense_temp =
+            std::dynamic_pointer_cast<phi::DenseTensor>(temp.impl());
+        auto dist_tensor = std::make_shared<phi::distributed::DistTensor>(
+            dense_temp,
+            dense_temp->meta(),
+            std::make_shared<
+                phi::distributed::auto_parallel::TensorDistAttr>());
+        temp.set_impl(dist_tensor);
+        buffer_[slot_id][rank] = temp;
+#endif
       } else {
         PADDLE_THROW(paddle::platform::errors::Fatal(
             "Only Support DENSE_TENSOR, SPARSE_COO_TENSOR, SPARSE_CSR_TENSOR "
@@ -178,6 +199,10 @@ void GradTensorHolder::add(size_t slot_id,
                                                         &buffer_values);
         }
       }
+#ifdef PADDLE_WITH_DISTRIBUTE
+    } else if (t.is_dist_tensor()) {
+      buffer_tensor = add_ad_func(t, buffer_tensor);
+#endif
     } else {
       // TODO(jiabin): Support Other TensorBase later
       // TODO(zhanlve): Replace SelectedRowsAddTensor with add_dygraph_function
 
@@ -297,13 +297,15 @@ void InitDistTensorWithTensor(
   if (place == src.place()) {
     std::shared_ptr<phi::DenseTensor> tensor =
         std::static_pointer_cast<phi::DenseTensor>(src.impl());
-    self->tensor.set_impl(std::make_shared<DistTensor>(tensor, dist_attr));
+    self->tensor.set_impl(
+        std::make_shared<DistTensor>(tensor, tensor->meta(), dist_attr));
     VLOG(4) << "Same place, do ShareDataWith for DistTensor.";
   } else {
     std::shared_ptr<phi::DenseTensor> tensor =
         std::static_pointer_cast<phi::DenseTensor>(
             src.copy_to(place, true).impl());
-    self->tensor.set_impl(std::make_shared<DistTensor>(tensor, dist_attr));
+    self->tensor.set_impl(
+        std::make_shared<DistTensor>(tensor, tensor->meta(), dist_attr));
     VLOG(4) << "Different place, do TensorCopy for DistTensor.";
   }
   if (src.get_autograd_meta()) {
 
@@ -249,6 +249,27 @@ static PyObject* tensor_method_numpy(TensorObject* self,
                            place,
                            dense_tensor->Holder()->ptr(),
                            dense_tensor->Holder()->size());
+#ifdef PADDLE_WITH_DISTRIBUTE
+    } else if (self->tensor.is_dist_tensor()) {
+      // TODO(chenweihang): deal with DistTensor as local DenseTensor now,
+      // if the local DenseTensor is shard or partial, do gather or reduce?
+      VLOG(6) << "Getting DistTensor's numpy value";
+      auto* dist_tensor =
+          static_cast<phi::distributed::DistTensor*>(self->tensor.impl().get());
+      auto& dense_tensor = dist_tensor->value();
+      cpu_tensor.set_meta(dense_tensor.meta());
+      // deep copy
+      auto tmp_allocation_ptr =
+          memory::Alloc(cpu_place, dense_tensor.Holder()->size());
+      cpu_tensor.ResetHolder(std::shared_ptr<phi::Allocation>(
+          tmp_allocation_ptr.release(), tmp_allocation_ptr.get_deleter()));
+      // deep copy
+      paddle::memory::Copy(place,
+                           cpu_tensor.Holder()->ptr(),
+                           place,
+                           dense_tensor.Holder()->ptr(),
+                           dense_tensor.Holder()->size());
+#endif
     } else {
       VLOG(6) << "Getting DenseTensor's numpy value";
       auto dense_tensor =
@@ -290,6 +311,22 @@ static PyObject* tensor_method_numpy(TensorObject* self,
                                       dense_tensor->Holder()->ptr(),
                                       dense_tensor->Holder()->size(),
                                       kind);
+#ifdef PADDLE_WITH_DISTRIBUTE
+    } else if (self->tensor.is_dist_tensor()) {
+      VLOG(6) << "Getting DistTensor's numpy value";
+      auto* dist_tensor =
+          static_cast<phi::distributed::DistTensor*>(self->tensor.impl().get());
+      auto& dense_tensor = dist_tensor->value();
+      cpu_tensor.set_meta(dense_tensor.meta());
+      auto tmp_allocation_ptr =
+          memory::Alloc(cpu_place, dense_tensor.Holder()->size());
+      cpu_tensor.ResetHolder(std::shared_ptr<phi::Allocation>(
+          tmp_allocation_ptr.release(), tmp_allocation_ptr.get_deleter()));
+      paddle::platform::GpuMemcpySync(cpu_tensor.Holder()->ptr(),
+                                      dense_tensor.Holder()->ptr(),
+                                      dense_tensor.Holder()->size(),
+                                      kind);
+#endif
     } else {
       VLOG(6) << "Getting DenseTensor's numpy value";
       auto dense_tensor =
 
@@ -9,6 +9,9 @@ set(api_header_file ${CMAKE_SOURCE_DIR}/paddle/phi/api/include/api.h)
 set(api_source_file ${CMAKE_SOURCE_DIR}/paddle/phi/api/lib/api.cc)
 set(api_header_file_tmp ${api_header_file}.tmp)
 set(api_source_file_tmp ${api_source_file}.tmp)
+# dist forward api file
+set(dist_api_gen_file
+    ${CMAKE_SOURCE_DIR}/paddle/phi/api/yaml/generator/dist_api_gen.py)
 
 # backward api file
 set(bw_api_gen_file
@@ -21,6 +24,9 @@ set(bw_api_header_file
 set(bw_api_source_file ${CMAKE_SOURCE_DIR}/paddle/phi/api/lib/backward_api.cc)
 set(bw_api_header_file_tmp ${bw_api_header_file}.tmp)
 set(bw_api_source_file_tmp ${bw_api_source_file}.tmp)
+# dist backward api file
+set(dist_bw_api_gen_file
+    ${CMAKE_SOURCE_DIR}/paddle/phi/api/yaml/generator/dist_bw_api_gen.py)
 
 # dygraph(intermediate) api file
 set(im_api_gen_file
@@ -124,19 +130,37 @@ endif()
 
 execute_process(COMMAND ${PYTHON_EXECUTABLE} -m pip install pyyaml)
 
-# generate forward api
-execute_process(
-  COMMAND
-    ${PYTHON_EXECUTABLE} ${api_gen_file} --api_yaml_path ${api_yaml_file}
-    ${legacy_api_yaml_file} --api_header_path ${api_header_file_tmp}
-    --api_source_path ${api_source_file_tmp})
+if(WITH_DISTRIBUTE)
+  # generate dist forward api
+  execute_process(
+    COMMAND
+      ${PYTHON_EXECUTABLE} ${dist_api_gen_file} --api_yaml_path
+      ${api_yaml_file} ${legacy_api_yaml_file} --api_header_path
+      ${api_header_file_tmp} --api_source_path ${api_source_file_tmp})
 
-# generate backward api
-execute_process(
-  COMMAND
-    ${PYTHON_EXECUTABLE} ${bw_api_gen_file} --backward_yaml_path
-    ${bw_api_yaml_file} ${legacy_bw_api_yaml_file} --backward_header_path
-    ${bw_api_header_file_tmp} --backward_source_path ${bw_api_source_file_tmp})
+  # generate dist backward api
+  execute_process(
+    COMMAND
+      ${PYTHON_EXECUTABLE} ${dist_bw_api_gen_file} --backward_yaml_path
+      ${bw_api_yaml_file} ${legacy_bw_api_yaml_file} --backward_header_path
+      ${bw_api_header_file_tmp} --backward_source_path
+      ${bw_api_source_file_tmp})
+else()
+  # generate forward api
+  execute_process(
+    COMMAND
+      ${PYTHON_EXECUTABLE} ${api_gen_file} --api_yaml_path ${api_yaml_file}
+      ${legacy_api_yaml_file} --api_header_path ${api_header_file_tmp}
+      --api_source_path ${api_source_file_tmp})
+
+  # generate backward api
+  execute_process(
+    COMMAND
+      ${PYTHON_EXECUTABLE} ${bw_api_gen_file} --backward_yaml_path
+      ${bw_api_yaml_file} ${legacy_bw_api_yaml_file} --backward_header_path
+      ${bw_api_header_file_tmp} --backward_source_path
+      ${bw_api_source_file_tmp})
+endif()
 
 # generate fused_op api
 execute_process(
 
@@ -19,6 +19,13 @@ limitations under the License. */
 
 DECLARE_bool(use_stride_kernel);
 
+#include "glog/logging.h"
+
+#ifdef PADDLE_WITH_DISTRIBUTE
+#include "paddle/phi/core/distributed/auto_parallel/dist_attr.h"
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+#endif
+
 namespace paddle {
 namespace experimental {
 
@@ -475,5 +482,26 @@ void TransStride(phi::DeviceContext* dev_ctx,
                  phi::SelectedRows* from,
                  phi::SelectedRows* to) {}
 
+#ifdef PADDLE_WITH_DISTRIBUTE
+/* ------------------ for auto parallel ----------------------- */
+
+phi::distributed::DistTensor* SetKernelDistOutput(Tensor* out) {
+  if (out) {
+    // TODO(chenweihang): now all dist case are nullptr
+    if (out->impl() == nullptr) {
+      auto dense_t = std::make_shared<phi::DenseTensor>();
+      // TODO(chenweihang): polish code, dist_attr is null now
+      auto dist_attr =
+          std::make_shared<phi::distributed::auto_parallel::TensorDistAttr>();
+      auto dist_t = std::make_shared<phi::distributed::DistTensor>(
+          dense_t, phi::DenseTensorMeta(), dist_attr);
+      out->set_impl(dist_t);
+    }
+    return static_cast<phi::distributed::DistTensor*>(out->impl().get());
+  }
+  return nullptr;
+}
+#endif
+
 }  // namespace experimental
 }  // namespace paddle
@@ -24,6 +24,12 @@ limitations under the License. */
 #include "paddle/phi/core/sparse_csr_tensor.h"
 #include "paddle/phi/core/string_tensor.h"
 
+namespace phi {
+namespace distributed {
+class DistTensor;
+}  // namespace distributed
+}  // namespace phi
+
 namespace paddle {
 namespace experimental {
 
@@ -127,5 +133,11 @@ void TransStride(phi::DeviceContext* dev_ctx,
                  phi::SelectedRows* from,
                  phi::SelectedRows* to);
 
+#ifdef PADDLE_WITH_DISTRIBUTE
+/* ------------------ for auto parallel ----------------------- */
+
+phi::distributed::DistTensor* SetKernelDistOutput(Tensor* out);
+#endif
+
 }  // namespace experimental
 }  // namespace paddle
@@ -27,6 +27,10 @@ limitations under the License. */
 #include "paddle/phi/kernels/cast_kernel.h"
 #include "paddle/phi/kernels/contiguous_kernel.h"
 #include "paddle/phi/kernels/transfer_layout_kernel.h"
+#ifdef PADDLE_WITH_DISTRIBUTE
+#include "paddle/phi/core/distributed/auto_parallel/dist_tensor.h"
+#endif
+
 DECLARE_bool(use_stride_kernel);
 
 namespace paddle {
@@ -567,5 +571,46 @@ void TransDataBackend(const phi::SelectedRows* tensor,
   }
 }
 
+#ifdef PADDLE_WITH_DISTRIBUTE
+/* ------------------ for auto parallel ----------------------- */
+
+std::shared_ptr<phi::distributed::DistTensor> PrepareDataForDistTensor(
+    const Tensor& input,
+    const phi::TensorArgDef& target_args_def,
+    const TransformFlag& transform_flag,
+    bool is_stride_kernel) {
+  const auto& tensor_in = input.impl();
+  if (tensor_in) {
+    phi::distributed::DistTensor* dist_tensor =
+        static_cast<phi::distributed::DistTensor*>(tensor_in.get());
+    phi::DenseTensor& dense_tensor = *(dist_tensor->mutable_value());
+    if (!transform_flag.NeedTransform() || !dense_tensor.initialized() ||
+        (!NeedTransformPlace(
+             dense_tensor.place(), target_args_def.backend, transform_flag) &&
+         !NeedTransformDataType(
+             dense_tensor.dtype(), target_args_def.dtype, transform_flag) &&
+         !NeedTransformLayout(dense_tensor.layout(),
+                              target_args_def.layout,
+                              dense_tensor.place(),
+                              transform_flag) &&
+         !NeedTransform2Contiguous(is_stride_kernel,
+                                   dense_tensor.meta().is_contiguous()))) {
+      return std::static_pointer_cast<phi::distributed::DistTensor>(tensor_in);
+    }
+    phi::DenseTensor out = TransformData(
+        &dense_tensor, target_args_def, transform_flag, is_stride_kernel);
+    // TODO(chenweihang): The global meta in DistTensor is not changed,
+    // but the local meta in DenseTensor maybe changed, such as layout
+    // change(NCHW->NHWC), so the new DistTensor's meta maybe not unified.
+    VLOG(6) << "PrepareDataForDistTensor return transformed dist tensor";
+    return std::make_shared<phi::distributed::DistTensor>(
+        std::make_shared<phi::DenseTensor>(std::move(out)),
+        dist_tensor->meta(),
+        dist_tensor->dist_attr());
+  }
+  return nullptr;
+}
+#endif
+
 }  // namespace experimental
 }  // namespace paddle