Update base for Update on "Arm backend: Add 16A8W support and test for add operation"

Add 16A8W quantization support and comprehensive tests for the add operation in ExecutorTorch ARM backend targeting Ethos U55 and U85 NPUs.

This follows the pattern established for linear operations, extending int16 support to add operations with hardware-specific testing.

Changes:
- Add INT16 dtype validation support in op_add.py
- Add test_add_tensor_16a8w_tosa_INT test function with U55/U85 pipeline support
- Add U55 and U85 specific 16A8W tests with proper xfail decorators
- Fix U55/U85 test parameter usage (remove unsupported tosa_extensions, clean quantizer function calls)
- Update xfail reasons to consistent 'Vela compilation fails with Invalid arguments' pattern

Differential Revision: [D80510463](https://our.internmc.facebook.com/intern/diff/D80510463)

cc digantdesai freddan80 per zingo oscarandersson8218

[ghstack-poisoned]

Author: Ninja91

backends/arm/operators/op_view.py

@@ -44,7 +44,13 @@ def define_node(
         validate_valid_dtype(
             self.target,
             [inputs[0], output],
-            [ts.DType.INT8, ts.DType.INT32, ts.DType.FP32, ts.DType.BOOL],
+            [
+                ts.DType.INT8,
+                ts.DType.INT16,
+                ts.DType.INT32,
+                ts.DType.FP32,
+                ts.DType.BOOL,
+            ],
             output.tosa_spec,
         )
 

backends/arm/quantizer/arm_quantizer.py

@@ -225,86 +225,6 @@ def get_symmetric_a16w8_quantization_config(
     return quantization_config
 
 
-@functools.lru_cache
-def get_symmetric_a16w8_quantization_config(
-    is_per_channel: bool = True,
-    is_qat: bool = False,
-    is_dynamic: bool = False,
-    weight_qmin: int = -127,
-    weight_qmax: int = 127,
-):
-    """
-    16A8W quantization config: 16-bit activations, 8-bit weights.
-
-    This configuration provides better accuracy than 8A8W while maintaining
-    reasonable memory usage through 8-bit weights.
-
-    Args:
-        is_per_channel: Whether to use per-channel quantization for weights
-        is_qat: Whether this is for Quantization Aware Training
-        is_dynamic: Whether to use dynamic quantization
-        weight_qmin: Minimum quantization value for weights
-        weight_qmax: Maximum quantization value for weights
-
-    Returns:
-        QuantizationConfig with 16-bit activations and 8-bit weights
-    """
-    extra_args: Dict[str, Any] = {"eps": 2**-12}
-
-    # Setup observer/fake-quant for 16-bit activations
-    if is_qat:
-        if is_dynamic:
-            act_observer_or_fake_quant_ctr = FakeQuantize
-            dynamic_quant_observer = MovingAverageMinMaxObserver.with_args(
-                averaging_constant=1
-            )
-            extra_args["observer"] = dynamic_quant_observer
-        else:
-            act_observer_or_fake_quant_ctr = FusedMovingAvgObsFakeQuantize  # type: ignore[assignment]
-    else:
-        if is_dynamic:
-            act_observer_or_fake_quant_ctr = PlaceholderObserver  # type: ignore[assignment]
-        else:
-            # HistogramObserver works well for 16-bit range
-            act_observer_or_fake_quant_ctr = HistogramObserver  # type: ignore[assignment]
-
-    # 16-bit activation quantization spec
-    act_quantization_spec = QuantizationSpec(
-        dtype=torch.int16,
-        quant_min=torch.iinfo(torch.int16).min,  # -32768
-        quant_max=torch.iinfo(torch.int16).max,  # 32767
-        qscheme=torch.per_tensor_symmetric,
-        is_dynamic=is_dynamic,
-        observer_or_fake_quant_ctr=act_observer_or_fake_quant_ctr.with_args(
-            **extra_args,
-        ),
-    )
-
-    # Instead of reconstructing quantization_config, just clone and update as needed
-    # Clone the quantization_config from get_symmetric_quantization_config and update activation spec
-    base_config = get_symmetric_quantization_config(
-        is_per_channel=is_per_channel,
-        is_qat=is_qat,
-        is_dynamic=is_dynamic,
-    )
-    # Replace activation quantization spec with 16-bit version
-    if is_dynamic:
-        quantization_config = QuantizationConfig(
-            act_quantization_spec,  # 16-bit input activations
-            None,
-            base_config.weight,  # 8-bit weights from base config
-            None,
-        )
-    else:
-        quantization_config = QuantizationConfig(
-            act_quantization_spec,  # 16-bit input activations
-            act_quantization_spec,  # 16-bit output activations
-            base_config.weight,  # 8-bit weights from base config
-            None,
-        )
-    return quantization_config
-
-
 NodeFilterType = Callable[[Node], bool]
 """Type for a Node Filter used by annotators. A Node filter is a function that takes
     a Node and returns whether the node should be annotated or not.

backends/arm/scripts/utils.sh

@@ -84,7 +84,7 @@ function check_os_support() {
     # Linux on arm64/aarch64
     # Darwin on arm64/aarch64
     if [[ "${ARCH}" == "aarch64" ]] || [[ "${ARCH}" == "arm64" ]]; then
-        if [[ "${OS}" != "Darwin" ]] || [[ "${OS}" != "Linux" ]]; then
+        if [[ "${OS}" != "Darwin" ]] && [[ "${OS}" != "Linux" ]]; then
             echo "Error: Only Linux and Darwin are supported on arm64"
             exit 1
         fi

backends/arm/test/ops/test_linear.py

@@ -9,9 +9,12 @@
 from typing import Tuple
 
 import pytest
-
 import torch
-from executorch.backends.arm.test import common
+from executorch.backends.arm.quantizer.arm_quantizer import (
+    get_symmetric_a16w8_quantization_config,
+    TOSAQuantizer,
+)
+from executorch.backends.arm.test import common, conftest
 
 from executorch.backends.arm.test.tester.test_pipeline import (
     EthosU55PipelineINT,
@@ -20,6 +23,8 @@
     TosaPipelineINT,
     VgfPipeline,
 )
+from executorch.backends.arm.tosa_specification import TosaSpecification
+from executorch.backends.xnnpack.test.tester import Quantize
 
 aten_op = "torch.ops.aten.linear.default"
 
@@ -143,7 +148,6 @@ def test_linear_tosa_FP(test_data: torch.Tensor):
     pipeline.run()
 
 
-@pytest.mark.flaky(reruns=5)  # TODO: Investigate flakyness.
 @common.parametrize("test_data", test_data_rank1_INT | test_data_rank4_INT)
 def test_linear_tosa_INT(test_data: torch.Tensor):
     test_data, out_features, has_bias, per_channel_quantization = test_data()
@@ -243,3 +247,64 @@ def test_linear_vgf_INT(test_data: torch.Tensor):
         per_channel_quantization=per_channel_quantization,
     )
     pipeline.run()
+
+
+def get_symmetric_a16w8_linear_quantizer(
+    u55_config=False, per_channel_quantization=False
+):
+    tosa_version = conftest.get_option("tosa_version")
+    tosa_profiles = {
+        "1.0": TosaSpecification.create_from_string("TOSA-1.0+INT+int16"),
+    }
+
+    quantizer = TOSAQuantizer(tosa_profiles[tosa_version])
+    quantizer.set_global(
+        get_symmetric_a16w8_quantization_config(is_per_channel=per_channel_quantization)
+    )
+    quantizer.set_module_type(
+        torch.nn.Linear,
+        get_symmetric_a16w8_quantization_config(
+            is_per_channel=per_channel_quantization
+        ),
+    )
+
+    return Quantize(
+        quantizer,
+        get_symmetric_a16w8_quantization_config(
+            is_per_channel=per_channel_quantization
+        ),
+    )
+
+
+@common.parametrize("test_data", test_data_rank1_INT | test_data_rank4_INT)
+@pytest.mark.xfail(
+    reason="missing int16 linear ops support; fails at TOSA reference model run with Invalid TOSA graph"
+)
+def test_linear_16a8w_tosa_INT(test_data: torch.Tensor):
+    """Test linear operation with 16A8W quantization (16-bit activations, 8-bit weights)"""
+    test_data, out_features, has_bias, per_channel_quantization = test_data()
+    in_features = test_data.shape[-1]
+
+    # Create pipeline with custom 16A8W quantization config
+    pipeline = TosaPipelineINT[input_t1](
+        Linear(
+            in_features=in_features,
+            out_features=out_features,
+            bias=has_bias,
+        ),
+        (test_data,),
+        aten_op,
+        exir_op=[],
+        per_channel_quantization=per_channel_quantization,
+        use_to_edge_transform_and_lower=True,
+        tosa_extensions=["int16"],
+    )
+
+    pipeline.change_args(
+        "quantize",
+        get_symmetric_a16w8_linear_quantizer(
+            per_channel_quantization=per_channel_quantization
+        ),
+    )
+    # Run the pipeline
+    pipeline.run()

backends/cadence/aot/ref_implementations.py

@@ -6,6 +6,8 @@
 
 # pyre-strict
 
+from typing import Optional
+
 import torch
 from executorch.exir.scalar_type import ScalarType
 from torch.library import impl, Library
@@ -177,6 +179,54 @@ def quantized_add(
     )
 
 
+@impl(m, "quantized_linear")
+def quantized_linear(
+    src: torch.Tensor,
+    weight: torch.Tensor,
+    bias: torch.Tensor,
+    in_zero_point: int,
+    weight_zero_point: torch.Tensor,
+    out_multiplier: torch.Tensor,
+    out_shift: torch.Tensor,
+    out_zero_point: int,
+    offset: Optional[torch.Tensor],
+) -> torch.Tensor:
+    """
+    Quantized linear (transposed matmul) operation.
+
+    Args:
+        - src (Tensor): The activations tensor
+        - weight (Tensor): The weight tensor
+        - bias (Tensor): The bias tensor
+        - in_zero_point (int): The quantized mapping of zero for the input
+        - weight_zero_point (Tensor): The quantized mapping of zero for the weight
+        - out_multiplier (Tensor): The multiplier used to scale the output
+        - out_shift (Tensor): The shift used to scale the output
+        - out_zero_point (int): The quantized mapping of zero for the output
+        - offset (Tensor): Unused
+    """
+    out_scale = -out_multiplier * (1 / (1 << 31)) * (2 ** out_shift[0])
+
+    N, K = weight.shape
+
+    leading_dims = src.shape[:-1]
+    src = src.view(-1, K)
+
+    dtype = src.dtype
+    supported_dtypes = [torch.int8, torch.uint8, torch.int32]
+    if dtype not in supported_dtypes:
+        raise ValueError(
+            f"Unsupported dtype to quantize to. Supported dtypes must be one of {supported_dtypes}"
+        )
+
+    out = torch.nn.functional.linear(
+        src - in_zero_point, weight - weight_zero_point, bias
+    )
+    return quantize_per_tensor(
+        out, out_scale, out_zero_point, -128, 127, dtype
+    ).reshape(*leading_dims, N)
+
+
 @impl(m, "requantize")
 def requantize(
     input: torch.Tensor,

backends/cadence/aot/tests/test_ref_implementations.py

@@ -5,15 +5,17 @@
 # LICENSE file in the root directory of this source tree.
 
 # pyre-strict
-
+import typing
 import unittest
 
+import numpy as np
 import torch
 
 from executorch.backends.cadence.aot.ref_implementations import (
     dequantize_per_tensor,
     quantize_per_tensor,
     quantized_add,
+    quantized_linear,
 )
 from executorch.backends.cadence.aot.typing_stubs import expand
 
@@ -138,3 +140,102 @@ def test_quantized_add(
             torch.equal(output, expected_output),
             f"Values don't match in {name}: got {output}, expected {expected_output}",
         )
+
+    @expand(
+        [
+            # Test case 1: 1x2 input, 1x2 weight (1 output feature)
+            (
+                torch.Size([1, 2]),  # src_shape: 1 sample, 2 input features
+                torch.Size([1, 2]),  # weight_shape: 1 output feature, 2 input features
+                0,  # in_zero_point
+                torch.tensor([0, 0], dtype=torch.int8),  # weight_zero_point
+                torch.tensor(
+                    [1073741824], dtype=torch.int32
+                ),  # out_multiplier (0.5 * 2^31)
+                torch.tensor([0], dtype=torch.int8),  # out_shift
+                0,  # out_zero_point
+                torch.tensor([[-2]], dtype=torch.int8),  # expected_output
+            ),
+            # Test case 2: 1x3 input, 2x3 weight (2 output features)
+            (
+                torch.Size([1, 3]),  # src_shape: 1 sample, 3 input features
+                torch.Size([2, 3]),  # weight_shape: 2 output features, 3 input features
+                0,  # in_zero_point
+                torch.tensor([0, 0, 0], dtype=torch.int8),  # weight_zero_point
+                torch.tensor(
+                    [1073741824], dtype=torch.int32
+                ),  # out_multiplier (0.5 * 2^31)
+                torch.tensor([0], dtype=torch.int8),  # out_shift
+                0,  # out_zero_point
+                torch.tensor([[-10, -30]], dtype=torch.int8),  # expected_output
+            ),
+            # Test case 3: Batch case with different dimensions
+            (
+                torch.Size([1, 2, 2]),  # src_shape: batch=1, seq=2, features=2
+                torch.Size([3, 2]),  # weight_shape: 3 output features, 2 input features
+                0,  # in_zero_point
+                torch.tensor([0, 0], dtype=torch.int8),  # weight_zero_point
+                torch.tensor(
+                    [1073741824], dtype=torch.int32
+                ),  # out_multiplier (0.5 * 2^31)
+                torch.tensor([0], dtype=torch.int8),  # out_shift
+                0,  # out_zero_point
+                torch.tensor(
+                    [[[-2, -8, -14], [-6, -28, -50]]], dtype=torch.int8
+                ),  # expected_output
+            ),
+            # Test case 4: Non-zero zero points
+            (
+                torch.Size([1, 2]),  # src_shape: 1 sample, 2 input features
+                torch.Size([2, 2]),  # weight_shape: 2 output feature, 1 input feature
+                2,  # in_zero_point
+                torch.tensor([1, 1], dtype=torch.int8),  # weight_zero_point
+                torch.tensor(
+                    [268435456], dtype=torch.int32
+                ),  # out_multiplier (1.0 * 2^31)
+                torch.tensor([0]),  # out_shift
+                1,  # out_zero_point
+                torch.tensor([[-15, 25]], dtype=torch.int8),  # expected_output
+            ),
+        ]
+    )
+    def test_quantized_linear(
+        self,
+        src_shape: torch.Size,
+        weight_shape: torch.Size,
+        in_zero_point: int,
+        weight_zero_point: torch.Tensor,
+        out_multiplier: torch.Tensor,
+        out_shift: torch.Tensor,
+        out_zero_point: int,
+        expected_output: torch.Tensor,
+    ) -> None:
+        src = (
+            torch.arange(np.product(src_shape))
+            .reshape(src_shape)
+            .to(expected_output.dtype)
+        )
+        weight = (
+            torch.arange(np.product(weight_shape))
+            .reshape(weight_shape)
+            .to(expected_output.dtype)
+        )
+        bias = torch.arange(weight_shape[0]).to(expected_output.dtype)
+        output = quantized_linear(
+            src,
+            weight,
+            bias,
+            in_zero_point,
+            weight_zero_point,
+            out_multiplier,
+            out_shift,
+            out_zero_point,
+            typing.cast(torch.Tensor, None),
+        )
+
+        self.assertTrue(output.dtype == expected_output.dtype, "Dtype mismatch")
+
+        self.assertTrue(
+            torch.equal(output, expected_output),
+            f"Values don't match: got {output}, expected {expected_output}",
+        )