can not share buffer if tensor not contiguous (#65050)

wanghuancoder · web-flow · commit 4b423e9f19d5 · 2024-06-19T12:57:01.000+08:00
* can not share buffer if tensor not contiguous

* refine

* refine

* refine

* for test

* refine
diff --git a/paddle/fluid/pybind/eager_method.cc b/paddle/fluid/pybind/eager_method.cc
@@ -1013,6 +1013,11 @@ static PyObject* tensor__share_buffer_to(TensorObject* self,
     auto* src_tensor =
         static_cast<phi::distributed::DistTensor*>(self->tensor.impl().get())
             ->unsafe_mutable_value();
+    if (!src_tensor->meta().is_contiguous()) {
+      PADDLE_THROW(platform::errors::InvalidArgument(
+          "Tensor %s is not contiguous, cannot call share_buffer_to on it.",
+          self->tensor.name()));
+    }
     if (!dst_ptr->defined()) {
       dst_ptr->set_impl(std::make_shared<phi::distributed::DistTensor>());
     }
@@ -1024,6 +1029,11 @@ static PyObject* tensor__share_buffer_to(TensorObject* self,
   } else {
     auto* src_tensor =
         static_cast<phi::DenseTensor*>(self->tensor.impl().get());
+    if (!src_tensor->meta().is_contiguous()) {
+      PADDLE_THROW(platform::errors::InvalidArgument(
+          "Tensor %s is not contiguous, cannot call share_buffer_to on it.",
+          self->tensor.name()));
+    }
     if (!dst_ptr->defined()) {
       dst_ptr->set_impl(std::make_shared<phi::DenseTensor>());
     }
diff --git a/paddle/fluid/pybind/tensor.cc b/paddle/fluid/pybind/tensor.cc
@@ -705,6 +705,11 @@ void BindTensor(pybind11::module &m) {  // NOLINT
       .def("_share_buffer_with",
            [](phi::DenseTensor &self, const phi::DenseTensor src,
               py::tuple t) {
+              if (!src.meta().is_contiguous()) {
+                PADDLE_THROW(platform::errors::InvalidArgument(
+                    "Tensor is not contiguous, cannot call "
+                    "share_buffer_with on it."));
+              }
              auto *cuda_ipc_allocation =
                  dynamic_cast<memory::allocation::CudaIpcAllocation *>(
                      src.Holder().get());
diff --git a/python/paddle/distributed/fleet/recompute/recompute.py b/python/paddle/distributed/fleet/recompute/recompute.py
@@ -345,15 +345,20 @@ def inner_pack(inner_x):
                 if holder_list[unpack_counter - 1]() is None:
                     return
 
-                tmp_tensor = core.eager.Tensor(
-                    inner_x.dtype,
-                    inner_x.shape,
-                    inner_x.name + "cpy",
-                    core.VarDesc.VarType.LOD_TENSOR,
-                    inner_x.persistable,
-                )
-                inner_x._share_buffer_to(tmp_tensor)
-                storage[holder_list[unpack_counter - 1]()] = tmp_tensor
+                if inner_x.is_contiguous():
+                    tmp_tensor = core.eager.Tensor(
+                        inner_x.dtype,
+                        inner_x.shape,
+                        inner_x.name + "cpy",
+                        core.VarDesc.VarType.LOD_TENSOR,
+                        inner_x.persistable,
+                    )
+                    inner_x._share_buffer_to(tmp_tensor)
+                    storage[holder_list[unpack_counter - 1]()] = tmp_tensor
+                else:
+                    storage[
+                        holder_list[unpack_counter - 1]()
+                    ] = inner_x.contiguous()
                 return
 
             def inner_unpack(inner_x):
