Skip to content

Commit 4c9beee

Browse files
SaoirseARMSaoirseARM
authored
Arm backend: Add TOSA dialect data layout node visitors (#20789)
Updates - CONCAT - RESHAPE - TILE - TRANSPOSE - REVERSE Adds to the aten->TOSA pass Signed-off-by: Saoirse Stewart <saoirse.stewart@arm.com>
1 parent f8b7f62 commit 4c9beee

11 files changed

Lines changed: 254 additions & 17 deletions

‎backends/arm/_passes/aten_to_tosa_data_layout.py‎

Lines changed: 98 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,98 @@
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 collections.abc import Sequence
7+
from typing import cast
8+
9+
from executorch.backends.transforms.aten_to_dialect_pass import (
10+
AtenToDialectPass,
11+
DialectNodeSpec,
12+
)
13+
from executorch.exir.dialects._ops import ops as exir_ops
14+
from torch.fx import Node
15+
16+
17+
def _get_arg(node: Node, index: int, name: str, default=None):
18+
if len(node.args) > index:
19+
return node.args[index]
20+
return node.kwargs.get(name, default)
21+
22+
23+
def _normalize_dim(dim: int, rank: int) -> int:
24+
return (dim + rank) % rank
25+
26+
27+
def _input_rank(node: Node) -> int:
28+
input_node = cast(Node, node.args[0])
29+
return len(input_node.meta["val"].shape)
30+
31+
32+
def _rewrite_cat(node: Node, pass_: AtenToDialectPass) -> DialectNodeSpec:
33+
tensors = cast(Sequence[Node], node.args[0])
34+
dim = _get_arg(node, 1, "dim", 0)
35+
first_tensor = tensors[0]
36+
axis = _normalize_dim(cast(int, dim), len(first_tensor.meta["val"].shape))
37+
return DialectNodeSpec(
38+
exir_ops.backend.tosa.CONCAT.default,
39+
(tensors,),
40+
{"axis": axis},
41+
)
42+
43+
44+
def _rewrite_view_copy(node: Node, pass_: AtenToDialectPass) -> DialectNodeSpec:
45+
return DialectNodeSpec(
46+
exir_ops.backend.tosa.RESHAPE.default,
47+
node.args,
48+
dict(node.kwargs),
49+
)
50+
51+
52+
def _rewrite_repeat(node: Node, pass_: AtenToDialectPass) -> DialectNodeSpec:
53+
return DialectNodeSpec(
54+
exir_ops.backend.tosa.TILE.default,
55+
node.args,
56+
dict(node.kwargs),
57+
)
58+
59+
60+
def _rewrite_permute_copy(node: Node, pass_: AtenToDialectPass) -> DialectNodeSpec:
61+
permutation = list(cast(Sequence[int], _get_arg(node, 1, "dims")))
62+
rank = _input_rank(node)
63+
permutation = [_normalize_dim(dim, rank) for dim in permutation]
64+
return DialectNodeSpec(
65+
exir_ops.backend.tosa.TRANSPOSE.default,
66+
(node.args[0], permutation),
67+
{},
68+
)
69+
70+
71+
def _rewrite_flip(node: Node, pass_: AtenToDialectPass) -> DialectNodeSpec | None:
72+
dims = list(cast(Sequence[int], _get_arg(node, 1, "dims")))
73+
if len(dims) != 1:
74+
return None
75+
76+
return DialectNodeSpec(
77+
exir_ops.backend.tosa.REVERSE.default,
78+
(node.args[0],),
79+
{"axis": _normalize_dim(dims[0], _input_rank(node))},
80+
)
81+
82+
83+
def rewrite_data_layout_operator(
84+
node: Node, pass_: AtenToDialectPass
85+
) -> DialectNodeSpec | None:
86+
match node.target:
87+
case exir_ops.edge.aten.cat.default:
88+
return _rewrite_cat(node, pass_)
89+
case exir_ops.edge.aten.view_copy.default:
90+
return _rewrite_view_copy(node, pass_)
91+
case exir_ops.edge.aten.repeat.default:
92+
return _rewrite_repeat(node, pass_)
93+
case exir_ops.edge.aten.permute_copy.default:
94+
return _rewrite_permute_copy(node, pass_)
95+
case exir_ops.edge.aten.flip.default:
96+
return _rewrite_flip(node, pass_)
97+
case _:
98+
return None

‎backends/arm/_passes/exir_to_tosa_pass.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,9 @@
99
from executorch.backends.arm._passes.aten_to_tosa_activation_functions import (
1010
get_activation_replacement,
1111
)
12+
from executorch.backends.arm._passes.aten_to_tosa_data_layout import (
13+
rewrite_data_layout_operator,
14+
)
1215
from executorch.backends.arm._passes.aten_to_tosa_tensor_operators import (
1316
rewrite_argmax,
1417
rewrite_binary_operator,
@@ -118,3 +121,16 @@ def _get_activation_replacement(
118121
node: Node, pass_: AtenToDialectPass
119122
) -> DialectNodeSpec | None:
120123
return get_activation_replacement(node, pass_)
124+
125+
126+
@register_dialect_substitutions(
127+
exir_ops.edge.aten.cat.default,
128+
exir_ops.edge.aten.flip.default,
129+
exir_ops.edge.aten.permute_copy.default,
130+
exir_ops.edge.aten.repeat.default,
131+
exir_ops.edge.aten.view_copy.default,
132+
)
133+
def _get_data_layout_replacement(
134+
node: Node, pass_: AtenToDialectPass
135+
) -> DialectNodeSpec | None:
136+
return rewrite_data_layout_operator(node, pass_)

‎backends/arm/_passes/insert_data_layout_casts_pass.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ class InsertDataLayoutCastsPass(ArmOpTargetedPass):
3535

3636
_concat_ops = {
3737
exir_ops.edge.aten.cat.default,
38-
exir_ops.edge.aten.concatenate.default,
38+
exir_ops.backend.tosa.CONCAT.default,
3939
}
4040
_single_input_ops = {
4141
exir_ops.edge.aten.constant_pad_nd.default,

‎backends/arm/operators/__init__.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,7 @@
1414
op_amax,
1515
op_amin,
1616
op_any,
17-
op_cat,
1817
op_cond_if,
19-
op_flip,
20-
op_permute,
21-
op_repeat,
2218
op_sum,
2319
op_to_dim_order_copy,
2420
op_tosa_abs,
@@ -32,6 +28,7 @@
3228
op_tosa_ceil,
3329
op_tosa_clamp,
3430
op_tosa_clz,
31+
op_tosa_concat,
3532
op_tosa_conv2d,
3633
op_tosa_conv2d_block_scaled,
3734
op_tosa_conv3d,
@@ -61,7 +58,9 @@
6158
op_tosa_pow,
6259
op_tosa_reciprocal,
6360
op_tosa_rescale,
61+
op_tosa_reshape,
6462
op_tosa_resize,
63+
op_tosa_reverse,
6564
op_tosa_rshift_tensor,
6665
op_tosa_rsqrt,
6766
op_tosa_scatter,
@@ -72,8 +71,9 @@
7271
op_tosa_sub,
7372
op_tosa_table,
7473
op_tosa_tanh,
74+
op_tosa_tile,
75+
op_tosa_transpose,
7576
op_tosa_transpose_conv2d,
76-
op_view,
7777
op_where,
7878
op_while,
7979
)

backends/arm/operators/op_cat.py renamed to backends/arm/operators/op_tosa_concat.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323

2424
@register_node_visitor
2525
class CatVisitor(NodeVisitor):
26-
target = "aten.cat.default"
26+
target = "tosa.CONCAT.default"
2727

2828
def __init__(self, *args):
2929
super().__init__(*args)
@@ -58,9 +58,7 @@ def define_node(
5858
self.tosa_spec,
5959
)
6060

61-
dim = 0 if len(inputs) < 2 else inputs[1].number
62-
rank = len(output.shape)
63-
dim = (dim + rank) % rank
61+
dim = node.kwargs["axis"]
6462

6563
attr = ts.TosaSerializerAttribute()
6664
attr.ConcatAttribute(dim)

backends/arm/operators/op_view.py renamed to backends/arm/operators/op_tosa_reshape.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323

2424
@register_node_visitor
2525
class ViewVisitor(NodeVisitor):
26-
target = "aten.view_copy.default"
26+
target = "tosa.RESHAPE.default"
2727

2828
def __init__(self, *args):
2929
super().__init__(*args)

‎backends/arm/operators/op_tosa_reverse.py‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
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 Any
7+
8+
import torch
9+
import tosa_serializer as ts
10+
11+
from executorch.backends.arm.operators.node_visitor import (
12+
NodeVisitor,
13+
register_node_visitor,
14+
)
15+
from executorch.backends.arm.operators.operator_validation_utils import (
16+
validate_num_inputs,
17+
validate_same_dtype,
18+
validate_valid_dtype,
19+
)
20+
from executorch.backends.arm.tosa.mapping import TosaArg
21+
22+
23+
@register_node_visitor
24+
class TosaReverseVisitor(NodeVisitor):
25+
target = "tosa.REVERSE.default"
26+
27+
def define_node(
28+
self,
29+
node: torch.fx.Node,
30+
tosa_graph: Any,
31+
inputs: list[TosaArg],
32+
output: TosaArg,
33+
) -> None:
34+
supported_dtypes = [ts.DType.BOOL]
35+
if self.tosa_spec.support_integer():
36+
supported_dtypes.extend([ts.DType.INT8, ts.DType.INT16, ts.DType.INT32])
37+
if self.tosa_spec.support_float():
38+
supported_dtypes.extend([ts.DType.FP16, ts.DType.FP32])
39+
if self.tosa_spec.support_extension("bf16"):
40+
supported_dtypes.append(ts.DType.BF16)
41+
if self.tosa_spec.support_extension("fp8e4m3"):
42+
supported_dtypes.append(ts.DType.FP8E4M3)
43+
if self.tosa_spec.support_extension("fp8e5m2"):
44+
supported_dtypes.append(ts.DType.FP8E5M2)
45+
46+
validate_num_inputs(self.target, inputs, 1)
47+
validate_same_dtype(self.target, [inputs[0], output], ts)
48+
validate_valid_dtype(
49+
self.target,
50+
[inputs[0], output],
51+
supported_dtypes,
52+
self.tosa_spec,
53+
)
54+
55+
attr = ts.TosaSerializerAttribute()
56+
attr.ReverseAttribute(node.kwargs["axis"])
57+
self._serialize_operator(
58+
node,
59+
tosa_graph,
60+
ts.Op.REVERSE,
61+
[inputs[0].name],
62+
[output.name],
63+
attr,
64+
)

backends/arm/operators/op_repeat.py renamed to backends/arm/operators/op_tosa_tile.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323

2424
@register_node_visitor
2525
class RepeatVisitor(NodeVisitor):
26-
target = "aten.repeat.default"
26+
target = "tosa.TILE.default"
2727

2828
def __init__(self, *args):
2929
super().__init__(*args)

backends/arm/operators/op_permute.py renamed to backends/arm/operators/op_tosa_transpose.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,7 @@
2424

2525
@register_node_visitor
2626
class PermuteVisitor(NodeVisitor):
27-
target = "aten.permute_copy.default"
27+
target = "tosa.TRANSPOSE.default"
2828

2929
def __init__(self, *args):
3030
super().__init__(*args)

‎backends/arm/test/misc/test_tosa_data_layout_visitors.py‎

Lines changed: 61 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,61 @@
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 types import SimpleNamespace
7+
from typing import Any, cast
8+
9+
import pytest
10+
import tosa_serializer as ts
11+
from executorch.backends.arm.operators.op_tosa_reverse import TosaReverseVisitor
12+
from executorch.backends.arm.tosa.mapping import TosaArg
13+
from executorch.backends.arm.tosa.specification import TosaSpecification
14+
from torch.fx import Node
15+
16+
17+
class CapturingTosaGraph:
18+
def __init__(self) -> None:
19+
self.operators: list[tuple[Any, tuple[Any, ...], tuple[Any, ...], Any, Any]] = (
20+
[]
21+
)
22+
23+
def addOperator(self, op, inputs, outputs, attributes=None, location=None):
24+
self.operators.append((op, tuple(inputs), tuple(outputs), attributes, location))
25+
26+
27+
def _tensor_arg(name: str, dtype) -> SimpleNamespace:
28+
return SimpleNamespace(name=name, dtype=dtype)
29+
30+
31+
def test_reverse_visitor_emits_tosa_reverse() -> None:
32+
visitor = TosaReverseVisitor(TosaSpecification.create_from_string("TOSA-1.1+FP"))
33+
tosa_graph = CapturingTosaGraph()
34+
35+
visitor.define_node(
36+
cast(Node, SimpleNamespace(kwargs={"axis": 1})),
37+
tosa_graph,
38+
[cast(TosaArg, _tensor_arg("input", ts.DType.FP32))],
39+
cast(TosaArg, _tensor_arg("output", ts.DType.FP32)),
40+
)
41+
42+
assert len(tosa_graph.operators) == 1
43+
op, inputs, outputs, _attributes, _location = tosa_graph.operators[0]
44+
assert op == ts.Op.REVERSE
45+
assert inputs == ("input",)
46+
assert outputs == ("output",)
47+
48+
49+
def test_reverse_visitor_rejects_bfloat16_without_extension() -> None:
50+
visitor = TosaReverseVisitor(TosaSpecification.create_from_string("TOSA-1.1+FP"))
51+
tosa_graph = CapturingTosaGraph()
52+
53+
with pytest.raises(ValueError):
54+
visitor.define_node(
55+
cast(Node, SimpleNamespace(kwargs={"axis": 0})),
56+
tosa_graph,
57+
[cast(TosaArg, _tensor_arg("input", ts.DType.BF16))],
58+
cast(TosaArg, _tensor_arg("output", ts.DType.BF16)),
59+
)
60+
61+
assert not tosa_graph.operators

0 commit comments

Comments
 (0)