s/AutoDispatchBelowAutograd/AutoDispatchBelowInplaceOrView. (#56657)

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

Test Plan: Imported from OSS

Reviewed By: albanD

Differential Revision: D27931526

Pulled By: ailzhang

fbshipit-source-id: 3af718df3435e2b0b30bc62070dbdc5aeeecdfb4

aten/src/ATen/TensorIndexing.cpp

@@ -45,7 +45,7 @@ static inline void set_item(const Tensor& self, ArrayRef<TensorIndex> indices, c
   Tensor value;
 
   {
-    at::AutoDispatchBelowAutograd guard;
+    at::AutoDispatchBelowInplaceOrView guard;
     // TODO: This qint special case looks very suspicious...
     if (isQIntType(self.scalar_type())) {
       value = at::indexing::scalarToTensor(v, device(kCPU).dtype(kFloat), at::Device(kCPU));

aten/src/ATen/TracerMode.h

@@ -8,12 +8,12 @@
 //
 // Historically, tracing function was controlled by two switches:
 //
-// - `AutoDispatchBelowAutograd` guard
+// - `AutoDispatchBelowInplaceOrView` guard
 //
 //    Tracing function used to be script-generated inside `VariableType_*.cpp`
 //    kernels, sharing the same `Autograd` dispatch key with autograd function.
 //    Therefore, before tracing function was moved out of VariableType,
-//    `AutoDispatchBelowAutograd` guard can also disable tracing as a side effect
+//    `AutoDispatchBelowInplaceOrView` guard can also disable tracing as a side effect
 //    of disabling `Autograd` dispatching.
 //
 // - `setTracingState()` API in `torch/csrc/jit/frontend/tracer.h`
@@ -42,7 +42,7 @@
 //
 // - `tracer::impl::NoTracerDispatchMode` guard
 //
-//    It's used to cover the old semantics of `AutoDispatchBelowAutograd` after
+//    It's used to cover the old semantics of `AutoDispatchBelowInplaceOrView` after
 //    tracing was moved out of VariableType.
 //
 // Before tracing function was moved out of VariableType, tracing was enabled
@@ -51,7 +51,7 @@
 //    1) `TracingState` object in TLS != null;
 //       - Either inside the execution scope of `tracer::trace()`, or
 //       - Eagerly called `setTracingState()` with non-null object.
-//    2) Not inside `AutoDispatchBelowAutograd` scope;
+//    2) Not inside `AutoDispatchBelowInplaceOrView` scope;
 //
 // After:
 //
@@ -69,10 +69,10 @@
 //   `setTracingState()` Python/C++ APIs (and other APIs calling it) so that
 //   these two can be unified.
 //
-// - `AutoDispatchBelowAutograd` v.s. `tracer::impl::NoTracerDispatchMode`
+// - `AutoDispatchBelowInplaceOrView` v.s. `tracer::impl::NoTracerDispatchMode`
 //
 //   We don't need to always set both guards together to keep semantics
-//   unchanged. For the follow use cases of `AutoDispatchBelowAutograd` we don't
+//   unchanged. For the follow use cases of `AutoDispatchBelowInplaceOrView` we don't
 //   need set the new tracer guard:
 //
 //   * Script-generated VariableType kernels. The guard is not necessary as
@@ -99,7 +99,7 @@
 //   * Some manually maintained functions, e.g.:
 //     `torch/csrc/autograd/VariableTypeManual.cpp`.
 //     Set the new guard if it's not obvious whether `setTracingState(null)`
-//     has been called before it reaches the `AutoDispatchBelowAutograd` guard.
+//     has been called before it reaches the `AutoDispatchBelowInplaceOrView` guard.
 //
 //   We might need tweak the usage of the new guard to optimize/fix things.
 //   It should only affect the correctness of tracing function, because the

aten/src/ATen/core/LegacyTypeDispatch.h

@@ -72,9 +72,9 @@ struct TORCH_API AutoNonVariableTypeMode {
   AutoNonVariableTypeMode(bool enabled = true) :
     autograd_guard_(c10::autograd_dispatch_keyset) {
     TORCH_WARN_ONCE("AutoNonVariableTypeMode is deprecated and will be removed in 1.10 release. "
-        "For kernel implementations please use AutoDispatchBelowAutograd instead, "
+        "For kernel implementations please use AutoDispatchBelowInplaceOrView instead, "
         "If you are looking for a user facing API to enable running your inference-only "
-        "workload, please use c10::InferenceMode. Using AutoDispatchBelowAutogradMode in user code "
+        "workload, please use c10::InferenceMode. Using AutoDispatchBelowInplaceOrView in user code "
         "is under risk of producing silent wrong result in some edge cases. "
         "See Note [AutoDispatchBelowAutograd] for more details.");
     TORCH_INTERNAL_ASSERT(enabled);
@@ -84,12 +84,16 @@ struct TORCH_API AutoNonVariableTypeMode {
   c10::impl::ExcludeDispatchKeyGuard autograd_guard_;
 };
 
-// Note this guard is used in VariableType kernels for functional ops
-// as well as InplaceOrView kernels for inplace/view ops to enforce the
-// invariant:
-//   Once you are in VariableType/InplaceOrView kernel for an op,
-//   you never go back to a kernel on same dispatch key until
-//   you finish the current op.
+/* Note [AutoDispatchBelowInplaceOrView]
+ * AutoDispatchBelowInplaceOrView is equivalent to AutoNonVariableTypeMode
+ * before we split inplace & view ops out of VariableType kernel.
+ * Note this guard is used in VariableType kernels for functional ops
+ * as well as InplaceOrView kernels for inplace/view ops to enforce the
+ * Invariant:
+ *   Once you are in VariableType/InplaceOrView kernel for an op,
+ *   you never go back to a kernel on same dispatch key until
+ *   you finish the current op.
+ */
 struct TORCH_API AutoDispatchBelowInplaceOrView {
   AutoDispatchBelowInplaceOrView() :
     dispatch_key_guard_(c10::autograd_dispatch_keyset_with_InplaceOrView) {

aten/src/ATen/cpp_custom_type_hack.h

@@ -82,7 +82,7 @@ template <typename T>
   "https://pytorch.org/tutorials/advanced/torch_script_custom_classes.html")]]
 Tensor create(std::unique_ptr<T> ptr, TensorOptions options) {
   // None of this should trace, so turn off Tracer dispatching
-  at::AutoDispatchBelowAutograd guard;  // TODO: remove
+  at::AutoDispatchBelowInplaceOrView guard;  // TODO: remove
   at::tracer::impl::NoTracerDispatchMode tracer_guard;
 
   // We store this instance away in a Tensor and register a deleter function

aten/src/ATen/native/ConvolutionMM2d.cpp

@@ -270,7 +270,7 @@ void slow_conv2d_backward_out_cpu_template(
   const int64_t batch_size = input.size(0);
   at::parallel_for(0, batch_size, 0, [&](int64_t start, int64_t end) {
     NoGradGuard no_grad;
-    AutoDispatchBelowAutograd non_variable_type_mode;
+    AutoDispatchBelowInplaceOrView non_variable_type_mode;
     for (int64_t t = start; t < end; t++) {
       Tensor grad_input_t = grad_input[t];
       Tensor grad_output_t = grad_output[t];
@@ -449,7 +449,7 @@ std::tuple<Tensor&, Tensor&, Tensor&> slow_conv2d_forward_out_cpu(const Tensor&
 
   at::parallel_for(0, batch_size, 0, [&](int64_t start, int64_t end) {
     NoGradGuard no_grad;
-    AutoDispatchBelowAutograd non_variable_type_mode;
+    AutoDispatchBelowInplaceOrView non_variable_type_mode;
     for (int64_t t = start; t < end; t++) {
       Tensor input_t = input[t].unsqueeze(0);
       Tensor output_t = output[t];

aten/src/ATen/native/ConvolutionMM3d.cpp

@@ -370,7 +370,7 @@ void slow_conv3d_backward_out_cpu_template(
   const int64_t batch_size = input.size(0);
   at::parallel_for(
       0, batch_size, CONV3D_GRAIN_SALT, [&](int64_t start, int64_t end) {
-        AutoDispatchBelowAutograd non_variable_type_mode;
+        AutoDispatchBelowInplaceOrView non_variable_type_mode;
         for (int64_t t = start; t < end; t++) {
           Tensor grad_input_t = grad_input[t];
           Tensor grad_output_t = grad_output_contiguous[t];
@@ -597,7 +597,7 @@ std::tuple<Tensor&, Tensor&, Tensor&> slow_conv3d_forward_out_cpu(const Tensor&
 
   at::parallel_for(
       0, batch_size, CONV3D_GRAIN_SALT, [&](int64_t start, int64_t end) {
-        AutoDispatchBelowAutograd non_variable_type_mode;
+        AutoDispatchBelowInplaceOrView non_variable_type_mode;
         for (int64_t t = start; t < end; t++) {
           Tensor input_t = input[t];
           Tensor output_t = output[t];

aten/src/ATen/templates/Functions.cpp

@@ -71,7 +71,7 @@ void noopDelete(void*) {}
 } // namespace detail
 
 Tensor TensorMaker::make_tensor() {
-  AutoDispatchBelowAutograd guard{}; // TODO: Remove.
+  AutoDispatchBelowInplaceOrView guard{}; // TODO: Remove.
   tracer::impl::NoTracerDispatchMode tracer_guard{};
 
   check_size_nonnegative(sizes_);

c10/core/DispatchKey.h

@@ -172,7 +172,7 @@ enum class DispatchKey : uint8_t {
   //     // inplace/view ops called through `at::` inside your backend
   //     // kernel will dispatch to InplaceOrView kernels and do a lot
   //     // of extra work.
-  //     at::AutoDispatchBelowInplaceOrView guard(true);
+  //     at::AutoDispatchBelowInplaceOrView guard;
   //     at::redispatch::my_functional_op(...);
   //   }
   // }

torch/csrc/autograd/python_variable_indexing.cpp

@@ -92,7 +92,7 @@ static inline Variable valueToTensor(c10::TensorOptions options, PyObject* value
   if (THPVariable_Check(value)) {
     return THPVariable_Unpack(value);
   }
-  at::AutoDispatchBelowAutograd guard;  // TODO: remove
+  at::AutoDispatchBelowInplaceOrView guard;  // TODO: remove
   at::tracer::impl::NoTracerDispatchMode tracer_guard;
   if (THPUtils_checkLong(value) || PyBool_Check(value)) {
     return at::indexing::scalarToTensor(Scalar(THPUtils_unpackLong(value)), options, device);

torch/csrc/jit/codegen/cuda/executor.cpp

@@ -186,7 +186,7 @@ at::Tensor inferAndAlloc(
   } else {
     c10::IntArrayRef isizes(sizes);
     // Non Variable type guard for empty_cuda call
-    at::AutoDispatchBelowAutograd non_variable_type_mode;
+    at::AutoDispatchBelowInplaceOrView non_variable_type_mode;
     return at::native::empty_cuda(
         isizes, at_type, c10::nullopt, options.device, c10::nullopt);
   }
@@ -410,7 +410,7 @@ std::vector<at::Tensor> FusionExecutor::runFusion(
   if (executor_entry && executor_entry->init) {
     {
       // context manager to disable auto grad for `empty_cuda` calls later;
-      at::AutoDispatchBelowAutograd non_variable_type_mode;
+      at::AutoDispatchBelowInplaceOrView non_variable_type_mode;
       // take the short-cut for launch if we see a recorded input set again;
       launch_params = executor_entry->launch_params;
       for (size_t i = 0; i < executor_entry->output_sizes.size(); i++) {

torch/csrc/jit/frontend/tracer.cpp

@@ -868,7 +868,7 @@ autograd::Variable getSizeOf(const autograd::Variable& var, int64_t dim) {
   Variable size_var;
   {
     // Make sure this scalar to tensor isn't traced!
-    at::AutoDispatchBelowAutograd guard;
+    at::AutoDispatchBelowInplaceOrView guard;
     size_var = scalar_to_tensor(at::Scalar(var.size(dim)));
   }
   auto* value = getValueTrace(var);

torch/csrc/utils/python_arg_parser.cpp

@@ -1092,7 +1092,7 @@ at::Tensor PythonArgs::tensor_slow(int i) {
     throw TypeError("expected Tensor as argument %d, but got %s", i,
         Py_TYPE(obj)->tp_name);
   }
-  at::AutoDispatchBelowAutograd guard;  // TODO: remove
+  at::AutoDispatchBelowInplaceOrView guard;  // TODO: remove
   at::tracer::impl::NoTracerDispatchMode tracer_guard;
 
   at::Tensor tensor = scalar_to_tensor(scalar);

torch/csrc/utils/tensor_new.cpp

@@ -256,7 +256,7 @@ Tensor internal_new_from_data(
   // here.
   Tensor tensor;
   {
-    at::AutoDispatchBelowAutograd guard;  // TODO: remove
+    at::AutoDispatchBelowInplaceOrView guard;  // TODO: remove
     at::tracer::impl::NoTracerDispatchMode tracer_guard;
     tensor = at::empty(sizes, at::initialTensorOptions().dtype(inferred_scalar_type).pinned_memory(pin_memory));
     recursive_store(