torch 2.6.0 upgrade patch

All patches to make ai-edge-torch working with torch 2.6.0

1. `capture_pre_autograd_graph` is removed in torch 2.6.0. Replace this with `export_for_training` in PT2E testing. Tutorial updates may come in later CLs.
2. `_decomp_table_to_post_autograd_aten` (added in 2.5.0) is removed in 2.6.0. For now there is no built-in decomp table to transform FX into the state we expect as in 2.4.0. Therefore a new core-aten based decomp regitry system is added in this CL to make decomp behavior expected in 2.4.0, 2.5.0, and 2.6.0. Tests are passing in 2.5.0 and 2.6.0 but impact on converted model performance need more investigation.
3. 2.6.0 introduce `_assert_tensor_metadata` op to guard custom ops in exported graph. This will break pattern matching and lowering. Add a `fx_infra` with `safe_run_decompositions` to exclude this op and all the other unexpected export behavior.

PiperOrigin-RevId: 718173217

Author: chunnienc

ai_edge_torch/__init__.py

@@ -12,14 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-
 from ai_edge_torch._config import config
 from ai_edge_torch._convert.converter import convert
 from ai_edge_torch._convert.converter import signature
 from ai_edge_torch._convert.to_channel_last_io import to_channel_last_io
 from ai_edge_torch.model import Model
 from ai_edge_torch.version import __version__
 
+
 def load(path: str) -> Model:
   """Imports an ai_edge_torch model from disk.
 

ai_edge_torch/_convert/conversion.py

@@ -17,7 +17,7 @@
 from typing import Any, Literal, Optional, Union
 
 import ai_edge_torch
-from ai_edge_torch import fx_pass_base
+from ai_edge_torch import fx_infra
 from ai_edge_torch import lowertools
 from ai_edge_torch import model
 from ai_edge_torch._convert import fx_passes
@@ -53,7 +53,7 @@ def _run_convert_passes(
         fx_passes.CanonicalizePass(),
     ]
 
-  exported_program = fx_pass_base.run_passes(exported_program, passes)
+  exported_program = fx_infra.run_passes(exported_program, passes)
   return exported_program
 
 
@@ -125,14 +125,10 @@ def export(**kwargs):
     else:
       exported_program = torch.export.export(**kwargs, strict=True)
 
-    if hasattr(torch._decomp, "_decomp_table_to_post_autograd_aten"):
-      # Available after torch 2.5.0: `_decomp_table_to_post_autograd_aten` is a
-      # stop-gap table which replicates the old behaviour of post-dispatch IR.
-      # This could help ensure the collection of aten ops remaining still as the
-      # implementation of torch.export changes.
-      exported_program = exported_program.run_decompositions(
-          torch._decomp._decomp_table_to_post_autograd_aten()
-      )
+    exported_program = fx_infra.safe_run_decompositions(
+        exported_program,
+        fx_infra.decomp.pre_convert_decomp(),
+    )
     return exported_program
 
   exported_programs: torch.export.ExportedProgram = [

ai_edge_torch/_convert/fx_passes/__init__.py

@@ -20,4 +20,4 @@
 from ai_edge_torch._convert.fx_passes.inject_mlir_debuginfo_pass import InjectMlirDebuginfoPass
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import OptimizeLayoutTransposesPass
 from ai_edge_torch._convert.fx_passes.remove_non_user_outputs_pass import RemoveNonUserOutputsPass
-from ai_edge_torch.fx_pass_base import CanonicalizePass
+from ai_edge_torch.fx_infra import CanonicalizePass

ai_edge_torch/_convert/fx_passes/build_aten_composite_pass.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 
 from typing import Any, Callable
-from ai_edge_torch import fx_pass_base
+from ai_edge_torch import fx_infra
 from ai_edge_torch import lowertools
 import torch
 import torch.utils._pytree as pytree
@@ -25,6 +25,9 @@
 
 
 def _register_composite_builder(op):
+  # Remove op from pre_convert_decomp to keep this in the decomposed graph.
+  fx_infra.decomp.remove_pre_convert_decomp(op)
+
   def inner(func):
     if isinstance(op, torch._ops.OpOverloadPacket):
       for overload in op.overloads():
@@ -276,7 +279,7 @@ def embedding(*args, **kwargs):
   node.target = embedding
 
 
-class BuildAtenCompositePass(fx_pass_base.PassBase):
+class BuildAtenCompositePass(fx_infra.PassBase):
 
   def call(self, graph_module: torch.fx.GraphModule):
     for node in graph_module.graph.nodes:
@@ -285,4 +288,4 @@ def call(self, graph_module: torch.fx.GraphModule):
 
     graph_module.graph.lint()
     graph_module.recompile()
-    return fx_pass_base.PassResult(graph_module, True)
+    return fx_infra.PassResult(graph_module, True)

ai_edge_torch/_convert/fx_passes/build_interpolate_composite_pass.py

@@ -16,7 +16,7 @@
 
 import functools
 
-from ai_edge_torch import fx_pass_base
+from ai_edge_torch import fx_infra
 from ai_edge_torch.hlfb import mark_pattern
 from ai_edge_torch.hlfb.mark_pattern import pattern as pattern_module
 import torch
@@ -41,7 +41,7 @@ def _get_upsample_bilinear2d_pattern():
           x, scale_factor=2, mode="bilinear", align_corners=False
       ),
       export_args=(torch.rand(1, 3, 100, 100),),
-      decomp_table=_INTERPOLATE_DECOMPOSITIONS,
+      extra_decomp_table=_INTERPOLATE_DECOMPOSITIONS,
   )
 
   @pattern.register_attr_builder
@@ -65,7 +65,7 @@ def _get_upsample_bilinear2d_align_corners_pattern():
           x, scale_factor=2, mode="bilinear", align_corners=True
       ),
       export_args=(torch.rand(1, 3, 100, 100),),
-      decomp_table=_INTERPOLATE_DECOMPOSITIONS,
+      extra_decomp_table=_INTERPOLATE_DECOMPOSITIONS,
   )
 
   @pattern.register_attr_builder
@@ -89,7 +89,7 @@ def _get_interpolate_nearest2d_pattern():
           x, scale_factor=2, mode="nearest"
       ),
       export_args=(torch.rand(1, 3, 100, 100),),
-      decomp_table=_INTERPOLATE_DECOMPOSITIONS,
+      extra_decomp_table=_INTERPOLATE_DECOMPOSITIONS,
   )
 
   @pattern.register_attr_builder
@@ -104,7 +104,7 @@ def attr_builder(pattern, graph_module, internal_match):
   return pattern
 
 
-class BuildInterpolateCompositePass(fx_pass_base.ExportedProgramPassBase):
+class BuildInterpolateCompositePass(fx_infra.ExportedProgramPassBase):
 
   def __init__(self):
     super().__init__()
@@ -115,11 +115,9 @@ def __init__(self):
     ]
 
   def call(self, exported_program: torch.export.ExportedProgram):
-    exported_program = fx_pass_base.run_passes(
-        exported_program, [fx_pass_base.CanonicalizePass()]
-    )
-    exported_program = exported_program.run_decompositions(
-        _INTERPOLATE_DECOMPOSITIONS
+    exported_program = fx_infra.safe_run_decompositions(
+        exported_program,
+        _INTERPOLATE_DECOMPOSITIONS,
     )
 
     graph_module = exported_program.graph_module
@@ -128,4 +126,4 @@ def call(self, exported_program: torch.export.ExportedProgram):
 
     graph_module.graph.lint()
     graph_module.recompile()
-    return fx_pass_base.ExportedProgramPassResult(exported_program, True)
+    return fx_infra.ExportedProgramPassResult(exported_program, True)

ai_edge_torch/_convert/fx_passes/inject_mlir_debuginfo_pass.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from ai_edge_torch import fx_pass_base
+from ai_edge_torch import fx_infra
 from ai_edge_torch import lowertools
 import torch
 import torch.utils._pytree as pytree
@@ -61,7 +61,7 @@ def debuginfo_writer(*args, **kwargs):
   node.target = debuginfo_writer
 
 
-class InjectMlirDebuginfoPass(fx_pass_base.PassBase):
+class InjectMlirDebuginfoPass(fx_infra.PassBase):
   """DEPRECATED: Debuginfo is injected automatically by odml_torch."""
 
   def call(self, graph_module: torch.fx.GraphModule):
@@ -70,4 +70,4 @@ def call(self, graph_module: torch.fx.GraphModule):
 
     graph_module.graph.lint()
     graph_module.recompile()
-    return fx_pass_base.PassResult(graph_module, True)
+    return fx_infra.PassResult(graph_module, True)

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/__init__.py

@@ -13,4 +13,5 @@
 # limitations under the License.
 # ==============================================================================
 
+from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import _decomp_registry
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass.pass_body import OptimizeLayoutTransposesPass  # NOQA

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/_decomp_registry.py

@@ -0,0 +1,33 @@
+# Copyright 2025 The AI Edge Torch Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Remove decompositions for ops to keep in layout optimization."""
+from ai_edge_torch import fx_infra
+import torch
+
+__all__ = []
+
+aten = torch.ops.aten
+
+_OPS_TO_KEEP = [
+    aten.conv2d,
+    aten.max_pool2d,
+    aten._softmax.default,
+    aten.group_norm.default,
+    aten.native_group_norm.default,
+    aten.reflection_pad2d.default,
+]
+
+for op in _OPS_TO_KEEP:
+  fx_infra.decomp.remove_pre_convert_decomp(op)

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/op_func_registry.py

