Arm backend: Add aten.flip support via TOSA REVERSE

backends/arm/_passes/__init__.py

@@ -53,6 +53,7 @@
 from .decompose_embedding_pass import DecomposeEmbeddingPass  # noqa  # noqa
 from .decompose_erfinv_pass import DecomposeErfinvPass  # noqa
 from .decompose_expm1_pass import DecomposeExpm1Pass  # noqa
+from .decompose_flip_pass import DecomposeFlipPass  # noqa
 from .decompose_floor_divide_pass import DecomposeFloorDividePass  # noqa
 from .decompose_gelu_pass import DecomposeGeluPass  # noqa
 from .decompose_glu_pass import DecomposeGluPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -58,6 +58,7 @@
     DecomposeEmbeddingPass,
     DecomposeErfinvPass,
     DecomposeExpm1Pass,
+    DecomposeFlipPass,
     DecomposeFloorDividePass,
     DecomposeGeluPass,
     DecomposeGluPass,
@@ -538,6 +539,7 @@ def _tosa_pipeline(
                 PromoteBoolOperandsPass(),
                 DecomposeSinhPass(),
                 DecomposeSignPass(),
+                DecomposeFlipPass(),
                 DecomposeFloorDividePass(),
                 DecomposeGeluPass(),
                 DecomposeAddSubAlphaPass(),

backends/arm/_passes/decompose_flip_pass.py

@@ -0,0 +1,35 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Set, Type
+
+from executorch.backends.arm._passes import ArmOpTargetedPass
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass
+
+
+class DecomposeFlipPass(ArmOpTargetedPass):
+    """Decompose a multi-axis ``aten.flip`` into a chain of single-axis flips.
+
+    TOSA ``REVERSE`` reverses a single ``axis``; a flip over N dims is the
+    composition of N independent single-axis reversals. Each single-axis
+    ``aten.flip`` is lowered to one ``REVERSE`` by ``FlipVisitor``. ``flip`` is a
+    shared-qparam data-movement op, so the chained flips inherit the original
+    node's quantization parameters via ``meta``.
+
+    """
+
+    _passes_required_after: Set[Type[ExportPass]] = set()
+    target_ops = {exir_ops.edge.aten.flip.default}
+
+    def call_operator(self, op, args, kwargs, meta, updated=False):
+        if op not in self.target_ops or len(args[1]) == 1:
+            return super().call_operator(op, args, kwargs, meta, updated)
+
+        # Chain one single-axis flip per dim; empty dims falls through as a no-op.
+        out = args[0]
+        for dim in args[1]:
+            out = super().call_operator(op, (out, [dim]), kwargs, meta, updated=True)
+        return out

backends/arm/operator_support/tosa_profile_supported_op_lists.py

@@ -60,6 +60,7 @@
     exir_ops.edge.aten.linear.default,
     exir_ops.edge.aten.split_with_sizes_copy.default,
     exir_ops.edge.aten.split_copy.Tensor,
+    exir_ops.edge.aten.flip.default,
     exir_ops.edge.aten.floor.default,
     exir_ops.edge.aten.full.default,
     exir_ops.edge.aten.full_like.default,
@@ -187,6 +188,7 @@
     exir_ops.edge.aten.linear.default,
     exir_ops.edge.aten.split_with_sizes_copy.default,
     exir_ops.edge.aten.split_copy.Tensor,
+    exir_ops.edge.aten.flip.default,
     exir_ops.edge.aten.floor.default,
     exir_ops.edge.aten.full.default,
     exir_ops.edge.aten.full_like.default,

backends/arm/operators/__init__.py

@@ -16,6 +16,7 @@
     op_any,
     op_cat,
     op_cond_if,
+    op_flip,
     op_permute,
     op_repeat,
     op_sum,

backends/arm/operators/op_flip.py

@@ -0,0 +1,77 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+from typing import Any, List
+
+import torch
+
+import tosa_serializer as ts
+
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.operators.operator_validation_utils import (
+    validate_num_inputs,
+    validate_same_dtype,
+    validate_valid_dtype,
+)
+from executorch.backends.arm.tosa.mapping import TosaArg
+
+
+@register_node_visitor
+class FlipVisitor(NodeVisitor):
+    target = "aten.flip.default"
+
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        tosa_graph: Any,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+        supported_dtypes = [ts.DType.BOOL]
+        if self.tosa_spec.support_integer():
+            supported_dtypes.extend([ts.DType.INT8, ts.DType.INT16, ts.DType.INT32])
+        if self.tosa_spec.support_float():
+            supported_dtypes.extend([ts.DType.FP16, ts.DType.FP32])
+        if self.tosa_spec.support_extension("bf16"):
+            supported_dtypes.append(ts.DType.BF16)
+        if self.tosa_spec.support_extension("fp8e4m3"):
+            supported_dtypes.append(ts.DType.FP8E4M3)
+        if self.tosa_spec.support_extension("fp8e5m2"):
+            supported_dtypes.append(ts.DType.FP8E5M2)
+
+        validate_num_inputs(self.target, inputs, 2)
+        validate_same_dtype(self.target, [inputs[0], output], ts)
+        validate_valid_dtype(
+            self.target,
+            [inputs[0], output],
+            supported_dtypes,
+            self.tosa_spec,
+        )
+
+        dims = inputs[1].special
+        if len(dims) != 1:
+            raise ValueError(
+                f"{self.target} expects a single flip dim after DecomposeFlipPass, "
+                f"got {dims}"
+            )
+        axis = dims[0] % len(inputs[0].shape)
+
+        attr = ts.TosaSerializerAttribute()
+        attr.ReverseAttribute(axis)
+        self._serialize_operator(
+            node,
+            tosa_graph,
+            ts.Op.REVERSE,
+            [inputs[0].name],
+            [output.name],
+            attr,
+        )

backends/arm/test/ops/test_flip.py

@@ -0,0 +1,226 @@
+# Copyright 2026 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+from typing import Tuple
+
+import torch
+
+from executorch.backends.arm.quantizer.arm_quantizer import (
+    get_symmetric_a16w8_quantization_config,
+)
+from executorch.backends.arm.test import common
+from executorch.backends.arm.test.tester.test_pipeline import (
+    EthosU85PipelineINT,
+    OpNotSupportedPipeline,
+    TosaPipelineFP,
+    TosaPipelineINT,
+    VgfPipeline,
+)
+
+aten_op = "torch.ops.aten.flip.default"
+exir_op = "executorch_exir_dialects_edge__ops_aten_flip_default"
+
+input_t1 = Tuple[torch.Tensor]  # Input x
+
+# (input tensor, dims) — single-axis, negative-axis, and multi-axis cases.
+test_data_suite = {
+    "rank1": lambda: (torch.rand(10), [0]),
+    "rank2_dim0": lambda: (torch.rand(4, 5), [0]),
+    "rank2_dim1_neg": lambda: (torch.rand(4, 5), [-1]),
+    "rank3_multi": lambda: (torch.rand(2, 3, 4), [0, 2]),
+    "rank4_dim2": lambda: (torch.rand(2, 3, 4, 5), [2]),
+    "rank4_multi": lambda: (torch.rand(2, 3, 4, 5), [1, 3]),
+    "rank4_all_neg": lambda: (torch.rand(2, 3, 4, 5), [-4, -3, -2, -1]),
+}
+
+test_data_suite_bf16 = {
+    "rank4_multi_bf16": lambda: (
+        torch.rand(2, 3, 4, 5, dtype=torch.bfloat16),
+        [1, 3],
+    ),
+}
+
+# The fp8 cases compare the lowered output against eager torch.flip. Eager
+# torch.flip on a float8 CPU tensor raises NotImplementedError ("flip_cpu not
+# implemented for Float8_*") when the flipped dims include the last axis, so
+# these cases flip outer axes only to keep that comparison runnable.
+test_data_suite_fp8 = {
+    "rank4_multi_fp8e4m3": lambda: (
+        torch.rand(2, 3, 4, 5).to(torch.float8_e4m3fn),
+        [0, 2],
+        "fp8e4m3",
+    ),
+    "rank4_multi_fp8e5m2": lambda: (
+        torch.rand(2, 3, 4, 5).to(torch.float8_e5m2),
+        [0, 2],
+        "fp8e5m2",
+    ),
+}
+
+# U55 has no REVERSE, so flip must not be delegated there.
+test_data_suite_u55_reject = {
+    "rank4_dim2": lambda: (torch.rand(2, 3, 4, 5), [2]),
+}
+
+# flip over no dims is the identity;
+# the partitioner must prune the trivial partition instead of delegating it.
+test_data_suite_empty = {
+    "empty_dims": lambda: (torch.rand(4, 5), []),
+}
+
+
+class Flip(torch.nn.Module):
+    def __init__(self, dims):
+        super().__init__()
+        self.dims = dims
+
+    def forward(self, x: torch.Tensor):
+        return torch.flip(x, self.dims)
+
+
+class FlipBool(torch.nn.Module):
+    def forward(self, x: torch.Tensor):
+        return torch.flip(x > 0, [1])
+
+
+@common.parametrize("test_data", test_data_suite)
+def test_flip_tosa_FP(test_data):
+    data, dims = test_data()
+    pipeline = TosaPipelineFP[input_t1](Flip(dims), (data,), aten_op, exir_op)
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite)
+def test_flip_tosa_FP_reverse_count(test_data):
+    """Each flip dim must lower to exactly one TOSA REVERSE."""
+    data, dims = test_data()
+    pipeline = TosaPipelineFP[input_t1](Flip(dims), (data,), aten_op, exir_op)
+    pipeline.count_tosa_ops({"REVERSE": len(dims)})
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite_bf16)
+def test_flip_tosa_FP_bf16(test_data):
+    data, dims = test_data()
+    pipeline = TosaPipelineFP[input_t1](
+        Flip(dims), (data,), aten_op, exir_op, tosa_extensions=["bf16"]
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite_fp8)
+def test_flip_tosa_FP_fp8(test_data):
+    data, dims, tosa_extension = test_data()
+    pipeline = TosaPipelineFP[input_t1](
+        Flip(dims),
+        (data,),
+        aten_op,
+        exir_op,
+        tosa_extensions=[tosa_extension],
+    )
+    pipeline.count_tosa_ops({"REVERSE": len(dims)})
+    pipeline.run()
+
+
+def test_flip_bool_tosa_FP():
+    """Flip on a bool mask (REVERSE supports bool)."""
+    pipeline = TosaPipelineFP[input_t1](
+        FlipBool(), (torch.randn(2, 4),), aten_op, exir_op
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite)
+def test_flip_tosa_INT(test_data):
+    data, dims = test_data()
+    pipeline = TosaPipelineINT[input_t1](Flip(dims), (data,), aten_op, exir_op)
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite)
+def test_flip_16a8w_tosa_INT(test_data):
+    """16A8W quantization (16-bit activations, 8-bit weights)."""
+    data, dims = test_data()
+    pipeline = TosaPipelineINT[input_t1](
+        Flip(dims),
+        (data,),
+        aten_op,
+        exir_op=[],
+        per_channel_quantization=False,
+        use_to_edge_transform_and_lower=True,
+        tosa_extensions=["int16"],
+    )
+    pipeline.quantizer.set_global(
+        get_symmetric_a16w8_quantization_config(is_per_channel=False)
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite_u55_reject)
+def test_flip_u55_INT_not_delegated(test_data):
+    data, dims = test_data()
+    pipeline = OpNotSupportedPipeline[input_t1](
+        Flip(dims),
+        (data,),
+        non_delegated_ops={exir_op: 1},
+        quantize=True,
+        u55_subset=True,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite_empty)
+def test_flip_empty_dims_not_delegated(test_data):
+    """Flip over no dims is the identity."""
+    data, dims = test_data()
+    pipeline = OpNotSupportedPipeline[input_t1](
+        Flip(dims),
+        (data,),
+        non_delegated_ops={exir_op: 1},
+        quantize=True,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite)
+@common.XfailIfNoCorstone320
+def test_flip_u85_INT(test_data):
+    data, dims = test_data()
+    pipeline = EthosU85PipelineINT[input_t1](
+        Flip(dims),
+        (data,),
+        aten_ops=[],
+        exir_ops=[],
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite)
+@common.SkipIfNoModelConverter
+def test_flip_vgf_no_quant(test_data):
+    data, dims = test_data()
+    pipeline = VgfPipeline[input_t1](
+        Flip(dims),
+        (data,),
+        aten_op,
+        exir_op,
+        quantize=False,
+    )
+    pipeline.run()
+
+
+@common.parametrize("test_data", test_data_suite)
+@common.SkipIfNoModelConverter
+def test_flip_vgf_quant(test_data):
+    data, dims = test_data()
+    pipeline = VgfPipeline[input_t1](
+        Flip(dims),
+        (data,),
+        aten_op,
+        exir_op,
+        quantize=True,
+    )
+    pipeline.run()

backends/arm/test/targets.bzl

@@ -38,6 +38,7 @@ def define_arm_tests():
         "ops/test_to_copy.py",
         "ops/test_exp.py",
         "ops/test_fft.py",
+        "ops/test_flip.py",
         "ops/test_reciprocal.py",
         "ops/test_mean_dim.py",
         "ops/test_var.py",