Fix Tensor clear to properly clear slots (#143203)

Fixes a bug introduced in https://github.com/pytorch/pytorch/pull/137267

While the test ensures the finalizer did run to make sure things are cleared, the objects are not properly collected by the gc due to the faulty tp_clear implementation. So, while the finalizer did run, the object was still alive.
Fixing this by giving tp_clear the same treatment as tp_traverse and tp_dealloc on Tensor: make it a unique function that handles the full subclass hierarchy in one place.

Pull Request resolved: https://github.com/pytorch/pytorch/pull/143203
Approved by: https://github.com/ezyang, https://github.com/colesbury
ghstack dependencies: #143202
This commit is contained in:
albanD
2024-12-13 14:26:46 -05:00
committed by PyTorch MergeBot
parent 8741d72e3c
commit 70be7900bb
2 changed files with 34 additions and 2 deletions

View File

@ -1828,9 +1828,12 @@ static int THPVariable_subclass_clear(THPVariable* self) {
// C++ owns the PyObject.
// See https://github.com/pytorch/pytorch/pull/75933 for more discussion.
if (isResurrectable((THPVariable*)self)) {
if (isResurrectable(self)) {
return 0;
}
// First clear Tensor specific things
Py_CLEAR(self->backward_hooks);
Py_CLEAR(self->post_accumulate_grad_hooks);
const auto& tensor = THPVariable_Unpack(self);
@ -1891,13 +1894,34 @@ static int THPVariable_subclass_clear(THPVariable* self) {
}
}
}
TORCH_INTERNAL_ASSERT(!isResurrectable((THPVariable*)self));
TORCH_INTERNAL_ASSERT(!isResurrectable(self));
{
// MapAllocator can take significant time to release large tensors;
// release the GIL here to avoid impacting main thread perf.
pybind11::gil_scoped_release no_gil;
self->cdata = MaybeOwned<Variable>();
}
// Since we override the basic subtype_clear from CPython, we need a crappy
// version here just like for traverse and dealloc
// Clear all slots until we get to the base Tensor class
PyTypeObject* type = Py_TYPE((PyObject*)self);
PyTypeObject* base = type;
while (base != &THPVariableType) {
if (Py_SIZE(base))
clear_slots(base, (PyObject*)self);
base = base->tp_base;
TORCH_INTERNAL_ASSERT(base);
}
// Assume we never have managed dict for Tensors as we don't set the flag on
// the base class
if (C10_LIKELY(type->tp_dictoffset)) {
PyObject** dictptr = _PyObject_GetDictPtr((PyObject*)self);
if (dictptr && *dictptr)
Py_CLEAR(*dictptr);
}
return 0;
}