Enable int16 for op permute

Differential Revision: D84948536

Pull Request resolved: pytorch#15256

Author: 3l1

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,
         )
 

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()

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",