Arm backend: Add tests for automatic list of supported ops (#21053)

Add tests for VGF backend 




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

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

Author: wwwind

backends/arm/test/ops/test_add.py

@@ -507,3 +507,44 @@ def test_add_tensor_u85_INT_16a8w(test_data: input_t1):
         get_symmetric_a16w8_quantization_config(is_per_channel=per_channel_quantization)
     )
     pipeline.run()
+
+
+aten_op_scalar = "torch.ops.aten.add.Scalar"
+exir_op_scalar = "executorch_exir_dialects_edge__ops_aten_add_Scalar"
+
+
+class AddScalar(torch.nn.Module):
+    def forward(self, x: torch.Tensor):
+        return torch.ops.aten.add.Scalar(x, 1.5)
+
+
+@common.parametrize(
+    "test_data", Add.test_data | Add.test_data_fp16 | Add.test_data_bf16
+)
+@common.SkipIfNoModelConverter
+def test_add_scalar_vgf_no_quant(test_data: input_t1):
+    pipeline = VgfPipeline[input_t1](
+        AddScalar(),
+        test_data(),
+        aten_op_scalar,
+        exir_op_scalar,
+        run_on_vulkan_runtime=True,
+        quantize=False,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", Add.test_data)
+@common.SkipIfNoModelConverter
+def test_add_scalar_vgf_quant(test_data: input_t1):
+    # Quantized scalar add is lowered through the Tensor overload.
+    # We keep this as add.Tensor rather than add.Scalar.
+    pipeline = VgfPipeline[input_t1](
+        AddScalar(),
+        test_data(),
+        aten_op,
+        exir_op,
+        run_on_vulkan_runtime=True,
+        quantize=True,
+    )
+    pipeline.run()

backends/arm/test/ops/test_div.py

@@ -8,6 +8,7 @@
 
 from typing import Optional, Tuple, Union
 
+import pytest
 import torch
 from executorch.backends.arm.test import common
 
@@ -89,6 +90,12 @@ def forward(
             return torch.div(input=input_, other=other_, rounding_mode=rounding_mode)
 
 
+# We need to get this op directly to have coverage
+class DivScalar(torch.nn.Module):
+    def forward(self, input_: torch.Tensor):
+        return torch.ops.aten.div.Scalar(input_, 2.0)
+
+
 @common.parametrize("test_data", test_data_suite)
 def test_div_tensor_tosa_FP(test_data: Tuple):
     pipeline = TosaPipelineFP[input_t1](Div(), test_data(), aten_op, exir_op)
@@ -149,3 +156,54 @@ def test_div_tensor_vgf_quant(test_data: Tuple):
         quantize=True,
     )
     pipeline.run()
+
+
+aten_op_scalar = "torch.ops.aten.div.Scalar"
+exir_op_scalar = "executorch_exir_dialects_edge__ops_aten_div_Scalar"
+
+test_data_suite_scalar = {
+    "op_div_scalar_rank1_rand": lambda: (torch.rand(5) + 1.0,),
+    "op_div_scalar_rank4_rand": lambda: (torch.rand(5, 10, 25, 20) + 1.0,),
+}
+
+
+@common.parametrize("test_data", test_data_suite_scalar)
+@common.SkipIfNoModelConverter
+def test_div_scalar_vgf_no_quant(test_data: input_t1):
+    """Test Tensor / Scalar division (VGF FP)."""
+    pipeline = VgfPipeline[input_t1](
+        DivScalar(),
+        test_data(),
+        aten_op_scalar,
+        exir_op_scalar,
+        quantize=False,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite_scalar)
+@pytest.mark.xfail(
+    reason=(
+        "Quantized div.Scalar is rewritten to div.Tensor after the initial "
+        "quantizer support match, so it misses DecomposeDivPass and remains "
+        "outside the VGF delegate. MLETORCH-2366"
+    ),
+    strict=True,
+)
+@common.SkipIfNoModelConverter
+def test_div_scalar_vgf_quant(test_data: input_t1):
+    """Test Tensor / Scalar division (VGF INT).
+
+    The quantized scalar path is rewritten before the ATen and should not be
+    checked as div.Scalar. We keep expected op lists empty, matching the
+    existing quantized div.Tensor test.
+
+    """
+    pipeline = VgfPipeline[input_t1](
+        DivScalar(),
+        test_data(),
+        [],
+        [],
+        quantize=True,
+    )
+    pipeline.run()

backends/arm/test/ops/test_index_tensor.py

@@ -12,6 +12,7 @@
     OpNotSupportedPipeline,
     TosaPipelineFP,
     TosaPipelineINT,
+    VgfPipeline,
 )
 
 
