Arm backend: Add tosa-dialect op for SLICE (pytorch#18093)

For the upcoming shape-extension, convolution will insert slices with
symints are arguments. To make this work we need to move logic out of
the slice-visitor and into a pass that rewrites aten.slice.Tensor to a
tosa.SLICE. The biggest difference is that tosa.SLICE operates on lists
as start and size arguments. This simplifies the conversion of symbolic
expressions to tosa-shape operations.



cc @digantdesai @freddan80 @per @zingo @mansnils @Sebastian-Larsson
@robell

Signed-off-by: Oscar Andersson <oscar.andersson@arm.com>
Co-authored-by: Per Åstrand <per.astrand@arm.com>

Author: oscarandersson8218

backends/arm/_passes/__init__.py

@@ -136,6 +136,7 @@
 from .rewrite_le_lt_to_ge_gt_pass import RewriteLeLtToGeGtPass  # noqa
 from .rewrite_matmul import RewriteMatmulPass  # noqa
 from .rewrite_pad import RewritePadPass  # noqa
+from .rewrite_slice import RewriteSlicePass  # noqa
 from .rewrite_upsample import RewriteUpsamplePass  # noqa
 from .scalars_to_attribute_pass import ScalarsToAttributePass  # noqa
 from .size_adjust_input_pass import SizeAdjustInputPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -121,6 +121,7 @@
     RewriteLeLtToGeGtPass,
     RewriteMatmulPass,
     RewritePadPass,
+    RewriteSlicePass,
     RewriteUpsamplePass,
     ScalarsToAttributePass,
     SizeAdjustInputPass,
@@ -374,6 +375,7 @@ def _tosa_pipeline(
                 RewriteConvPass(exported_program),
                 RewriteMatmulPass(),
                 RewritePadPass(),
+                RewriteSlicePass(),
             ]
         )
 

backends/arm/_passes/rewrite_slice.py

@@ -0,0 +1,74 @@
+# 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 ArmPass
+
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass, ProxyValue
+
+from torch import SymInt
+
+
+class RewriteSlicePass(ArmPass):
+    """Rewrite slice operations with step of 1 to TOSA slice operators."""
+
+    _passes_required_after: Set[Type[ExportPass]] = set()
+
+    def _fixup_start(self, start, input_shape, dim) -> int:
+        """Convert negative and out-of-bounds start indices to valid positive
+        indices.
+        """
+        idx = start or 0
+        if idx < 0:
+            idx = start % input_shape[dim]
+        if idx >= input_shape[dim]:
+            idx = input_shape[dim] - 1
+        return idx
+
+    def call_operator(self, op, args, kwargs, meta, updated=False) -> ProxyValue:
+        if op not in (exir_ops.edge.aten.slice_copy.Tensor,):
+            return super().call_operator(op, args, kwargs, meta, updated)
+
+        if len(args) == 5 and args[4] != 1:
+            raise ValueError(
+                f"Only slice with 4 arguments and step of 1 is supported, got {len(args)} arguments and step {args[4]}"
+            )
+        input, dim, start, end = args[:4]
+        input_shape = input.data.shape
+        start_index = self._fixup_start(start, input_shape, dim)
+
+        start_list = [0] * len(input_shape)
+        start_list[dim] = start_index
+        size_list = list(meta.data["val"].shape)
+
+        if any(isinstance(dim, SymInt) for dim in start_list):
+            starts = start_list
+        else:
+            starts = super().call_shape_operator(  # type: ignore[assignment]
+                exir_ops.backend.tosa.CONST_SHAPE.default,
+                (start_list,),
+                {},
+                meta,
+                True,
+            )
+        if any(isinstance(dim, SymInt) for dim in size_list):
+            sizes = size_list
+        else:
+            sizes = super().call_shape_operator(  # type: ignore[assignment]
+                exir_ops.backend.tosa.CONST_SHAPE.default,
+                (size_list,),
+                {},
+                meta,
+                True,
+            )
+
+        return super().call_operator(
+            exir_ops.backend.tosa.SLICE.default,
+            (input, starts, sizes),
+            kwargs,
+            meta,
+            updated=True,
+        )

