Skip to content

Commit 112eb50

Browse files
vacu9708vacu9708
authored
Arm backend: Add aten.flip support via TOSA REVERSE (#20592)
## Summary TOSA `REVERSE` reverses a single axis, while `aten.flip` can take several. A multi-axis flip is a chain of single-axis reversals, so: - `DecomposeFlipPass` rewrites a multi-axis flip into single-axis flips (one dim passes through, zero dims is a no-op). - The quantizer annotates it as a shared-qspec op. The decompose runs after Q/DQ folding, so the chained flips inherit the shared quant params from the node metadata. - `FlipVisitor` lowers each single-axis flip to one `REVERSE` (`axis = dim % rank`). ```mermaid graph TD src["aten.flip(x, [d0, d1, d2])"] src -->|DecomposeFlipPass| F0 subgraph flips["FX graph"] direction LR F0["aten.flip [d0]"] --> F1["aten.flip [d1]"] --> F2["aten.flip [d2]"] end subgraph reverses["Serialized TOSA"] direction LR R0["REVERSE axis=d0"] --> R1["REVERSE axis=d1"] --> R2["REVERSE axis=d2"] end F0 -.FlipVisitor.-> R0 F1 -.FlipVisitor.-> R1 F2 -.FlipVisitor.-> R2 ``` ## Changes I referred to existing ops for each new file: `op_permute.py` for `op_flip.py`, `decompose_strided_slice_copy_pass` for `decompose_flip_pass.py`, and `test_permute.py` for `test_flip.py`. | File | Change | | --- | --- | | `_passes/decompose_flip_pass.py` (new) | `DecomposeFlipPass` | | `operators/op_flip.py` (new) | `FlipVisitor` | | `_passes/__init__.py` | export `DecomposeFlipPass` | | `_passes/arm_pass_manager.py` | register `DecomposeFlipPass` in the post-fold decompose stage | | `operators/__init__.py` | register `op_flip` | | `operator_support/tosa_profile_supported_op_lists.py` | add `aten.flip.default` to the INT and FP support lists | | `test/ops/test_flip.py` (new) | Test cases | ## Testing ``` $ pytest backends/arm/test/ops/test_flip.py -q 33 passed, 14 skipped, 7 xfailed $ lintrunner # new and changed files ok No lint issues. ``` | Test | Dtype / target | Ranks / dims | What it verifies | Result | | --- | --- | --- | --- | --- | | `tosa_FP` | FP32 | 1–4; single, negative, multi-axis | output matches `torch.flip` | pass | | `tosa_FP_reverse_count` | FP32 | 1–4; single, negative, multi-axis | one TOSA `REVERSE` per flipped dim | pass | | `tosa_FP_bf16` | bf16 | rank 4; multi-axis | output matches `torch.flip` | pass | | `tosa_FP_fp8` | fp8 e4m3 / e5m2 | rank 4; multi-axis | output matches `torch.flip` | pass | | `bool_tosa_FP` | bool | rank 2; single-axis | output matches `torch.flip` on a bool mask | pass | | `tosa_INT` | INT8 | 1–4; single, negative, multi-axis | output matches `torch.flip` | pass | | `16a8w_tosa_INT` | INT16 (16a8w) | 1–4; single, negative, multi-axis | output matches `torch.flip` | pass | | `u55_INT_not_delegated` | INT8 / U55 | rank 4; single-axis | flip is rejected and runs on CPU (U55 has no `REVERSE`) | pass | | `u85_INT` | INT8 / Ethos-U85 | 1–4; single, negative, multi-axis | runs on Ethos-U85 | xfail locally (no FVP); runs in CI | | `vgf_no_quant`, `vgf_quant` | FP & INT / VGF | 1–4; single, negative, multi-axis | runs on the VGF backend | skip locally (no converter); runs in CI | cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani Signed-off-by: Youngsik Yang <vacu9708@gmail.com>
1 parent 572f102 commit 112eb50

9 files changed

Lines changed: 354 additions & 0 deletions

File tree

‎backends/arm/_passes/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -53,6 +53,7 @@
5353
from .decompose_embedding_pass import DecomposeEmbeddingPass # noqa # noqa
5454
from .decompose_erfinv_pass import DecomposeErfinvPass # noqa
5555
from .decompose_expm1_pass import DecomposeExpm1Pass # noqa
56+
from .decompose_flip_pass import DecomposeFlipPass # noqa
5657
from .decompose_floor_divide_pass import DecomposeFloorDividePass # noqa
5758
from .decompose_gelu_pass import DecomposeGeluPass # noqa
5859
from .decompose_glu_pass import DecomposeGluPass # noqa

‎backends/arm/_passes/arm_pass_manager.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,7 @@
5757
DecomposeEmbeddingPass,
5858
DecomposeErfinvPass,
5959
DecomposeExpm1Pass,
60+
DecomposeFlipPass,
6061
DecomposeFloorDividePass,
6162
DecomposeGeluPass,
6263
DecomposeGluPass,
@@ -535,6 +536,7 @@ def _tosa_pipeline(
535536
PromoteBoolOperandsPass(),
536537
DecomposeSinhPass(),
537538
DecomposeSignPass(),
539+
DecomposeFlipPass(),
538540
DecomposeFloorDividePass(),
539541
DecomposeGeluPass(),
540542
DecomposeAddSubAlphaPass(),

‎backends/arm/_passes/decompose_flip_pass.py‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
# Copyright 2026 Arm Limited and/or its affiliates.
2+
#
3+
# This source code is licensed under the BSD-style license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
from typing import Set, Type
7+
8+
from executorch.backends.arm._passes import ArmOpTargetedPass
9+
from executorch.exir.dialects._ops import ops as exir_ops
10+
from executorch.exir.pass_base import ExportPass
11+
12+
13+
class DecomposeFlipPass(ArmOpTargetedPass):
14+
"""Decompose a multi-axis ``aten.flip`` into a chain of single-axis flips.
15+
16+
TOSA ``REVERSE`` reverses a single ``axis``; a flip over N dims is the
17+
composition of N independent single-axis reversals. Each single-axis
18+
``aten.flip`` is lowered to one ``REVERSE`` by ``FlipVisitor``. ``flip`` is a
19+
shared-qparam data-movement op, so the chained flips inherit the original
20+
node's quantization parameters via ``meta``.
21+
22+
"""
23+
24+
_passes_required_after: Set[Type[ExportPass]] = set()
25+
target_ops = {exir_ops.edge.aten.flip.default}
26+
27+
def call_operator(self, op, args, kwargs, meta, updated=False):
28+
if op not in self.target_ops or len(args[1]) == 1:
29+
return super().call_operator(op, args, kwargs, meta, updated)
30+
31+
# Chain one single-axis flip per dim; empty dims falls through as a no-op.
32+
out = args[0]
33+
for dim in args[1]:
34+
out = super().call_operator(op, (out, [dim]), kwargs, meta, updated=True)
35+
return out

‎backends/arm/operator_support/tosa_profile_supported_op_lists.py‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@
6060
exir_ops.edge.aten.linear.default,
6161
exir_ops.edge.aten.split_with_sizes_copy.default,
6262
exir_ops.edge.aten.split_copy.Tensor,
63+
exir_ops.edge.aten.flip.default,
6364
exir_ops.edge.aten.floor.default,
6465
exir_ops.edge.aten.full.default,
6566
exir_ops.edge.aten.full_like.default,
@@ -187,6 +188,7 @@
187188
exir_ops.edge.aten.linear.default,
188189
exir_ops.edge.aten.split_with_sizes_copy.default,
189190
exir_ops.edge.aten.split_copy.Tensor,
191+
exir_ops.edge.aten.flip.default,
190192
exir_ops.edge.aten.floor.default,
191193
exir_ops.edge.aten.full.default,
192194
exir_ops.edge.aten.full_like.default,

‎backends/arm/operators/__init__.py‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
op_any,
1717
op_cat,
1818
op_cond_if,
19+
op_flip,
1920
op_permute,
2021
op_repeat,
2122
op_sum,

‎backends/arm/operators/op_flip.py‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
# Copyright 2026 Arm Limited and/or its affiliates.
2+
#
3+
# This source code is licensed under the BSD-style license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
7+
from typing import Any, List
8+
9+
import torch
10+
11+
import tosa_serializer as ts
12+
13+
from executorch.backends.arm.operators.node_visitor import (
14+
NodeVisitor,
15+
register_node_visitor,
16+
)
17+
from executorch.backends.arm.operators.operator_validation_utils import (
18+
validate_num_inputs,
19+
validate_same_dtype,
20+
validate_valid_dtype,
21+
)
22+
from executorch.backends.arm.tosa.mapping import TosaArg
23+
24+
25+
@register_node_visitor
26+
class FlipVisitor(NodeVisitor):
27+
target = "aten.flip.default"
28+
29+
def __init__(self, *args):
30+
super().__init__(*args)
31+
32+
def define_node(
33+
self,
34+
node: torch.fx.Node,
35+
tosa_graph: Any,
36+
inputs: List[TosaArg],
37+
output: TosaArg,
38+
) -> None:
39+
supported_dtypes = [ts.DType.BOOL]
40+
if self.tosa_spec.support_integer():
41+
supported_dtypes.extend([ts.DType.INT8, ts.DType.INT16, ts.DType.INT32])
42+
if self.tosa_spec.support_float():
43+
supported_dtypes.extend([ts.DType.FP16, ts.DType.FP32])
44+
if self.tosa_spec.support_extension("bf16"):
45+
supported_dtypes.append(ts.DType.BF16)
46+
if self.tosa_spec.support_extension("fp8e4m3"):
47+
supported_dtypes.append(ts.DType.FP8E4M3)
48+
if self.tosa_spec.support_extension("fp8e5m2"):
49+
supported_dtypes.append(ts.DType.FP8E5M2)
50+
51+
validate_num_inputs(self.target, inputs, 2)
52+
validate_same_dtype(self.target, [inputs[0], output], ts)
53+
validate_valid_dtype(
54+
self.target,
55+
[inputs[0], output],
56+
supported_dtypes,
57+
self.tosa_spec,
58+
)
59+
60+
dims = inputs[1].special
61+
if len(dims) != 1:
62+
raise ValueError(
63+
f"{self.target} expects a single flip dim after DecomposeFlipPass, "
64+
f"got {dims}"
65+
)
66+
axis = dims[0] % len(inputs[0].shape)
67+
68+
attr = ts.TosaSerializerAttribute()
69+
attr.ReverseAttribute(axis)
70+
self._serialize_operator(
71+
node,
72+
tosa_graph,
73+
ts.Op.REVERSE,
74+
[inputs[0].name],
75+
[output.name],
76+
attr,
77+
)

‎backends/arm/test/ops/test_flip.py‎

Lines changed: 226 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,226 @@
1+
# Copyright 2026 Arm Limited and/or its affiliates.
2+
#
3+
# This source code is licensed under the BSD-style license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
6+
7+
from typing import Tuple
8+
9+
import torch
10+
11+
from executorch.backends.arm.quantizer.arm_quantizer import (
12+
get_symmetric_a16w8_quantization_config,
13+
)
14+
from executorch.backends.arm.test import common
15+
from executorch.backends.arm.test.tester.test_pipeline import (
16+
EthosU85PipelineINT,
17+
OpNotSupportedPipeline,
18+
TosaPipelineFP,
19+
TosaPipelineINT,
20+
VgfPipeline,
21+
)
22+
23+
aten_op = "torch.ops.aten.flip.default"
24+
exir_op = "executorch_exir_dialects_edge__ops_aten_flip_default"
25+
26+
input_t1 = Tuple[torch.Tensor] # Input x
27+
28+
# (input tensor, dims) — single-axis, negative-axis, and multi-axis cases.
29+
test_data_suite = {
30+
"rank1": lambda: (torch.rand(10), [0]),
31+
"rank2_dim0": lambda: (torch.rand(4, 5), [0]),
32+
"rank2_dim1_neg": lambda: (torch.rand(4, 5), [-1]),
33+
"rank3_multi": lambda: (torch.rand(2, 3, 4), [0, 2]),
34+
"rank4_dim2": lambda: (torch.rand(2, 3, 4, 5), [2]),
35+
"rank4_multi": lambda: (torch.rand(2, 3, 4, 5), [1, 3]),
36+
"rank4_all_neg": lambda: (torch.rand(2, 3, 4, 5), [-4, -3, -2, -1]),
37+
}
38+
39+
test_data_suite_bf16 = {
40+
"rank4_multi_bf16": lambda: (
41+
torch.rand(2, 3, 4, 5, dtype=torch.bfloat16),
42+
[1, 3],
43+
),
44+
}
45+
46+
# The fp8 cases compare the lowered output against eager torch.flip. Eager
47+
# torch.flip on a float8 CPU tensor raises NotImplementedError ("flip_cpu not
48+
# implemented for Float8_*") when the flipped dims include the last axis, so
49+
# these cases flip outer axes only to keep that comparison runnable.
50+
test_data_suite_fp8 = {
51+
"rank4_multi_fp8e4m3": lambda: (
52+
torch.rand(2, 3, 4, 5).to(torch.float8_e4m3fn),
53+
[0, 2],
54+
"fp8e4m3",
55+
),
56+
"rank4_multi_fp8e5m2": lambda: (
57+
torch.rand(2, 3, 4, 5).to(torch.float8_e5m2),
58+
[0, 2],
59+
"fp8e5m2",
60+
),
61+
}
62+
63+
# U55 has no REVERSE, so flip must not be delegated there.
64+
test_data_suite_u55_reject = {
65+
"rank4_dim2": lambda: (torch.rand(2, 3, 4, 5), [2]),
66+
}
67+
68+
# flip over no dims is the identity;
69+
# the partitioner must prune the trivial partition instead of delegating it.
70+
test_data_suite_empty = {
71+
"empty_dims": lambda: (torch.rand(4, 5), []),
72+
}
73+
74+
75+
class Flip(torch.nn.Module):
76+
def __init__(self, dims):
77+
super().__init__()
78+
self.dims = dims
79+
80+
def forward(self, x: torch.Tensor):
81+
return torch.flip(x, self.dims)
82+
83+
84+
class FlipBool(torch.nn.Module):
85+
def forward(self, x: torch.Tensor):
86+
return torch.flip(x > 0, [1])
87+
88+
89+
@common.parametrize("test_data", test_data_suite)
90+
def test_flip_tosa_FP(test_data):
91+
data, dims = test_data()
92+
pipeline = TosaPipelineFP[input_t1](Flip(dims), (data,), aten_op, exir_op)
93+
pipeline.run()
94+
95+
96+
@common.parametrize("test_data", test_data_suite)
97+
def test_flip_tosa_FP_reverse_count(test_data):
98+
"""Each flip dim must lower to exactly one TOSA REVERSE."""
99+
data, dims = test_data()
100+
pipeline = TosaPipelineFP[input_t1](Flip(dims), (data,), aten_op, exir_op)
101+
pipeline.count_tosa_ops({"REVERSE": len(dims)})
102+
pipeline.run()
103+
104+
105+
@common.parametrize("test_data", test_data_suite_bf16)
106+
def test_flip_tosa_FP_bf16(test_data):
107+
data, dims = test_data()
108+
pipeline = TosaPipelineFP[input_t1](
109+
Flip(dims), (data,), aten_op, exir_op, tosa_extensions=["bf16"]
110+
)
111+
pipeline.run()
112+
113+
114+
@common.parametrize("test_data", test_data_suite_fp8)
115+
def test_flip_tosa_FP_fp8(test_data):
116+
data, dims, tosa_extension = test_data()
117+
pipeline = TosaPipelineFP[input_t1](
118+
Flip(dims),
119+
(data,),
120+
aten_op,
121+
exir_op,
122+
tosa_extensions=[tosa_extension],
123+
)
124+
pipeline.count_tosa_ops({"REVERSE": len(dims)})
125+
pipeline.run()
126+
127+
128+
def test_flip_bool_tosa_FP():
129+
"""Flip on a bool mask (REVERSE supports bool)."""
130+
pipeline = TosaPipelineFP[input_t1](
131+
FlipBool(), (torch.randn(2, 4),), aten_op, exir_op
132+
)
133+
pipeline.run()
134+
135+
136+
@common.parametrize("test_data", test_data_suite)
137+
def test_flip_tosa_INT(test_data):
138+
data, dims = test_data()
139+
pipeline = TosaPipelineINT[input_t1](Flip(dims), (data,), aten_op, exir_op)
140+
pipeline.run()
141+
142+
143+
@common.parametrize("test_data", test_data_suite)
144+
def test_flip_16a8w_tosa_INT(test_data):
145+
"""16A8W quantization (16-bit activations, 8-bit weights)."""
146+
data, dims = test_data()
147+
pipeline = TosaPipelineINT[input_t1](
148+
Flip(dims),
149+
(data,),
150+
aten_op,
151+
exir_op=[],
152+
per_channel_quantization=False,
153+
use_to_edge_transform_and_lower=True,
154+
tosa_extensions=["int16"],
155+
)
156+
pipeline.quantizer.set_global(
157+
get_symmetric_a16w8_quantization_config(is_per_channel=False)
158+
)
159+
pipeline.run()
160+
161+
162+
@common.parametrize("test_data", test_data_suite_u55_reject)
163+
def test_flip_u55_INT_not_delegated(test_data):
164+
data, dims = test_data()
165+
pipeline = OpNotSupportedPipeline[input_t1](
166+
Flip(dims),
167+
(data,),
168+
non_delegated_ops={exir_op: 1},
169+
quantize=True,
170+
u55_subset=True,
171+
)
172+
pipeline.run()
173+
174+
175+
@common.parametrize("test_data", test_data_suite_empty)
176+
def test_flip_empty_dims_not_delegated(test_data):
177+
"""Flip over no dims is the identity."""
178+
data, dims = test_data()
179+
pipeline = OpNotSupportedPipeline[input_t1](
180+
Flip(dims),
181+
(data,),
182+
non_delegated_ops={exir_op: 1},
183+
quantize=True,
184+
)
185+
pipeline.run()
186+
187+
188+
@common.parametrize("test_data", test_data_suite)
189+
@common.XfailIfNoCorstone320
190+
def test_flip_u85_INT(test_data):
191+
data, dims = test_data()
192+
pipeline = EthosU85PipelineINT[input_t1](
193+
Flip(dims),
194+
(data,),
195+
aten_ops=[],
196+
exir_ops=[],
197+
)
198+
pipeline.run()
199+
200+
201+
@common.parametrize("test_data", test_data_suite)
202+
@common.SkipIfNoModelConverter
203+
def test_flip_vgf_no_quant(test_data):
204+
data, dims = test_data()
205+
pipeline = VgfPipeline[input_t1](
206+
Flip(dims),
207+
(data,),
208+
aten_op,
209+
exir_op,
210+
quantize=False,
211+
)
212+
pipeline.run()
213+
214+
215+
@common.parametrize("test_data", test_data_suite)
216+
@common.SkipIfNoModelConverter
217+
def test_flip_vgf_quant(test_data):
218+
data, dims = test_data()
219+
pipeline = VgfPipeline[input_t1](
220+
Flip(dims),
221+
(data,),
222+
aten_op,
223+
exir_op,
224+
quantize=True,
225+
)
226+
pipeline.run()

‎backends/arm/test/targets.bzl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ def define_arm_tests():
3838
"ops/test_to_copy.py",
3939
"ops/test_exp.py",
4040
"ops/test_fft.py",
41+
"ops/test_flip.py",
4142
"ops/test_reciprocal.py",
4243
"ops/test_mean_dim.py",
4344
"ops/test_var.py",

0 commit comments

Comments
 (0)