@@ -587,3 +588,37 @@ def test_index_tensor_u55_INT_not_delegated(test_data: input_params):
             quantize=True,
             u55_subset=True,
         ).run()
+
+
+@common.parametrize(
+    "test_data",
+    IndexTensor.test_data_fp | IndexTensor.test_data_bf16 | IndexTensor.test_data_fp16,
+)
+@common.SkipIfNoModelConverter
+def test_index_tensor_vgf_no_quant(test_data: input_params):
+    test_input = test_data
+    pipeline = VgfPipeline[input_params](
+        IndexTensor(),
+        test_input,
+        IndexTensorTestCommon.aten_op,
+        IndexTensorTestCommon.exir_op,
+        atol=IndexTensorTestCommon.atol,
+        rtol=IndexTensorTestCommon.rtol,
+        tosa_extensions=["bf16"],
+        quantize=False,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", IndexTensor.test_data_int | IndexTensor.test_data_fp)
+@common.SkipIfNoModelConverter
+def test_index_tensor_vgf_quant(test_data: input_params):
+    test_input = test_data
+    pipeline = VgfPipeline[input_params](
+        IndexTensor(),
+        test_input,
+        IndexTensorTestCommon.aten_op,
+        IndexTensorTestCommon.exir_op,
+        quantize=True,
+    )
+    pipeline.run()

backends/arm/test/ops/test_mean_dim.py

@@ -4,7 +4,6 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
-from typing import Callable
 
 import pytest
 import torch
@@ -351,43 +350,43 @@ def test_mean_dim_vgf_quant(test_data):
     pipeline = VgfPipeline[input_t](
         MeanDim(dim, keep_dim),
         (test_data_val,),
+        [],  # Quantized mean.dim may be decomposed to sum, avgpool, or both.
         [],
         symmetric_io_quantization=True,
         quantize=True,
     )
     pipeline.run()
 
 
-mean_input_t = tuple[torch.Tensor, bool]
-
-
 class MeanDefault(torch.nn.Module):
-    def forward(self, tensor: torch.Tensor, keepdim: bool):
-        return tensor.mean()
+    torch_op = "torch.ops.aten.mean.default"
+    exir_op = "executorch_exir_dialects_edge__ops_aten_mean_default"
 
-    test_data_suite: dict[str, Callable[[], mean_input_t]] = {
-        "rank_2": lambda: (
-            torch.rand(1, 2),
-            False,
-        ),
-        "rank_2_keepdim": lambda: (torch.rand(5, 5), True),
-        "rank_4": lambda: (torch.rand(5, 1, 10, 1), False),
+    test_data_suite = {
+        "rank_1": lambda: (torch.rand(7),),
+        "rank_2": lambda: (torch.rand(2, 5),),
+        "rank_3": lambda: (torch.rand(2, 5, 7),),
+        "rank_4": lambda: (torch.rand(1, 5, 7, 3),),
     }
 
+    def forward(self, x: torch.Tensor):
+        return torch.mean(x)
+
 
 @common.parametrize("test_data", MeanDefault.test_data_suite)
 def test_mean_tosa_FP(test_data):
-    pipeline = TosaPipelineFP[mean_input_t](
+    pipeline = TosaPipelineFP[input_t](
         MeanDefault(),
         test_data(),
-        [],  # Might be sum, avgpool, or both
+        MeanDefault.torch_op,
+        MeanDefault.exir_op,
     )
     pipeline.run()
 
 
 @common.parametrize("test_data", MeanDefault.test_data_suite)
 def test_mean_tosa_INT(test_data):
-    pipeline = TosaPipelineINT[mean_input_t](
+    pipeline = TosaPipelineINT[input_t](
         MeanDefault(),
         test_data(),
         [],  # Might be sum, avgpool, or both
@@ -396,6 +395,33 @@ def test_mean_tosa_INT(test_data):
     pipeline.run()
 
 
+@common.parametrize("test_data", MeanDefault.test_data_suite)
+@common.SkipIfNoModelConverter
+def test_mean_vgf_no_quant(test_data):
+    pipeline = VgfPipeline[input_t](
+        MeanDefault(),
+        test_data(),
+        MeanDefault.torch_op,
+        MeanDefault.exir_op,
+        quantize=False,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", MeanDefault.test_data_suite)
+@common.SkipIfNoModelConverter
+def test_mean_vgf_quant(test_data):
+    pipeline = VgfPipeline[input_t](
+        MeanDefault(),
+        test_data(),
+        [],  # Quantized mean.default may be decomposed to sum, avgpool, or both.
+        [],
+        symmetric_io_quantization=True,
+        quantize=True,
+    )
+    pipeline.run()
+
+
 a16w8_mean_test_parameters = {
     "rank_2_keepdim": lambda: (torch.rand(7, 3), (0, 1), True),
     "rank_3_keepdim": lambda: (torch.rand(5, 7, 3), (0, 1, 2), True),

backends/arm/test/ops/test_mul.py

@@ -20,7 +20,13 @@
 )
 
 input_t1 = Tuple[torch.Tensor, torch.Tensor]  # Input x
+input_t_scalar = Tuple[torch.Tensor]
+
 aten_op = "torch.ops.aten.mul.Tensor"
+exir_op = "executorch_exir_dialects_edge__ops_aten_mul_Tensor"
+aten_op_scalar = "torch.ops.aten.mul.Scalar"
+exir_op_scalar = "executorch_exir_dialects_edge__ops_aten_mul_Scalar"
+
 
 test_data_suite = {
     # (test_name, input, other,) See torch.mul() for info
@@ -117,6 +123,11 @@
     ),
 }
 
+test_data_suite_scalar = {
+    "op_mul_scalar_rank1_rand": lambda: (torch.rand(5) * 3.7,),
+    "op_mul_scalar_rank4_randn": lambda: (torch.randn(1, 10, 25, 20),),
+}
+
 
 class Mul(torch.nn.Module):
 
@@ -128,6 +139,11 @@ def forward(
         return input_ * other_
 
 
+class MulScalar(torch.nn.Module):
+    def forward(self, input_: torch.Tensor):
+        return torch.ops.aten.mul.Scalar(input_, 2.0)
+
+
 @common.parametrize(
     "test_data", test_data_suite | test_data_suite_bf16 | test_data_suite_fp16
 )
@@ -365,3 +381,36 @@ def test_mul_tensor_16a8w_u85_INT(test_data: input_t1):
         a16w8_quantization=True,
     )
     pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite_scalar)
+@common.SkipIfNoModelConverter
+def test_mul_scalar_vgf_no_quant(test_data: input_t_scalar):
+    """Test Tensor * Scalar multiplication (VGF FP)."""
+    pipeline = VgfPipeline[input_t_scalar](
+        MulScalar(),
+        test_data(),
+        aten_op_scalar,
+        exir_op_scalar,
+        quantize=False,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite_scalar)
+@common.SkipIfNoModelConverter
+def test_mul_scalar_vgf_quant(test_data: input_t_scalar):
+    """Test Tensor * Scalar multiplication (VGF INT).
+
+    The quantized scalar path is lowered through the Tensor overload, so use
+    mul.Tensor expectations rather than mul.Scalar.
+
+    """
+    pipeline = VgfPipeline[input_t_scalar](
+        MulScalar(),
+        test_data(),
+        aten_op,
+        exir_op,
+        quantize=True,
+    )
+    pipeline.run()