Use view name instead of view_copy name for functional inverses (pytorch#117056)

Ex: `unsqueeze_copy_inverse()` -> `unsqueeze_inverse()`
Pull Request resolved: pytorch#117056
Approved by: https://github.com/bdhirsh

Author: jbschlosser

aten/src/ATen/FunctionalInverses.cpp

@@ -115,10 +115,10 @@ Tensor _test_check_tensor(const Tensor& self) {
 
 namespace at::functionalization {
 
-// view_copy ops must have a functional inverse registered
-Tensor FunctionalInverses::_test_autograd_multiple_dispatch_view_copy_inverse(const at::Tensor& base, const at::Tensor& mutated_view, InverseReturnMode inverse_return_mode) {
+// view ops must have a functional inverse registered
+Tensor FunctionalInverses::_test_autograd_multiple_dispatch_view_inverse(const at::Tensor& base, const at::Tensor& mutated_view, InverseReturnMode inverse_return_mode) {
     TORCH_INTERNAL_ASSERT(false,
-    "Attempted to call _test_autograd_multiple_dispatch_view_copy_inverse() during the functionalization pass. ",
+    "Attempted to call _test_autograd_multiple_dispatch_view_inverse() during the functionalization pass. ",
     "This function is for testing only and should never be called.");
     return Tensor();
 }

aten/src/ATen/native/TestOps.cpp

@@ -25,8 +25,8 @@ struct FunctionalInverses {
 // NB: These are not generated! They're manually implemented in the template.
 // TODO: Change codegen to generate these. See the following link:
 // https://github.com/pytorch/pytorch/blob/main/torchgen/model.py#L2583-L2585
-static at::Tensor chunk_copy_inverse(const at::Tensor & base, const at::Tensor & mutated_view, InverseReturnMode inverse_return_mode, int64_t mutated_view_idx, int chunks, int dim);
-static at::Tensor narrow_copy_inverse(const at::Tensor & base, const at::Tensor & mutated_view, InverseReturnMode inverse_return_mode, int dim, c10::SymInt start, c10::SymInt length);
+static at::Tensor chunk_inverse(const at::Tensor & base, const at::Tensor & mutated_view, InverseReturnMode inverse_return_mode, int64_t mutated_view_idx, int chunks, int dim);
+static at::Tensor narrow_inverse(const at::Tensor & base, const at::Tensor & mutated_view, InverseReturnMode inverse_return_mode, int dim, c10::SymInt start, c10::SymInt length);
 
 };
 }

aten/src/ATen/templates/FunctionalInverses.h

@@ -158,7 +158,7 @@
 
 REVERSE_VIEW_DISPATCH = CodeTemplate(
     """\
-at::functionalization::FunctionalInverses::${reverse_name}(${unpacked_args})"""
+${reverse_name}(${unpacked_args})"""
 )
 
 MULTI_OUTPUT_VIEW_ITERATION = CodeTemplate(
@@ -429,8 +429,10 @@ def emit_view_lambda(
         *updated_args[1:],
     ]
 
+    from torchgen.api.functionalization import reverse_name
+
     reverse_replay_view_call = REVERSE_VIEW_DISPATCH.substitute(
-        reverse_name=inverse_view_name(f),
+        reverse_name=reverse_name(f, include_namespace=True),
         unpacked_args=reverse_unpacked_args,
     )
     reverse_replay_view_func = REVERSE_REPLAY_VIEW_LAMBDA_FUNC.substitute(

tools/autograd/gen_inplace_or_view_type.py

@@ -17,6 +17,7 @@
     BaseTy,
     BaseType,
     FunctionSchema,
+    NativeFunction,
     NativeFunctionsViewGroup,
 )
 
@@ -99,16 +100,7 @@ def name(
         # since we always plumb the runtime "reapply_views" argument into the reverse function.
         assert is_reverse
     if is_reverse:
-        # for the reverse: the name of the inverse function always involves "view_copy",
-        # and we plumb the "reapply_views" flag into that function.
-        # (We could avoid doing that, but that would require writing out twice as many view inverse functions).
-        assert g.view_copy is not None
-        api_name = g.view_copy.func.name.unambiguous_name()
-        # in the reverse case, we codegen both the call-sites (which need the full namespace) and the declarations (which don't)
-        if include_namespace:
-            return f"at::functionalization::FunctionalInverses::{api_name}_inverse"
-        else:
-            return f"{api_name}_inverse"
+        return reverse_name(g.view, include_namespace)
     # in the forward case, we just directly call into the at::_ops API (so we always need the namespace)
     assert include_namespace
     assert g.view_copy is not None
@@ -120,6 +112,18 @@ def name(
     return f"at::_ops::{api_name}::call"
 
 
+def reverse_name(f: NativeFunction, include_namespace: bool) -> str:
+    # for the reverse: we plumb the "reapply_views" flag into that function and support
+    # both copy and non-copy variants. (We could avoid doing that, but that would require
+    # writing out twice as many view inverse functions).
+    api_name = f.func.name.unambiguous_name()
+    # in the reverse case, we codegen both the call-sites (which need the full namespace) and the declarations (which don't)
+    if include_namespace:
+        return f"at::functionalization::FunctionalInverses::{api_name}_inverse"
+    else:
+        return f"{api_name}_inverse"
+
+
 def capture_arguments(func: FunctionSchema, *, is_reverse: bool) -> List[Binding]:
     # capture arguments include all arguments except `self`.
     # Importantly, they don't include any C++ reference types (or else we'll get a dangling reference in the capture),

torchgen/api/functionalization.py

@@ -289,16 +289,14 @@ class ViewInverseSignature:
     g: NativeFunctionsViewGroup
 
     def name(self) -> str:
-        assert self.g.view_copy is not None
-        return functionalization.name(self.g, is_reverse=True, include_namespace=False)
+        return functionalization.reverse_name(self.g.view, include_namespace=False)
 
     def decl(self) -> str:
-        assert self.g.view_copy is not None
-        return_type = functionalization.returns_type(self.g.view_copy.func)
+        return_type = functionalization.returns_type(self.g.view.func)
         decls = [
             a.decl()
             for a in functionalization.inner_arguments(
-                self.g.view_copy.func, is_reverse=True
+                self.g.view.func, is_reverse=True
             )
         ]
         return f"static {return_type.cpp_type()} {self.name()}({', '.join(decls)});"

torchgen/api/types/signatures.py

@@ -686,8 +686,8 @@ def gen_functionalization_view_inverse_declaration(
     def emit_decl_helper(g: NativeFunctionsViewGroup) -> Optional[str]:
         if g.view.has_composite_implicit_autograd_kernel:
             return None
-        view_copy_inverse_sig = ViewInverseSignature(g)
-        return view_copy_inverse_sig.decl()
+        view_inverse_sig = ViewInverseSignature(g)
+        return view_inverse_sig.decl()
 
     return emit_decl_helper(g)