From ecf2b059b3bb89f61bdac32326bd5dac5e01ecf8 Mon Sep 17 00:00:00 2001 From: Saoirse Stewart Date: Mon, 22 Jun 2026 14:01:10 +0100 Subject: [PATCH] Arm backend: Add TOSA dialect data layout node visitors Signed-off-by: Saoirse Stewart --- .../arm/_passes/aten_to_tosa_data_layout.py | 98 +++++++++++++++++++ backends/arm/_passes/exir_to_tosa_pass.py | 16 +++ .../_passes/insert_data_layout_casts_pass.py | 2 +- backends/arm/operators/__init__.py | 10 +- .../{op_cat.py => op_tosa_concat.py} | 6 +- .../{op_view.py => op_tosa_reshape.py} | 2 +- backends/arm/operators/op_tosa_reverse.py | 64 ++++++++++++ .../{op_repeat.py => op_tosa_tile.py} | 2 +- .../{op_permute.py => op_tosa_transpose.py} | 2 +- .../misc/test_tosa_data_layout_visitors.py | 61 ++++++++++++ .../test_decompose_dynamic_full_pass.py | 8 +- 11 files changed, 254 insertions(+), 17 deletions(-) create mode 100644 backends/arm/_passes/aten_to_tosa_data_layout.py rename backends/arm/operators/{op_cat.py => op_tosa_concat.py} (93%) rename backends/arm/operators/{op_view.py => op_tosa_reshape.py} (98%) create mode 100644 backends/arm/operators/op_tosa_reverse.py rename backends/arm/operators/{op_repeat.py => op_tosa_tile.py} (98%) rename backends/arm/operators/{op_permute.py => op_tosa_transpose.py} (98%) create mode 100644 backends/arm/test/misc/test_tosa_data_layout_visitors.py diff --git a/backends/arm/_passes/aten_to_tosa_data_layout.py b/backends/arm/_passes/aten_to_tosa_data_layout.py new file mode 100644 index 00000000000..025dfeeec9e --- /dev/null +++ b/backends/arm/_passes/aten_to_tosa_data_layout.py @@ -0,0 +1,98 @@ +# 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 collections.abc import Sequence +from typing import cast + +from executorch.backends.transforms.aten_to_dialect_pass import ( + AtenToDialectPass, + DialectNodeSpec, +) +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx import Node + + +def _get_arg(node: Node, index: int, name: str, default=None): + if len(node.args) > index: + return node.args[index] + return node.kwargs.get(name, default) + + +def _normalize_dim(dim: int, rank: int) -> int: + return (dim + rank) % rank + + +def _input_rank(node: Node) -> int: + input_node = cast(Node, node.args[0]) + return len(input_node.meta["val"].shape) + + +def _rewrite_cat(node: Node, pass_: AtenToDialectPass) -> DialectNodeSpec: + tensors = cast(Sequence[Node], node.args[0]) + dim = _get_arg(node, 1, "dim", 0) + first_tensor = tensors[0] + axis = _normalize_dim(cast(int, dim), len(first_tensor.meta["val"].shape)) + return DialectNodeSpec( + exir_ops.backend.tosa.CONCAT.default, + (tensors,), + {"axis": axis}, + ) + + +def _rewrite_view_copy(node: Node, pass_: AtenToDialectPass) -> DialectNodeSpec: + return DialectNodeSpec( + exir_ops.backend.tosa.RESHAPE.default, + node.args, + dict(node.kwargs), + ) + + +def _rewrite_repeat(node: Node, pass_: AtenToDialectPass) -> DialectNodeSpec: + return DialectNodeSpec( + exir_ops.backend.tosa.TILE.default, + node.args, + dict(node.kwargs), + ) + + +def _rewrite_permute_copy(node: Node, pass_: AtenToDialectPass) -> DialectNodeSpec: + permutation = list(cast(Sequence[int], _get_arg(node, 1, "dims"))) + rank = _input_rank(node) + permutation = [_normalize_dim(dim, rank) for dim in permutation] + return DialectNodeSpec( + exir_ops.backend.tosa.TRANSPOSE.default, + (node.args[0], permutation), + {}, + ) + + +def _rewrite_flip(node: Node, pass_: AtenToDialectPass) -> DialectNodeSpec | None: + dims = list(cast(Sequence[int], _get_arg(node, 1, "dims"))) + if len(dims) != 1: + return None + + return DialectNodeSpec( + exir_ops.backend.tosa.REVERSE.default, + (node.args[0],), + {"axis": _normalize_dim(dims[0], _input_rank(node))}, + ) + + +def rewrite_data_layout_operator( + node: Node, pass_: AtenToDialectPass +) -> DialectNodeSpec | None: + match node.target: + case exir_ops.edge.aten.cat.default: + return _rewrite_cat(node, pass_) + case exir_ops.edge.aten.view_copy.default: + return _rewrite_view_copy(node, pass_) + case exir_ops.edge.aten.repeat.default: + return _rewrite_repeat(node, pass_) + case exir_ops.edge.aten.permute_copy.default: + return _rewrite_permute_copy(node, pass_) + case exir_ops.edge.aten.flip.default: + return _rewrite_flip(node, pass_) + case _: + return None diff --git a/backends/arm/_passes/exir_to_tosa_pass.py b/backends/arm/_passes/exir_to_tosa_pass.py index 865a8b47cd7..942f88d517d 100644 --- a/backends/arm/_passes/exir_to_tosa_pass.py +++ b/backends/arm/_passes/exir_to_tosa_pass.py @@ -9,6 +9,9 @@ from executorch.backends.arm._passes.aten_to_tosa_activation_functions import ( get_activation_replacement, ) +from executorch.backends.arm._passes.aten_to_tosa_data_layout import ( + rewrite_data_layout_operator, +) from executorch.backends.arm._passes.aten_to_tosa_tensor_operators import ( rewrite_argmax, rewrite_binary_operator, @@ -118,3 +121,16 @@ def _get_activation_replacement( node: Node, pass_: AtenToDialectPass ) -> DialectNodeSpec | None: return get_activation_replacement(node, pass_) + + +@register_dialect_substitutions( + exir_ops.edge.aten.cat.default, + exir_ops.edge.aten.flip.default, + exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.aten.repeat.default, + exir_ops.edge.aten.view_copy.default, +) +def _get_data_layout_replacement( + node: Node, pass_: AtenToDialectPass +) -> DialectNodeSpec | None: + return rewrite_data_layout_operator(node, pass_) diff --git a/backends/arm/_passes/insert_data_layout_casts_pass.py b/backends/arm/_passes/insert_data_layout_casts_pass.py index 07a2d186895..39714c05147 100644 --- a/backends/arm/_passes/insert_data_layout_casts_pass.py +++ b/backends/arm/_passes/insert_data_layout_casts_pass.py @@ -35,7 +35,7 @@ class InsertDataLayoutCastsPass(ArmOpTargetedPass): _concat_ops = { exir_ops.edge.aten.cat.default, - exir_ops.edge.aten.concatenate.default, + exir_ops.backend.tosa.CONCAT.default, } _single_input_ops = { exir_ops.edge.aten.constant_pad_nd.default, diff --git a/backends/arm/operators/__init__.py b/backends/arm/operators/__init__.py index 6abc3bfb83e..15cb14eddf1 100644 --- a/backends/arm/operators/__init__.py +++ b/backends/arm/operators/__init__.py @@ -14,11 +14,7 @@ op_amax, op_amin, op_any, - op_cat, op_cond_if, - op_flip, - op_permute, - op_repeat, op_sum, op_to_dim_order_copy, op_tosa_abs, @@ -32,6 +28,7 @@ op_tosa_ceil, op_tosa_clamp, op_tosa_clz, + op_tosa_concat, op_tosa_conv2d, op_tosa_conv2d_block_scaled, op_tosa_conv3d, @@ -61,7 +58,9 @@ op_tosa_pow, op_tosa_reciprocal, op_tosa_rescale, + op_tosa_reshape, op_tosa_resize, + op_tosa_reverse, op_tosa_rshift_tensor, op_tosa_rsqrt, op_tosa_scatter, @@ -72,8 +71,9 @@ op_tosa_sub, op_tosa_table, op_tosa_tanh, + op_tosa_tile, + op_tosa_transpose, op_tosa_transpose_conv2d, - op_view, op_where, op_while, ) diff --git a/backends/arm/operators/op_cat.py b/backends/arm/operators/op_tosa_concat.py similarity index 93% rename from backends/arm/operators/op_cat.py rename to backends/arm/operators/op_tosa_concat.py index 97ea651cb12..ba61793106c 100644 --- a/backends/arm/operators/op_cat.py +++ b/backends/arm/operators/op_tosa_concat.py @@ -23,7 +23,7 @@ @register_node_visitor class CatVisitor(NodeVisitor): - target = "aten.cat.default" + target = "tosa.CONCAT.default" def __init__(self, *args): super().__init__(*args) @@ -58,9 +58,7 @@ def define_node( self.tosa_spec, ) - dim = 0 if len(inputs) < 2 else inputs[1].number - rank = len(output.shape) - dim = (dim + rank) % rank + dim = node.kwargs["axis"] attr = ts.TosaSerializerAttribute() attr.ConcatAttribute(dim) diff --git a/backends/arm/operators/op_view.py b/backends/arm/operators/op_tosa_reshape.py similarity index 98% rename from backends/arm/operators/op_view.py rename to backends/arm/operators/op_tosa_reshape.py index 6d399b65801..aef67f64ee5 100644 --- a/backends/arm/operators/op_view.py +++ b/backends/arm/operators/op_tosa_reshape.py @@ -23,7 +23,7 @@ @register_node_visitor class ViewVisitor(NodeVisitor): - target = "aten.view_copy.default" + target = "tosa.RESHAPE.default" def __init__(self, *args): super().__init__(*args) diff --git a/backends/arm/operators/op_tosa_reverse.py b/backends/arm/operators/op_tosa_reverse.py new file mode 100644 index 00000000000..01a91fd05ed --- /dev/null +++ b/backends/arm/operators/op_tosa_reverse.py @@ -0,0 +1,64 @@ +# 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 + +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 TosaReverseVisitor(NodeVisitor): + target = "tosa.REVERSE.default" + + 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, 1) + validate_same_dtype(self.target, [inputs[0], output], ts) + validate_valid_dtype( + self.target, + [inputs[0], output], + supported_dtypes, + self.tosa_spec, + ) + + attr = ts.TosaSerializerAttribute() + attr.ReverseAttribute(node.kwargs["axis"]) + self._serialize_operator( + node, + tosa_graph, + ts.Op.REVERSE, + [inputs[0].name], + [output.name], + attr, + ) diff --git a/backends/arm/operators/op_repeat.py b/backends/arm/operators/op_tosa_tile.py similarity index 98% rename from backends/arm/operators/op_repeat.py rename to backends/arm/operators/op_tosa_tile.py index f990dbef64b..5bf30d9929f 100644 --- a/backends/arm/operators/op_repeat.py +++ b/backends/arm/operators/op_tosa_tile.py @@ -23,7 +23,7 @@ @register_node_visitor class RepeatVisitor(NodeVisitor): - target = "aten.repeat.default" + target = "tosa.TILE.default" def __init__(self, *args): super().__init__(*args) diff --git a/backends/arm/operators/op_permute.py b/backends/arm/operators/op_tosa_transpose.py similarity index 98% rename from backends/arm/operators/op_permute.py rename to backends/arm/operators/op_tosa_transpose.py index 2418131af3e..bfc867de543 100644 --- a/backends/arm/operators/op_permute.py +++ b/backends/arm/operators/op_tosa_transpose.py @@ -24,7 +24,7 @@ @register_node_visitor class PermuteVisitor(NodeVisitor): - target = "aten.permute_copy.default" + target = "tosa.TRANSPOSE.default" def __init__(self, *args): super().__init__(*args) diff --git a/backends/arm/test/misc/test_tosa_data_layout_visitors.py b/backends/arm/test/misc/test_tosa_data_layout_visitors.py new file mode 100644 index 00000000000..c69aa92e0bf --- /dev/null +++ b/backends/arm/test/misc/test_tosa_data_layout_visitors.py @@ -0,0 +1,61 @@ +# 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 types import SimpleNamespace +from typing import Any, cast + +import pytest +import tosa_serializer as ts +from executorch.backends.arm.operators.op_tosa_reverse import TosaReverseVisitor +from executorch.backends.arm.tosa.mapping import TosaArg +from executorch.backends.arm.tosa.specification import TosaSpecification +from torch.fx import Node + + +class CapturingTosaGraph: + def __init__(self) -> None: + self.operators: list[tuple[Any, tuple[Any, ...], tuple[Any, ...], Any, Any]] = ( + [] + ) + + def addOperator(self, op, inputs, outputs, attributes=None, location=None): + self.operators.append((op, tuple(inputs), tuple(outputs), attributes, location)) + + +def _tensor_arg(name: str, dtype) -> SimpleNamespace: + return SimpleNamespace(name=name, dtype=dtype) + + +def test_reverse_visitor_emits_tosa_reverse() -> None: + visitor = TosaReverseVisitor(TosaSpecification.create_from_string("TOSA-1.1+FP")) + tosa_graph = CapturingTosaGraph() + + visitor.define_node( + cast(Node, SimpleNamespace(kwargs={"axis": 1})), + tosa_graph, + [cast(TosaArg, _tensor_arg("input", ts.DType.FP32))], + cast(TosaArg, _tensor_arg("output", ts.DType.FP32)), + ) + + assert len(tosa_graph.operators) == 1 + op, inputs, outputs, _attributes, _location = tosa_graph.operators[0] + assert op == ts.Op.REVERSE + assert inputs == ("input",) + assert outputs == ("output",) + + +def test_reverse_visitor_rejects_bfloat16_without_extension() -> None: + visitor = TosaReverseVisitor(TosaSpecification.create_from_string("TOSA-1.1+FP")) + tosa_graph = CapturingTosaGraph() + + with pytest.raises(ValueError): + visitor.define_node( + cast(Node, SimpleNamespace(kwargs={"axis": 0})), + tosa_graph, + [cast(TosaArg, _tensor_arg("input", ts.DType.BF16))], + cast(TosaArg, _tensor_arg("output", ts.DType.BF16)), + ) + + assert not tosa_graph.operators diff --git a/backends/arm/test/passes/test_decompose_dynamic_full_pass.py b/backends/arm/test/passes/test_decompose_dynamic_full_pass.py index d8b56cac291..123914e4c9e 100644 --- a/backends/arm/test/passes/test_decompose_dynamic_full_pass.py +++ b/backends/arm/test/passes/test_decompose_dynamic_full_pass.py @@ -140,17 +140,17 @@ def test_backend_pipeline_decomposes_dynamic_full_like() -> None: if node.op == "call_function" and node.target == exir_ops.edge.aten.full_like.default ] - repeat_nodes = [ + tile_nodes = [ node for node in graph_module.graph.nodes if node.op == "call_function" - and node.target == exir_ops.edge.aten.repeat.default + and node.target == exir_ops.backend.tosa.TILE.default ] assert not full_nodes assert not full_like_nodes - assert len(repeat_nodes) == 1 - assert repeat_nodes[0].args[1][1] == 3 + assert len(tile_nodes) == 1 + assert tile_nodes[0].args[1][1] == 3 def test_decompose_dynamic_full_leaves_static_full_unchanged() -> None: