fix recompute memory leak (#63420)

deepllz · web-flow · commit e67841a43f8c · 2024-04-12T15:31:16.000+08:00
* fix recompute memory leak

* polish code
diff --git a/paddle/fluid/eager/tensor_wrapper.h b/paddle/fluid/eager/tensor_wrapper.h
@@ -142,33 +142,6 @@ class TensorWrapper {
     }
   }
 
-#ifndef PADDLE_NO_PYTHON
-  TensorWrapper(const TensorWrapper& other) {
-    no_need_buffer_ = other.no_need_buffer_;
-    intermidiate_tensor_ = other.intermidiate_tensor_;
-    weak_grad_node_ = other.weak_grad_node_;
-    inplace_version_snapshot_ = other.inplace_version_snapshot_;
-    packed_value_ = other.packed_value_;
-    unpack_hook_ = other.unpack_hook_;
-    if (packed_value_) {
-      packed_value_->inc_ref();
-    }
-  }
-
-  TensorWrapper& operator=(const TensorWrapper& other) {
-    no_need_buffer_ = other.no_need_buffer_;
-    intermidiate_tensor_ = other.intermidiate_tensor_;
-    weak_grad_node_ = other.weak_grad_node_;
-    inplace_version_snapshot_ = other.inplace_version_snapshot_;
-    packed_value_ = other.packed_value_;
-    unpack_hook_ = other.unpack_hook_;
-    if (packed_value_) {
-      packed_value_->inc_ref();
-    }
-    return *this;
-  }
-#endif
-
   paddle::Tensor recover() {
     VLOG(6) << "Recover tensor: " << intermidiate_tensor_.name()
             << " for wrapper";
diff --git a/paddle/fluid/pybind/eager_utils.cc b/paddle/fluid/pybind/eager_utils.cc
@@ -2478,6 +2478,7 @@ PyObjectHolder::PyObjectHolder(PyObject* ptr) { ptr_ = ptr; }
 
 PyObjectHolder::~PyObjectHolder() {  // NOLINT
   ::pybind11::gil_scoped_acquire gil;
+  // NOTE(deepllz): ptr_ is owned by this object, so release it in destructor.
   Py_XDECREF(ptr_);
 }
 
@@ -2512,7 +2513,10 @@ std::shared_ptr<egr::PyObjectHolderBase> PackHook::operator()(
   bool grad_tmp = egr::Controller::Instance().HasGrad();
   egr::Controller::Instance().SetHasGrad(false);
   ::pybind11::gil_scoped_acquire gil;
-  auto args = PyTuple_New(1);
+  PyObject* args = PyTuple_New(1);
+  PADDLE_ENFORCE_NOT_NULL(args,
+                          paddle::platform::errors::External(
+                              pybind11::detail::error_string().c_str()));
   PyTuple_SET_ITEM(args, 0, paddle::pybind::ToPyObject(tensor));
   PyObject* ret = PyObject_Call(hook_, args, nullptr);
   PADDLE_ENFORCE_NOT_NULL(ret,
@@ -2527,7 +2531,10 @@ void* PackHook::operator()(void* py_tensor) {
   bool grad_tmp = egr::Controller::Instance().HasGrad();
   egr::Controller::Instance().SetHasGrad(false);
   ::pybind11::gil_scoped_acquire gil;
-  auto args = PyTuple_New(1);
+  PyObject* args = PyTuple_New(1);
+  PADDLE_ENFORCE_NOT_NULL(args,
+                          paddle::platform::errors::External(
+                              pybind11::detail::error_string().c_str()));
   Py_INCREF(reinterpret_cast<PyObject*>(py_tensor));
   PyTuple_SET_ITEM(args, 0, reinterpret_cast<PyObject*>(py_tensor));
   PyObject* ret = PyObject_Call(hook_, args, nullptr);
@@ -2551,13 +2558,20 @@ paddle::Tensor UnPackHook::operator()(
   bool grad_tmp = egr::Controller::Instance().HasGrad();
   egr::Controller::Instance().SetHasGrad(false);
   ::pybind11::gil_scoped_acquire gil;
-  auto args = PyTuple_New(1);
-  Py_INCREF(reinterpret_cast<PyObject*>(packed_value->get()));
-  PyTuple_SET_ITEM(args, 0, reinterpret_cast<PyObject*>(packed_value->get()));
+  PyObject* args = PyTuple_New(1);
+  PADDLE_ENFORCE_NOT_NULL(args,
+                          paddle::platform::errors::External(
+                              pybind11::detail::error_string().c_str()));
+  PyObject* py_packed_value = reinterpret_cast<PyObject*>(packed_value->get());
+  Py_INCREF(py_packed_value);
+  PyTuple_SET_ITEM(args, 0, py_packed_value);
   PyObject* ret = PyObject_Call(hook_, args, nullptr);
   PADDLE_ENFORCE_NOT_NULL(ret,
                           paddle::platform::errors::External(
                               pybind11::detail::error_string().c_str()));
+  // NOTE(deepllz): tupledealloc will cause the reference count of the objects
+  // in it to be decremented by one, so no need to call
+  // Py_XDECREF(py_packed_value)
   Py_XDECREF(args);
   egr::Controller::Instance().SetHasGrad(grad_tmp);
 
@@ -2576,7 +2590,10 @@ void* UnPackHook::operator()(void* packed_value, void* other) {
   bool grad_tmp = egr::Controller::Instance().HasGrad();
   egr::Controller::Instance().SetHasGrad(false);
   ::pybind11::gil_scoped_acquire gil;
-  auto args = PyTuple_New(1);
+  PyObject* args = PyTuple_New(1);
+  PADDLE_ENFORCE_NOT_NULL(args,
+                          paddle::platform::errors::External(
+                              pybind11::detail::error_string().c_str()));
   Py_INCREF(reinterpret_cast<PyObject*>(packed_value));
   PyTuple_SET_ITEM(args, 0, reinterpret_cast<PyObject*>(packed_value));
   PyObject* ret = PyObject_Call(hook_, args, nullptr);
