Delete type_ field from TensorImpl, replaced with backend_/scalar_typ… (pytorch#9787)

Summary:
…e_/is_variable_

The basic game plan is to stop accessing the type_ field directly,
and instead using the stored backend_, scalar_type_ and
is_variable_ to look up the appropriate Type from Context.
Storage of backend_ and scalar_type_ are new.

At some future point in time, I'd like to look at this code
carefully to see if I can get everything in this codepath inlining.
I didn't do it in this patch because there are circular include
problems making things difficult.

Some other details:

- Added Device::backend() which does what it says on the tin

- SparseTensorImpl is temporarily hard-coded to root in at::Context
  for the appropriate context.  If/when we put this in shared code,
  we'll have to break this dep too, but for now it should be OK.

- There's a stupid problem with globalContext() deadlocking if
  you didn't actually initialize it before loading libtorch.so
  (which is bringing along the variable hooks).  I didn't fix
  it in this PR; it's tracked in pytorch#9784

Signed-off-by: Edward Z. Yang <[email protected]>
Pull Request resolved: pytorch#9787

Reviewed By: cpuhrsch

Differential Revision: D8980971

Pulled By: ezyang

fbshipit-source-id: 2b4d867abfdc3999a836a220c638c109053145a8

Author: ezyang

aten/src/ATen/Context.cpp

@@ -37,8 +37,11 @@ Context::Context()
   Type::registerCPU(this);
 }
 
+// NB: Ensure that globalContext is initialized before we load
+// variable hooks, otherwise we will deadlock.  Regardless, the
+// deadlock is bad, and being tracked at https://github.com/pytorch/pytorch/issues/9784
+static Context globalContext_;
 Context & globalContext() {
-  static Context globalContext_;
   return globalContext_;
 }
 

aten/src/ATen/SparseTensorImpl.cpp

