fix contiguous warning bug (#65619)

wanghuancoder · web-flow · commit 986939a2016d · 2024-07-02T16:26:28.000+08:00
diff --git a/paddle/fluid/pybind/eager_method.cc b/paddle/fluid/pybind/eager_method.cc
@@ -1046,6 +1046,50 @@ static PyObject* tensor__share_buffer_to(TensorObject* self,
   EAGER_CATCH_AND_THROW_RETURN_NULL
 }
 
+static PyObject* tensor__unsafe_share_buffer_to(TensorObject* self,
+                                                PyObject* args,
+                                                PyObject* kwargs) {
+  EAGER_TRY
+  paddle::Tensor* dst_ptr =
+      &(reinterpret_cast<TensorObject*>(PyTuple_GET_ITEM(args, 0))->tensor);
+  if (!self->tensor.initialized()) {
+    if (self->tensor.numel() == 0) {
+      // Do nothing for 0-size Tensor
+      Py_RETURN_NONE;
+    } else {
+      PADDLE_THROW(platform::errors::InvalidArgument(
+          "Tensor %s has not been initialized! please initialize "
+          "src tensor before share_buffer_with to other.",
+          self->tensor.name()));
+    }
+  }
+  if (self->tensor.is_dist_tensor()) {
+    auto* src_tensor =
+        static_cast<phi::distributed::DistTensor*>(self->tensor.impl().get())
+            ->unsafe_mutable_value();
+    if (!dst_ptr->defined()) {
+      dst_ptr->set_impl(std::make_shared<phi::distributed::DistTensor>());
+    }
+    auto dst_tensor =
+        static_cast<phi::distributed::DistTensor*>(dst_ptr->impl().get())
+            ->unsafe_mutable_value();
+    dst_tensor->ShareBufferWith(*src_tensor);
+    dst_tensor->ShareDataTypeWith(*src_tensor);
+  } else {
+    auto* src_tensor =
+        static_cast<phi::DenseTensor*>(self->tensor.impl().get());
+    if (!dst_ptr->defined()) {
+      dst_ptr->set_impl(std::make_shared<phi::DenseTensor>());
+    }
+    auto dst_tensor = static_cast<phi::DenseTensor*>(dst_ptr->impl().get());
+    dst_tensor->ShareBufferWith(*src_tensor);
+    dst_tensor->ShareDataTypeWith(*src_tensor);
+  }
+  RETURN_PY_NONE
+
+  EAGER_CATCH_AND_THROW_RETURN_NULL
+}
+
 static PyObject* tensor__is_shared_buffer_with(TensorObject* self,
                                                PyObject* args,
                                                PyObject* kwargs) {
@@ -3288,6 +3332,10 @@ PyMethodDef variable_methods[] = {  // NOLINT
      (PyCFunction)(void (*)())tensor__share_buffer_to,
      METH_VARARGS | METH_KEYWORDS,
      nullptr},
+    {"_unsafe_share_buffer_to",
+     (PyCFunction)(void (*)())tensor__unsafe_share_buffer_to,
+     METH_VARARGS | METH_KEYWORDS,
+     nullptr},
     {"_is_shared_buffer_with",
      (PyCFunction)(void (*)())tensor__is_shared_buffer_with,
      METH_VARARGS | METH_KEYWORDS,
diff --git a/python/paddle/distributed/fleet/recompute/recompute.py b/python/paddle/distributed/fleet/recompute/recompute.py
@@ -347,33 +347,28 @@ def inner_pack(inner_x):
                 if holder_list[unpack_counter - 1]() is None:
                     return
 
-                if inner_x.is_contiguous():
-                    if inner_x.is_dist():
-                        # TODO(jeff41404): it seems better to use `tmp_tensor = core.eager.Tensor(inner_x)`,
-                        # but other errors will be triggered during the current period, and can be modified after resolution
-                        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.process_mesh,
-                            inner_x.placements,
-                        )
-                    else:
-                        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_dist():
+                    # TODO(jeff41404): it seems better to use `tmp_tensor = core.eager.Tensor(inner_x)`,
+                    # but other errors will be triggered during the current period, and can be modified after resolution
+                    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.process_mesh,
+                        inner_x.placements,
+                    )
                 else:
-                    storage[
-                        holder_list[unpack_counter - 1]()
-                    ] = inner_x.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._unsafe_share_buffer_to(tmp_tensor)
+                storage[holder_list[unpack_counter - 1]()] = tmp_tensor
                 return
 
             def inner_unpack(inner_x):
