From f4f913b963e187060a4d30c90deb07de89729f51 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=C3=A5ns=20Nilsson?= Date: Fri, 21 Aug 2026 13:10:22 +0200 Subject: [PATCH 1/2] Arm backend: Guard argmax and argmin int32 propagation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Previously, once an argmax or argmin index fit int32, the pass sent every consumer through the int32 path. Later arithmetic could then overflow even though the raw index itself was safe. Keep unsafe direct consumers on the original int64 value. When a safe int32 prefix reaches an unsafe consumer, insert an int64 boundary. This allows preceding safe operations to remain delegatable while keeping calculations correct and model output dtypes unchanged. Cover safe and overflowing paths in both ATen and Edge graphs. Signed-off-by: Måns Nilsson Change-Id: I6849f1e33e17ab9661cf25002923e180755e0a46 --- backends/arm/README.md | 9 +- .../convert_int64_output_ops_to_int32.py | 377 ++++++++++++++++-- .../test_convert_int64_output_ops_to_int32.py | 104 ++++- 3 files changed, 459 insertions(+), 31 deletions(-) diff --git a/backends/arm/README.md b/backends/arm/README.md index 28bbb8f8351..ce0c49919e4 100644 --- a/backends/arm/README.md +++ b/backends/arm/README.md @@ -360,11 +360,16 @@ List of model specific and optional passes: - Supported Ops: - torch.ops.aten.to.\[dtype|dtype_layout\] - exir_ops.edge.dim_order_ops.\_to_dim_order_copy.default - 2. Post-process argmax outputs: - - Inserts an int64->int32 cast after the argmax operations that produce int64 outputs: + 2. Post-process argmax and argmin outputs: + - Converts only downstream paths whose statically inferred values remain + within the int32 range. + - Leaves unsafe direct consumers on int64 and inserts int64 boundary + casts where converted paths reach unsafe consumers or model outputs. - Supported Ops: - torch.ops.aten.argmax.default - exir_ops.edge.aten.argmax.default + - torch.ops.aten.argmin.default + - exir_ops.edge.aten.argmin.default - Example usage: - (Functionality 1) backends/arm/test/models/stable_diffusion/test_T5EncoderModel.py - (Functionality 2) backends/arm/test/models/stable_diffusion/test_CLIPTextModelWithProjection.py diff --git a/backends/arm/_passes/convert_int64_output_ops_to_int32.py b/backends/arm/_passes/convert_int64_output_ops_to_int32.py index 061ffd3a4a6..c269d4aa396 100644 --- a/backends/arm/_passes/convert_int64_output_ops_to_int32.py +++ b/backends/arm/_passes/convert_int64_output_ops_to_int32.py @@ -5,7 +5,7 @@ import logging -from typing import cast, Literal, Set, Type +from typing import Any, cast, Dict, Literal, Optional, Set, Tuple, Type import torch from executorch.backends.arm._passes import ArmPass @@ -16,6 +16,7 @@ ) from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass, PassResult +from torch._subclasses.fake_tensor import FakeTensor logger = logging.getLogger(__name__) @@ -31,7 +32,13 @@ class ConvertInt64OutputOpsToInt32Pass(ArmPass): 2. other types -> int64: rewrites the cast to produce int32 instead of int64. 3. torch.argmax() / torch.argmin() - insert an int64->int32 cast after the argmax/argmin node + insert an int64->int32 cast only along downstream paths whose values + are proven to remain within the int32 range. Other paths keep the + original int64 value or receive an int32->int64 boundary cast. + + Argmax and argmin are currently the only bounded-index sources. Range + propagation from those sources recognizes a separate allowlist of safe + shape and arithmetic operations. Future extensions may include other operators that return int64 outputs by default, rewriting them or inserting an int64 -> int32 cast to yield int32 @@ -63,15 +70,15 @@ def __init__( ) self.on_overflow = on_overflow - def _is_int32_range_safe(self, node: torch.fx.Node) -> bool: - """Return True if the argmax/argmin index output fits in int32.""" + def _index_range(self, node: torch.fx.Node) -> Tuple[int, int]: + """Return the inclusive output range of an argmax/argmin node.""" input_tensor = get_first_fake_tensor(cast(torch.fx.Node, node.args[0])) dim = node.args[1] if len(node.args) > 1 and node.args[1] is not None else None if dim is None: size = input_tensor.numel() else: size = input_tensor.shape[cast(int, dim)] - return size <= self._INT32_MAX + return 0, int(size) - 1 aten_cast_ops = ( torch.ops.aten.to.dtype, @@ -85,8 +92,31 @@ def _is_int32_range_safe(self, node: torch.fx.Node) -> bool: aten_argmin_ops = (torch.ops.aten.argmin.default,) edge_argmin_ops = (exir_ops.edge.aten.argmin.default,) - aten_ops = aten_cast_ops + aten_argmax_ops + aten_argmin_ops - edge_ops = edge_cast_ops + edge_argmax_ops + edge_argmin_ops + aten_bounded_index_ops = aten_argmax_ops + aten_argmin_ops + edge_bounded_index_ops = edge_argmax_ops + edge_argmin_ops + + aten_index_relay_ops = ( + torch.ops.aten.unsqueeze.default, + torch.ops.aten.expand.default, + torch.ops.aten.view.default, + ) + edge_index_relay_ops = ( + exir_ops.edge.aten.unsqueeze_copy.default, + exir_ops.edge.aten.expand_copy.default, + exir_ops.edge.aten.view_copy.default, + ) + aten_index_add_ops = (torch.ops.aten.add.Tensor,) + edge_index_add_ops = (exir_ops.edge.aten.add.Tensor,) + aten_index_sub_ops = (torch.ops.aten.sub.Tensor,) + edge_index_sub_ops = (exir_ops.edge.aten.sub.Tensor,) + aten_index_mul_ops = (torch.ops.aten.mul.Tensor,) + edge_index_mul_ops = (exir_ops.edge.aten.mul.Tensor,) + + aten_index_binary_ops = aten_index_add_ops + aten_index_sub_ops + aten_index_mul_ops + edge_index_binary_ops = edge_index_add_ops + edge_index_sub_ops + edge_index_mul_ops + + aten_ops = aten_cast_ops + aten_bounded_index_ops + edge_ops = edge_cast_ops + edge_bounded_index_ops # dtype is specified in args cast_ops_args = ( @@ -135,25 +165,312 @@ def _convert_casting_operators(self, node: torch.fx.Node): f" {input_dtype}->torch.int32 defined in {node.meta.get('stack_trace','[no stack trace found]')}" ) - def _cast_int64_output_to_int32(self, node: torch.fx.Node, graph: torch.fx.Graph): - output_tensor = node - to_copy_op = self._get_decomposition(node.target) - with graph.inserting_after(node): - cast_after = create_node( + def _range_fits_int32(self, value_range: Tuple[int, int]) -> bool: + return -self._INT32_MAX - 1 <= value_range[0] and ( + value_range[1] <= self._INT32_MAX + ) + + def _index_size_fits_int32_policy(self, index_range: Tuple[int, int]) -> bool: + """Apply the established source-dimension overflow policy. + + Args: + index_range (tuple): Inclusive index range. + + Returns: + bool: True when the source dimension is accepted. + + """ + return index_range[1] < self._INT32_MAX + + @staticmethod + def _scalar_int(value: Any) -> Optional[int]: + if isinstance(value, int): + return value + if ( + isinstance(value, torch.Tensor) + and not isinstance(value, FakeTensor) + and value.numel() == 1 + and not value.dtype.is_floating_point + and not value.dtype.is_complex + ): + return int(value.item()) + return None + + def _constant_range( + self, value: Any, graph_module: torch.fx.GraphModule + ) -> Optional[Tuple[int, int]]: + scalar = self._scalar_int(value) + if scalar is not None: + return scalar, scalar + if not isinstance(value, torch.fx.Node): + return None + + constant = None + if value.op == "get_attr" and isinstance(value.target, str): + constant = getattr(graph_module, value.target, None) + elif value.op == "placeholder" and isinstance(value.target, str): + buffer_name = value.target.removeprefix("_lifted") + if buffer_name != value.target: + try: + constant = graph_module.get_buffer(buffer_name) + except AttributeError: + pass + + scalar = self._scalar_int(constant) + return None if scalar is None else (scalar, scalar) + + def _operand_range( + self, + value: Any, + ranges: Dict[torch.fx.Node, Tuple[int, int]], + graph_module: torch.fx.GraphModule, + ) -> Optional[Tuple[int, int]]: + if isinstance(value, torch.fx.Node) and value in ranges: + return ranges[value] + return self._constant_range(value, graph_module) + + @staticmethod + def _scale_range(value_range: Tuple[int, int], scale: int) -> Tuple[int, int]: + values = value_range[0] * scale, value_range[1] * scale + return min(values), max(values) + + def _infer_binary_range( + self, + node: torch.fx.Node, + ranges: Dict[torch.fx.Node, Tuple[int, int]], + graph_module: torch.fx.GraphModule, + ) -> Optional[Tuple[int, int]]: + if get_first_fake_tensor(node).dtype != torch.int64 or len(node.args) < 2: + return None + lhs = self._operand_range(node.args[0], ranges, graph_module) + rhs = self._operand_range(node.args[1], ranges, graph_module) + if lhs is None or rhs is None: + return None + if not self._range_fits_int32(lhs) or not self._range_fits_int32(rhs): + return None + + if node.target in self.aten_index_mul_ops + self.edge_index_mul_ops: + products = ( + lhs[0] * rhs[0], + lhs[0] * rhs[1], + lhs[1] * rhs[0], + lhs[1] * rhs[1], + ) + result = min(products), max(products) + else: + alpha = self._scalar_int(node.kwargs.get("alpha", 1)) + if alpha is None: + return None + rhs = self._scale_range(rhs, alpha) + if node.target in self.aten_index_add_ops + self.edge_index_add_ops: + result = lhs[0] + rhs[0], lhs[1] + rhs[1] + elif node.target in self.aten_index_sub_ops + self.edge_index_sub_ops: + result = lhs[0] - rhs[1], lhs[1] - rhs[0] + else: + return None + return result if self._range_fits_int32(result) else None + + def _infer_safe_int32_range( + self, + node: torch.fx.Node, + ranges: Dict[torch.fx.Node, Tuple[int, int]], + graph_module: torch.fx.GraphModule, + ) -> Optional[Tuple[int, int]]: + if node.target in self.aten_index_relay_ops + self.edge_index_relay_ops: + if get_first_fake_tensor(node).dtype != torch.int64: + return None + return ranges.get(cast(torch.fx.Node, node.args[0])) + return self._infer_binary_range(node, ranges, graph_module) + + def _is_safe_widening_consumer( + self, + node: torch.fx.Node, + ranges: Dict[torch.fx.Node, Tuple[int, int]], + ) -> bool: + """Return whether a node safely ends integer range propagation. + + Args: + node (torch.fx.Node): Candidate downstream consumer. + ranges (dict): Proven ranges. + + Returns: + bool: True when the node safely produces a non-int64 value. + + """ + output = node.meta.get("val") + if not isinstance(output, torch.Tensor) or output.dtype == torch.int64: + return False + int64_inputs = [ + input_node + for input_node in node.all_input_nodes + if isinstance(input_node.meta.get("val"), torch.Tensor) + and input_node.meta["val"].dtype == torch.int64 + ] + if not int64_inputs or not all(node in ranges for node in int64_inputs): + return False + return node.target in ( + self.aten_cast_ops + + self.edge_cast_ops + + self.aten_index_add_ops + + self.edge_index_add_ops + + self.aten_index_sub_ops + + self.edge_index_sub_ops + + self.aten_index_mul_ops + + self.edge_index_mul_ops + ) + + def _find_safe_index_consumers( + self, + graph_module: torch.fx.GraphModule, + source: torch.fx.Node, + source_range: Tuple[int, int], + ) -> Tuple[Dict[torch.fx.Node, Tuple[int, int]], Set[torch.fx.Node]]: + """Collect consumers proven safe for the int32 index path. + + Args: + graph_module (torch.fx.GraphModule): Graph containing the source. + source (torch.fx.Node): Bounded int64 index source. + source_range (tuple): Inclusive source range. + + Returns: + tuple: Proven ranges and safe consumers. + + """ + ranges = {source: source_range} + safe_consumers: Set[torch.fx.Node] = set() + for node in graph_module.graph.nodes: + if node.op != "call_function": + continue + inferred_range = self._infer_safe_int32_range(node, ranges, graph_module) + if inferred_range is not None: + ranges[node] = inferred_range + safe_consumers.add(node) + elif self._is_safe_widening_consumer(node, ranges): + safe_consumers.add(node) + return ranges, safe_consumers + + @staticmethod + def _insert_int64_boundary( + graph: torch.fx.Graph, + node: torch.fx.Node, + to_copy_op, + boundaries: Dict[torch.fx.Node, torch.fx.Node], + ) -> torch.fx.Node: + if node not in boundaries: + with graph.inserting_after(node): + boundaries[node] = create_node( + graph, + to_copy_op, + args=(node,), + kwargs={"dtype": torch.int64}, + ) + return boundaries[node] + + def _cast_safe_scalar_constants_to_int32( + self, + graph_module: torch.fx.GraphModule, + safe_consumers: Set[torch.fx.Node], + ranges: Dict[torch.fx.Node, Tuple[int, int]], + to_copy_op, + ) -> None: + """Cast int64 scalar constants used by safe binary consumers. + + Args: + graph_module (torch.fx.GraphModule): Graph being transformed. + safe_consumers (set): Consumers on int32 paths. + ranges (dict): Proven ranges. + to_copy_op (Any): Dialect-specific operator used for casts. + + """ + graph = graph_module.graph + constant_casts: Dict[torch.fx.Node, torch.fx.Node] = {} + for consumer in safe_consumers: + if consumer.target not in ( + self.aten_index_binary_ops + self.edge_index_binary_ops + ): + continue + for input_node in consumer.all_input_nodes: + if input_node in ranges: + continue + input_value = input_node.meta.get("val") + if ( + not isinstance(input_value, torch.Tensor) + or input_value.dtype != torch.int64 + or self._constant_range(input_node, graph_module) is None + ): + continue + if input_node not in constant_casts: + with graph.inserting_after(input_node): + constant_casts[input_node] = create_node( + graph, + to_copy_op, + args=(input_node,), + kwargs={"dtype": torch.int32}, + ) + consumer.replace_input_with(input_node, constant_casts[input_node]) + + def _cast_safe_index_paths_to_int32( + self, + graph_module: torch.fx.GraphModule, + source: torch.fx.Node, + source_range: Tuple[int, int], + to_copy_op, + ) -> bool: + """Convert proven-safe paths from a bounded index source to int32. + + The caller identifies the bounded source and supplies its inclusive + value range. Direct consumers that cannot be proven safe retain the + original int64 source. An int64 boundary cast is inserted when an + unproven consumer follows an intermediate converted to int32. + + Args: + graph_module (torch.fx.GraphModule): Graph containing the source. + source (torch.fx.Node): Int64 node with a statically known range. + source_range (Tuple[int, int]): Inclusive minimum and maximum. + to_copy_op (Any): Dialect-specific operator used for casts. + + Returns: + bool: True when at least one path is converted to int32. + + """ + ranges, safe_consumers = self._find_safe_index_consumers( + graph_module, source, source_range + ) + if not safe_consumers: + return False + + graph = graph_module.graph + original_users = {node: list(node.users) for node in ranges} + with graph.inserting_after(source): + cast_to_int32 = create_node( graph, to_copy_op, - args=(output_tensor,), - kwargs={ - "dtype": torch.int32, - }, + args=(source,), + kwargs={"dtype": torch.int32}, ) - users = [user for user in node.users if user != cast_after] + + self._cast_safe_scalar_constants_to_int32( + graph_module, safe_consumers, ranges, to_copy_op + ) + + boundaries: Dict[torch.fx.Node, torch.fx.Node] = {} + for node, users in original_users.items(): for user in users: - user.replace_input_with(output_tensor, cast_after) - logger.warning( - f"Inserting a casting node {cast_after.name} after {node.name} to cast int64 output" - f" to int32 for {node.name} defined in {node.meta.get('stack_trace','[no stack trace found]')}" - ) + if user in safe_consumers: + if node is source: + user.replace_input_with(source, cast_to_int32) + elif node is not source: + boundary = self._insert_int64_boundary( + graph, node, to_copy_op, boundaries + ) + user.replace_input_with(node, boundary) + + logger.warning( + f"Inserting a casting node {cast_to_int32.name} after " + f"{source.name} for range-safe index consumers defined in " + f"{source.meta.get('stack_trace','[no stack trace found]')}" + ) + return True def call(self, graph_module: torch.fx.GraphModule): modified = False @@ -170,12 +487,10 @@ def call(self, graph_module: torch.fx.GraphModule): if node.target in self.aten_cast_ops + self.edge_cast_ops: self._convert_casting_operators(node) elif node.target in ( - self.aten_argmax_ops - + self.edge_argmax_ops - + self.aten_argmin_ops - + self.edge_argmin_ops + self.aten_bounded_index_ops + self.edge_bounded_index_ops ): - if not self._is_int32_range_safe(node): + index_range = self._index_range(node) + if not self._index_size_fits_int32_policy(index_range): msg = ( f"{node.target} reduces over more than {self._INT32_MAX} elements; " f"the int64 index cannot be safely cast to int32." @@ -185,7 +500,13 @@ def call(self, graph_module: torch.fx.GraphModule): if self.on_overflow == "warn": logger.warning(msg) continue - self._cast_int64_output_to_int32(node, graph) + if not self._cast_safe_index_paths_to_int32( + graph_module, + node, + index_range, + self._get_decomposition(node.target), + ): + continue else: raise RuntimeError(f"Unexpected target {node.target} in {node.name}") diff --git a/backends/arm/test/passes/test_convert_int64_output_ops_to_int32.py b/backends/arm/test/passes/test_convert_int64_output_ops_to_int32.py index f64b17297ca..a76ba2efb98 100644 --- a/backends/arm/test/passes/test_convert_int64_output_ops_to_int32.py +++ b/backends/arm/test/passes/test_convert_int64_output_ops_to_int32.py @@ -1,4 +1,4 @@ -# Copyright 2025 Arm Limited and/or its affiliates. +# Copyright 2025-2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -12,6 +12,8 @@ from executorch.backends.arm.test import common from executorch.backends.arm.test.tester.test_pipeline import TosaPipelineFP +from executorch.exir import EdgeCompileConfig, to_edge +from executorch.exir.dialects._ops import ops as exir_ops from torch.fx import Graph, GraphModule input_t1 = Tuple[torch.Tensor] # Input x @@ -119,6 +121,106 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: pipeline.run() +@pytest.mark.parametrize( + "arg_op", + [torch.argmax, torch.argmin], + ids=["argmax", "argmin"], +) +def test_arg_op_safe_edge_scalar_constant_is_cast_to_int32(arg_op): + class SafeScalarArithmetic(torch.nn.Module): + def forward(self, x: torch.Tensor): + return arg_op(x, dim=1) * 10 + + module = SafeScalarArithmetic() + test_input = torch.randn(2, 8) + exported_program = to_edge( + torch.export.export(module, (test_input,)), + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ).exported_program() + + result = ConvertInt64OutputOpsToInt32Pass().call(exported_program.graph_module) + + mul = next( + node + for node in result.graph_module.graph.nodes + if node.target == exir_ops.edge.aten.mul.Tensor + ) + assert mul.args[0].meta["val"].dtype == torch.int32 + assert mul.args[1].meta["val"].dtype == torch.int32 + + actual = result.graph_module(torch.tensor(10), test_input)[0] + expected = module(test_input) + assert actual.dtype == torch.int64 + torch.testing.assert_close(actual, expected) + + +@pytest.mark.parametrize( + "arg_op", + [torch.argmax, torch.argmin], + ids=["argmax", "argmin"], +) +@pytest.mark.parametrize("use_edge_ops", [False, True], ids=["aten", "edge"]) +def test_arg_op_unsafe_arithmetic_stays_int64(arg_op, use_edge_ops: bool): + class UnsafeArithmetic(torch.nn.Module): + def forward(self, x: torch.Tensor): + indices = arg_op(x, dim=1).unsqueeze(-1) + return indices * indices, indices + + module = UnsafeArithmetic() + test_input = torch.zeros(1, 50001) + test_input[0, -1] = 1 if arg_op is torch.argmax else -1 + exported_program = torch.export.export(module, (test_input,)) + if use_edge_ops: + exported_program = to_edge( + exported_program, + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ).exported_program() + + result = ConvertInt64OutputOpsToInt32Pass().call(exported_program.graph_module) + + mul_targets = { + torch.ops.aten.mul.Tensor, + exir_ops.edge.aten.mul.Tensor, + } + relay_targets = { + torch.ops.aten.unsqueeze.default, + exir_ops.edge.aten.unsqueeze_copy.default, + } + mul = next( + node for node in result.graph_module.graph.nodes if node.target in mul_targets + ) + relay = next( + node for node in result.graph_module.graph.nodes if node.target in relay_targets + ) + assert mul.args[0].meta["val"].dtype == torch.int64 + assert mul.args[1].meta["val"].dtype == torch.int64 + assert relay.args[0].meta["val"].dtype == torch.int32 + + expected = module(test_input) + actual = result.graph_module(test_input) + assert actual[0].item() == 2_500_000_000 + assert actual[1].dtype == torch.int64 + for actual_output, expected_output in zip(actual, expected, strict=True): + torch.testing.assert_close(actual_output, expected_output) + + +@pytest.mark.parametrize( + "arg_op", + [torch.argmax, torch.argmin], + ids=["argmax", "argmin"], +) +def test_arg_op_direct_output_is_unchanged(arg_op): + class DirectOutput(torch.nn.Module): + def forward(self, x: torch.Tensor): + return arg_op(x, dim=1) + + exported_program = torch.export.export(DirectOutput(), (torch.randn(2, 8),)) + result = ConvertInt64OutputOpsToInt32Pass().call(exported_program.graph_module) + + assert not result.modified + assert result.graph_module(torch.randn(2, 8))[0].dtype == torch.int64 + + ############################################################## ## Test on_overflow range check for argmax/argmin ## ############################################################## From 7105b90f5cd5cd70bd95b311b438d8adcf24e1f4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=C3=A5ns=20Nilsson?= Date: Wed, 26 Aug 2026 13:54:30 +0200 Subject: [PATCH 2/2] Arm backend: Allow int32 argmax and argmin indices for index.Tensor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Måns Nilsson Change-Id: I4cafb654c315f77dab4b80a20de669cb0c161610 --- backends/arm/_passes/convert_int64_output_ops_to_int32.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/backends/arm/_passes/convert_int64_output_ops_to_int32.py b/backends/arm/_passes/convert_int64_output_ops_to_int32.py index c269d4aa396..e9b3a5b96fd 100644 --- a/backends/arm/_passes/convert_int64_output_ops_to_int32.py +++ b/backends/arm/_passes/convert_int64_output_ops_to_int32.py @@ -94,6 +94,8 @@ def _index_range(self, node: torch.fx.Node) -> Tuple[int, int]: aten_bounded_index_ops = aten_argmax_ops + aten_argmin_ops edge_bounded_index_ops = edge_argmax_ops + edge_argmin_ops + aten_index_consumer_ops = (torch.ops.aten.index.Tensor,) + edge_index_consumer_ops = (exir_ops.edge.aten.index.Tensor,) aten_index_relay_ops = ( torch.ops.aten.unsqueeze.default, @@ -311,6 +313,8 @@ def _is_safe_widening_consumer( return node.target in ( self.aten_cast_ops + self.edge_cast_ops + + self.aten_index_consumer_ops + + self.edge_index_consumer_ops + self.aten_index_add_ops + self.edge_index_add_ops + self.aten_index_sub_ops