@@ -18,14 +18,14 @@ namespace at {
 // tensor and a [0] size values tensor for such an empty tensor.  However,
 // we don't currently support zero-size dimensions, so we can't actually
 // do this; so we just allocate zero-size tensors for everything.
-SparseTensorImpl::SparseTensorImpl(Type * type)
-    : TensorImpl(type, nullptr)
+SparseTensorImpl::SparseTensorImpl(at::Backend backend, at::ScalarType scalar_type)
+    : TensorImpl(backend, scalar_type, nullptr, false)
     , size_{0}
     , sparseDims_(1)
     , denseDims_(0)
-    , indices_(type->toDense().toScalarType(ScalarType::Long).tensor())
-    , values_(type->toDense().tensor()) {
-      AT_ASSERT(type->is_sparse());
+    , indices_(globalContext().getTypeOpt(toDense(backend), ScalarType::Long)->tensor())
+    , values_(globalContext().getTypeOpt(toDense(backend), scalar_type)->tensor()) {
+      AT_ASSERT(backend == Backend::SparseCPU || backend == Backend::SparseCUDA);
     }
 
 IntList SparseTensorImpl::sizes() const {

aten/src/ATen/SparseTensorImpl.h

@@ -48,7 +48,7 @@ struct AT_API SparseTensorImpl : public TensorImpl {
 
 public:
   // Public for now...
-  explicit SparseTensorImpl(Type * type);
+  explicit SparseTensorImpl(at::Backend, at::ScalarType);
 
   int64_t nnz() const { return nnz_; }
   int64_t sparseDims() const { return sparseDims_; }

aten/src/ATen/TensorImpl.cpp

@@ -2,10 +2,23 @@
 
 #include <ATen/Tensor.h>
 #include <ATen/optional.h>
+#include <ATen/Context.h>
+
+#include <ATen/detail/VariableHooksInterface.h>
 
 #include <TH/THTensor.hpp>
 
 namespace at {
+
+Type& TensorImpl::type() const {
+  Type* base_type = &globalContext().getType(backend_, scalar_type_);
+  if (is_variable_) {
+    return detail::getVariableHooks().getVariableType(*base_type);
+  } else {
+    return *base_type;
+  }
+}
+
 Tensor& TensorImpl::grad() {
   AT_ERROR("grad is not implemented for Tensor");
 }

aten/src/ATen/TensorImpl.h

@@ -18,16 +18,18 @@ struct Tensor;
 
 namespace at {
 struct AT_API TensorImpl : public Retainable {
-  explicit TensorImpl(Type * type, THTensor * tensor)
-  : type_(type), tensor(tensor) {}
+  explicit TensorImpl(Backend backend, ScalarType scalar_type, THTensor * tensor, bool is_variable)
+  : backend_(backend), scalar_type_(scalar_type), is_variable_(is_variable), tensor(tensor) {}
 
   virtual ~TensorImpl();
 
   virtual void release_resources() override;
 
-  Type & type() const {
-    return *type_;
-  }
+  // The implementation of this method will have to be hoisted out and
+  // hooked in, so that Caffe2 doesn't need to know about Context
+  // TODO: This really really needs to be inlined.
+  Type & type() const;
+
   const char * toString() const;
   virtual IntList sizes() const;
   virtual IntList strides() const;
@@ -91,8 +93,12 @@ struct AT_API TensorImpl : public Retainable {
   virtual void set_data(Tensor new_data);
 
 protected:
+  Backend backend_;
+  // INVARIANT: When storage is non-null, this scalar type must
+  // agree with the scalar type in storage
+  ScalarType scalar_type_;
+  bool is_variable_ = false;
   bool is_wrapped_number_ = false;
-  Type * type_;
 public:
   THTensor * tensor;
 };

aten/src/ATen/UndefinedTensor.cpp

@@ -6,7 +6,7 @@ namespace at {
 
 // should this use the globalContext?  Can it get a context passed in somehow?
 UndefinedTensor::UndefinedTensor()
-: TensorImpl(&(globalContext().getType(Backend::Undefined,ScalarType::Undefined)), nullptr) {
+: TensorImpl(Backend::Undefined, ScalarType::Undefined, nullptr, /* is variable */ false) {
 }
 
 IntList UndefinedTensor::sizes() const {

aten/src/ATen/detail/VariableHooksInterface.h

@@ -3,6 +3,7 @@
 #include <ATen/Registry.h>
 #include <ATen/Error.h>
 #include <ATen/ScalarType.h>
+#include <ATen/Type.h>
 
 namespace at {
   class Context;
@@ -25,6 +26,10 @@ struct AT_API VariableHooksInterface {
   // squelch -Werror=non-virtual-dtor
   virtual ~VariableHooksInterface() {}
 
+  virtual Type& getVariableType(const at::Type& baseType) const {
+    AT_ERROR("cannot getVariableType without libtorch");
+  }
+
   virtual void registerVariableTypeFor(Context*, Backend backend, ScalarType scalar_type) const {
     // no-op if Variable not available; it'll get handled (if at all) when
     // libtorch.so gets loaded

aten/src/ATen/native/sparse/SparseTensor.cpp

@@ -63,7 +63,7 @@ SparseTensor new_sparse(const SparseType& dtype) {
   AT_ASSERT(!dtype.is_variable());
   AT_ASSERT(dtype.is_sparse());
   // TODO: Hmm... this const_cast business seems a bit dodgy
-  return SparseTensor(new SparseTensorImpl(const_cast<SparseType*>(&dtype)), /* retain */ false);
+  return SparseTensor(new SparseTensorImpl(dtype.backend(), dtype.scalarType()), /* retain */ false);
 }
 
 /*** Helper methods ***/

aten/src/ATen/templates/TensorDerived.cpp

@@ -21,7 +21,7 @@ namespace detail {
 }
 
 ${Tensor}::${Tensor}(${THTensor} * tensor)
-: TensorImpl(&globalContext().getType(Backend::${Backend},ScalarType::${ScalarName}), tensor)
+: TensorImpl(Backend::${Backend}, ScalarType::${ScalarName}, tensor, /* is variable */ false)
 {}
 
 ${TensorDenseOrSparse}

aten/src/TH/THTensor.hpp

@@ -56,6 +56,10 @@ struct THTensor
       return sizes_.size();
     }
 
+    at::ScalarType scalar_type() const {
+      return storage_->scalar_type;
+    }
+
     ptrdiff_t storage_offset() const {
       return storage_offset_;
     }

torch/csrc/autograd/aten_variable_hooks.cpp

@@ -6,6 +6,7 @@ namespace torch { namespace autograd {
 struct VariableHooks : public at::VariableHooksInterface {
   VariableHooks(at::VariableHooksArgs) {}
   void registerVariableTypeFor(at::Context*, at::Backend, at::ScalarType) const override;
+  at::Type& getVariableType(const at::Type&) const override;
 };
 
 // Sigh, the registry doesn't support namespaces :(
@@ -20,4 +21,8 @@ void VariableHooks::registerVariableTypeFor(at::Context* context, at::Backend ba
   register_variable_type_for(baseType);
 }
 
+at::Type& VariableHooks::getVariableType(const at::Type& baseType) const {
+  return *VariableType::getType(baseType);
+}
+
 }} // torch::autograd

torch/csrc/autograd/variable.cpp

@@ -22,7 +22,7 @@
 namespace torch {
 namespace autograd {
 Variable::Impl::Impl(at::Tensor data, bool requires_grad, Edge gradient_edge)
-    : TensorImpl(VariableType::getType(data), nullptr),
+    : TensorImpl(data.type().backend(), data.type().scalarType(), nullptr, /* is variable */ true),
       data_(std::move(data)),
       grad_fn_(std::move(gradient_edge.function)),
       requires_grad_(false),
@@ -118,7 +118,9 @@ void Variable::Impl::backward(
 
 void Variable::Impl::set_data(Tensor new_data) {
   if (new_data.type() != data_.type()) {
-    type_ = VariableType::getType(new_data.type());
+    scalar_type_ = new_data.type().scalarType();
+    backend_ = new_data.type().backend();
+    is_variable_ = true;
     // Clear grad_accumulator if it exists, since it stores the old type info.
     grad_accumulator_.reset();
   }

torch/csrc/autograd/variable.h

@@ -327,9 +327,6 @@ struct Variable::Impl : public at::TensorImpl {
   /// Reset all expensive fields to free up resources
   void release_resources() override;
 
-  // Make this field public so we can access it from `Variable`.
-  using at::TensorImpl::type_;
-
   std::string name;
   at::Tensor data_;
 

torch/csrc/jit/interpreter.cpp

@@ -337,7 +337,7 @@ struct PreprocessGraph {
 struct ContainerTensor : public at::TensorImpl {
 public:
   ContainerTensor()
-  : TensorImpl(&(at::globalContext().getType(at::Backend::Undefined,at::ScalarType::Undefined)), nullptr) {}
+  : TensorImpl(at::Backend::Undefined,at::ScalarType::Undefined, nullptr, /* is_variable */ false) {}
 
   virtual ~ContainerTensor() = default;
   virtual at::IntList sizes() const override {