@@ -20,6 +20,7 @@
 class OpFuncRegistry(dict):
 
   def register(self, op):
+
     ops = utils.flatten_torch_op_overloads(op)
 
     def inner(func):

ai_edge_torch/_convert/fx_passes/optimize_layout_transposes_pass/pass_body.py

@@ -18,7 +18,7 @@
 import os
 from typing import Union
 
-from ai_edge_torch import fx_pass_base
+from ai_edge_torch import fx_infra
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_check  # NOQA
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_mark  # NOQA
 from ai_edge_torch._convert.fx_passes.optimize_layout_transposes_pass import layout_partitioners  # NOQA
@@ -30,7 +30,7 @@
 TransposeFunc = Union[utils.tensor_to_nchw, utils.tensor_to_nhwc]
 
 
-class OptimizeLayoutTransposesPass(fx_pass_base.ExportedProgramPassBase):
+class OptimizeLayoutTransposesPass(fx_infra.ExportedProgramPassBase):
 
   def get_source_meta(self, node: torch.fx.Node):
     keys = ["stack_trace", "nn_module_stack", "source_fn_stack", "from_node"]
@@ -300,4 +300,4 @@ def call(self, exported_program: torch.export.ExportedProgram):
     # Mark const node again for debugging
     self.mark_const_nodes(exported_program)
 
-    return fx_pass_base.ExportedProgramPassResult(exported_program, True)
+    return fx_infra.ExportedProgramPassResult(exported_program, True)