Move pyobj_ to TensorImpl (#18225)

Summary:
Currently, `THPVariable_Wrap(…)` and `THPVariable_NewWithVar(…)` depend on the existence of `pyobj_` in the autograd metadata of a Variable to convert the Variable to a Python tensor. However, after the Variable/Tensor merge, there will be Variables that don't contain autograd metadata, and to allow the conversion from non-autograd-meta Variable to a Python tensor we need to store the `pyobj_` outside of autograd metadata and in a place where it will always be available.

This PR makes it possible by moving `pyobj_` into TensorImpl, so that `THPVariable_Wrap(…)` and `THPVariable_NewWithVar(…)` can always access a Variable's `pyobj_` and convert the Variable to a Python tensor.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/18225

Differential Revision: D14562616

Pulled By: yf225

fbshipit-source-id: 18d4aaace70eee6120abaf9276036d1f8f51b18d
This commit is contained in:
Will Feng 2019-03-23 12:47:15 -07:00 committed by Facebook Github Bot
parent 5860fa5dcf
commit 32d0e7e339
4 changed files with 19 additions and 5 deletions

View file

@ -15,6 +15,7 @@
#include <c10/util/Optional.h>
#include <c10/util/Flags.h>
#include <c10/util/Logging.h>
#include <c10/util/python_stub.h>
// A global boolean variable to control whether we free memory when a Tensor
// is shrinked to a smaller size. As a result, a Tensor is always going to
@ -860,6 +861,14 @@ struct C10_API TensorImpl : public c10::intrusive_ptr_target {
return impl;
}
inline void set_pyobj(PyObject* pyobj) noexcept {
pyobj_ = pyobj;
}
inline PyObject* pyobj() const noexcept {
return pyobj_;
}
private:
// As an optimization, get_device handles the typical CUDA Tensor case and
// calls get_device_slow if the tensor stores its device somewhere else
@ -1368,6 +1377,8 @@ protected:
// at a time).
std::unique_ptr<c10::AutogradMetaInterface> autograd_meta_ = nullptr;
PyObject* pyobj_ = nullptr; // weak reference
// We could save a word or two by combining the SmallVector structs,
// since their size is redundant, and if we need to overflow the buffer space
// we could keep the two pointers together. However, that would require
@ -1461,10 +1472,11 @@ protected:
// numel
// data type pointer
// autograd metadata pointer
// PyObject pointer
// miscellaneous bitfield
//
static_assert(sizeof(void*) != sizeof(int64_t) || // if 64-bit...
sizeof(TensorImpl) == sizeof(int64_t) * 25,
sizeof(TensorImpl) == sizeof(int64_t) * 26,
"You changed the size of TensorImpl on 64-bit arch."
"See Note [TensorImpl size constraints] on how to proceed.");

4
c10/util/python_stub.h Normal file
View file

@ -0,0 +1,4 @@
#pragma once
struct _object;
using PyObject = _object;

View file

@ -28,7 +28,6 @@ Variable::Impl::Impl(at::Tensor data, std::unique_ptr<Variable::AutogradMeta> au
autograd_meta->requires_grad_ = false;
autograd_meta->is_view_ = false;
autograd_meta->output_nr_ = gradient_edge.input_nr;
autograd_meta->pyobj_ = nullptr;
// set_requires_grad also checks error conditions.
autograd_meta->set_requires_grad(requires_grad, this);

View file

@ -348,7 +348,6 @@ struct TORCH_API Variable::AutogradMeta : public c10::AutogradMetaInterface {
// We use this to make sure we can setup the backwards trace
// correctly when this variable is passed to another function.
uint32_t output_nr_;
PyObject* pyobj_ = nullptr; // weak reference
// Mutex to ensure that concurrent read operations that modify internal
// state are still thread-safe. Used by grad_fn() and
@ -756,11 +755,11 @@ inline const std::string& Variable::name() const noexcept {
}
inline void Variable::set_pyobj(PyObject* pyobj) noexcept {
get_autograd_meta()->pyobj_ = pyobj;
get()->set_pyobj(pyobj);
}
inline PyObject* Variable::pyobj() const noexcept {
return get_autograd_meta()->pyobj_;
return get()->pyobj();
}
inline Variable::AutogradMeta* Variable::get_autograd_meta() const noexcept {