Update

ezyang · ezyang · commit a0038dc37ce3 · 2025-07-21T11:32:27.000-07:00
[ghstack-poisoned]
diff --git a/torch/_functorch/_aot_autograd/collect_metadata_analysis.py b/torch/_functorch/_aot_autograd/collect_metadata_analysis.py
@@ -31,8 +31,8 @@
 from .descriptors import (
     AOTInput,
     AOTOutput,
-    InputMutationTangentAOTInput,
-    IntermediateBaseTangentAOTInput,
+    InputMutationAOTOutput,
+    IntermediateBaseAOTOutput,
     PlainAOTOutput,
     TangentAOTInput,
 )
@@ -602,8 +602,13 @@ def inner(*flat_args):
                                 new_out_idx
                             )
                             intermediate_bases.append(o._base)
+                            # NB: The desc we picked here is guaranteed to be
+                            # synchronized with the one in
+                            # graph_capture_wrappers.py because we
+                            # SPECIFICALLY notated this output as
+                            # alias_of_intermediate_save_as_output
                             intermediate_bases_descs.append(
-                                IntermediateBaseTangentAOTInput(desc)
+                                TangentAOTInput(IntermediateBaseAOTOutput(desc))
                             )
             elif (
                 # See https://github.com/pytorch/pytorch/issues/100348 for this case.
@@ -741,7 +746,7 @@ def _is_subclass_mutated_input_tangent_always_subclass(inp):
                     and not _is_subclass_mutated_input_tangent_always_subclass(inp)
                     else inp
                 ),
-                InputMutationTangentAOTInput(inp_desc),
+                TangentAOTInput(InputMutationAOTOutput(inp_desc)),
             )
             for inp, inp_desc, info in zip(flat_f_args, flat_f_args_descs, input_info)
             if info.mutation_type == MutationType.MUTATED_OUT_GRAPH
diff --git a/torch/_functorch/_aot_autograd/descriptors.py b/torch/_functorch/_aot_autograd/descriptors.py
@@ -82,7 +82,15 @@ def wrapped_graph(*args):
 
 The full set of descriptors that are possible to annotate a graph can be
 understood by looking at the overall ordering of transformations in
-AOTAutograd.  At a high level, AOTAutograd is structured as a series of
+AOTAutograd.  First, we present a highly stylized account which is semantically
+equivalent to what AOTAutograd does, but doesn't follow the implementation.  Then,
+we describe how the implementation works.
+
+
+
+
+
+At a high level, AOTAutograd is structured as a series of
 wrappers on the original user function, which are composed together to form
 the final function to trace.  As a result of this, AOTAutograd ends up first
 building the full AOTInputs for a function to be traced, and then in reverse
@@ -103,18 +111,18 @@ def wrapped_graph(*args):
 
 Precompile dedupe           (remove dupes)                      (nothing)
 
-Precompile synthetic base   SyntheticBaseAOTInput               AliasedArgWithMetadataMutation...
+Precompile synthetic base   SyntheticBaseAOTInput               MetadataMutationAOTOutput
                             ViewBaseAOTInput
 
 Forward metadata trace      PlainAOTOutput                      (n/a)
-                            AliasedArgWithMetadataMutation...
+                            MetadataMutationAOTOutput
 
 Prepare for autograd        (nothing)                           InputMutationAOTOutput
                                                                 IntermediateBaseAOTOutput
 
 Create joint                TangentAOTInput                     GradAOTOutput
-                            IntermediateBaseTangentAOTInput
-                            InputMutationTangentAOTInput
+                            above w/ InputMutationAOTOutput
+                            above w/ IntermediateBaseAOTOutput
 
 Precompile subclass         SubclassGetAttrAOTInput et al.      SubclassGetAttrAOTOutput et al.
 
@@ -132,9 +140,10 @@ def wrapped_graph(*args):
   - [IN] Precompile synthetic base: SyntheticBaseAOTInput, ViewBaseAOTInput
   - Forward metadata trace (mini output desc propagation)
     - [OUT] Original output convention: PlainAOTOutput
-    - [OUT] Precompile synthetic base: AliasedArgWithMetadataMutationAOTOutput
+    - [OUT] Precompile synthetic base: MetadataMutationAOTOutput
   - [IN] Prepare for autograd: (nothing)
-  - [IN] Create joint: TangentAOTInput, IntermediateBaseTangentAOTInput, InputMutationTangentAOTInput
+  - [IN] Create joint: TangentAOTInput (potentially w/
+         IntermediateBaseAOTOutput, InputMutationAOTOutput)
   - [IN] Precompile subclass: SubclassGetAttrAOTInput et al.
   - [IN] Effect tokens: ForwardTokenAOTInput, BackwardTokenAOTInput
 - Trigger a trace with the modified inputs on the wrapper
@@ -144,7 +153,7 @@ def wrapped_graph(*args):
   - [OUT] Precompile subclass: SubclassGetAttrAOTOutput et al.
   - [OUT] Create joint: GradAOTOutput
   - [OUT] Prepare for autograd: InputMutationAOTOutput, IntermediateBaseAOTOutput
