Move backward/set_data to Type-based dispatch.

Summary: Pull Request resolved: pytorch#11440

Differential Revision: D9736565

Pulled By: gchanan

fbshipit-source-id: 1e66f54f1c87084f37c0b014030f0d6d2f8dfaee

Author: gchanan

aten/src/ATen/SparseTensorImpl.cpp

@@ -29,7 +29,7 @@ namespace {
 // This means that we allocate a [1,0] size indices tensor and a [0] size
 // values tensor for such an empty tensor.
 SparseTensorImpl::SparseTensorImpl(at::TensorTypeId type_id, at::ScalarType scalar_type)
-    : TensorImpl(type_id, scalar_type, false)
+    : TensorImpl(type_id, scalar_type, nullptr, false)
     , size_{0}
     , sparseDims_(1)
     , denseDims_(0)

aten/src/ATen/TensorImpl.cpp

@@ -1,6 +1,5 @@
 #include <ATen/TensorImpl.h>
 
-#include <ATen/Type.h>
 #include <ATen/core/optional.h>
 #include <ATen/core/Backend.h>
 #include <ATen/core/WrapDimMinimal.h>
@@ -18,31 +17,12 @@ const Tensor& TensorImpl::grad() const {
   AT_ERROR("grad is not implemented for Tensor");
 }
 
-void TensorImpl::backward(
-    at::optional<Tensor> gradient,
-    bool keep_graph,
-    bool create_graph) {
-  AT_ERROR("backward is not implemented for Tensor");
-}
-
-void TensorImpl::set_data(Tensor new_data) {
-  AT_ERROR("set_type is not implemented for Tensor");
-}
-
-void Tensor::backward(
-    at::optional<Tensor> gradient,
-    bool keep_graph,
-    bool create_graph) {
-  tensor_impl_->backward(std::move(gradient), keep_graph, create_graph);
-}
-
-TensorImpl::TensorImpl(TensorTypeId type_id, ScalarType scalar_type, bool is_variable)
+TensorImpl::TensorImpl(TensorTypeId type_id, ScalarType scalar_type, Allocator *allocator, bool is_variable)
     : TensorImpl({}, type_id, scalar_type, is_variable) {
   // UndefinedTensors and SparseTensors don't have storages.
   if (type_id != UndefinedTensorId() && scalar_type != ScalarType::Undefined
       && type_id != SparseCPUTensorId() && type_id != SparseCUDATensorId()) {
-    auto type = &globalLegacyTypeDispatch().getNonVariableType(tensorTypeIdToBackend(type_id), scalar_type);
-    storage_ = type->storage(true);
+    storage_ = Storage(scalar_type, 0, allocator, true);
   }
 }
 

aten/src/ATen/TensorImpl.h

@@ -22,7 +22,7 @@ struct Tensor;
 namespace at {
 struct AT_API TensorImpl : public c10::intrusive_ptr_target {
   TensorImpl() = delete;
-  TensorImpl(TensorTypeId type_id, ScalarType scalar_type, bool is_variable);
+  TensorImpl(TensorTypeId type_id, ScalarType scalar_type, Allocator *allocator, bool is_variable);
   TensorImpl(Storage&& storage, TensorTypeId type_id, bool is_variable);
 
   virtual void release_resources() override;
@@ -90,13 +90,6 @@ struct AT_API TensorImpl : public c10::intrusive_ptr_target {
   virtual Tensor& grad();
   virtual const Tensor& grad() const;
 
-  virtual void backward(
-      at::optional<Tensor> gradient,
-      bool keep_graph,
-      bool create_graph);
-
-  virtual void set_data(Tensor new_data);
-
   // TODO: make these protected
   // Note: storage->size() may be greater than the recorded size
   // of a tensor

aten/src/ATen/UndefinedTensor.cpp

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

aten/src/ATen/function_wrapper.py

@@ -331,17 +331,17 @@ def __init__(self, reason):
 
 ALLOC_NOARGS_WRAP = {
     'THTensor*': 'c10::make_intrusive<TensorImpl, UndefinedTensor>'
-                 '(${Backend}TensorId(), ScalarType::${ScalarName}, false).release()',
+                 '(${Backend}TensorId(), ScalarType::${ScalarName}, allocator(), false).release()',
     'THBoolTensor*': 'c10::make_intrusive<TensorImpl, UndefinedTensor>'
-                     '(${Backend}TensorId(), ScalarType::Byte, false).release()',
+                     '(${Backend}TensorId(), ScalarType::Byte, allocator(), false).release()',
     'THIndexTensor*': 'c10::make_intrusive<TensorImpl, UndefinedTensor>'
-                      '(${Backend}TensorId(), ScalarType::Long, false).release()',
+                      '(${Backend}TensorId(), ScalarType::Long, allocator(), false).release()',
     'THIntegerTensor*': 'c10::make_intrusive<TensorImpl, UndefinedTensor>'
-                        '(${Backend}TensorId(), ScalarType::Int, false).release()',
+                        '(${Backend}TensorId(), ScalarType::Int, allocator(), false).release()',
     'THDenseTensor*': 'c10::make_intrusive<TensorImpl, UndefinedTensor>'
-                      '(${Backend}TensorId(), ScalarType::${ScalarName}, false).release()',
+                      '(${Backend}TensorId(), ScalarType::${ScalarName}, allocator(), false).release()',
     'THDenseIndexTensor*': 'c10::make_intrusive<TensorImpl, UndefinedTensor>'
-                           '(${Backend}TensorId(), ScalarType::Long, false).release()'
+                           '(${Backend}TensorId(), ScalarType::Long, allocator(), false).release()'
 }
 
 ALLOC_WRAP = {

aten/src/ATen/templates/Tensor.h

@@ -244,9 +244,7 @@ struct AT_API Tensor {
     return tensor_impl_->grad();
   }
 
-  void set_data(Tensor new_data) {
-    tensor_impl_->set_data(new_data);
-  }
+  void set_data(Tensor new_data);
 
   /// Computes the gradient of current tensor w.r.t. graph leaves.
   void backward(

aten/src/ATen/templates/TensorMethods.h

@@ -83,6 +83,17 @@ inline Tensor Tensor::to(Device device, bool non_blocking) const {
   return detail::to(*this, options().device(device), non_blocking);
 }
 
+inline void Tensor::backward(
+    at::optional<Tensor> gradient,
+    bool keep_graph,
+    bool create_graph) {
+  type().backward(*this, std::move(gradient), keep_graph, create_graph);
+}
+
+inline void Tensor::set_data(Tensor new_data) {
+  type().set_data(*this, new_data);
+}
+
 // all static inline to allow for inlining of the non-dynamic part of dispatch
 ${tensor_method_definitions}
 

aten/src/ATen/templates/Type.h

@@ -104,6 +104,9 @@ struct AT_API Type {
   virtual Tensor & s_copy_(Tensor & self, const Tensor & src, bool non_blocking) const = 0;
   virtual Tensor & _s_copy_from(const Tensor & self, Tensor & dst, bool non_blocking) const = 0;
 
+  virtual void backward(Tensor & self, at::optional<Tensor> gradient, bool keep_graph, bool create_graph) const = 0;
+  virtual void set_data(Tensor & self, Tensor new_data) const = 0;
+
   virtual Tensor tensorFromBlob(void * data, IntList sizes, const std::function<void(void*)> & deleter=noop_deleter) const = 0;
   virtual Tensor tensorFromBlob(void * data, IntList sizes, IntList strides, const std::function<void(void*)> & deleter=noop_deleter) const = 0;
   virtual Tensor tensorWithAllocator(IntList sizes, Allocator* allocator) const = 0;

aten/src/ATen/templates/TypeDefault.cpp

@@ -40,6 +40,14 @@ Tensor TypeDefault::copy(const Tensor & src, bool non_blocking) const {
   }
 }
 
+void TypeDefault::backward(Tensor & self, at::optional<Tensor> gradient, bool keep_graph, bool create_graph) const {
+  AT_ERROR("backward is not implemented for Tensor");
+}
+
+void TypeDefault::set_data(Tensor & self, Tensor new_data) const {
+  AT_ERROR("set_data is not implemented for Tensor");
+}
+
 Type & TypeDefault::toBackend(Backend b) const {
   return at::globalContext().getNonVariableType(b,scalarType());
 }

aten/src/ATen/templates/TypeDefault.h

@@ -28,6 +28,9 @@ struct AT_API TypeDefault : public Type {
   Tensor copy(const Tensor & src, bool non_blocking=false) const override;
   Tensor & copy_(Tensor & self, const Tensor & src, bool non_blocking=false) const override;
 
+  void backward(Tensor & self, at::optional<Tensor> gradient, bool keep_graph, bool create_graph) const override;
+  void set_data(Tensor & self, Tensor new_data) const override;
+
   Tensor tensorFromBlob(void * data, IntList sizes, const std::function<void(void*)> & deleter=noop_deleter) const override;
   Tensor tensorFromBlob(void * data, IntList sizes, IntList strides, const std::function<void(void*)> & deleter=noop_deleter) const override;
   Tensor tensorWithAllocator(IntList sizes, Allocator* allocator) const override;

tools/autograd/templates/VariableType.cpp

@@ -378,6 +378,13 @@ static bool isFloatingPoint(ScalarType s) {
   return s == kFloat || s == kDouble || s == kHalf;
 }
 
+void VariableType::backward(Tensor & self, at::optional<Tensor> gradient, bool keep_graph, bool create_graph) const {
+  as_variable_ref(self).backward(gradient, keep_graph, create_graph);
+}
+
+void VariableType::set_data(Tensor & self, Tensor new_data) const {
+  as_variable_ref(self).set_data(new_data);
+}
 Tensor & VariableType::s_copy_(Tensor & self, const Tensor & src, bool non_blocking) const {
   jit::Node* node = nullptr;
   if(torch::jit::tracer::isTracing()) {

tools/autograd/templates/VariableType.h

@@ -58,6 +58,10 @@ struct TORCH_API VariableType final : public at::TypeDefault {
 
   Tensor & s_copy_(Tensor & self, const Tensor & src, bool non_blocking) const override;
   Tensor & _s_copy_from(const Tensor & self, Tensor & dst, bool non_blocking) const override;
+
+  void backward(Tensor & self, at::optional<Tensor> gradient, bool keep_graph, bool create_graph) const override;
+  void set_data(Tensor & self, Tensor new_data) const override;
+
   ${type_derived_method_declarations}
 
 private:

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(data.type().type_id(), data.type().scalarType(), /* is variable */ true),
+    : TensorImpl(data.type().type_id(), data.type().scalarType(), data.type().allocator(), /* is variable */ true),
       data_(std::move(data)),
       grad_fn_(std::move(gradient_edge.function)),
       requires_grad_(false),

torch/csrc/autograd/variable.h

@@ -187,6 +187,12 @@ struct TORCH_API Variable : public at::Tensor {
   /// this. If this `Variable` is a view, throws an `std::runtime_error()`.
   void detach_();
 
+  /// Computes the gradient of current tensor w.r.t. graph leaves.
+  void backward(at::optional<Tensor> gradient, bool keep_graph, bool create_graph) const;
+
+  /// Sets the type of the Variable.
+  void set_data(Tensor new_data) const;
+
   /// Set the gradient edge -- i.e. `grad_fn` and `input_nr` -- of the
   /// `Variable`.
   /// NOTE: This will always set the `grad_fn`, even if this is a leaf variable,
@@ -324,14 +330,12 @@ struct TORCH_API Variable::Impl : public at::TensorImpl {
   Variable detach() const;
   void detach_();
 
-  /// Sets the type of the Variable.
-  void set_data(Tensor new_data) override;
+  void set_data(Tensor new_data);
 
-  /// Computes the gradient of current tensor w.r.t. graph leaves.
   void backward(
       at::optional<at::Tensor> gradient,
       bool keep_graph,
-      bool create_graph) override;
+      bool create_graph);
 
   /// Reset all expensive fields to free up resources
   void release_resources() override;
@@ -500,6 +504,14 @@ inline void Variable::detach_() {
   get()->detach_();
 }
 
+inline void Variable::backward(at::optional<Tensor> gradient, bool keep_graph, bool create_graph) const {
+  get()->backward(gradient, keep_graph, create_graph);
+}
+
+inline void Variable::set_data(Tensor new_data) const {
+  get()->set_data(new_data);
+}
+
 inline void Variable::set_gradient_edge(Edge edge) noexcept {
   get()->grad_fn_ = std::move(edge.function);
   get()->output_nr_ = edge.input_nr;

torch/csrc/jit/interpreter.cpp

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