Arm backend: Add TOSA dialect data layout node visitors

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

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_)

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,

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,
 )

backends/arm/operators/op_cat.py → 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)

backends/arm/operators/op_view.py → 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)

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,
+        )

backends/arm/operators/op_repeat.py → 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)

backends/arm/operators/op_permute.py → 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)

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

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: