mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
a3b7dd7b78
Summary: When default hooks are set, they are pushed onto a stack. When nesting context-manager, only the inner-most hooks will be applied. There is special care needed to update the TLS code. See also https://github.com/pytorch/pytorch/issues/70940 (i.e. do we need to be storing the enabled flag as well?) Fixes https://github.com/pytorch/pytorch/issues/70134 Pull Request resolved: https://github.com/pytorch/pytorch/pull/70932 Reviewed By: mruberry Differential Revision: D33530370 Pulled By: albanD fbshipit-source-id: 3197d585d77563f36c175d3949115a0776b309f4
79 lines
3.1 KiB
C++
79 lines
3.1 KiB
C++
#include <torch/csrc/autograd/python_saved_variable_hooks.h>
|
|
#include <ATen/SavedTensorHooks.h>
|
|
|
|
#include <torch/csrc/THP.h>
|
|
|
|
namespace py = pybind11;
|
|
|
|
namespace torch { namespace autograd {
|
|
PySavedVariableHooks::PySavedVariableHooks(py::function &pack_hook, py::function &unpack_hook) :
|
|
// steals the reference (we will decref ourselves)
|
|
pack_hook_(pack_hook.release().ptr()),
|
|
unpack_hook_(unpack_hook.release().ptr()) {}
|
|
|
|
void PySavedVariableHooks::call_pack_hook(const at::Tensor &tensor) {
|
|
py::gil_scoped_acquire acquire;
|
|
auto pack_hook = py::reinterpret_borrow<py::function>(pack_hook_);
|
|
auto wrapped = THPVariable_Wrap(tensor);
|
|
py::object obj = py::reinterpret_steal<py::object>(wrapped);
|
|
py::object packed = pack_hook(obj);
|
|
data_ = packed.release().ptr();
|
|
// pack_hook, obj are decrefed on exit
|
|
// wrapped and packed had their references stolen
|
|
// pack_hook_ and data_ will be manually decrefed when the saved variable is released
|
|
}
|
|
|
|
at::Tensor PySavedVariableHooks::call_unpack_hook() {
|
|
py::gil_scoped_acquire acquire;
|
|
auto unpack_hook = py::reinterpret_borrow<py::function>(unpack_hook_);
|
|
py::object obj = py::cast<py::object>(data_);
|
|
py::object res = unpack_hook(obj);
|
|
PyObject* ptr = res.ptr();
|
|
TORCH_CHECK_TYPE(THPVariable_Check(ptr), "Output of saved tensor unpack_hook expected to be a Tensor but got result of type ", THPUtils_typename(ptr));
|
|
return THPVariable_Unpack(ptr);
|
|
// unpack_hook, obj and res are decrefed on exit
|
|
// ptr is only alive as long as res is
|
|
// unpack_hook_ will be manually decrefed when the saved variable is released
|
|
}
|
|
|
|
PySavedVariableHooks::~PySavedVariableHooks() {
|
|
// If python is already dead, leak the wrapped python objects
|
|
if (Py_IsInitialized()) {
|
|
py::gil_scoped_acquire gil;
|
|
Py_XDECREF(pack_hook_);
|
|
Py_XDECREF(unpack_hook_);
|
|
Py_XDECREF(data_);
|
|
}
|
|
}
|
|
|
|
void PyDefaultSavedVariableHooks::push_hooks(py::function &pack_hook, py::function &unpack_hook) {
|
|
at::SavedTensorDefaultHooks::enable();
|
|
at::SavedTensorDefaultHooks::push_hooks(pack_hook.release().ptr(), unpack_hook.release().ptr());
|
|
}
|
|
|
|
void PyDefaultSavedVariableHooks::pop_hooks() {
|
|
PyObject *pack_hook(nullptr), *unpack_hook(nullptr);
|
|
std::tie(pack_hook, unpack_hook) = at::SavedTensorDefaultHooks::get_hooks();
|
|
TORCH_INTERNAL_ASSERT(pack_hook != nullptr && unpack_hook != nullptr);
|
|
if (Py_IsInitialized()) {
|
|
py::gil_scoped_acquire gil;
|
|
Py_XDECREF(pack_hook);
|
|
Py_XDECREF(unpack_hook);
|
|
}
|
|
at::SavedTensorDefaultHooks::pop_hooks();
|
|
}
|
|
|
|
std::unique_ptr<SavedVariableHooks> PyDefaultSavedVariableHooks::get_hooks() {
|
|
PyObject *pack_hook(nullptr), *unpack_hook(nullptr);
|
|
std::tie(pack_hook, unpack_hook) = at::SavedTensorDefaultHooks::get_hooks();
|
|
if (!pack_hook || !unpack_hook) {
|
|
return nullptr;
|
|
}
|
|
py::gil_scoped_acquire gil;
|
|
py::function pack_hook_ = py::reinterpret_borrow<py::function>(pack_hook);
|
|
py::function unpack_hook_ = py::reinterpret_borrow<py::function>(unpack_hook);
|
|
return std::make_unique<PySavedVariableHooks>(pack_hook_, unpack_hook_);
|
|
}
|
|
|
|
}}
|