From e2a5b203fccf78717c5171e89e2f9d258c487fe2 Mon Sep 17 00:00:00 2001 From: Will Feng Date: Mon, 11 Feb 2019 12:48:17 -0800 Subject: [PATCH] Enforce same input tensor storage in VariableType functions (#16305) Summary: In VariableType.cpp, when a function modifies its input tensors, it should only change the input tensors' storage data in-place, and should never change the input tensors' storage pointers. This PR adds checks for this, and also fixes functions that fail this test. This is part of the Variable/Tensor merge work (https://github.com/pytorch/pytorch/issues/13638). Pull Request resolved: https://github.com/pytorch/pytorch/pull/16305 Differential Revision: D13897855 Pulled By: yf225 fbshipit-source-id: 0c4fc7eb530d30db88037b1f0981f6f8454d3b79 --- aten/src/ATen/SparseTensorImpl.cpp | 3 + aten/src/ATen/SparseTensorImpl.h | 1 + aten/src/ATen/core/Tensor.h | 3 + aten/src/ATen/native/BatchLinearAlgebra.cpp | 5 +- aten/src/ATen/templates/Tensor.h | 3 + c10/core/TensorImpl.cpp | 4 + c10/core/TensorImpl.h | 7 +- c10/core/UndefinedTensorImpl.cpp | 4 + c10/core/UndefinedTensorImpl.h | 1 + tools/autograd/gen_variable_type.py | 97 ++++++++++++++++++++- torch/csrc/autograd/variable.cpp | 4 + torch/csrc/autograd/variable.h | 1 + 12 files changed, 130 insertions(+), 3 deletions(-) diff --git a/aten/src/ATen/SparseTensorImpl.cpp b/aten/src/ATen/SparseTensorImpl.cpp index 016a3e7e2cf..d7b80e0901a 100644 --- a/aten/src/ATen/SparseTensorImpl.cpp +++ b/aten/src/ATen/SparseTensorImpl.cpp @@ -72,6 +72,9 @@ TensorImpl* SparseTensorImpl::maybe_zero_dim(bool condition_when_zero_dim) { " changing dimensionality via maybe_zero_dim"); return this; } +bool SparseTensorImpl::has_storage() const { + return false; +} const Storage& SparseTensorImpl::storage() const { AT_ERROR("sparse tensors do not have storage"); } diff --git a/aten/src/ATen/SparseTensorImpl.h b/aten/src/ATen/SparseTensorImpl.h index 9da294fab4e..a5de5a1b548 100644 --- a/aten/src/ATen/SparseTensorImpl.h +++ b/aten/src/ATen/SparseTensorImpl.h @@ -51,6 +51,7 @@ public: int64_t dim() const override; TensorImpl* maybe_zero_dim(bool condition_when_zero_dim) override; + bool has_storage() const override; const Storage& storage() const override; int64_t storage_offset() const override; diff --git a/aten/src/ATen/core/Tensor.h b/aten/src/ATen/core/Tensor.h index a110df7f18c..266315586f6 100644 --- a/aten/src/ATen/core/Tensor.h +++ b/aten/src/ATen/core/Tensor.h @@ -184,6 +184,9 @@ class CAFFE2_API Tensor { ScalarType scalar_type() const { return typeMetaToScalarType(impl_->dtype()); } + bool has_storage() const { + return defined() && impl_->has_storage(); + } const Storage& storage() const { return impl_->storage(); } diff --git a/aten/src/ATen/native/BatchLinearAlgebra.cpp b/aten/src/ATen/native/BatchLinearAlgebra.cpp index 4cb57dc9e78..e067878bd77 100644 --- a/aten/src/ATen/native/BatchLinearAlgebra.cpp +++ b/aten/src/ATen/native/BatchLinearAlgebra.cpp @@ -175,7 +175,10 @@ std::tuple gesv_out(Tensor& solution, Tensor& lu, const Tensor& AT_CHECK(self.dim() == 2 && A.dim() == 2, "torch.gesv() with the `out` keyword does not support batching. " "b.dim() (", self.dim(), ") and A.dim() (", A.dim(), ") must both be 2."); - std::tie(solution, lu) = at::_gesv_helper(self, A); + Tensor solution_tmp, lu_tmp; + std::tie(solution_tmp, lu_tmp) = at::_gesv_helper(self, A); + solution.resize_as_(solution_tmp).copy_(solution_tmp); + lu.resize_as_(lu_tmp).copy_(lu_tmp); return std::tuple(solution, lu); } diff --git a/aten/src/ATen/templates/Tensor.h b/aten/src/ATen/templates/Tensor.h index aa031fe036a..3d99b8aaf32 100644 --- a/aten/src/ATen/templates/Tensor.h +++ b/aten/src/ATen/templates/Tensor.h @@ -184,6 +184,9 @@ class CAFFE2_API Tensor { ScalarType scalar_type() const { return typeMetaToScalarType(impl_->dtype()); } + bool has_storage() const { + return defined() && impl_->has_storage(); + } const Storage& storage() const { return impl_->storage(); } diff --git a/c10/core/TensorImpl.cpp b/c10/core/TensorImpl.cpp index 33d677f390a..446d206c24b 100644 --- a/c10/core/TensorImpl.cpp +++ b/c10/core/TensorImpl.cpp @@ -110,6 +110,10 @@ TensorImpl* TensorImpl::maybe_zero_dim(bool condition_when_zero_dim) { return this; } +bool TensorImpl::has_storage() const { + return storage_; +} + const Storage& TensorImpl::storage() const { return storage_; } diff --git a/c10/core/TensorImpl.h b/c10/core/TensorImpl.h index bd44ed53487..f97798d3338 100644 --- a/c10/core/TensorImpl.h +++ b/c10/core/TensorImpl.h @@ -292,7 +292,12 @@ struct C10_API TensorImpl : public c10::intrusive_ptr_target { virtual int64_t dim() const; /** - * Return the underyling storage of a Tensor. Multiple tensors may share + * True if this tensor has storage. See storage() for details. + */ + virtual bool has_storage() const; + + /** + * Return the underlying storage of a Tensor. Multiple tensors may share * a single storage. A Storage is an impoverished, Tensor-like class * which supports far less operations than Tensor. * diff --git a/c10/core/UndefinedTensorImpl.cpp b/c10/core/UndefinedTensorImpl.cpp index 106e44165a9..badc5ec7e46 100644 --- a/c10/core/UndefinedTensorImpl.cpp +++ b/c10/core/UndefinedTensorImpl.cpp @@ -24,6 +24,10 @@ int64_t UndefinedTensorImpl::dim() const { AT_ERROR("dim() called on undefined Tensor"); } +bool UndefinedTensorImpl::has_storage() const { + AT_ERROR("has_storage() called on undefined Tensor"); +} + const Storage& UndefinedTensorImpl::storage() const { AT_ERROR("storage() called on undefined Tensor"); } diff --git a/c10/core/UndefinedTensorImpl.h b/c10/core/UndefinedTensorImpl.h index 8ea4ba9b042..9f1cb93c10e 100644 --- a/c10/core/UndefinedTensorImpl.h +++ b/c10/core/UndefinedTensorImpl.h @@ -22,6 +22,7 @@ struct C10_API UndefinedTensorImpl final : public TensorImpl { int64_t size(int64_t d) const override; int64_t stride(int64_t d) const override; int64_t dim() const override; + bool has_storage() const override; const Storage& storage() const override; int64_t storage_offset() const override; private: diff --git a/tools/autograd/gen_variable_type.py b/tools/autograd/gen_variable_type.py index 289ac3a78df..2453352605f 100644 --- a/tools/autograd/gen_variable_type.py +++ b/tools/autograd/gen_variable_type.py @@ -29,7 +29,6 @@ from .utils import CodeTemplate, nested_dict, write, uninplace_api_name from .gen_autograd import VIEW_FUNCTIONS from .gen_autograd_functions import uses_single_grad - # These functions are written manually in templates/VariableType.cpp MANUAL_IMPLEMENTATIONS = { 'resize_', 'resize_as_', 'detach', 'detach_', 's_copy_', '_s_copy_from' @@ -80,6 +79,68 @@ DONT_REQUIRE_DERIVATIVE = { '_coalesced_', } +# NOTE [ Invariant: TensorImpl and Storage Pointer Equality ] +# +# When a function modifies its input tensors (via inplace or out-variants), +# it should never change the the input tensors' underlying c10::TensorImpl pointers +# or c10::Storage pointers. +# +# The following code templates implement the checks for this invariant: +SAVE_TENSOR_STORAGE = CodeTemplate("""\ +c10::optional ${tensor_name}_storage_saved = + ${tensor_name}.has_storage() ? c10::optional(${tensor_name}.storage()) : c10::nullopt; +""") + +ENFORCE_SAME_TENSOR_STORAGE = CodeTemplate("""\ +if (${tensor_name}_storage_saved.has_value()) + AT_ASSERT(${tensor_name}_storage_saved.value().is_alias_of(${tensor_name}.storage())); +""") + +SAVE_TENSORLIST_STORAGE = CodeTemplate("""\ +std::vector> ${tensorlist_name}_storage_saved(${tensorlist_name}.size()); +for (Tensor tensor : ${tensorlist_name}) + ${tensorlist_name}_storage_saved.push_back( + tensor.has_storage() ? c10::optional(tensor.storage()) : c10::nullopt); +""") + +ENFORCE_SAME_TENSORLIST_STORAGE = CodeTemplate("""\ +for (size_t i=0; i<${tensorlist_name}.size(); i++) { + if (${tensorlist_name}_storage_saved[i].has_value()) + AT_ASSERT(${tensorlist_name}_storage_saved[i].value().is_alias_of(${tensorlist_name}[i].storage())); +} +""") + +SAVE_TENSOR_IMPL = CodeTemplate("""\ +c10::intrusive_ptr ${tensor_name}_impl_saved; +if (${tensor_name}.defined()) ${tensor_name}_impl_saved = ${tensor_name}.getIntrusivePtr(); +""") + +ENFORCE_SAME_TENSOR_IMPL = CodeTemplate("""\ +if (${tensor_name}_impl_saved) AT_ASSERT(${tensor_name}_impl_saved == ${tensor_name}.getIntrusivePtr()); +""") + +SAVE_TENSORLIST_IMPL = CodeTemplate("""\ +std::vector> ${tensorlist_name}_impl_saved(${tensorlist_name}.size()); +for (size_t i=0; i<${tensorlist_name}.size(); i++) + if (${tensorlist_name}[i].defined()) ${tensorlist_name}_impl_saved[i] = ${tensorlist_name}[i].getIntrusivePtr(); +""") + +ENFORCE_SAME_TENSORLIST_IMPL = CodeTemplate("""\ +for (size_t i=0; i<${tensorlist_name}.size(); i++) { + if (${tensorlist_name}_impl_saved[i]) + AT_ASSERT(${tensorlist_name}_impl_saved[i] == ${tensorlist_name}[i].getIntrusivePtr()); +} +""") + +# The following list contains functions that we don't enforce the invariant on. +DONT_ENFORCE_SAME_TENSOR_IMPL_OR_STORAGE = { + # These functions are expected to change impl or storage of input tensors + '_th_set_', '_cudnn_rnn_flatten_weight', + # TODO: Fix these functions to update input tensor in-place + 'tril_', 'triu_', +} +# END CHECKS FOR [ Invariant: TensorImpl and Storage Pointer Equality ] + METHOD_DECLARATION = CodeTemplate("""\ ${return_type} ${method_prefix_derived}${api_name}(${type_method_formals}) const override; """) @@ -189,6 +250,12 @@ if (tracer_state) { } """) +RUN_ONLY_IN_DEBUG_MODE = CodeTemplate("""\ +#ifndef NDEBUG +${statements} +#endif +""") + FACTORY_FUNCTION_NAMES = None @@ -608,6 +675,29 @@ def emit_body(declaration): else: return 'as_variable({})'.format(call), [] + def enforce_same_tensorimpl_and_storage(env, call): + save_ptrs_stmts = [] + enforce_same_ptrs_stmts = [] + if declaration['name'] not in DONT_ENFORCE_SAME_TENSOR_IMPL_OR_STORAGE: + for arg in env.get('unpacked_args', []): + simple_type = env['unpacked_args_simple_type'][arg] + if simple_type == 'TensorList': + save_ptrs_stmts += [SAVE_TENSORLIST_STORAGE.substitute(tensorlist_name=arg), + SAVE_TENSORLIST_IMPL.substitute(tensorlist_name=arg)] + enforce_same_ptrs_stmts += [ENFORCE_SAME_TENSORLIST_STORAGE.substitute(tensorlist_name=arg), + ENFORCE_SAME_TENSORLIST_IMPL.substitute(tensorlist_name=arg)] + elif simple_type == 'Tensor': + save_ptrs_stmts += [SAVE_TENSOR_STORAGE.substitute(tensor_name=arg), + SAVE_TENSOR_IMPL.substitute(tensor_name=arg)] + enforce_same_ptrs_stmts += [ENFORCE_SAME_TENSOR_STORAGE.substitute(tensor_name=arg), + ENFORCE_SAME_TENSOR_IMPL.substitute(tensor_name=arg)] + assert (save_ptrs_stmts and enforce_same_ptrs_stmts) or (not save_ptrs_stmts and not enforce_same_ptrs_stmts) + if save_ptrs_stmts and enforce_same_ptrs_stmts: + call = RUN_ONLY_IN_DEBUG_MODE.substitute(statements=save_ptrs_stmts) + \ + call + \ + RUN_ONLY_IN_DEBUG_MODE.substitute(statements=enforce_same_ptrs_stmts) + return call + def emit_call(env): combined = nested_dict(env, declaration) extra_wrapping_stmts = [] @@ -634,6 +724,7 @@ def emit_body(declaration): call = call + ';' for stmt in extra_wrapping_stmts: call += '\n' + stmt + call = enforce_same_tensorimpl_and_storage(env, call) return call def tie_return_values(): @@ -725,9 +816,11 @@ def unpack_args(env, declaration): body = [] unpacked_args = [] + unpacked_args_simple_type = {} for i, arg in enumerate(declaration['arguments']): if not requires_unpack(arg): unpacked_args.append(arg['name']) + unpacked_args_simple_type[arg['name']] = arg['simple_type'] continue dynamic_type = arg['dynamic_type'] @@ -749,8 +842,10 @@ def unpack_args(env, declaration): body.append(UNPACK_OPTIONS.substitute(arg_name=arg['name'])) unpacked_args.append(arg['name'] + '_') + unpacked_args_simple_type[arg['name'] + '_'] = arg['simple_type'] env['unpacked_args'] = unpacked_args + env['unpacked_args_simple_type'] = unpacked_args_simple_type return body diff --git a/torch/csrc/autograd/variable.cpp b/torch/csrc/autograd/variable.cpp index 3da736dadc3..c79e7650714 100644 --- a/torch/csrc/autograd/variable.cpp +++ b/torch/csrc/autograd/variable.cpp @@ -92,6 +92,10 @@ void* Variable::Impl::slow_data() const { return data_.unsafeGetTensorImpl()->slow_data(); } +bool Variable::Impl::has_storage() const { + return data_.has_storage(); +} + const at::Storage& Variable::Impl::storage() const { return data_.storage(); } diff --git a/torch/csrc/autograd/variable.h b/torch/csrc/autograd/variable.h index ab632c4ca4d..7eea2b04246 100644 --- a/torch/csrc/autograd/variable.h +++ b/torch/csrc/autograd/variable.h @@ -407,6 +407,7 @@ struct TORCH_API Variable::Impl : public at::TensorImpl { void set_storage_offset(int64_t storage_offset) override; int64_t dim() const override; + bool has_storage() const override; const at::Storage& storage() const override; void* slow_data() const override;