[AutoParallel] Eager method support autoparallel3 (#58476)

wanghuancoder · web-flow · commit 3852ae03aa2c · 2023-10-31T15:05:54.000+08:00
* PHI copy support auto parallel
diff --git a/paddle/fluid/pybind/eager_method.cc b/paddle/fluid/pybind/eager_method.cc
@@ -2962,7 +2962,8 @@ static PyObject* tensor_method_strides(TensorObject* self,
                                        PyObject* kwargs) {
   EAGER_TRY
   std::vector<int64_t> value;
-  if (!self->tensor.defined() || !self->tensor.is_dense_tensor()) {
+  if (!self->tensor.defined() ||
+      (!self->tensor.is_dense_tensor() && !self->tensor.is_dist_tensor())) {
     return ToPyObject(value);
   }
   auto stride = self->tensor.strides();
@@ -3002,20 +3003,24 @@ static PyObject* tensor_contiguous(TensorObject* self,
                                    PyObject* args,
                                    PyObject* kwargs) {
   EAGER_TRY
-  if (self->tensor.is_dense_tensor()) {
-    auto dense_tensor =
-        std::dynamic_pointer_cast<phi::DenseTensor>(self->tensor.impl());
+  if (self->tensor.is_dense_tensor() || self->tensor.is_dist_tensor()) {
+    phi::DenseTensor* dense_tensor = nullptr;
+    if (self->tensor.is_dist_tensor()) {
+      dense_tensor =
+          static_cast<phi::distributed::DistTensor*>(self->tensor.impl().get())
+              ->unsafe_mutable_value();
+    } else {
+      dense_tensor = static_cast<phi::DenseTensor*>(self->tensor.impl().get());
+    }
     if (dense_tensor->meta().is_contiguous()) {
       Py_INCREF(self);
       return reinterpret_cast<PyObject*>(self);
     } else {
       eager_gil_scoped_release guard;
-      self->tensor.set_impl(std::make_shared<phi::DenseTensor>(std::move(
-          paddle::experimental::Trans2Contiguous(*(dense_tensor.get())))));
+      *dense_tensor = paddle::experimental::Trans2Contiguous(*dense_tensor);
       Py_INCREF(self);
       return reinterpret_cast<PyObject*>(self);
     }
-
   } else {
     Py_INCREF(self);
     return reinterpret_cast<PyObject*>(self);
@@ -3050,6 +3055,11 @@ static PyObject* tensor_is_contiguous(TensorObject* self,
     auto dense_tensor =
         std::dynamic_pointer_cast<phi::DenseTensor>(self->tensor.impl());
     return ToPyObject(dense_tensor->meta().is_contiguous());
+  } else if (self->tensor.is_dist_tensor()) {
+    auto dense_tensor = std::dynamic_pointer_cast<phi::distributed::DistTensor>(
+                            self->tensor.impl())
+                            ->unsafe_mutable_value();
+    return ToPyObject(dense_tensor->meta().is_contiguous());
   } else {
     return ToPyObject(true);
   }
@@ -3074,19 +3084,27 @@ static PyObject* tensor_method__uva(TensorObject* self,
                                     PyObject* kwargs) {
   EAGER_TRY
   VLOG(4) << "Running in tensor_method__uva.";
-  PADDLE_ENFORCE_EQ(self->tensor.is_dense_tensor(),
-                    true,
-                    platform::errors::InvalidArgument(
-                        "Unified virtual addressing only support "
-                        "DenseTensor currently."));
+  PADDLE_ENFORCE_EQ(
+      self->tensor.is_dense_tensor() || self->tensor.is_dist_tensor(),
+      true,
+      platform::errors::InvalidArgument(
+          "Unified virtual addressing only support "
+          "DenseTensor and DistTensor currently."));
   PADDLE_ENFORCE_EQ(platform::is_cpu_place(self->tensor.place()),
                     true,
                     platform::errors::InvalidArgument(
                         "Unified virtual addressing only support "
                         "CPU Tensor currently."));
   int device_id = pybind::CastPyArg2AttrLong(PyTuple_GET_ITEM(args, 0), 0);
-  auto* self_tensor = static_cast<phi::DenseTensor*>(self->tensor.impl().get());
-  tensor_uva(self_tensor, device_id);
+  phi::DenseTensor* dense_tensor = nullptr;
+  if (self->tensor.is_dist_tensor()) {
+    dense_tensor =
+        static_cast<phi::distributed::DistTensor*>(self->tensor.impl().get())
+            ->unsafe_mutable_value();
+  } else {
+    dense_tensor = static_cast<phi::DenseTensor*>(self->tensor.impl().get());
+  }
+  tensor_uva(dense_tensor, device_id);
 
   RETURN_PY_NONE
 
diff --git a/paddle/phi/api/lib/tensor.cc b/paddle/phi/api/lib/tensor.cc
@@ -113,6 +113,10 @@ std::vector<int64_t> Tensor::shape() const {
 const phi::DDim &Tensor::strides() const {
   if (is_dense_tensor()) {
     return static_cast<phi::DenseTensor *>(impl_.get())->strides();
+  } else if (is_dist_tensor()) {
+    return static_cast<phi::distributed::DistTensor *>(impl_.get())
+        ->unsafe_mutable_value()
+        ->strides();
   } else {
     PADDLE_THROW(phi::errors::Unimplemented(
         "Only support strides operation on DenseTensor now."));
diff --git a/python/paddle/distributed/auto_parallel/api.py b/python/paddle/distributed/auto_parallel/api.py
@@ -136,6 +136,10 @@ def shard_tensor(
             >>> print(d_tensor)
 
     """
+    if place is None:
+        place = paddle.framework._current_expected_place()
+    place = paddle.framework._get_paddle_place(place)
+
     # 1. create dense tensor
     # `paddle.to_tensor` supports both dynamic and static mode
     tensor = paddle.to_tensor(
@@ -154,7 +158,7 @@ def shard_tensor(
                 tensor, dist_attr=dist_attr, **tensor.__dict__
             )
         else:
-            return paddle.Tensor(tensor, dist_attr=dist_attr)
+            return paddle.Tensor(tensor, dist_attr=dist_attr, place=place)
     else:
         # TODO(zhiqiu): we need to refine the static shard_tensor
         return shard_tensor_static(
diff --git a/test/auto_parallel/test_semi_auto_parallel_functional_in_single_card.py b/test/auto_parallel/test_semi_auto_parallel_functional_in_single_card.py
@@ -14,6 +14,8 @@
 
 import unittest
 
+import numpy as np
+
 import paddle
 import paddle.distributed as dist
 
@@ -84,6 +86,30 @@ def test_tensor__is_shared_buffer_with(self):
         dist_tensor._share_buffer_to(to)
         self.assertTrue(dist_tensor._is_shared_buffer_with(to))
 
+    def test_tensor_strides(self):
+        mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+        dense_tensor = paddle.randn([10, 20])
+        dense_tensor = dense_tensor.reshape([20, 10])
+        dist_tensor = dist.shard_tensor(
+            dense_tensor,
+            dist_attr=dist.DistAttr(mesh=mesh, sharding_specs=[None, None]),
+        )
+        strides = dist_tensor.get_strides()
+        is_contiguous = dist_tensor.is_contiguous()
+        dist_tensor = dist_tensor.contiguous()
+
+    def test_tensor_uva(self):
+        mesh = dist.ProcessMesh([0, 1], dim_names=["x"])
+        place = paddle.CPUPlace()
+        np_value = np.random.random(size=[10, 30]).astype('float32')
+        dense_tensor = paddle.to_tensor(np_value, place=place)
+        dist_tensor = dist.shard_tensor(
+            dense_tensor,
+            place=place,
+            dist_attr=dist.DistAttr(mesh=mesh, sharding_specs=[None, None]),
+        )
+        dist_tensor._uva()
+
 
 if __name__ == "__main__":
     unittest.main()
