diff --git a/backends/arm/operators/op_permute.py b/backends/arm/operators/op_permute.py index 92cc2b37479..3d7a3569882 100644 --- a/backends/arm/operators/op_permute.py +++ b/backends/arm/operators/op_permute.py @@ -117,7 +117,7 @@ def define_node( validate_valid_dtype( self.target, [inputs[0], output], - [ts.DType.INT8, ts.DType.INT32, ts.DType.FP32], + [ts.DType.INT8, ts.DType.INT16, ts.DType.INT32, ts.DType.FP32], output.tosa_spec, ) diff --git a/backends/arm/test/ops/test_permute.py b/backends/arm/test/ops/test_permute.py index 6fd8555b56b..aa2f49b5e53 100644 --- a/backends/arm/test/ops/test_permute.py +++ b/backends/arm/test/ops/test_permute.py @@ -9,8 +9,11 @@ from typing import Tuple 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, @@ -19,7 +22,8 @@ TosaPipelineINT, VgfPipeline, ) -from torchvision.ops import Permute +from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.xnnpack.test.tester import Quantize input_t1 = Tuple[torch.Tensor] # Input x @@ -42,10 +46,10 @@ class SimplePermute(torch.nn.Module): def __init__(self, dims: list[int]): super().__init__() - self.permute = Permute(dims=dims) + self.dims = dims def forward(self, x): - return self.permute(x) + return torch.permute(x, self.dims) @common.parametrize("test_data", test_data_suite) @@ -128,3 +132,98 @@ def test_permute_vgf_INT(test_data): tosa_version="TOSA-1.0+INT", ) pipeline.run() + + +def get_symmetric_a16w8_permute_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_permute_int16_tosa_INT(test_data: torch.Tensor): + """Test permute operation with int16 quantization""" + test_data, dims = test_data() + pipeline = TosaPipelineINT[input_t1]( + SimplePermute(dims=dims), + (test_data,), + aten_op, + exir_op=[], + per_channel_quantization=False, + use_to_edge_transform_and_lower=True, + tosa_extensions=["int16"], + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_permute_quantizer(per_channel_quantization=False), + ) + # Run the pipeline + pipeline.run() + + +test_data_suite_exact = { + x: test_data_suite[x] for x in test_data_suite if x != "rank_4_3" +} + + +@common.parametrize("test_data", test_data_suite_exact) +@common.XfailIfNoCorstone300 +def test_permute_int16_u55_INT16(test_data: torch.Tensor): + """Test permute operation with int16 quantization on U55""" + test_data, dims = test_data() + pipeline = EthosU55PipelineINT[input_t1]( + SimplePermute(dims=dims), + (test_data,), + aten_op, + exir_ops=[], + per_channel_quantization=True, + use_to_edge_transform_and_lower=True, + atol=1e-02, + rtol=1e-02, + run_on_fvp=True, + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_permute_quantizer(per_channel_quantization=False), + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_suite) +@common.XfailIfNoCorstone320 +def test_permute_int16_u85_INT16(test_data: torch.Tensor): + """Test permute operation with int16 quantization on U85""" + test_data, dims = test_data() + pipeline = EthosU85PipelineINT[input_t1]( + SimplePermute(dims=dims), + (test_data,), + aten_op, + exir_ops=[], + use_to_edge_transform_and_lower=True, + atol=1e-03, + rtol=1e-03, + run_on_fvp=True, + ) + + pipeline.change_args( + "quantize", + get_symmetric_a16w8_permute_quantizer(per_channel_quantization=False), + ) + pipeline.run() diff --git a/backends/arm/test/targets.bzl b/backends/arm/test/targets.bzl index 5fdd1c3d827..b94ea23c6c1 100644 --- a/backends/arm/test/targets.bzl +++ b/backends/arm/test/targets.bzl @@ -20,6 +20,7 @@ def define_arm_tests(): "ops/test_cat.py", "ops/test_linear.py", "ops/test_mul.py", + "ops/test_permute.py", "ops/test_slice.py", "ops/test_sigmoid.py", "ops/test_sub.py",