From 4034974d7524127efd16f1e101958c4dd973e6a5 Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Wed, 12 Aug 2026 11:46:28 -0700 Subject: [PATCH 1/2] Avoid no-op decomposition retraces for AOTI --- backends/aoti/aoti_backend.py | 9 +++++-- exir/graph_module.py | 28 +++++++++++++++++++- exir/program/_program.py | 9 +++++-- exir/tests/test_graph_module.py | 47 +++++++++++++++++++++++++++++++++ 4 files changed, 88 insertions(+), 5 deletions(-) create mode 100644 exir/tests/test_graph_module.py diff --git a/backends/aoti/aoti_backend.py b/backends/aoti/aoti_backend.py index 85eb0b1cc00..09a1ab7afa8 100644 --- a/backends/aoti/aoti_backend.py +++ b/backends/aoti/aoti_backend.py @@ -20,6 +20,7 @@ from executorch.exir._warnings import experimental from executorch.exir.backend.backend_details import ExportedProgram, PreprocessResult from executorch.exir.backend.compile_spec_schema import CompileSpec +from executorch.exir.graph_module import contains_any_op from torch._inductor.codegen.cpp_wrapper_cpu import CppWrapperCpu from torch.export.passes import move_to_device_pass @@ -248,8 +249,12 @@ def preprocess( else: custom_pass(device_edge_program.graph_module) - # Run decompositions if any - if decomposition_table: + # ``run_decompositions`` retraces the complete ExportedProgram even + # when none of the table's operators occur in the graph. Large CUDA + # models make that no-op expensive, so only run it when it can apply. + if contains_any_op( + device_edge_program.graph_module, decomposition_table.keys() + ): device_edge_program = device_edge_program.run_decompositions( decomposition_table ) diff --git a/exir/graph_module.py b/exir/graph_module.py index cc85a596bfb..7120efc86a6 100644 --- a/exir/graph_module.py +++ b/exir/graph_module.py @@ -8,7 +8,7 @@ # pyre-strict from types import FunctionType as function -from typing import Callable, Dict, List, Tuple, Union +from typing import Any, Callable, Collection, Dict, List, Tuple, Union import torch from torch._ops import HigherOrderOperator @@ -151,3 +151,29 @@ def bfs_trace_with_node_process( for _, submodule, _ in get_control_flow_submodules(current_graph_module) ] queue.extend(control_flow_submodules) + + +def contains_any_op( + graph_module: torch.fx.GraphModule, + ops: Collection[Any], +) -> bool: + """Return whether ``graph_module`` or a control-flow subgraph uses ``ops``.""" + if not ops: + return False + + queue = [graph_module] + while queue: + current_graph_module = queue.pop(0) + for node in current_graph_module.graph.nodes: + # EdgeOpOverload wraps the original ATen overload in ``_op``. + # Decomposition tables are keyed by the latter. + target = node.target + if not isinstance(target, torch._ops.OpOverload): + target = getattr(target, "_op", target) + if node.op == "call_function" and target in ops: + return True + queue.extend( + submodule + for _, submodule, _ in get_control_flow_submodules(current_graph_module) + ) + return False diff --git a/exir/program/_program.py b/exir/program/_program.py index 0e63266e663..dd4c9e47a3e 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -33,7 +33,7 @@ from executorch.exir.emit import emit_program, EmitterOutput from executorch.exir.emit._emitter import _DelegateDebugIdentifierMap from executorch.exir.error import ExportError -from executorch.exir.graph_module import get_control_flow_submodules +from executorch.exir.graph_module import contains_any_op, get_control_flow_submodules from executorch.exir.operator.convert import _pybind_schema_to_native_schema from executorch.exir.operator.util import _QUANT_PRIMITIVES from executorch.exir.pass_base import PassBase @@ -1166,7 +1166,12 @@ def _gen_edge_manager_for_partitioners( if table.pop(op, None) is not None: ops_needing_preservation.append(op) - program = program.run_decompositions(table) + # The initial ``run_decompositions({})`` above has already + # functionalized the graph. This second call only applies + # operator decompositions from the remaining table, so it is + # safe to skip when none of those targets occur in the graph. + if contains_any_op(program.graph_module, table.keys()): + program = program.run_decompositions(table) final_ops_to_preserve.update(ops_needing_preservation) else: # EDGE_DO_NOT_DECOMP path for the partitioner diff --git a/exir/tests/test_graph_module.py b/exir/tests/test_graph_module.py new file mode 100644 index 00000000000..e2186b7f45d --- /dev/null +++ b/exir/tests/test_graph_module.py @@ -0,0 +1,47 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +import torch +from executorch.exir.graph_module import contains_any_op + + +class TestContainsAnyOp(unittest.TestCase): + @staticmethod + def _add_graph() -> tuple[torch.fx.GraphModule, torch.fx.Node]: + graph = torch.fx.Graph() + lhs = graph.placeholder("lhs") + rhs = graph.placeholder("rhs") + add = graph.call_function(torch.ops.aten.add.Tensor, (lhs, rhs)) + graph.output(add) + return torch.fx.GraphModule({}, graph), add + + def test_matches_aten_op(self) -> None: + graph_module, _ = self._add_graph() + + self.assertTrue( + contains_any_op(graph_module, {torch.ops.aten.add.Tensor}) + ) + self.assertFalse( + contains_any_op(graph_module, {torch.ops.aten.mul.Tensor}) + ) + + def test_matches_wrapped_edge_op(self) -> None: + graph_module, add = self._add_graph() + + class EdgeOp: + _op = torch.ops.aten.add.Tensor + + add.target = EdgeOp() + + self.assertTrue( + contains_any_op(graph_module, {torch.ops.aten.add.Tensor}) + ) + + +if __name__ == "__main__": + unittest.main() From a81d8ccc8b5b5dafae3acd7db1cf0bb0d9191daa Mon Sep 17 00:00:00 2001 From: gasoonjia Date: Wed, 12 Aug 2026 14:45:17 -0700 Subject: [PATCH 2/2] Remove redundant CUDA floor-div rewrite pass --- backends/cuda/BUCK | 1 - backends/cuda/cuda_backend.py | 5 +- .../cuda/passes/replace_int64_floordiv.py | 152 ------------ .../tests/test_replace_int64_floordiv.py | 216 ------------------ exir/tests/test_graph_module.py | 12 +- 5 files changed, 4 insertions(+), 382 deletions(-) delete mode 100644 backends/cuda/passes/replace_int64_floordiv.py delete mode 100644 backends/cuda/passes/tests/test_replace_int64_floordiv.py diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index 1fbbf61be3c..746b8caa7ad 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -77,7 +77,6 @@ fbcode_target( srcs = [ "passes/__init__.py", "passes/move_cond_predicate_to_cpu.py", - "passes/replace_int64_floordiv.py", ], visibility = [ "//executorch/backends/cuda/...", diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 1f6a5df9ea2..9232dd57324 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -23,9 +23,6 @@ from executorch.backends.cuda.passes.move_cond_predicate_to_cpu import ( MoveCondPredicateToCpuPass, ) -from executorch.backends.cuda.passes.replace_int64_floordiv import ( - ReplaceInt64FloorDivWithFloatPass, -) from executorch.backends.cuda.triton.replacement_pass import ( ReplaceEdgeOpWithTritonOpPass, ) @@ -568,7 +565,7 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any] f"Invalid triton_kernel_mode: {mode}. Expected 'ON' or 'OFF'." ) triton_kernel_mode = mode - passes = [MoveCondPredicateToCpuPass(), ReplaceInt64FloorDivWithFloatPass()] + passes = [MoveCondPredicateToCpuPass()] if triton_kernel_mode == "ON": passes.append(ReplaceEdgeOpWithTritonOpPass()) return passes diff --git a/backends/cuda/passes/replace_int64_floordiv.py b/backends/cuda/passes/replace_int64_floordiv.py deleted file mode 100644 index 85cd201416e..00000000000 --- a/backends/cuda/passes/replace_int64_floordiv.py +++ /dev/null @@ -1,152 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -""" -Graph Transformation Pass for Integer Floor-Division Replacement. - -Rewrites integer (int64/int32) floor-division into a float64-domain floor to -work around a torch-2.12 AOTInductor/Inductor CUDA miscompile: - - floor_divide(a, b) -> floor(a.to(float64) / b.to(float64)).to(orig_int_dtype) -""" - -import logging - -import torch -from executorch.exir.dialects._ops import ops as exir_ops -from torch.fx import GraphModule, Node -from torch.fx.passes.infra.pass_base import PassBase, PassResult - -logger = logging.getLogger(__name__) - -# NOTE: Integer dtypes we rewrite. float64 (53-bit mantissa) is for -# |value| < 2**53, which covers models' index ranges but not enough -# for extreme large numbers. -_INT_DTYPES = (torch.int64, torch.int32) - -# Edge ops that perform a floor-rounded integer division. -_FLOOR_DIVIDE_OP = exir_ops.edge.aten.floor_divide.default -_DIV_MODE_OPS = ( - exir_ops.edge.aten.div.Tensor_mode, - exir_ops.edge.aten.div.Scalar_mode, -) - - -class ReplaceInt64FloorDivWithFloatPass(PassBase): - # Work around a torch-2.12 AOTInductor/Inductor CUDA miscompile of integer - # (int64) floor-division: fused/broadcast int64 floor_divide is mis-lowered - # (truncation instead of floor; cross-division term bleed under dynamic shapes). - # TODO(gasoonjia): remove this pass once the upstream issue solved. - # Upstream issue: https://github.com/pytorch/pytorch/issues/186164 - """ - Pass to rewrite integer floor-division into a float64-domain floor. - - Matches ``floor_divide.default`` and the floor-mode ``div.Tensor_mode`` / - ``div.Scalar_mode`` overloads on integer operands, and replaces each with - ``floor(a.to(float64) / b.to(float64)).to(orig_int_dtype)`` built from edge - dialect ops. Float floor-division and non-integer nodes are left untouched. - """ - - def __init__(self): - super().__init__() - self._replacement_count = 0 - - def call(self, graph_module: GraphModule) -> PassResult: - self._replacement_count = 0 - modified = False - - for node in graph_module.graph.nodes: - if not self._should_replace_node(node): - continue - try: - self._replace_node(graph_module, node) - modified = True - self._replacement_count += 1 - except Exception as e: - logger.warning(f"Failed to rewrite floor-div node {node.name}: {e}") - # Continue with other nodes even if one fails. - - if modified: - graph_module.recompile() - - logger.info( - f"Rewrote {self._replacement_count} integer floor-division nodes " - f"into float64-domain floor" - ) - - return PassResult(graph_module, modified) - - @staticmethod - def _node_dtype(node: Node): - val = node.meta.get("val", None) - if isinstance(val, torch.Tensor): - return val.dtype - return None - - @staticmethod - def _rounding_mode(node: Node): - if "rounding_mode" in node.kwargs: - return node.kwargs["rounding_mode"] - # Trailing positional arg: div(self, other, rounding_mode) - if len(node.args) > 2: - return node.args[2] - return None - - def _should_replace_node(self, node: Node) -> bool: - if node.op != "call_function": - return False - - if node.target == _FLOOR_DIVIDE_OP: - pass - elif node.target in _DIV_MODE_OPS: - if self._rounding_mode(node) != "floor": - return False - else: - return False - - # Only rewrite when the result is an integer tensor. Guard meta access: - # a node may lack meta["val"]; skip conservatively if so. - out_dtype = self._node_dtype(node) - if out_dtype not in _INT_DTYPES: - return False - - return True - - def _replace_node(self, graph_module: GraphModule, node: Node) -> None: - orig_dtype = self._node_dtype(node) - a = node.args[0] - b = node.args[1] - - graph = graph_module.graph - with graph.inserting_before(node): - a_f = graph.call_function( - exir_ops.edge.aten._to_copy.default, - args=(a,), - kwargs={"dtype": torch.float64}, - ) - if isinstance(b, Node): - b_f = graph.call_function( - exir_ops.edge.aten._to_copy.default, - args=(b,), - kwargs={"dtype": torch.float64}, - ) - q = graph.call_function(exir_ops.edge.aten.div.Tensor, args=(a_f, b_f)) - else: - # Python-scalar divisor: stays bit-exact, no cast needed for b. - q = graph.call_function( - exir_ops.edge.aten.div.Scalar, args=(a_f, float(b)) - ) - fl = graph.call_function(exir_ops.edge.aten.floor.default, args=(q,)) - new_node = graph.call_function( - exir_ops.edge.aten._to_copy.default, - args=(fl,), - kwargs={"dtype": orig_dtype}, - ) - - new_node.meta = node.meta.copy() - - node.replace_all_uses_with(new_node) - graph.erase_node(node) diff --git a/backends/cuda/passes/tests/test_replace_int64_floordiv.py b/backends/cuda/passes/tests/test_replace_int64_floordiv.py deleted file mode 100644 index 9632611890b..00000000000 --- a/backends/cuda/passes/tests/test_replace_int64_floordiv.py +++ /dev/null @@ -1,216 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -import unittest - -import torch -from backends.cuda.passes.replace_int64_floordiv import ( - ReplaceInt64FloorDivWithFloatPass, -) -from executorch.exir import to_edge -from executorch.exir.dialects._ops import ops as exir_ops -from torch.export import export - - -_INT_DIV_OPS = ( - exir_ops.edge.aten.floor_divide.default, - exir_ops.edge.aten.div.Tensor_mode, - exir_ops.edge.aten.div.Scalar_mode, -) - - -def _count_int_floordiv(graph_module) -> int: - """Count integer floor-division nodes remaining in the graph.""" - n = 0 - for node in graph_module.graph.nodes: - if node.op != "call_function" or node.target not in _INT_DIV_OPS: - continue - if node.target in ( - exir_ops.edge.aten.div.Tensor_mode, - exir_ops.edge.aten.div.Scalar_mode, - ): - rmode = node.kwargs.get("rounding_mode", None) - if rmode != "floor": - continue - val = node.meta.get("val", None) - if isinstance(val, torch.Tensor) and val.dtype in ( - torch.int64, - torch.int32, - ): - n += 1 - return n - - -class TestReplaceInt64FloorDivWithFloatPass(unittest.TestCase): - """Test the ReplaceInt64FloorDivWithFloatPass transformation pass.""" - - def _edge_gm(self, module, inputs): - ep = to_edge(export(module, inputs, strict=True)) - return ep, ep.exported_program().graph_module - - def test_tensor_tensor_floordiv_rewritten(self): - """int64 a // b (tensor/tensor), including negative numerators.""" - - class M(torch.nn.Module): - def forward(self, a, b): - return a // b - - a = torch.tensor([-5, 7, -8, 9, -1, 0], dtype=torch.long) - b = torch.tensor([2, 3, 4, 5, 3, 7], dtype=torch.long) - ep, gm = self._edge_gm(M().eval(), (a, b)) - - self.assertGreater(_count_int_floordiv(gm), 0) - ReplaceInt64FloorDivWithFloatPass()(gm) - self.assertEqual(_count_int_floordiv(gm), 0) - - out = ep.exported_program().module()(a, b) - self.assertEqual(out.dtype, torch.int64) - self.assertTrue(torch.equal(out, a // b)) - - def test_scalar_divisor_floordiv_rewritten(self): - """int64 a // 3 (scalar divisor lifted to a 0-d tensor constant).""" - - class M(torch.nn.Module): - def forward(self, a): - return a // 3 - - a = torch.tensor([-5, 7, -8, 9, -1, 0], dtype=torch.long) - ep, gm = self._edge_gm(M().eval(), (a,)) - - self.assertGreater(_count_int_floordiv(gm), 0) - ReplaceInt64FloorDivWithFloatPass()(gm) - self.assertEqual(_count_int_floordiv(gm), 0) - - out = ep.exported_program().module()(a) - self.assertTrue(torch.equal(out, a // 3)) - - def test_div_rounding_mode_floor_rewritten(self): - """torch.div(..., rounding_mode='floor') on int64 is rewritten.""" - - class M(torch.nn.Module): - def forward(self, a, b): - return torch.div(a, b, rounding_mode="floor") - - a = torch.tensor([-5, 7, -8, 9], dtype=torch.long) - b = torch.tensor([2, 3, 4, 5], dtype=torch.long) - ep, gm = self._edge_gm(M().eval(), (a, b)) - - self.assertGreater(_count_int_floordiv(gm), 0) - ReplaceInt64FloorDivWithFloatPass()(gm) - self.assertEqual(_count_int_floordiv(gm), 0) - - out = ep.exported_program().module()(a, b) - self.assertTrue(torch.equal(out, torch.div(a, b, rounding_mode="floor"))) - - def test_int32_floordiv_rewritten(self): - """int32 floor-division is also rewritten and stays int32.""" - - class M(torch.nn.Module): - def forward(self, a, b): - return a // b - - a = torch.tensor([-5, 7, -8, 9], dtype=torch.int32) - b = torch.tensor([2, 3, 4, 5], dtype=torch.int32) - ep, gm = self._edge_gm(M().eval(), (a, b)) - - self.assertGreater(_count_int_floordiv(gm), 0) - ReplaceInt64FloorDivWithFloatPass()(gm) - self.assertEqual(_count_int_floordiv(gm), 0) - - out = ep.exported_program().module()(a, b) - self.assertEqual(out.dtype, torch.int32) - self.assertTrue(torch.equal(out, a // b)) - - def test_float_division_untouched(self): - """Real float division must not be rewritten.""" - - class M(torch.nn.Module): - def forward(self, a, b): - return a / b - - a = torch.tensor([1.0, 2.0, 3.0]) - b = torch.tensor([2.0, 3.0, 4.0]) - ep, gm = self._edge_gm(M().eval(), (a, b)) - - before = [n.target for n in gm.graph.nodes if n.op == "call_function"] - result = ReplaceInt64FloorDivWithFloatPass()(gm) - self.assertFalse(result.modified) - after = [n.target for n in gm.graph.nodes if n.op == "call_function"] - self.assertEqual(before, after) - - def test_trunc_rounding_mode_untouched(self): - """div with rounding_mode='trunc' must not be rewritten.""" - - class M(torch.nn.Module): - def forward(self, a, b): - return torch.div(a, b, rounding_mode="trunc") - - a = torch.tensor([-5, 7, -8, 9], dtype=torch.long) - b = torch.tensor([2, 3, 4, 5], dtype=torch.long) - ep, gm = self._edge_gm(M().eval(), (a, b)) - - result = ReplaceInt64FloorDivWithFloatPass()(gm) - self.assertFalse(result.modified) - - def test_floor_divide_default_branch(self): - """Exercise the floor_divide.default match/rewrite branch. - - This pin lowers ``//`` to ``div.Tensor_mode``; floor_divide.default does - not appear naturally, so we synthesize it by retargeting a node. - """ - - class M(torch.nn.Module): - def forward(self, a, b): - return a // b - - a = torch.tensor([-5, 7, -8, 9], dtype=torch.long) - b = torch.tensor([2, 3, 4, 5], dtype=torch.long) - ep, gm = self._edge_gm(M().eval(), (a, b)) - - # Retarget the div.Tensor_mode node to floor_divide.default. - for node in list(gm.graph.nodes): - if node.target == exir_ops.edge.aten.div.Tensor_mode: - with gm.graph.inserting_before(node): - new = gm.graph.call_function( - exir_ops.edge.aten.floor_divide.default, args=node.args - ) - new.meta = node.meta.copy() - node.replace_all_uses_with(new) - gm.graph.erase_node(node) - gm.recompile() - - self.assertGreater(_count_int_floordiv(gm), 0) - ReplaceInt64FloorDivWithFloatPass()(gm) - self.assertEqual(_count_int_floordiv(gm), 0) - - out = ep.exported_program().module()(a, b) - self.assertTrue(torch.equal(out, a // b)) - - def test_ring_buffer_mask_analog(self): - """gemma4_31b sliding-window analog: negative numerators + scalar divisor.""" - - class M(torch.nn.Module): - def forward(self, input_pos): - buf_size = 8 - seq_len = input_pos.shape[0] - total_written = input_pos[0] + seq_len - j = torch.arange(buf_size, dtype=torch.long) - wraps = (total_written - 1 - j) // buf_size - return j + wraps * buf_size - - input_pos = torch.arange(3, dtype=torch.long) - ep, gm = self._edge_gm(M().eval(), (input_pos,)) - - ReplaceInt64FloorDivWithFloatPass()(gm) - self.assertEqual(_count_int_floordiv(gm), 0) - - out = ep.exported_program().module()(input_pos) - ref = M()(input_pos) - self.assertTrue(torch.equal(out, ref)) - - -if __name__ == "__main__": - unittest.main() diff --git a/exir/tests/test_graph_module.py b/exir/tests/test_graph_module.py index e2186b7f45d..560c35e747d 100644 --- a/exir/tests/test_graph_module.py +++ b/exir/tests/test_graph_module.py @@ -23,12 +23,8 @@ def _add_graph() -> tuple[torch.fx.GraphModule, torch.fx.Node]: def test_matches_aten_op(self) -> None: graph_module, _ = self._add_graph() - self.assertTrue( - contains_any_op(graph_module, {torch.ops.aten.add.Tensor}) - ) - self.assertFalse( - contains_any_op(graph_module, {torch.ops.aten.mul.Tensor}) - ) + self.assertTrue(contains_any_op(graph_module, {torch.ops.aten.add.Tensor})) + self.assertFalse(contains_any_op(graph_module, {torch.ops.aten.mul.Tensor})) def test_matches_wrapped_edge_op(self) -> None: graph_module, add = self._add_graph() @@ -38,9 +34,7 @@ class EdgeOp: add.target = EdgeOp() - self.assertTrue( - contains_any_op(graph_module, {torch.ops.aten.add.Tensor}) - ) + self.assertTrue(contains_any_op(graph_module, {torch.ops.aten.add.Tensor})) if __name__ == "__main__":