Arm backend: Add missing tests/fixes for U55/U85. (pytorch#23052)

This PR adds fixes and tests, so the tolerance in the script for
automatic list of supported ops is zero.
The script is a proper coverage tool now.

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

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

Author: wwwind

backends/arm/scripts/docgen/generate_op_support.py

@@ -103,7 +103,7 @@ class BackendConfig:
         # operator is unsupported. They should still count as static evidence.
         infrastructure_xfail_markers=frozenset({"XfailIfNoCorstone300"}),
         filter_u55_unsupported_ops=True,
-        max_missing_profile_cells=4,
+        max_missing_profile_cells=0,
     ),
     "u85": BackendConfig(
         key="u85",
@@ -119,7 +119,7 @@ class BackendConfig:
         # A missing Corstone-320 FVP is infrastructure, not evidence that the
         # operator itself is unsupported.
         infrastructure_xfail_markers=frozenset({"XfailIfNoCorstone320"}),
-        max_missing_profile_cells=3,
+        max_missing_profile_cells=0,
     ),
 }
 
@@ -355,6 +355,7 @@ def _activate_backend(backend: str | BackendConfig) -> BackendConfig:
     "torch.ops.aten.native_group_norm.default",
     "torch.ops.aten.native_layer_norm.default",
     "torch.ops.aten.embedding.default",
+    "torch.ops.aten.pow.Tensor_Tensor",
 }
 
 
@@ -421,6 +422,12 @@ def _activate_backend(backend: str | BackendConfig) -> BackendConfig:
     ): {
         "INT": {"torch.ops.aten.index_select.default"},
     },
+    (
+        "backends/arm/test/ops/test_index_tensor.py",
+        "test_index_tensor_u55_INT_constant",
+    ): {
+        "INT": {"torch.ops.aten.index.Tensor"},
+    },
     (
         "backends/arm/test/ops/test_silu.py",
         "test_silu_u55_INT",
@@ -433,11 +440,19 @@ def _activate_backend(backend: str | BackendConfig) -> BackendConfig:
     ): {
         "INT": {"torch.ops.aten.unfold_copy.default"},
     },
+    (
+        "backends/arm/test/ops/test_pow.py",
+        "test_pow_tensor_tensor_u55_INT",
+    ): {
+        "INT": {"torch.ops.aten.pow.Tensor_Tensor"},
+    },
 }
 
 # Existing U85 runtime tests below intentionally suppress direct ATen/Edge
 # assertions because quantization/decomposition changes the graph. They still
 # provide positive runtime coverage for the exported operator.
+# The test for embedding op has pipeline.pop_stage("check.aten"),
+# that is why we need to set it explicitly.
 U85_EXPLICIT_BACKEND_COVERAGE: dict[tuple[str, str], dict[str, set[str]]] = {
     (
         "backends/arm/test/ops/test_div_tensor_mode.py",
@@ -451,6 +466,18 @@ def _activate_backend(backend: str | BackendConfig) -> BackendConfig:
     ): {
         "INT": {"torch.ops.aten.silu.default"},
     },
+    (
+        "backends/arm/test/ops/test_embedding.py",
+        "test_embedding_u85_INT",
+    ): {
+        "INT": {"torch.ops.aten.embedding.default"},
+    },
+    (
+        "backends/arm/test/ops/test_pow.py",
+        "test_pow_tensor_tensor_u85_INT",
+    ): {
+        "INT": {"torch.ops.aten.pow.Tensor_Tensor"},
+    },
 }
 
 

backends/arm/test/ops/test_argmax.py

@@ -8,6 +8,8 @@
 import torch
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.test_pipeline import (
+    EthosU55PipelineINT,
+    EthosU85PipelineINT,
     OpNotSupportedPipeline,
     TosaPipelineFP,
     TosaPipelineINT,
@@ -65,6 +67,21 @@ def forward(self, x: torch.Tensor):
         ),
     }
 
+    test_data_u55_int: dict[str, Tuple[input_t, int]] = {
+        "rank_1_dim_0_int8": lambda: (
+            (torch.randint(-128, 127, (10,), dtype=torch.int8),),
+            0,
+        ),
+        "rank_2_dim_1_int8": lambda: (
+            (torch.randint(-128, 127, (2, 5), dtype=torch.int8),),
+            1,
+        ),
+        "rank_4_dim_3_int8": lambda: (
+            (torch.randint(-128, 127, (1, 3, 4, 5), dtype=torch.int8),),
+            3,
+        ),
+    }
+
 
 class ArgmaxAll(torch.nn.Module):
     def forward(self, x: torch.Tensor):
@@ -137,6 +154,32 @@ def test_argmax_tosa_INT(test_data: Tuple[input_t, int]):
     pipeline.run()
 
 
