Arm backend: Fix index_put broadcasted indices in Arm pass (#17247)

- Compute W from broadcasted index shapes in RewriteIndexPutPas
- Add broadcast_indices_mismatch test for index_put broadcasting

Change-Id: I1a85889754bc46807bbc3fc41846ae89ac48abd1
Signed-off-by: Rob Elliott <Robert.Elliott@arm.com>

Author: robell

backends/arm/_passes/rewrite_index_put_pass.py

@@ -6,6 +6,8 @@
 import math
 from typing import Any, Iterable, List, Sequence, Set, Type
 
+import torch
+
 from executorch.backends.arm._passes import ArmPass
 from executorch.backends.arm._passes.fuse_view_copy_transform_pass import (
     FuseViewCopyTransformPass,
@@ -189,9 +191,17 @@ def call_operator(self, op, args, kwargs, meta):
         processed_indices = self._expand_none_indices(
             source_shape, indices, plain_meta, full_op
         )
+        index_shapes = [tuple(idx.data.shape) for idx in processed_indices]
+        try:
+            broadcast_shape = torch.broadcast_shapes(*index_shapes)
+        except Exception as exc:
+            raise RuntimeError(
+                "RewriteIndexPutPass: failed to broadcast index shapes %s: %s"
+                % (index_shapes, exc)
+            ) from exc
 
         N, K, W, C = calculate_tosa_values(
-            list(processed_indices[0].data.shape),
+            list(broadcast_shape),
             [idx.node for idx in processed_indices],
             source_shape,
         )
@@ -204,6 +214,13 @@ def call_operator(self, op, args, kwargs, meta):
             full_op,
             plain_meta,
         )
+        idx_shape = list(indices_reshaped.data.shape)
+        idx_numel = math.prod(idx_shape)
+        if idx_numel != N * W:
+            raise RuntimeError(
+                "RewriteIndexPutPass: flat index numel (%s) does not match expected N*W (%s)"
+                % (idx_numel, N * W)
+            )
 
         # Scatter expects a 3D layout; flatten everything into [N, K, C].
         reshape_indices = super().call_operator(

backends/arm/test/ops/test_index_put.py

@@ -147,6 +147,19 @@
         ),
         0,  # used for u55 tests to config n_expected_delgates, only 1 when accumulate is True
     ),
+    "broadcast_indices_mismatch": (
+        lambda: (
+            torch.zeros((1, 2, 3, 4), dtype=torch.float32),
+            (
+                torch.tensor([0], dtype=torch.int64),
+                torch.tensor([0, 1], dtype=torch.int64),
+                torch.tensor([1], dtype=torch.int64),
+            ),
+            torch.randn((2, 4), dtype=torch.float32),
+            False,
+        ),
+        0,
+    ),
 }
 test_data_suite_bf16 = {
     "rank3_rand_bf16": (