Arm backend: Add aten.select to "view" operator support check (#12672)

Select decomposes into squeeze, which in turn becomes a view. Therefore,
perform the same check on select operators as view operators.

Signed-off-by: Sebastian Larsson <[email protected]>

Author: Sebastian-Larsson

backends/arm/operator_support/ethos_u55_support.py

@@ -6,15 +6,16 @@
 # pyre-unsafe
 
 import typing
+from typing import cast
 
 import torch
 import torch.fx as fx
+
 from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
 from executorch.backends.arm._passes.insert_table_ops import TableOps
 from executorch.backends.arm.operators.op_permute import transform_permutation_vector
 from executorch.backends.arm.tosa_utils import tosa_shape
 from executorch.exir.backend.utils import WhyNoPartitionReporter
-
 from executorch.exir.dialects._ops import ops as exir_ops
 from torch.fx.passes.operator_support import OperatorSupportBase
 
@@ -212,19 +213,51 @@ def is_node_supported(
         Returns:
             False if the operator is not support and True if it is supported.
         """
-        if not node.target == exir_ops.edge.aten.view_copy.default:
+        # Select decomposes into squeeze, which in turn becomes a view. Therefore,
+        # perform the same check on select operators as view operators.
+        if node.target not in (
+            exir_ops.edge.aten.view_copy.default,
+            exir_ops.edge.aten.select.int,
+            exir_ops.edge.aten.select_copy.int,
+        ):
             return True
 
-        shape = list(get_first_fake_tensor(node).shape)
+        if node.target in (
+            exir_ops.edge.aten.select.int,
+            exir_ops.edge.aten.select_copy.int,
+        ):
+            input_node, dim, index = cast(tuple[fx.Node, int, int], node.args)
+
+            shape = input_node.meta["val"].shape
+            rank = len(shape)
+            if not -rank <= dim < rank:
+                raise IndexError(
+                    f"Dim {dim} is outside of the range for tensor '{node.target}' of "
+                    f"rank {rank}"
+                )
+            dim = dim % rank
+
+            size = shape[dim]
+            if not -size <= index < size:
+                raise IndexError(
+                    f"Index {index} is outside of the range for dim {dim} with size "
+                    f"{size} for tensor {node.target}"
+                )
+            index = index % size
+
+            # Shape after squeeze. This may get converted into a view which may become
+            # a transpose. This is why we're checking select.
+            squeezed_shape = shape[:dim] + shape[dim + 1 :]
+            shape = squeezed_shape
+        else:
+            shape = list(get_first_fake_tensor(node).shape)
+
         dtype = _try_determine_dtype(node)
-        permutation = list(typing.cast(list[int], node.args[1]))
 
         rank = len(shape)
         if rank > 4:
             if dtype == torch.int32:
-                self.reporter.report_reject(
-                    node, f"No support for {permutation=} in int32."
-                )
+                self.reporter.report_reject(node, "No support for rank > 4 in int32.")
                 return False
 
         if dtype in (torch.int8, torch.int16):

backends/arm/test/ops/test_select.py

@@ -13,6 +13,7 @@
 from executorch.backends.arm.test.tester.test_pipeline import (
     EthosU55PipelineBI,
     EthosU85PipelineBI,
+    OpNotSupportedPipeline,
     TosaPipelineBI,
     TosaPipelineMI,
 )
@@ -32,6 +33,10 @@
     "select3d_0_dim_1_index": lambda: (torch.arange(-16, 16, 0.2), 0, 1),
 }
 
+test_data_not_delegated = {
+    "select3d_large_after_squeeze": lambda: (torch.rand(3, 64, 3, 49, 32), 0, 0),
+}
+
 aten_op_copy = "torch.ops.aten.select_copy.int"
 aten_op_int = "torch.ops.aten.select.int"
 
@@ -129,6 +134,19 @@ def test_select_int_u55_BI(test_data: Tuple):
     pipeline.run()
 
 
+@common.parametrize("test_data", test_data_not_delegated)
+def test_select_int_u55_BI_not_delegated(test_data: Tuple):
+    pipeline = OpNotSupportedPipeline[input_t1](
+        SelectInt(),
+        test_data(),
+        {aten_op_copy: 0},
+        n_expected_delegates=0,
+        quantize=True,
+        u55_subset=True,
+    )
+    pipeline.run()
+
+
 @common.parametrize("test_data", test_data_suite, x_fails)
 @common.XfailIfNoCorstone320
 def test_select_int_u85_BI_copy(test_data: Tuple):