From 94cf213dc15f5c7f08f8d01f16f7ce0a0a9be536 Mon Sep 17 00:00:00 2001
From: DesmonDay <908660116@qq.com>
Date: Wed, 1 Feb 2023 10:41:57 +0000
Subject: [PATCH] fix reference count

---
 paddle/fluid/framework/dlpack_tensor.cc | 11 +++++++++--
 paddle/fluid/framework/dlpack_tensor.h  |  4 +++-
 paddle/fluid/pybind/tensor.cc           |  2 ++
 3 files changed, 14 insertions(+), 3 deletions(-)

diff --git a/paddle/fluid/framework/dlpack_tensor.cc b/paddle/fluid/framework/dlpack_tensor.cc
index 32c6e17143..419eb0e7fe 100644
--- a/paddle/fluid/framework/dlpack_tensor.cc
+++ b/paddle/fluid/framework/dlpack_tensor.cc
@@ -12,6 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 #include "paddle/fluid/framework/dlpack_tensor.h"
+#include "pybind11/pybind11.h"
 
 #include "paddle/fluid/framework/convert_utils.h"
 #include "paddle/fluid/framework/data_type.h"
@@ -134,8 +135,9 @@ struct DLDeviceVisitor
 };
 }  // namespace internal
 
-DLPackTensor::DLPackTensor(const phi::DenseTensor &tensor, LaneType lanes) {
+DLPackTensor::DLPackTensor(phi::DenseTensor &tensor, LaneType lanes) {
   // init data, data buffer
+  dt_ = &tensor;
   t_.data = const_cast<void *>(tensor.data());
 
   // init device, DLDevice type with device_type and device_id
@@ -188,12 +190,17 @@ DLPackTensor::DLPackTensor(const phi::DenseTensor &tensor, LaneType lanes) {
   tensor->dl_tensor = t_;
 
   tensor->deleter = [](DLManagedTensor *arg) {
+    phi::DenseTensor *tensor_ptr =
+        reinterpret_cast<phi::DenseTensor *>(arg->manager_ctx);
+    pybind11::handle tensor_handle = pybind11::cast(tensor_ptr);
+    tensor_handle.dec_ref();
+
     delete[] arg->dl_tensor.shape;
     delete[] arg->dl_tensor.strides;
     delete arg;
   };
 
-  tensor->manager_ctx = nullptr;
+  tensor->manager_ctx = dt_;
 
   return tensor;
 }
diff --git a/paddle/fluid/framework/dlpack_tensor.h b/paddle/fluid/framework/dlpack_tensor.h
index c6fca6707f..2fcd2511af 100644
--- a/paddle/fluid/framework/dlpack_tensor.h
+++ b/paddle/fluid/framework/dlpack_tensor.h
@@ -28,7 +28,7 @@ class DLPackTensor {
       std::remove_reference<decltype(::DLTensor::shape[0])>::type;  // int64_t
 
   // lanes is only used in CPU to enable vectorization
-  explicit DLPackTensor(const phi::DenseTensor& tensor, LaneType lanes = 1);
+  explicit DLPackTensor(phi::DenseTensor& tensor, LaneType lanes = 1);
 
   inline operator const ::DLTensor&() const { return t_; }
 
@@ -42,6 +42,8 @@ class DLPackTensor {
   // The shape in DLTensor is defined as int64_t*
   // Add this member to make TVMTensor init without heap allocation
   ShapeType shape_[DDim::kMaxRank];
+
+  phi::DenseTensor* dt_;
 };
 
 }  // namespace framework
diff --git a/paddle/fluid/pybind/tensor.cc b/paddle/fluid/pybind/tensor.cc
index 29a17d25d9..b58e697e69 100644
--- a/paddle/fluid/pybind/tensor.cc
+++ b/paddle/fluid/pybind/tensor.cc
@@ -474,6 +474,8 @@ void BindTensor(pybind11::module &m) {  // NOLINT
       .def("_to_dlpack",
            [](phi::DenseTensor &self) {
              DLPackTensor dlpack_tensor(self, 1);
+             pybind11::handle tensor_handle = pybind11::cast(&self);
+             tensor_handle.inc_ref();
              DLManagedTensor *dmt = dlpack_tensor.ToDLManagedTensor();
              auto capsule = pybind11::capsule(
                  static_cast<void *>(dmt), "dltensor", [](PyObject *ptr) {
-- 
GitLab