-  - [OUT] Precompile synthetic base: AliasedArgWithMetadataMutationAOTOutput
+  - [OUT] Precompile synthetic base: MetadataMutationAOTOutput
   - [OUT] Precompile dedupe: (nothing)
 
 The important thing to know about this list is that descriptors can only be
@@ -188,8 +197,6 @@ def wrapped_graph(*args):
   - Less important, can be eliminated by cloning inputs of graph:
     - ViewBaseAOTInput
     - SyntheticBaseAOTInput
-    - InputMutationTangentAOTInput
-    - IntermediateBaseTangentAOTInput
   - Non-tensor, mostly just ignore them:
     - DummyAOTInput
     - PhiloxForwardSeedAOTInput
@@ -206,7 +213,7 @@ def wrapped_graph(*args):
     - InputMutationAOTOutput
   - Less important, can be eliminated by cloning outputs of graph:
     - IntermediateBaseAOTOutput
-    - AliasedArgWithMetadataMutationAOTOutput
+    - MetadataMutationAOTOutput
   - Non-tensor, mostly just ignore them:
     - PhiloxUpdatedForwardOffsetAOTOutput
     - PhiloxUpdatedBackwardOffsetAOTOutput
@@ -393,19 +400,6 @@ def expr(self) -> str:
         return "__backward_token"
 
 
-@dataclasses.dataclass(frozen=True)
-class InputMutationTangentAOTInput(AOTInput):
-    """An input to the joint graph representing the tangent of a mutated input
-    (which is therefore implicitly an output)"""
-
-    base: AOTInput
-
-    def expr(self) -> str:
-        # these are not "real" functions, in that they can't actually be executed;
-        # they truly are new inputs, but they're parameterized by other sources
-        return f"__input_mutation_tangent({self.base.expr()})"
-
-
 # Technically the "output" here is redundant, tangents always correspond to
 # outputs
 @dataclasses.dataclass(frozen=True)
@@ -418,18 +412,6 @@ def expr(self) -> str:
         return f"__output_tangent({self.output.expr()})"
 
 
-@dataclasses.dataclass(frozen=True)
-class IntermediateBaseTangentAOTInput(AOTInput):
-    """An input to the joint graph representing the tangent of an
-    'intermediate base' output that was created to deal with aliasing
-    outputs."""
-
-    output: "AOTOutput"
-
-    def expr(self) -> str:
-        return f"__output_intermediate_base_tangent({self.output.expr()})"
-
-
 # ------------
 
 # AOTOutput
@@ -469,7 +451,7 @@ def expr(self) -> str:
 
 
 @dataclasses.dataclass(frozen=True)
-class AliasedArgWithMetadataMutationAOTOutput(AOTOutput):
+class MetadataMutationAOTOutput(AOTOutput):
     # TODO: we're not recording detailed information about this
 
     def expr(self) -> str:
diff --git a/torch/_functorch/_aot_autograd/input_output_analysis.py b/torch/_functorch/_aot_autograd/input_output_analysis.py
@@ -24,7 +24,7 @@
 from torch.fx.experimental.symbolic_shapes import is_concrete_int
 
 from .collect_metadata_analysis import coerce_tangent_and_suggest_memory_format
-from .descriptors import AOTInput, InputMutationTangentAOTInput
+from .descriptors import AOTInput, InputMutationAOTOutput, TangentAOTInput
 from .schemas import (
     BackwardSignature,
     GraphSignature,
@@ -250,7 +250,7 @@ def create_synthetic_base_metadata(
         # See Note [Tangents memory format]
         (
             coerce_tangent_and_suggest_memory_format(x),
-            InputMutationTangentAOTInput(x_desc),
+            TangentAOTInput(InputMutationAOTOutput(x_desc)),
         )
         for inner_idx, (x, x_desc) in enumerate(zip(inner_args, inner_args_desc))
         if input_infos[inner_idx].mutates_data and input_infos[inner_idx].requires_grad
diff --git a/torch/_functorch/_aot_autograd/runtime_wrappers.py b/torch/_functorch/_aot_autograd/runtime_wrappers.py
@@ -46,10 +46,10 @@
 from .. import config
 from .collect_metadata_analysis import run_functionalized_fw_and_collect_metadata
 from .descriptors import (
-    AliasedArgWithMetadataMutationAOTOutput,
     AOTInput,
     AOTOutput,
     DummyAOTInput,
+    MetadataMutationAOTOutput,
     SyntheticBaseAOTInput,
     ViewBaseAOTInput,
 )
@@ -1178,7 +1178,7 @@ def wrapped_flat_fn(*args):
                 return (*out, *aliased_args_with_metadata_mutations), (
                     *out_descs,
                     *(
-                        [AliasedArgWithMetadataMutationAOTOutput()]
+                        [MetadataMutationAOTOutput()]
                         * len(self.aliased_arg_idx_with_metadata_mutations)
                     ),
                 )
