Arm backend: Add 16x8 tests for VGF (pytorch#18839)

Introduce int16a8w for a couple of op tests
on VGF

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

Signed-off-by: Christoffer J.L <christoffer.johanssonlundqvist@arm.com>

Author: Christoffer-JL

backends/arm/test/ops/test_adaptive_avg_pool2d.py

@@ -7,6 +7,9 @@
 
 import torch
 
+from executorch.backends.arm.quantizer.arm_quantizer import (
+    get_symmetric_a16w8_quantization_config,
+)
 from executorch.backends.arm.test import common
 
 from executorch.backends.arm.test.tester.test_pipeline import (
@@ -244,3 +247,19 @@ def test_adaptive_avg_pool2d_vgf_quant(test_module):
         quantize=True,
     )
     pipeline.run()
+
+
+@common.parametrize("test_module", test_modules)
+@common.SkipIfNoModelConverter
+def test_adaptive_avg_pool2d_vgf_quant_a16w8(test_module):
+    model, input_tensor = test_module()
+    pipeline = VgfPipeline[input_t](
+        model,
+        input_tensor,
+        [],
+        exir_op,
+        quantize=True,
+        tosa_extensions=["int16"],
+    )
+    pipeline.quantizer.set_global(get_symmetric_a16w8_quantization_config())
+    pipeline.run()

backends/arm/test/ops/test_add.py

@@ -278,6 +278,22 @@ def test_add_tensor_vgf_quant(test_data: input_t1):
     pipeline.run()
 
 
+@common.parametrize("test_data", Add.test_data)
+@common.SkipIfNoModelConverter
+def test_add_tensor_vgf_quant_a16w8(test_data: input_t1):
+    pipeline = VgfPipeline[input_t1](
+        Add(),
+        test_data(),
+        aten_op,
+        exir_op,
+        run_on_vulkan_runtime=True,
+        quantize=True,
+        tosa_extensions=["int16"],
+    )
+    pipeline.quantizer.set_global(get_symmetric_a16w8_quantization_config())
+    pipeline.run()
+
+
 @common.parametrize("test_data", Add.test_data)
 def test_add_tensor_tosa_INT_16a8w(test_data: input_t1):
     """Test add operation with 16A8W quantization (16-bit activations, 8-bit

backends/arm/test/ops/test_amax.py

@@ -8,6 +8,9 @@
 
 import pytest
 import torch
+from executorch.backends.arm.quantizer.arm_quantizer import (
+    get_symmetric_a16w8_quantization_config,
+)
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.test_pipeline import (
     EthosU55PipelineINT,
@@ -263,6 +266,23 @@ def test_max_dim_vgf_quant_to_amax(test_data: Max.input_t):
     pipeline.run()
 
 
+@common.parametrize("test_data", Amax.test_data)
+@common.SkipIfNoModelConverter
+def test_amax_vgf_quant_a16w8(test_data: Amax.input_t):
+    data, dim, keep_dims = test_data()
+    module = Amax(dim, keep_dims)
+    pipeline = VgfPipeline[Max.input_t](
+        module,
+        data,
+        amax_aten_op,
+        amax_exir_op,
+        quantize=True,
+        tosa_extensions=["int16"],
+    )
+    pipeline.quantizer.set_global(get_symmetric_a16w8_quantization_config())
+    pipeline.run()
+
+
 @common.parametrize("test_data", Amax.test_data)
 def test_amax_tosa_INT_a16w8(test_data: Amax.input_t):
     """Test amax with 16A8W quantization for TOSA INT."""

backends/arm/test/ops/test_amin.py

@@ -9,6 +9,9 @@
 import pytest
 
 import torch
+from executorch.backends.arm.quantizer.arm_quantizer import (
+    get_symmetric_a16w8_quantization_config,
+)
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.test_pipeline import (
     EthosU55PipelineINT,
@@ -273,6 +276,21 @@ def test_min_dim_vgf_quant_to_amin(test_data: Min.input_t):
     pipeline.run()
 
 
+@common.parametrize("test_data", Amin.test_data)
+@common.SkipIfNoModelConverter
+def test_amin_vgf_quant_a16w8(test_data: Amin.input_t):
+    data, dim, keep_dims = test_data()
+    pipeline = VgfPipeline[Amin.input_t](
+        Amin(dim, keep_dims),
+        data,
+        amin_aten_op,
+        quantize=True,
+        tosa_extensions=["int16"],
+    )
+    pipeline.quantizer.set_global(get_symmetric_a16w8_quantization_config())
+    pipeline.run()
+
+
 @common.parametrize("test_data", Amin.test_data)
 def test_amin_tosa_INT_a16w8(test_data: Amin.input_t):
     """Test amin with 16A8W quantization for TOSA INT."""

backends/arm/test/ops/test_avg_pool2d.py

@@ -11,6 +11,9 @@
 import conftest
 
 import torch
+from executorch.backends.arm.quantizer.arm_quantizer import (
+    get_symmetric_a16w8_quantization_config,
+)
 
 from executorch.backends.arm.test import common
 
@@ -289,6 +292,22 @@ def test_avg_pool2d_vgf_quant(test_module):
     pipeline.run()
 
 
+@common.parametrize("test_module", test_modules)
+@common.SkipIfNoModelConverter
+def test_avg_pool2d_vgf_quant_a16w8(test_module):
+    model, input_tensor = test_module()
+    pipeline = VgfPipeline[input_t](
+        model,
+        input_tensor,
+        aten_op,
+        exir_op,
+        quantize=True,
+        tosa_extensions=["int16"],
+    )
+    pipeline.quantizer.set_global(get_symmetric_a16w8_quantization_config())
+    pipeline.run()
+
+
 reject_modules = {
     "kernel_1x1_stride_1_pad_0": lambda: (AvgPool2d(1, 1, 0), torch.rand(2, 5, 5, 5)),
     "kernel_2x9_stride_1_pad_1": lambda: (

backends/arm/test/ops/test_clamp.py

@@ -7,6 +7,9 @@
 from typing import Tuple, Union
 
 import torch
+from executorch.backends.arm.quantizer.arm_quantizer import (
+    get_symmetric_a16w8_quantization_config,
+)
 
 from executorch.backends.arm.test import common
 
@@ -260,6 +263,25 @@ def test_clamp_vgf_quant(test_data):
 aten_op_tensor = [
     "torch.ops.aten.clamp.Tensor",
 ]
+
+
+@common.parametrize("test_data", test_data_suite)
+@common.SkipIfNoModelConverter
+def test_clamp_vgf_quant_a16w8(test_data):
+    input_tensor, min_val, max_val = test_data()
+    model = Clamp(min_val, max_val)
+    pipeline = VgfPipeline[input_t](
+        model,
+        (input_tensor,),
+        aten_op,
+        exir_op,
+        quantize=True,
+        tosa_extensions=["int16"],
+    )
+    pipeline.quantizer.set_global(get_symmetric_a16w8_quantization_config())
+    pipeline.run()
+
+
 exir_op_tensor = "executorch_exir_dialects_edge__ops_aten_clamp_Tensor"
 
 test_data_suite_tensor_FP = {

backends/arm/test/ops/test_constant_pad_nd.py

@@ -9,6 +9,9 @@
 
 import torch
 import torch.nn.functional as F
+from executorch.backends.arm.quantizer.arm_quantizer import (
+    get_symmetric_a16w8_quantization_config,
+)
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.test_pipeline import (
     TosaPipelineFP,
@@ -208,3 +211,19 @@ def test_constant_pad_nd_vgf_quant(test_data: Tuple):
         quantize=True,
     )
     pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite)
+@common.SkipIfNoModelConverter
+def test_constant_pad_nd_vgf_quant_a16w8(test_data: Tuple):
+    inp, padding, value, mode = test_data()
+    pipeline = VgfPipeline[input_t1](
+        ConstantPadND(padding, value, mode),
+        (inp,),
+        aten_op,
+        exir_op,
+        quantize=True,
+        tosa_extensions=["int16"],
+    )
+    pipeline.quantizer.set_global(get_symmetric_a16w8_quantization_config())
+    pipeline.run()

backends/arm/test/ops/test_conv2d.py

@@ -8,6 +8,7 @@
 
 import torch
 from executorch.backends.arm.quantizer.arm_quantizer import (
+    get_symmetric_a16w8_quantization_config,
     get_symmetric_a8w4_quantization_config,
 )
 from executorch.backends.arm.test import common
@@ -644,6 +645,25 @@ def test_convolution_2d_vgf_quant(test_data):
     pipeline.run()
 
 
+@common.parametrize("test_data", test_data_INT)
+@common.SkipIfNoModelConverter
+def test_convolution_2d_vgf_quant_a16w8(test_data):
+    model, per_channel_quantization = test_data()
+    pipeline = VgfPipeline[input_t](
+        model,
+        model.get_inputs(),
+        aten_op,
+        exir_op,
+        per_channel_quantization=per_channel_quantization,
+        quantize=True,
+        tosa_extensions=["int16"],
+    )
+    pipeline.quantizer.set_global(
+        get_symmetric_a16w8_quantization_config(is_per_channel=per_channel_quantization)
+    )
+    pipeline.run()
+
+
 reject_suite = {
     "large_stride": lambda: Conv2d(
         in_channels=1,

backends/arm/test/ops/test_conv3d.py

@@ -9,6 +9,7 @@
 import pytest
 import torch
 from executorch.backends.arm.quantizer.arm_quantizer import (
+    get_symmetric_a16w8_quantization_config,
     get_symmetric_a8w4_quantization_config,
 )
 from executorch.backends.arm.test import common
@@ -741,6 +742,26 @@ def test_convolution_3d_vgf_quant_a8w4(test_data):
     pipeline.run()
 
 
+@common.parametrize("test_data", test_data_INT16)
+@common.SkipIfNoModelConverter
+def test_convolution_3d_vgf_quant_a16w8(test_data):
+    model, per_channel_quantization = test_data()
+    pipeline = VgfPipeline[input_t](
+        model,
+        model.get_inputs(),
+        aten_op,
+        exir_op,
+        per_channel_quantization=per_channel_quantization,
+        quantize=True,
+        tosa_extensions=["int16"],
+        qtol=1,
+    )
+    pipeline.quantizer.set_global(
+        get_symmetric_a16w8_quantization_config(is_per_channel=per_channel_quantization)
+    )
+    pipeline.run()
+
+
 @common.SkipIfNoModelConverter
 def test_convolution_3d_vgf_no_quant_multi_op():
     """Ensure mixed Conv3d/Conv2d graphs keep correct spatial annotations."""

backends/arm/test/ops/test_depthwise_conv.py

@@ -9,6 +9,7 @@
 
 import torch
 from executorch.backends.arm.quantizer.arm_quantizer import (
+    get_symmetric_a16w8_quantization_config,
     get_symmetric_a8w4_quantization_config,
 )
 from executorch.backends.arm.test import common
@@ -338,6 +339,25 @@ def test_convolution_2d_vgf_quant_a8w4_depthwise(test_data):
     pipeline.run()
 
 
+@common.parametrize("test_data", test_data_conv2d_INT)
+@common.SkipIfNoModelConverter
+def test_convolution_2d_vgf_quant_a16w8_depthwise(test_data):
+    model, per_channel_quantization = test_data()
+    pipeline = VgfPipeline[input_t](
+        model,
+        model.get_inputs(),
+        aten_op=[],
+        exir_op=exir_op,
+        per_channel_quantization=per_channel_quantization,
+        quantize=True,
+        tosa_extensions=["int16"],
+    )
+    pipeline.quantizer.set_global(
+        get_symmetric_a16w8_quantization_config(is_per_channel=per_channel_quantization)
+    )
+    pipeline.run()
+
+
 @common.XfailIfNoCorstone300
 @common.parametrize("test_data", test_data_conv2d_INT)
 def test_convolution_2d_u55_INT_depthwise(test_data):