mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-22 06:11:27 +08:00
This PR is a bit more involved but effectively works to drastically simplify PyObjectSlot and PyInterpreter. 1) For PyObjectSlot we now use a global pyinterpreter since there only is one. From here we change all of the call sites to rely on this assumption. 2) We also remove the "tags" of the PyInterpreter by deprecating `PyInterpreterStatus`. For the reviewer, sadly it seems like `functorch/csrc/dim/dim.cpp` needed to get linted, so there is an unreadable amount of changes there. Fortunately, the only actual change in the file is as follows which just removes `getPyInterpreter()` from the `check_pyobj` call. ``` mpy::handle handle_from_tensor(Arena& A, TensorRef t) { - // fast case: tensor is live in python - std::optional<PyObject*> mb_obj = - t->unsafeGetTensorImpl()->pyobj_slot()->check_pyobj(getPyInterpreter(), /*ignore_hermetic_tls=*/false); - if (mb_obj.has_value() && !t->unsafeGetTensorImpl()->pyobj_slot()->owns_pyobj()) { - return *mb_obj; - } - return A.autorelease(mpy::object::checked_steal(THPVariable_Wrap(*t))); -} -} + // fast case: tensor is live in python + std::optional<PyObject*> mb_obj = + t->unsafeGetTensorImpl()->pyobj_slot()->check_pyobj( + /*ignore_hermetic_tls=*/false); + if (mb_obj.has_value() && + !t->unsafeGetTensorImpl()->pyobj_slot()->owns_pyobj()) { + return *mb_obj; + } + return A.autorelease(mpy::object::checked_steal(THPVariable_Wrap(*t))); +} ``` Pull Request resolved: https://github.com/pytorch/pytorch/pull/158427 Approved by: https://github.com/albanD
57 lines
1.8 KiB
C++
57 lines
1.8 KiB
C++
#include <c10/core/impl/PyObjectSlot.h>
|
|
|
|
namespace c10::impl {
|
|
|
|
PyObjectSlot::PyObjectSlot() : pyobj_interpreter_(nullptr), pyobj_(nullptr) {}
|
|
|
|
PyObjectSlot::~PyObjectSlot() {
|
|
maybe_destroy_pyobj();
|
|
}
|
|
|
|
void PyObjectSlot::maybe_destroy_pyobj() {
|
|
if (owns_pyobj()) {
|
|
TORCH_INTERNAL_ASSERT(pyobj_interpreter_ != nullptr);
|
|
TORCH_INTERNAL_ASSERT(pyobj_ != nullptr);
|
|
(*pyobj_interpreter_.load(std::memory_order_acquire))
|
|
->decref(_unchecked_untagged_pyobj(), /*has_pyobj_slot*/ true);
|
|
// NB: this destructor can only be entered when there are no
|
|
// references to this C++ object (obviously), NOR any references
|
|
// to the PyObject (if there are references to the PyObject,
|
|
// then the PyObject holds an owning reference to the tensor).
|
|
// So it is OK to clear pyobj_ here as it is impossible for it to
|
|
// be used again (modulo weak reference races)
|
|
pyobj_ = nullptr; // for safety
|
|
}
|
|
}
|
|
|
|
PyInterpreter* PyObjectSlot::pyobj_interpreter() {
|
|
return pyobj_interpreter_.load(std::memory_order_acquire);
|
|
}
|
|
|
|
PyObject* PyObjectSlot::_unchecked_untagged_pyobj() const {
|
|
// NOLINTNEXTLINE(performance-no-int-to-ptr)
|
|
return reinterpret_cast<PyObject*>(
|
|
reinterpret_cast<uintptr_t>(pyobj_) & ~0x1ULL);
|
|
}
|
|
|
|
PyInterpreter& PyObjectSlot::load_pyobj_interpreter() const {
|
|
auto interpreter = pyobj_interpreter_.load(std::memory_order_acquire);
|
|
if (interpreter) {
|
|
return *interpreter;
|
|
}
|
|
TORCH_CHECK(false, "cannot access PyObject for Tensor - no interpreter set");
|
|
}
|
|
|
|
bool PyObjectSlot::owns_pyobj() {
|
|
// NOLINTNEXTLINE(performance-no-int-to-ptr)
|
|
return reinterpret_cast<uintptr_t>(pyobj_) & 1;
|
|
}
|
|
|
|
void PyObjectSlot::set_owns_pyobj(bool b) {
|
|
// NOLINTNEXTLINE(performance-no-int-to-ptr)
|
|
pyobj_ = reinterpret_cast<PyObject*>(
|
|
reinterpret_cast<uintptr_t>(_unchecked_untagged_pyobj()) | b);
|
|
}
|
|
|
|
} // namespace c10::impl
|