backends/arm/operators/__init__.py

@@ -46,7 +46,6 @@
     op_rsqrt,
     op_sigmoid,
     op_sin,
-    op_slice,
     op_sub,
     op_sum,
     op_tanh,
@@ -61,6 +60,7 @@
     op_tosa_resize,
     op_tosa_scatter,
     op_tosa_shapes,
+    op_tosa_slice,
     op_tosa_table,
     op_tosa_transpose,
     op_tosa_transpose_conv2d,

backends/arm/operators/op_slice.py

@@ -0,0 +1,44 @@
+# 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 tosa_serializer as ts
+
+from executorch.backends.arm.operators.node_visitor import (
+    NodeVisitor,
+    register_node_visitor,
+)
+from executorch.backends.arm.tosa.mapping import TosaArg
+from torch.fx import Node
+
+
+@register_node_visitor
+class SliceVisitor(NodeVisitor):
+    target = "tosa.SLICE.default"
+
+    def __init__(self, *args):
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: Node,
+        tosa_graph: Any,
+        inputs: List[TosaArg],
+        output: TosaArg,
+    ) -> None:
+
+        input_node, starts, sizes = inputs
+        attr = ts.TosaSerializerAttribute()
+        attr.SliceAttribute()
+        self._serialize_operator(
+            node,
+            tosa_graph,
+            ts.Op.SLICE,
+            [input_node.name, starts.name, sizes.name],
+            [output.name],
+            attr,
+        )

backends/arm/operators/op_tosa_slice.py

@@ -14,6 +14,7 @@
     resize,
     scatter,
     shape_ops,
+    slice,
     table,
     transpose,
     transpose_conv2d,

backends/arm/tosa/dialect/__init__.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.
+
+import torch
+from executorch.backends.arm.tosa.dialect.lib import TosaValueError
+from executorch.backends.arm.tosa.dialect.ops_registration import register_fake_tosa_op
+
+from executorch.backends.arm.tosa.specification import (
+    get_context_spec,
+    TosaSpecification,
+)
+
+
+@register_fake_tosa_op(
+    "SLICE(Tensor input1, SymInt[] start, SymInt[] size) -> Tensor",  # schema
+    TosaSpecification.all_versions_and_profiles(),  # target TOSA specifications
+)
+def SLICE(a, start, size):
+    tosa_spec = get_context_spec()
+
+    # Rank validation
+    input_rank = a.dim()
+    if input_rank != len(start):
+        raise TosaValueError(
+            f"start list does not have the same rank {len(start)} as input {input_rank}"
+        )
+    if len(start) != len(size):
+        raise TosaValueError(
+            f"size list does not have the same rank {len(size)} as start list {len(start)}"
+        )
+
+    # Shape validation
+    for i in range(len(start)):
+        dim_start = start[i]
+        if dim_start < 0 or dim_start > a.shape[i]:
+            raise TosaValueError(
+                f"Expected start values between [0, {a.shape[i]}] but got {dim_start}"
+            )
+        dim_size = size[i]
+        if dim_size < 0 or dim_start + dim_size > a.shape[i]:
+            raise TosaValueError(
+                f"Expected start + size values between [0, {a.shape[i]}] but got {dim_start + dim_size}"
+            )
+
+    # Dtype validation
+    supported_dtypes = [torch.bool]
+    if tosa_spec.support_integer():
+        supported_dtypes += [torch.int8, torch.int16, torch.int32]
+    if tosa_spec.support_float():
+        supported_dtypes += [torch.float16, torch.float32]
+    if tosa_spec.support_extension("bf16"):
+        supported_dtypes += [torch.bfloat16]
+
+    if a.dtype not in supported_dtypes:
+        raise TosaValueError(
+            f"Unsupported dtype {a.dtype} for SLICE. Supported dtypes are {supported_dtypes}"
+        )
+
+    return torch.empty(size=size, dtype=a.dtype)