AutoDispatchBelowAutograd takes no arguments. (#56424)

Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/56424

Test Plan: Imported from OSS

Reviewed By: nikithamalgifb

Differential Revision: D27866607

Pulled By: ailzhang

fbshipit-source-id: b82cfb90af5bc7b4129266083fe31f8b335a5b41

aten/src/ATen/ScalarOps.cpp

@@ -31,7 +31,7 @@ Tensor& scalar_fill(Tensor& self, const Scalar& value) {
 
 Tensor scalar_tensor_static(const Scalar& s, c10::optional<ScalarType> dtype_opt, c10::optional<Device> device_opt) {
   at::tracer::impl::NoTracerDispatchMode tracer_guard;
-  at::AutoDispatchBelowAutograd non_var_type_mode(true);
+  at::AutoDispatchBelowAutograd mode;
   auto result = at::detail::empty_cpu({}, dtype_opt, c10::nullopt, device_opt, c10::nullopt, c10::nullopt);
   scalar_fill(result, s);
   return result;

aten/src/ATen/core/LegacyTypeDispatch.h

@@ -59,12 +59,8 @@ namespace at {
  * error saying "one of the variables needed for gradient computation has be modified."
  */
 struct TORCH_API AutoDispatchBelowAutograd {
-  // NB: The enabled parameter must ALWAYS be black, as Henry Ford used to say.
-  // TODO: Eliminate this parameter entirely
-  AutoDispatchBelowAutograd(bool enabled = true) :
+  AutoDispatchBelowAutograd() :
     autograd_guard_(c10::autograd_dispatch_keyset) {
-
-    TORCH_INTERNAL_ASSERT(enabled);
   }
 
   // disable all autograd dispatch keys

aten/src/ATen/core/boxing/impl/kernel_function_legacy_test.cpp

@@ -45,7 +45,7 @@ int64_t decrementKernel(const Tensor& tensor, int64_t input) {
 }
 
 void expectCallsIncrement(DispatchKey dispatch_key) {
-  at::AutoDispatchBelowAutograd non_var_type_mode(true);
+  at::AutoDispatchBelowAutograd mode;
 
   // assert that schema and cpu kernel are present
   auto op = c10::Dispatcher::singleton().findSchema({"_test::my_op", ""});
@@ -56,7 +56,7 @@ void expectCallsIncrement(DispatchKey dispatch_key) {
 }
 
 void expectCallsDecrement(DispatchKey dispatch_key) {
-  at::AutoDispatchBelowAutograd non_var_type_mode(true);
+  at::AutoDispatchBelowAutograd mode;
 
   // assert that schema and cpu kernel are present
   auto op = c10::Dispatcher::singleton().findSchema({"_test::my_op", ""});
@@ -911,7 +911,7 @@ std::string concatKernel(const Tensor& tensor1, std::string a, const std::string
 }
 
 void expectCallsConcatUnboxed(DispatchKey dispatch_key) {
-  at::AutoDispatchBelowAutograd non_var_type_mode(true);
+  at::AutoDispatchBelowAutograd mode;
 
   // assert that schema and cpu kernel are present
   auto op = c10::Dispatcher::singleton().findSchema({"_test::my_op", ""});

aten/src/ATen/core/boxing/impl/kernel_function_test.cpp

@@ -34,7 +34,7 @@ int64_t decrementKernel(const Tensor& tensor, int64_t input) {
 }
 
 void expectCallsIncrement(DispatchKey dispatch_key) {
-  at::AutoDispatchBelowAutograd non_var_type_mode(true);
+  at::AutoDispatchBelowAutograd mode;
 
   // assert that schema and cpu kernel are present
   auto op = c10::Dispatcher::singleton().findSchema({"_test::my_op", ""});
@@ -45,7 +45,7 @@ void expectCallsIncrement(DispatchKey dispatch_key) {
 }
 
 void expectCallsDecrement(DispatchKey dispatch_key) {
-  at::AutoDispatchBelowAutograd non_var_type_mode(true);
+  at::AutoDispatchBelowAutograd mode;
 
   // assert that schema and cpu kernel are present
   auto op = c10::Dispatcher::singleton().findSchema({"_test::my_op", ""});
@@ -653,7 +653,7 @@ std::string concatKernel(const Tensor& tensor1, std::string a, const std::string
 }
 
 void expectCallsConcatUnboxed(DispatchKey dispatch_key) {
-  at::AutoDispatchBelowAutograd non_var_type_mode(true);
+  at::AutoDispatchBelowAutograd mode;
 
   // assert that schema and cpu kernel are present
   auto op = c10::Dispatcher::singleton().findSchema({"_test::my_op", ""});
@@ -663,7 +663,7 @@ void expectCallsConcatUnboxed(DispatchKey dispatch_key) {
 }
 
 void expectCannotCallConcatBoxed(DispatchKey dispatch_key) {
-  at::AutoDispatchBelowAutograd non_var_type_mode(true);
+  at::AutoDispatchBelowAutograd mode;
 
   // assert that schema and cpu kernel are present
   auto op = c10::Dispatcher::singleton().findSchema({"_test::my_op", ""});

aten/src/ATen/core/boxing/impl/kernel_lambda_legacy_test.cpp

@@ -31,7 +31,7 @@ using std::unique_ptr;
 namespace {
 
 void expectCallsIncrement(DispatchKey dispatch_key) {
-  at::AutoDispatchBelowAutograd non_var_type_mode(true);
+  at::AutoDispatchBelowAutograd mode;
 
   // assert that schema and cpu kernel are present
   auto op = c10::Dispatcher::singleton().findSchema({"_test::my_op", ""});
@@ -839,7 +839,7 @@ TEST(OperatorRegistrationTest_LegacyLambdaBasedKernel, givenKernelWithOptionalIn
 }
 
 void expectCallsConcatUnboxed(DispatchKey dispatch_key) {
-  at::AutoDispatchBelowAutograd non_var_type_mode(true);
+  at::AutoDispatchBelowAutograd mode;
 
   // assert that schema and cpu kernel are present
   auto op = c10::Dispatcher::singleton().findSchema({"_test::my_op", ""});

aten/src/ATen/core/boxing/impl/kernel_lambda_test.cpp

@@ -21,7 +21,7 @@ using std::unique_ptr;
 namespace {
 
 void expectCallsIncrement(DispatchKey dispatch_key) {
-  at::AutoDispatchBelowAutograd non_var_type_mode(true);
+  at::AutoDispatchBelowAutograd mode;
 
   // assert that schema and cpu kernel are present
   auto op = c10::Dispatcher::singleton().findSchema({"_test::my_op", ""});
@@ -32,7 +32,7 @@ void expectCallsIncrement(DispatchKey dispatch_key) {
 }
 
 void expectCallsDecrement(DispatchKey dispatch_key) {
-  at::AutoDispatchBelowAutograd non_var_type_mode(true);
+  at::AutoDispatchBelowAutograd mode;
 
   // assert that schema and cpu kernel are present
   auto op = c10::Dispatcher::singleton().findSchema({"_test::my_op", ""});
@@ -564,7 +564,7 @@ TEST(OperatorRegistrationTest_LambdaBasedKernel, givenKernelWithOptionalInputs_w
 }
 
 void expectCallsConcatUnboxed(DispatchKey dispatch_key) {
-  at::AutoDispatchBelowAutograd non_var_type_mode(true);
+  at::AutoDispatchBelowAutograd mode;
 
   // assert that schema and cpu kernel are present
   auto op = c10::Dispatcher::singleton().findSchema({"_test::my_op", ""});

aten/src/ATen/core/boxing/impl/kernel_stackbased_test.cpp

@@ -43,7 +43,7 @@ void redispatchingKernel_with_DispatchKeySet(const OperatorHandle& op, c10::Disp
 }
 
 void expectCallsIncrement(c10::DispatchKeySet ks) {
-  at::AutoDispatchBelowAutograd non_var_type_mode(true);
+  at::AutoDispatchBelowAutograd mode;
 
   // assert that schema and cpu kernel are present
   auto op = c10::Dispatcher::singleton().findSchema({"_test::my_op", ""});
@@ -58,7 +58,7 @@ void expectCallsIncrement(DispatchKey dispatch_key) {
 }
 
 void expectCallsIncrementUnboxed(DispatchKey dispatch_key) {
-  at::AutoDispatchBelowAutograd non_var_type_mode(true);
+  at::AutoDispatchBelowAutograd mode;
 
   // assert that schema and cpu kernel are present
   auto op = c10::Dispatcher::singleton().findSchema({"_test::my_op", ""});
@@ -68,7 +68,7 @@ void expectCallsIncrementUnboxed(DispatchKey dispatch_key) {
 }
 
 void expectCallsDecrement(DispatchKey dispatch_key) {
-  at::AutoDispatchBelowAutograd non_var_type_mode(true);
+  at::AutoDispatchBelowAutograd mode;
 
   // assert that schema and cpu kernel are present
   auto op = c10::Dispatcher::singleton().findSchema({"_test::my_op", ""});

aten/src/ATen/core/boxing/impl/make_boxed_from_unboxed_functor_test.cpp

@@ -41,7 +41,7 @@ struct DecrementKernel final : OperatorKernel {
 };
 
 void expectCallsIncrement(DispatchKey dispatch_key) {
-  at::AutoDispatchBelowAutograd non_var_type_mode(true);
+  at::AutoDispatchBelowAutograd mode;
 
   // assert that schema and cpu kernel are present
   auto op = c10::Dispatcher::singleton().findSchema({"_test::my_op", ""});
@@ -52,7 +52,7 @@ void expectCallsIncrement(DispatchKey dispatch_key) {
 }
 
 void expectCallsDecrement(DispatchKey dispatch_key) {
-  at::AutoDispatchBelowAutograd non_var_type_mode(true);
+  at::AutoDispatchBelowAutograd mode;
 
   // assert that schema and cpu kernel are present
   auto op = c10::Dispatcher::singleton().findSchema({"_test::my_op", ""});
@@ -794,7 +794,7 @@ struct ConcatKernel final : OperatorKernel {
 };
 
 void expectCallsConcatUnboxed(DispatchKey dispatch_key) {
-  at::AutoDispatchBelowAutograd non_var_type_mode(true);
+  at::AutoDispatchBelowAutograd mode;
 
   // assert that schema and cpu kernel are present
   auto op = c10::Dispatcher::singleton().findSchema({"_test::my_op", ""});

aten/src/ATen/native/TensorFactories.cpp

@@ -619,7 +619,7 @@ Tensor scalar_tensor(const Scalar& s,
     // revert this to following:
     //   auto result = at::empty({}, options);
     at::tracer::impl::NoTracerDispatchMode tracer_guard;
-    at::AutoDispatchBelowAutograd non_var_type_mode(true);
+    at::AutoDispatchBelowAutograd mode;
     auto result = empty_cpu({}, optTypeMetaToScalarType(options.dtype_opt()), options.layout_opt(), options.device_opt(), options.pinned_memory_opt());
     at::native::fill_(result, s);
     return result;

aten/src/ATen/quantized/Quantizer.cpp

@@ -31,7 +31,7 @@ namespace {
 // Note: this is not a native function as Quantizer is not exposed to python yet
 QuantizerPtr Tensor::quantizer() const {
   // This is a terrible hack to emulate what VariableType is doing
-  at::AutoDispatchBelowAutograd non_var_type_mode(true);
+  at::AutoDispatchBelowAutograd mode;
   return get_qtensorimpl(*this)->quantizer();
 }
 

test/custom_operator/test_custom_ops.cpp

@@ -81,7 +81,7 @@ void get_autograd_operator_from_registry_and_execute() {
 }
 
 void get_autograd_operator_from_registry_and_execute_in_nograd_mode() {
-  at::AutoDispatchBelowAutograd _var_guard(true);
+  at::AutoDispatchBelowAutograd guard;
 
   torch::Tensor x = torch::randn({5,5}, torch::requires_grad());
   torch::Tensor y = torch::randn({5,5}, torch::requires_grad());

tools/autograd/templates/variable_factories.h

@@ -65,7 +65,7 @@ inline at::Tensor from_blob(
     const Deleter& deleter,
     const at::TensorOptions& options = at::TensorOptions()) {
   at::Tensor tensor = ([&]() {
-    at::AutoDispatchBelowAutograd non_var_type_mode(true);  // TODO: remove
+    at::AutoDispatchBelowAutograd guard;  // TODO: remove
     at::tracer::impl::NoTracerDispatchMode tracer_guard;
     return at::from_blob(data, sizes, strides, deleter, options.requires_grad(c10::nullopt));
   })();
@@ -83,7 +83,7 @@ inline at::Tensor from_blob(
     at::IntArrayRef strides,
     const at::TensorOptions& options = at::TensorOptions()) {
   at::Tensor tensor = ([&]() {
-    at::AutoDispatchBelowAutograd non_var_type_mode(true);  // TODO: remove
+    at::AutoDispatchBelowAutograd guard;  // TODO: remove
     at::tracer::impl::NoTracerDispatchMode tracer_guard;
     return at::from_blob(data, sizes, strides, options.requires_grad(c10::nullopt));
   })();
@@ -102,7 +102,7 @@ inline at::Tensor from_blob(
     const Deleter& deleter,
     const at::TensorOptions& options = at::TensorOptions()) {
   at::Tensor tensor = ([&]() {
-    at::AutoDispatchBelowAutograd non_var_type_mode(true);  // TODO: remove
+    at::AutoDispatchBelowAutograd guard;  // TODO: remove
     at::tracer::impl::NoTracerDispatchMode tracer_guard;
     return at::from_blob(data, sizes, deleter, options.requires_grad(c10::nullopt));
   })();
@@ -118,7 +118,7 @@ inline at::Tensor from_blob(
     at::IntArrayRef sizes,
     const at::TensorOptions& options = at::TensorOptions()) {
   at::Tensor tensor = ([&]() {
-    at::AutoDispatchBelowAutograd non_var_type_mode(true);  // TODO: remove
+    at::AutoDispatchBelowAutograd guard;  // TODO: remove
     at::tracer::impl::NoTracerDispatchMode tracer_guard;
     return at::from_blob(data, sizes, options.requires_grad(c10::nullopt));
   })();

torch/csrc/api/include/torch/detail/TensorDataContainer.h

@@ -137,7 +137,7 @@ AT_FORALL_COMPLEX_TYPES(TENSOR)
       sizes_({(int64_t)values.size()}), \
       scalar_type_(at::k##S), \
       type_(TensorDataContainerType::Tensor) { \
-    at::AutoDispatchBelowAutograd non_var_type_mode(true); \
+    at::AutoDispatchBelowAutograd mode; \
     if (scalar_type_ == at::kBool) { \
       tensor_ = at::tensor(values, at::TensorOptions().device(at::kCPU)); \
     } else { \
@@ -212,7 +212,7 @@ AT_FORALL_COMPLEX_TYPES(TENSOR)
     }
 
     if (is_scalar()) {
-      at::AutoDispatchBelowAutograd non_var_type_mode(true);
+      at::AutoDispatchBelowAutograd mode;
       return at::scalar_tensor(scalar_, options);
     } else if (is_init_list()) {
       // NOTE: Here we explicitly choose to initialize the tensor on CPU first,
@@ -222,7 +222,7 @@ AT_FORALL_COMPLEX_TYPES(TENSOR)
       // filling each element of it (which involves `N` CUDA kernel launches where
       // `N` is the number of the elements in the tensor).
       at::Tensor tensor = ([&]() {
-        at::AutoDispatchBelowAutograd non_var_type_mode(true);
+        at::AutoDispatchBelowAutograd mode;
         return at::empty(sizes_, options.device(at::kCPU));
       })();
       fill_tensor(tensor);

torch/csrc/autograd/VariableTypeManual.cpp

@@ -90,7 +90,7 @@ Tensor _fw_primal(const Tensor & self, int64_t level) {
     grad_fn->set_next_edges(collect_next_edges( self ));
   }
   auto tmp = ([&]() {
-    at::AutoDispatchBelowAutograd non_var_type_mode(true);
+    at::AutoDispatchBelowAutograd mode;
     return self_.alias();
   })();
   std::function<at::Tensor(const at::Tensor&)> func=nullptr;
@@ -131,7 +131,7 @@ Tensor & copy_(c10::DispatchKeySet ks, Tensor & self, const Tensor & src, bool n
     grad_fn->src_device = src.device();
   }
   {
-    at::AutoDispatchBelowAutograd non_var_type_mode(true);
+    at::AutoDispatchBelowAutograd mode;
     at::redispatch::copy_(ks & c10::after_autograd_keyset, self_, src_, non_blocking);
   }
   rebase_history(self , std::move(grad_fn));
@@ -166,7 +166,7 @@ Tensor& resize_(
     AT_ERROR("cannot resize variables that require grad");
   }
   {
-    at::AutoDispatchBelowAutograd non_var_type_mode(true);
+    at::AutoDispatchBelowAutograd mode;
     at::redispatch::resize_(ks & c10::after_autograd_keyset, self_, size, optional_memory_format);
   }
 
@@ -188,7 +188,7 @@ Tensor& resize_as_(
     AT_ERROR("cannot resize variables that require grad");
   }
   {
-    at::AutoDispatchBelowAutograd non_var_type_mode(true);
+    at::AutoDispatchBelowAutograd mode;
     at::redispatch::resize_as_(ks & c10::after_autograd_keyset, self_, the_template_, optional_memory_format);
   }
 

torch/csrc/autograd/functions/comm.cpp

@@ -124,7 +124,7 @@ variable_list Gather::apply(variable_list&& inputs) {
   // so no need for extra logic here
   at::Tensor variable;
   {
-    at::AutoDispatchBelowAutograd non_var_type_mode(true);
+    at::AutoDispatchBelowAutograd mode;
     // This is special logic for torch::cuda::gather!
     const auto destination_index =
         destination_device_.is_cpu() ? -1 : destination_device_.index();

torch/csrc/jit/passes/frozen_ops_to_mkldnn.cpp

@@ -441,7 +441,7 @@ jit::RegisterOperators reg_fut_ops({
           Tensor weight = pop(stack).toTensor();
           Tensor input = pop(stack).toTensor();
 
-          at::AutoDispatchBelowAutograd non_var_type_mode(true);
+          at::AutoDispatchBelowAutograd mode;
           // aten::convolution takes care of 0 dim case before calls into
           // backends
           if (input.size(0) == 0) {

torch/lib/c10d/ProcessGroupGloo.cpp

@@ -1057,7 +1057,7 @@ class AsyncSparseAllreduceWork : public ProcessGroupGloo::AsyncWork {
     //
     // The correct fix is to stop allocating tensors that are not variables,
     // but to conveniently do this c10d must depend on torch not ATen
-    at::AutoDispatchBelowAutograd _no_grad(true);
+    at::AutoDispatchBelowAutograd guard;
     auto input = tensors[0];
 
     // Perform local reduction if we have multiple inputs.