Quantize moveaxis/movedim so they delegate to Ethos-U (#20314)

backends/arm/quantizer/quantization_annotator.py

@@ -619,6 +619,10 @@ def _get_fixed_qparams_qspec(
     # dequant -> neg -> requant chain.
     torch.ops.aten.neg.default,
     torch.ops.aten.detach_copy.default,
+    torch.ops.aten.moveaxis.int,
+    torch.ops.aten.moveaxis.intlist,
+    torch.ops.aten.movedim.int,
+    torch.ops.aten.movedim.intlist,
 }
 
 # Dimname has been removed from upstream PyTorch, but there may be a window
@@ -630,6 +634,7 @@ def _get_fixed_qparams_qspec(
 if _transpose_dimname is not None:
     _one_to_one_shared_input_qspec.add(_transpose_dimname)
 
+
 _one_to_one_shared_input_or_input_act_qspec: set[OpOverload] = {
     torch.ops.aten.alias.default,
     torch.ops.aten.clone.default,

backends/arm/test/ops/test_permute.py

@@ -78,6 +78,12 @@ def forward(self, x):
         return torch.permute(x, self.dims)
 
 
+class SimpleMoveAxis(torch.nn.Module):
+
+    def forward(self, x):
+        return torch.moveaxis(x, 1, -1)
+
+
 @common.parametrize(
     "test_data", test_data_suite | test_data_suite_fp16 | test_data_suite_bf16
 )
@@ -118,6 +124,17 @@ def test_permute_u55_INT(test_data):
     pipeline.run()
 
 
+def test_moveaxis_u55_INT():
+    pipeline = EthosU55PipelineINT[input_t1](
+        SimpleMoveAxis(),
+        (torch.rand(1, 4, 5, 6),),
+        "torch.ops.aten.moveaxis.int",
+        exir_ops="executorch_exir_dialects_edge__ops_aten_permute_copy_default",
+        run_on_fvp=False,
+    )
+    pipeline.run()
+
+
 @common.parametrize("test_data", test_data_suite_u55_reject)
 def test_permute_u55_INT_not_delegated(test_data: torch.Tensor):
     test_data, dims = test_data()

backends/arm/test/quantizer/test_generic_annotater.py

@@ -7,7 +7,7 @@
 from typing import Any, Callable, Tuple
 
 import torch
-from executorch.backends.arm.quantizer import is_annotated
+from executorch.backends.arm.quantizer import is_annotated, quantization_annotator
 from executorch.backends.arm.test.tester.test_pipeline import TosaPipelineINT
 from executorch.backends.test.harness.stages import StageType
 
@@ -89,6 +89,17 @@ def test_transpose_tosa_INT():
     )
 
 
+def test_moveaxis_movedim_shared_qspec_annotations():
+    expected_ops = {
+        torch.ops.aten.moveaxis.int,
+        torch.ops.aten.moveaxis.intlist,
+        torch.ops.aten.movedim.int,
+        torch.ops.aten.movedim.intlist,
+    }
+
+    assert expected_ops <= quantization_annotator._one_to_one_shared_input_qspec
+
+
 def test_tile_tosa_INT():
     check_annotation(
         SingleOpModel(torch.tile, (torch.randn(4, 4),), dims=(2,)),