Arm backend: Fix int64 buffers before index gather (#20764)

Move CastInt64BuffersToInt32Pass before
DecomposeIndexTensorToGatherPass. This converts int64 get_attr index
buffers before index.Tensor is lowered to TOSA GATHER.

This matches Swin's relative-position lookup pattern:

  table[index:int64 buffer]
    -> table[index:int32 buffer]
    -> TOSA GATHER

Non-persistent buffers are also handled. They are present in
inputs_to_buffers, but their tensor values live in ExportedProgram
constants instead of the state_dict. This covers transformer buffers
such as position_ids.

Without this ordering, gather decomposition sees the original int64
buffer and fails with "index.Tensor requires index dtype must be int32".
Without the constants handling, non-persistent int64 buffers fail with a
KeyError while being cast.

Fixes the flow test:

  test_swin_v2_t[arm_tosa_fp-static_shapes-float32]

Signed-off-by: Zingo Andersen <Zingo.Andersen@arm.com>

Author: zingo

backends/arm/_passes/arm_pass_manager.py

@@ -547,6 +547,7 @@ def _tosa_pipeline(
                 DecomposeUnfoldToGatherPass(),
                 DecomposeEmbeddingPass(),
                 DecomposeIndexSelectToGatherPass(),
+                CastInt64BuffersToInt32Pass(exported_program),
                 DecomposeStridedSliceCopyPass(),
                 DecomposeSliceScatterPass(),
                 AccumulateIndexPutPass(),

backends/arm/_passes/cast_int64_pass.py

@@ -48,6 +48,34 @@ 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.
+            store = self.exported_program.constants
+        else:
+            logger.warning(
+                "Skipping int64 buffer %s (%s): value not found in state_dict "
+                "or constants",
+                node.name,
+                buffer_name,
+            )
+            return False
+
+        buffer = store[buffer_name]
+        self._assert_within_int32(buffer, node)
+        logger.warning(
+            f"Casting buffer {node.name} from torch.int64 to torch.int32"
+            f" defined in {node.meta.get('stack_trace','[no stack trace found]')}"
+        )
+        store[buffer_name] = buffer.to(torch.int32)
+        node.meta["val"] = node.meta["val"].to(torch.int32)
+        return True
+
     def _to_int32(self, graph_module: torch.fx.GraphModule) -> bool:
         modified = False
         for node in graph_module.graph.nodes:
@@ -61,19 +89,7 @@ def _to_int32(self, graph_module: torch.fx.GraphModule) -> bool:
             if fake_tensor.dtype != torch.int64:
                 continue
             if is_buffer(self.exported_program, node):
-                node.meta["val"] = fake_tensor.to(torch.int32)
-                buffer_name = self.exported_program.graph_signature.inputs_to_buffers[
-                    node.name
-                ]
-                buffer = self.exported_program.state_dict[buffer_name]
-                self._assert_within_int32(buffer, node)
-                logger.warning(
-                    f"Casting buffer {node.name} from torch.int64 to torch.int32"
-                    f" defined in {node.meta.get('stack_trace','[no stack trace found]')}"
-                )
-                buffer_int32 = buffer.to(torch.int32)
-                self.exported_program.state_dict[buffer_name] = buffer_int32
-                modified = True
+                modified |= self._cast_buffer_to_int32(node)
         return modified
 
     def call(self, graph_module: torch.fx.GraphModule):

backends/arm/test/ops/test_index_tensor.py

@@ -29,6 +29,33 @@ class IndexTensorTestCommon:
     AFTER = "AFTER"
 
 
+class IndexTensorInt64Buffer(torch.nn.Module):
+    """Swin-style indexing with an int64 get_attr index buffer."""
+
+    def __init__(self):
+        super().__init__()
+        self.register_buffer("index", torch.tensor([0, 2], dtype=torch.int64))
+
+    def forward(self, x: torch.Tensor):
+        return x[self.index]
+
+
+def test_index_tensor_tosa_FP_int64_buffer_index():
+    # This mirrors torchvision Swin relative_position_bias_table[index]. The
+    # int64 get_attr must be cast before index.Tensor is decomposed to GATHER:
+    #
+    #   x[index:int64 buffer] -> x[index:int32 buffer] -> TOSA GATHER
+    pipeline = TosaPipelineFP[Tuple[torch.Tensor]](
+        IndexTensorInt64Buffer(),
+        (torch.rand(4, 3),),
+        IndexTensorTestCommon.aten_op,
+        IndexTensorTestCommon.exir_op,
+        atol=IndexTensorTestCommon.atol,
+        rtol=IndexTensorTestCommon.rtol,
+    )
+    pipeline.run()
+
+
 input_params_slice = Tuple[torch.Tensor, int, int, str, Tuple[torch.Tensor]]
 input_params = Tuple[torch.Tensor, Tuple[torch.Tensor]]
 

backends/arm/test/passes/test_cast_int64_pass.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.
@@ -25,6 +25,21 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return x + 3
 
 
+class NonPersistentInt64BufferModel(torch.nn.Module):
+    test_data = {
+        "rand": (torch.rand(4),),
+    }
+
+    def __init__(self):
+        super().__init__()
+        self.register_buffer(
+            "position_ids", torch.arange(4, dtype=torch.int64), persistent=False
+        )
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x + self.get_buffer("position_ids").to(torch.float32)
+
+
 @common.parametrize("test_data", Int64Model.test_data)
 def test_cast_int64_buffers_to_int32_tosa_FP(test_data: input_t):
     module = Int64Model()
@@ -47,3 +62,20 @@ def test_cast_int64_buffers_to_int32_tosa_FP(test_data: input_t):
     ).exported_program()
     for state in exported_program.state_dict:
         assert exported_program.state_dict[state].dtype == torch.int32
+
+
+@common.parametrize("test_data", NonPersistentInt64BufferModel.test_data)
+def test_cast_non_persistent_int64_buffers_to_int32_tosa_FP(test_data: input_t):
+    module = NonPersistentInt64BufferModel()
+    pipeline = PassPipeline[input_t](
+        module,
+        test_data,
+        quantize=False,
+        passes_with_exported_program=[CastInt64BuffersToInt32Pass],
+    )
+    pipeline.run()
+
+    exported_program = pipeline.tester.get_artifact(
+        StageType.RUN_PASSES
+    ).exported_program()
+    assert exported_program.constants["position_ids"].dtype == torch.int32