diff --git a/backends/arm/operators/op_mul.py b/backends/arm/operators/op_mul.py index 7d9f6eac6aa..669dd506fc2 100644 --- a/backends/arm/operators/op_mul.py +++ b/backends/arm/operators/op_mul.py @@ -51,7 +51,7 @@ def define_node( validate_valid_dtype( self.target, [*inputs, output], - [ts.DType.INT8, ts.DType.INT32], + [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32], output.tosa_spec, ) diff --git a/backends/arm/test/ops/test_mul.py b/backends/arm/test/ops/test_mul.py index d8f9e947ce3..ebe3ed7f8f4 100644 --- a/backends/arm/test/ops/test_mul.py +++ b/backends/arm/test/ops/test_mul.py @@ -9,8 +9,12 @@ from typing import Tuple import torch +from executorch.backends.arm.quantizer.arm_quantizer import ( + get_symmetric_a16w8_quantization_config, + TOSAQuantizer, +) -from executorch.backends.arm.test import common +from executorch.backends.arm.test import common, conftest from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, EthosU85PipelineINT, @@ -18,6 +22,8 @@ TosaPipelineINT, VgfPipeline, ) +from executorch.backends.arm.tosa_specification import TosaSpecification +from executorch.backends.xnnpack.test.tester import Quantize input_t1 = Tuple[torch.Tensor, torch.Tensor] # Input x aten_op = "torch.ops.aten.mul.Tensor" @@ -268,3 +274,46 @@ def test_mul_tensor_vgf_INT_int32(test_data: torch.Tensor): ) pipeline.pop_stage("check.quant_nodes") pipeline.run() + + +def get_symmetric_a16w8_mul_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) + ) + + return Quantize( + quantizer, + get_symmetric_a16w8_quantization_config( + is_per_channel=per_channel_quantization + ), + ) + + +@common.parametrize("test_data", test_data_suite) +def test_mul_tensor_16a8w_tosa_INT(test_data: input_t1): + """Test mul operation with 16A8W quantization (16-bit activations, 8-bit weights)""" + per_channel_quantization = False + + pipeline = TosaPipelineINT[input_t1]( + Mul(), + 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_mul_quantizer( + per_channel_quantization=per_channel_quantization + ), + ) + pipeline.run() diff --git a/backends/arm/test/targets.bzl b/backends/arm/test/targets.bzl index 405f1bbf081..b438e556cca 100644 --- a/backends/arm/test/targets.bzl +++ b/backends/arm/test/targets.bzl @@ -16,6 +16,7 @@ def define_arm_tests(): "ops/test_add.py", "ops/test_avg_pool2d.py", "ops/test_linear.py", + "ops/test_mul.py", "ops/test_slice.py", "ops/test_sigmoid.py", "ops/test_tanh.py",