Arm backend: Fix lifted int64 index constants. (pytorch#24794)

CastInt64BuffersToInt32Pass previously handled buffers but not tensor
constants lifted by torch.export.

Extend the pass to convert lifted int64 constants used exclusively as
indexing operands, provided their values fit within int32.

Preserve the original dtype for constants that are also used as data or
other non-index operands.



cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils
@Sebastian-Larsson @robell @rascani

Signed-off-by: Elena Zhelezina <elena.zhelezina@arm.com>

Author: wwwind

backends/arm/_passes/cast_int64_pass.py

@@ -9,36 +9,94 @@
 
 import torch
 from executorch.backends.arm._passes.arm_pass import ArmPass
+from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
-from torch._export.utils import is_buffer
+from torch._export.utils import is_buffer, is_lifted_tensor_constant
 from torch.export import ExportedProgram
 
 logger = logging.getLogger(__name__)
 
 
 class CastInt64BuffersToInt32Pass(ArmPass):
-    """Cast int64 buffers to int32 if the int64 data is in int32 range."""
+    """Cast int64 buffers and lifted tensor constants to int32 when in range."""
 
     _passes_required_after: Set[Type[ExportPass]] = set()
 
     def __init__(self, exported_program: ExportedProgram, *args, **kwargs):
         super().__init__(*args, **kwargs)
         self.exported_program = exported_program
 
-    def should_run_pass(self, graph_module: torch.fx.GraphModule) -> bool:
-        for node in graph_module.graph.nodes:
-            if len(node.users) == 0:
-                continue
-            fake_tensor = node.meta.get("val")
-            if not isinstance(fake_tensor, torch._subclasses.fake_tensor.FakeTensor):
-                continue
-            if fake_tensor.dtype == torch.int64 and is_buffer(
-                self.exported_program, node
-            ):
+    @staticmethod
+    def _is_index_tensor_input(node: torch.fx.Node) -> bool:
+        """Return True only when a constant is used exclusively as an index.
+
+        ``node.users`` tracks consumer operations, not operand occurrences.
+        A consumer can use this node as both indexed data and an index.
+
+        """
+        index_ops = {
+            torch.ops.aten.index.Tensor,
+            exir_ops.edge.aten.index.Tensor,
+        }
+
+        def contains_node(value: object) -> bool:
+            if value is node:
                 return True
-        return False
+            if isinstance(value, (tuple, list)):
+                return any(contains_node(item) for item in value)
+            if isinstance(value, dict):
+                return any(contains_node(item) for item in value.values())
+            if isinstance(value, slice):
+                return any(
+                    contains_node(item)
+                    for item in (value.start, value.stop, value.step)
+                )
+            return False
+
+        if not node.users:
+            return False
+        for user in node.users:
+            if (
+                user.op != "call_function"
+                or user.target not in index_ops
+                or len(user.args) < 2
+                or not isinstance(user.args[1], (list, tuple))
+                or not any(index is node for index in user.args[1])
+            ):
+                return False
+
+            # args[1] is the index-list operand. All other uses are unsafe.
+            if any(
+                contains_node(argument)
+                for i, argument in enumerate(user.args)
+                if i != 1
+            ) or contains_node(user.kwargs):
+                return False
+        return True
+
+    def _is_cast_candidate(self, node: torch.fx.Node) -> bool:
+        """Select int64 buffers and index-only lifted tensor constants."""
+        if not node.users:
+            return False
+        fake_tensor = node.meta.get("val")
+        if not isinstance(fake_tensor, torch._subclasses.fake_tensor.FakeTensor):
+            return False
+        if fake_tensor.dtype != torch.int64:
+            return False
+        if is_buffer(self.exported_program, node):
+            return True
+        return is_lifted_tensor_constant(
+            self.exported_program, node
+        ) and self._is_index_tensor_input(node)
+
+    def should_run_pass(self, graph_module: torch.fx.GraphModule) -> bool:
+        return any(self._is_cast_candidate(node) for node in graph_module.graph.nodes)
 
     def _assert_within_int32(self, tensor: torch.Tensor, node: torch.fx.Node):
+        # An empty tensor has no values outside the int32 range.
+        if tensor.numel() == 0:
+            return
+
         if torch.min(tensor) < torch.iinfo(torch.int32).min:
             raise RuntimeError(
                 f"Node {node.name} has value < {torch.iinfo(torch.int32).min}"
@@ -48,44 +106,42 @@ def _assert_within_int32(self, tensor: torch.Tensor, node: torch.fx.Node):
                 f"Node {node.name} has value > {torch.iinfo(torch.int32).max}"
             )
 
-    def _cast_buffer_to_int32(self, node: torch.fx.Node) -> bool:
-        buffer_name = self.exported_program.graph_signature.inputs_to_buffers[node.name]
-        if buffer_name in self.exported_program.state_dict:
-            store = self.exported_program.state_dict
-        elif buffer_name in self.exported_program.constants:
-            # Non-persistent buffers are tracked by inputs_to_buffers, but their
-            # tensor values live in ExportedProgram.constants instead of the
-            # state_dict. Examples include transformer position_ids buffers.
+    def _cast_constant_to_int32(self, node: torch.fx.Node) -> bool:
+        signature = self.exported_program.graph_signature
+        if is_buffer(self.exported_program, node):
+            constant_name = signature.inputs_to_buffers[node.name]
+            # Persistent buffers live in state_dict, non-persistent ones in constants.
+            store = (
+                self.exported_program.state_dict
+                if constant_name in self.exported_program.state_dict
+                else self.exported_program.constants
+            )
+        elif is_lifted_tensor_constant(self.exported_program, node):
+            constant_name = signature.inputs_to_lifted_tensor_constants[node.name]
             store = self.exported_program.constants
         else:
+            return False
+
+        if constant_name not in store:
             logger.warning(
-                "Skipping int64 buffer %s (%s): value not found in state_dict "
-                "or constants",
+                "Skipping int64 constant %s (%s): value not found in its storage",
                 node.name,
-                buffer_name,
+                constant_name,
             )
             return False
 
-        buffer = store[buffer_name]
-        self._assert_within_int32(buffer, node)
-        store[buffer_name] = buffer.to(torch.int32)
+        value = store[constant_name]
+        # Do not truncate values outside the signed int32 range.
+        self._assert_within_int32(value, node)
+        store[constant_name] = value.to(torch.int32)
         node.meta["val"] = node.meta["val"].to(torch.int32)
         return True
 
     def _to_int32(self, graph_module: torch.fx.GraphModule) -> int:
         converted = 0
         for node in graph_module.graph.nodes:
-            if len(node.users) == 0:
-                continue
-            if "val" not in node.meta:
-                continue
-            fake_tensor = node.meta["val"]
-            if not isinstance(fake_tensor, torch._subclasses.fake_tensor.FakeTensor):
-                continue
-            if fake_tensor.dtype != torch.int64:
-                continue
-            if is_buffer(self.exported_program, node):
-                converted += self._cast_buffer_to_int32(node)
+            if self._is_cast_candidate(node):
+                converted += self._cast_constant_to_int32(node)
         return converted
 
     def call(self, graph_module: torch.fx.GraphModule):
@@ -94,7 +150,8 @@ def call(self, graph_module: torch.fx.GraphModule):
         if modified:
             graph_module = super().call(graph_module).graph_module
             logger.warning(
-                "CastInt64BuffersToInt32Pass: cast %d int64 buffer(s) to int32.",
+                "CastInt64BuffersToInt32Pass: cast %d int64 buffer(s) or "
+                "lifted tensor constant(s) to int32.",
                 converted,
             )
         return PassResult(graph_module, modified)

backends/arm/test/passes/test_cast_int64_pass.py

@@ -5,6 +5,7 @@
 
 from typing import Tuple
 
+import pytest
 import torch
 from executorch.backends.arm._passes import CastInt64BuffersToInt32Pass
 
@@ -79,3 +80,149 @@ def test_cast_non_persistent_int64_buffers_to_int32_tosa_FP(test_data: input_t):
         StageType.RUN_PASSES
     ).exported_program()
     assert exported_program.constants["position_ids"].dtype == torch.int32
+
+
+class LiftedInt64IndexModel(torch.nn.Module):
+    """A plain tensor attribute is lifted as CONSTANT_TENSOR, not BUFFER."""
+
+    def __init__(self):
+        super().__init__()
+        self.indices = torch.tensor([0, 2], dtype=torch.int64)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x[:, self.indices]
+
+
+def test_cast_lifted_int64_index_constant_to_int32():
+    """Lifted indices must reach gather as int32."""
+    inputs = (torch.arange(32, dtype=torch.float32).reshape(1, 4, 8),)
+    pipeline = PassPipeline[input_t](
+        LiftedInt64IndexModel(),
+        inputs,
+        quantize=False,
+        passes_with_exported_program=[CastInt64BuffersToInt32Pass],
+    )
+    # Also compares the output of the transformed Edge graph to the original.
+    pipeline.run()
+
+    ep = pipeline.tester.get_artifact(StageType.RUN_PASSES).exported_program()
+    lifted = ep.graph_signature.inputs_to_lifted_tensor_constants
+    assert len(lifted) == 1
+    placeholder_name, constant_name = next(iter(lifted.items()))
+    assert ep.constants[constant_name].dtype == torch.int32
+    placeholder = next(
+        node for node in ep.graph_module.graph.nodes if node.name == placeholder_name
+    )
+    assert placeholder.meta["val"].dtype == torch.int32
+
+
+@pytest.mark.parametrize(
+    "bad_index,expected_error",
+    [
+        (2**31, "value >"),
+        (-(2**31) - 1, "value <"),
+    ],
+)
+def test_lifted_int64_constant_out_of_int32_range_is_not_truncated(
+    bad_index: int, expected_error: str
+):
+    from executorch.exir import to_edge
+
+    inputs = (torch.randn(1, 4, 8),)
+    ep = to_edge(
+        torch.export.export(LiftedInt64IndexModel(), inputs)
+    ).exported_program()
+    lifted = ep.graph_signature.inputs_to_lifted_tensor_constants
+    assert len(lifted) == 1
+    _, constant_name = next(iter(lifted.items()))
+    # Modify after export: an out-of-bounds index need not be run during tracing.
+    ep.constants[constant_name] = torch.tensor([0, bad_index], dtype=torch.int64)
+
+    with pytest.raises(RuntimeError, match=expected_error):
+        CastInt64BuffersToInt32Pass(ep)(ep.graph_module)
+    assert ep.constants[constant_name].dtype == torch.int64
+
+
+def test_lifted_int64_constant_used_as_data_and_index_is_not_narrowed():
+    """The same lifted constant is indexed data and index."""
+
+    class SelfIndexModel(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.values = torch.tensor([1, 0], dtype=torch.int64)
+
+        def forward(self) -> torch.Tensor:
+            return self.values[self.values]
+
+    model = SelfIndexModel().eval()
+    expected = model()
+    ep = torch.export.export(model, ())
+    lifted = ep.graph_signature.inputs_to_lifted_tensor_constants
+    assert len(lifted) == 1
+    placeholder_name, constant_name = next(iter(lifted.items()))
+    node = next(n for n in ep.graph_module.graph.nodes if n.name == placeholder_name)
+    users = [user for user in node.users if user.target == torch.ops.aten.index.Tensor]
+    assert len(users) == 1
+    assert users[0].args[0] is node
+    assert any(index is node for index in users[0].args[1])
+
+    cast_pass = CastInt64BuffersToInt32Pass(ep)
+    assert not cast_pass._is_index_tensor_input(node)
+    assert not cast_pass.should_run_pass(ep.graph_module)
+    assert cast_pass._to_int32(ep.graph_module) == 0
+    assert ep.constants[constant_name].dtype == torch.int64
+    assert node.meta["val"].dtype == torch.int64
+    result = ep.module()()
+    assert result.dtype == torch.int64
+    torch.testing.assert_close(result, expected)
+
+
+def test_lifted_int64_index_rejects_other_operand_occurrences():
+    """The predicate must reject nested positional/keyword operand reuse."""
+    for location in ("indexed_data", "positional", "keyword"):
+        graph = torch.fx.Graph()
+        data = graph.placeholder("data")
+        index = graph.placeholder("index")
+        args = (data, [index])
+        kwargs = {}
+        if location == "indexed_data":
+            args = (index, [index])
+        elif location == "positional":
+            args += ({"nested": [index]},)
+        else:
+            kwargs["extra"] = {"nested": (index,)}
+        graph.call_function(torch.ops.aten.index.Tensor, args=args, kwargs=kwargs)
+        assert not CastInt64BuffersToInt32Pass._is_index_tensor_input(index)
+
+    graph = torch.fx.Graph()
+    data = graph.placeholder("data")
+    index = graph.placeholder("index")
+    graph.call_function(torch.ops.aten.index.Tensor, args=(data, [index]))
+    assert CastInt64BuffersToInt32Pass._is_index_tensor_input(index)
+
+
+def test_cast_empty_lifted_int64_index_constant_to_int32():
+    """Empty lifted indices are valid and must convert without reductions."""
+    from executorch.exir import to_edge
+
+    model = LiftedInt64IndexModel()
+    model.indices = torch.empty(0, dtype=torch.int64)
+    inputs = (torch.arange(32, dtype=torch.float32).reshape(1, 4, 8),)
+    ep = to_edge(torch.export.export(model, inputs)).exported_program()
+
+    lifted = ep.graph_signature.inputs_to_lifted_tensor_constants
+    assert len(lifted) == 1
+    placeholder_name, constant_name = next(iter(lifted.items()))
+    assert ep.constants[constant_name].numel() == 0
+    assert ep.constants[constant_name].dtype == torch.int64
+
+    cast_pass = CastInt64BuffersToInt32Pass(ep)
+    assert cast_pass.should_run_pass(ep.graph_module)
+    cast_pass(ep.graph_module)
+
+    assert ep.constants[constant_name].numel() == 0
+    assert ep.constants[constant_name].dtype == torch.int32
+    placeholder = next(
+        node for node in ep.graph_module.graph.nodes if node.name == placeholder_name
+    )
+    assert placeholder.meta["val"].dtype == torch.int32