[AutoParallel] tensor.cc support auto parallel (PaddlePaddle#58656)

* tensor.cc support auto parallel
SecretXV · Nov 28, 2023 · 84f3532 · 84f3532
1 parent f09e2d3
commit 84f3532
Showing 1 changed file with 26 additions and 7 deletions.
diff --git a/paddle/phi/api/lib/tensor.cc b/paddle/phi/api/lib/tensor.cc
@@ -115,11 +115,11 @@ const phi::DDim &Tensor::strides() const {
     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();
+        ->value()
+        .strides();
   } else {
     PADDLE_THROW(phi::errors::Unimplemented(
-        "Only support strides operation on DenseTensor now."));
+        "Only support strides operation on DenseTensor and DistTensor now."));
   }
 }
 
@@ -437,9 +437,16 @@ void Tensor::bump_inplace_version() {
     auto &inplace_version_counter =
         static_cast<phi::DenseTensor *>(impl_.get())->InplaceVersionCounter();
     inplace_version_counter.Bump();
+  } else if (is_dist_tensor()) {
+    auto &inplace_version_counter =
+        static_cast<phi::distributed::DistTensor *>(impl_.get())
+            ->unsafe_mutable_value()
+            ->InplaceVersionCounter();
+    inplace_version_counter.Bump();
   } else {
-    PADDLE_THROW(phi::errors::Unimplemented(
-        "bump_inplace_version is only supported on DenseTensor now."));
+    PADDLE_THROW(
+        phi::errors::Unimplemented("bump_inplace_version is only supported on "
+                                   "DenseTensor and DistTensor now."));
   }
 }
 
@@ -448,9 +455,15 @@ uint32_t Tensor::current_inplace_version() {
     auto &inplace_version_counter =
         static_cast<phi::DenseTensor *>(impl_.get())->InplaceVersionCounter();
     return inplace_version_counter.CurrentVersion();
+  } else if (is_dist_tensor()) {
+    auto &inplace_version_counter =
+        static_cast<phi::distributed::DistTensor *>(impl_.get())
+            ->unsafe_mutable_value()
+            ->InplaceVersionCounter();
+    return inplace_version_counter.CurrentVersion();
   } else {
-    LOG_FIRST_N(WARNING, 1)
-        << "current_inplace_version is only supported on DenseTensor now.";
+    LOG_FIRST_N(WARNING, 1) << "current_inplace_version is only supported on "
+                               "DenseTensor DistTensor now.";
   }
   return 0;
 }
@@ -461,6 +474,12 @@ void Tensor::reset_inplace_version(bool set_to_zero) {
       auto &inplace_version_counter =
           static_cast<phi::DenseTensor *>(impl_.get())->InplaceVersionCounter();
       inplace_version_counter.SetInplaceVersionToZero();
+    } else if (is_dist_tensor()) {
+      auto &inplace_version_counter =
+          static_cast<phi::distributed::DistTensor *>(impl_.get())
+              ->unsafe_mutable_value()
+              ->InplaceVersionCounter();
+      return inplace_version_counter.SetInplaceVersionToZero();
     }
   }
 }