Arm backend: Add index_tensor/index_select to unsupported u55 operators (pytorch#13176)

index_tensor and index_select both become a GATHER Tosa operators, which
is not support on U55.

Add test cases for U55 for these operators. The test cases verify that
the ops do not get delegated.

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

Author: Sebastian-Larsson

backends/arm/operator_support/ethos_u55_support.py

@@ -149,6 +149,8 @@ class EthosU55NotSupported(OperatorSupportBase):
         exir_ops.edge.aten.ne.Scalar,
         exir_ops.edge.aten.flip.default,  # REVERSE
         exir_ops.edge.aten.grid_sampler_2d,  # GATHER
+        exir_ops.edge.aten.index.Tensor,  # GATHER
+        exir_ops.edge.aten.index_select.default,  # GATHER
         exir_ops.edge.aten.scatter.src,
         exir_ops.edge.aten.scatter.value,
         exir_ops.edge.aten.select_scatter.default,

backends/arm/test/ops/test_index_select.py

@@ -12,6 +12,7 @@
 
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.test_pipeline import (
+    OpNotSupportedPipeline,
     TosaPipelineFP,
     TosaPipelineINT,
     VgfPipeline,
@@ -120,6 +121,20 @@ def test_index_select_tosa_INT_rand(test_data: input_params):
     pipeline.run()
 
 
+@pytest.mark.parametrize("test_data", list(test_data.values())[-1:])
+def test_index_select_u55_INT_not_delegated(test_data: input_params):
+    op, test_input = test_data
+
+    pipeline = OpNotSupportedPipeline[input_params](
+        op,
+        test_input,
+        {op.exir_op: 1},
+        quantize=True,
+        u55_subset=True,
+    )
+    pipeline.run()
+
+
 @pytest.mark.parametrize("test_data", list(test_data.values()))
 @common.SkipIfNoModelConverter
 def test_index_select_vgf_FP(test_data: input_params):

backends/arm/test/ops/test_index_tensor.py

@@ -10,6 +10,7 @@
 import torch
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.test_pipeline import (
+    OpNotSupportedPipeline,
     TosaPipelineFP,
     TosaPipelineINT,
 )
@@ -460,3 +461,18 @@ def test_index_tensor_tosa_INT_none(test_data: input_params):
                 IndexTensorTestCommon.exir_op,
             ).run()
         )
+
+
+@common.parametrize("test_data", IndexTensor.test_data)
+@common.XfailIfNoCorstone300
+def test_index_tensor_u55_INT_not_delegated(test_data: input_params):
+    """Ethos-U55 backend BI pipeline test for index.Tensor"""
+    test_input = test_data
+    with torch.no_grad():
+        OpNotSupportedPipeline[input_params](
+            IndexTensor(),
+            test_input,
+            {IndexTensorTestCommon.exir_op: 1},
+            quantize=True,
+            u55_subset=True,
+        ).run()