+@common.parametrize("test_data", Argmax.test_data_u55_int)
+@common.XfailIfNoCorstone300
+def test_argmax_u55_INT(test_data: Tuple[input_t, int]):
+    data, dim = test_data()
+    pipeline = EthosU55PipelineINT[input_t](
+        Argmax(dim),
+        data,
+        aten_op,
+        [exir_op, to_copy_exir_op],
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", Argmax.test_data_int)
+@common.XfailIfNoCorstone320
+def test_argmax_u85_INT(test_data: Tuple[input_t, int]):
+    data, dim = test_data()
+    pipeline = EthosU85PipelineINT[input_t](
+        Argmax(dim),
+        data,
+        aten_op,
+        [exir_op, to_copy_exir_op],
+    )
+    pipeline.run()
+
+
 @common.parametrize(
     "test_data",
     Argmax.test_data | Argmax.test_data_fp16 | Argmax.test_data_bf16,

backends/arm/test/ops/test_pow.py

@@ -356,6 +356,21 @@ def test_pow_tensor_tensor_u55_INT(
     pipeline.run()
 
 
+@common.parametrize("test_data", Pow_TensorTensor.test_data_quant)
+@common.XfailIfNoCorstone320
+def test_pow_tensor_tensor_u85_INT(
+    test_data: Pow_TensorTensor.input_t,
+):
+    pipeline = EthosU85PipelineINT[Pow_TensorTensor.input_t](
+        Pow_TensorTensorPositiveBase(),
+        test_data(),
+        POW_TENSOR_TENSOR_INT_ATEN_OPS,
+        POW_TENSOR_TENSOR_INT_EXIR_OPS,
+        qtol=4,
+    )
+    pipeline.run()
+
+
 def test_pow_tensor_tensor_int_rejects_base_lower_bound_quantized_to_zero():
     """A positive FP lower bound must stay positive after LOG input QDQ."""
     # abs(x) + 0.25 is structurally >= 0.25, but this calibration range makes

docs/source/backends/arm-ethos-u/U55_op_support.md

@@ -6,7 +6,7 @@ This page lists Ethos-U55-supported PyTorch APIs and the dtype and quantization
 
 `8x8` means 8-bit activations and 8-bit weights. `16x8` means 16-bit activations and 8-bit weights. `8x4` means 8-bit activations and 4-bit weights.
 
-Total supported PyTorch APIs: **101**.
+Total supported PyTorch APIs: **103**.
 
 | PyTorch API | Support profile | DType | Quantization mode |
 | --- | --- | --- | --- |
@@ -18,6 +18,7 @@ Total supported PyTorch APIs: **101**.
 | `torch.amax` | INT | `INT8` | 8x8 |
 | `torch.amin` | INT | `INT8` | 8x8 |
 | `torch.arange` | INT | `INT8` | 8x8 |
+| `torch.argmax` | INT | `INT8` | 8x8 |
 | `torch.as_strided_copy` | INT | `INT8` | 8x8 |
 | `torch.asin` | INT | `INT8` | 8x8 |
 | `torch.asinh` | INT | `INT8` | 8x8 |
@@ -101,6 +102,7 @@ Total supported PyTorch APIs: **101**.
 | `torch.t` / `torch.Tensor.t` | INT | `INT8` | 8x8 |
 | `torch.tan` | INT | `INT8` | 8x8 |
 | `torch.tanh` / `torch.nn.Tanh` | INT | `INT8`, `INT16` | 8x8, 16x8 |
+| `torch.Tensor.__getitem__` / `tensor indexing` | INT | `INT8` | 8x8 |
 | `torch.Tensor.__getitem__` / `tensor slicing` | INT | `INT8`, `INT16` | 8x8, 16x8 |
 | `torch.Tensor.copy_` | INT | `INT8` | 8x8 |
 | `torch.Tensor.expand` | INT | `INT8` | 8x8 |

docs/source/backends/arm-ethos-u/U85_op_support.md

@@ -6,7 +6,7 @@ This page lists Ethos-U85-supported PyTorch APIs and the dtype and quantization
 
 `8x8` means 8-bit activations and 8-bit weights. `16x8` means 16-bit activations and 8-bit weights. `8x4` means 8-bit activations and 4-bit weights.
 
-Total supported PyTorch APIs: **124**.
+Total supported PyTorch APIs: **126**.
 
 | PyTorch API | Support profile | DType | Quantization mode |
 | --- | --- | --- | --- |
@@ -20,6 +20,7 @@ Total supported PyTorch APIs: **124**.
 | `torch.amin` | INT | `INT8`, `INT16` | 8x8, 16x8 |
 | `torch.any` | INT | `INT8` | 8x8 |
 | `torch.arange` | INT | `INT8` | 8x8 |
+| `torch.argmax` | INT | `INT8` | 8x8 |
 | `torch.as_strided_copy` | INT | `INT8` | 8x8 |
 | `torch.asin` | INT | `INT8` | 8x8 |
 | `torch.asinh` | INT | `INT8` | 8x8 |
@@ -84,6 +85,7 @@ Total supported PyTorch APIs: **124**.
 | `torch.nn.AvgPool2d` / `torch.nn.functional.avg_pool2d` | INT | `INT8`, `INT16` | 8x8, 16x8 |
 | `torch.nn.Conv2d` / `torch.nn.functional.conv2d` | INT | `INT8`, `INT16`, `INT4` | 8x8, 8x4, 16x8 |
 | `torch.nn.ELU` / `torch.nn.functional.elu` | INT | `INT8` | 8x8 |
+| `torch.nn.Embedding` / `torch.nn.functional.embedding` | INT | `INT8` | 8x8 |
 | `torch.nn.functional.pad` | INT | `INT8` | 8x8 |
 | `torch.nn.GELU` / `torch.nn.functional.gelu` | INT | `INT8`, `INT16` | 8x8, 16x8 |
 | `torch.nn.Hardsigmoid` / `torch.nn.functional.hardsigmoid` | INT | `INT8` | 8x8 |