Arm backend: Fuse consecutive slices (pytorch#20839)

Fuse consecutive slices with static shapes.


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

Signed-off-by: Oscar Andersson <oscar.andersson@arm.com>

Author: oscarandersson8218

backends/arm/_passes/__init__.py

@@ -114,6 +114,7 @@
 from .fuse_batch_norm2d_pass import FuseBatchNorm2dPass  # noqa
 from .fuse_consecutive_concat_shapes import FuseConsecutiveConcatShapesPass  # noqa
 from .fuse_consecutive_rescales_pass import FuseConsecutiveRescalesPass  # noqa
+from .fuse_consecutive_slices_pass import FuseConsecutiveSlicesPass  # noqa
 from .fuse_constant_ops_pass import (  # noqa
     ComputeConstantOpsAOTPass,
     FuseConstantArgsPass,

backends/arm/_passes/arm_pass_manager.py

@@ -110,6 +110,7 @@
     FuseBatchNorm2dPass,
     FuseConsecutiveConcatShapesPass,
     FuseConsecutiveRescalesPass,
+    FuseConsecutiveSlicesPass,
     FuseConstantArgsPass,
     FuseDuplicateUsersPass,
     FuseEqualPlaceholdersPass,
@@ -630,6 +631,7 @@ def _tosa_pipeline(
                 RewriteHighRankSingletonPermutePass(),
                 DecomposePermuteForU55Pass(),
                 RewriteSlicePass(),
+                FuseConsecutiveSlicesPass(),
                 InsertConstShapesPass(),
             ]
         )

backends/arm/_passes/fuse_consecutive_slices_pass.py

@@ -0,0 +1,112 @@
+# 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 Sequence, Set, Type
+
+import torch
+from executorch.backends.arm._passes import ArmPass
+from executorch.backends.arm._passes.arm_pass_utils import create_node
+from executorch.backends.arm.tosa.mapping import TosaSpecialDtype
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import ExportPass, PassResult
+from torch.fx import GraphModule, Node
+
+_SLICE_INPUT = 0
+_SLICE_START = 1
+_SLICE_SIZE = 2
+
+
+class FuseConsecutiveSlicesPass(ArmPass):
+    """Fuse consecutive ``tosa.SLICE`` operations into one ``tosa.SLICE``.
+
+    For static slice operands, ``slice(slice(x, start_1, size_1), start_2,
+    size_2)`` is equivalent to ``slice(x, start_1 + start_2, size_2)``. The
+    pass rewrites the second slice to use the original input and removes the
+    first slice when it becomes dead. Dynamic shape operands are left unchanged.
+
+    """
+
+    _passes_required_after: Set[Type[ExportPass]] = set()
+
+    def call(self, graph_module: GraphModule) -> PassResult:
+        graph = graph_module.graph
+        modified = False
+
+        for node in list(graph.nodes):
+            if self._try_fuse_slice(node):
+                modified = True
+
+        if modified:
+            graph.eliminate_dead_code()
+            graph.lint()
+            graph_module.recompile()
+
+        return PassResult(graph_module, modified)
+
+    @staticmethod
+    def _is_slice(node: Node) -> bool:
+        return (
+            node.op == "call_function"
+            and node.target == exir_ops.backend.tosa.SLICE.default
+        )
+
+    @staticmethod
+    def _constant_shape(shape_arg) -> list[int] | None:
+        if isinstance(shape_arg, Node):
+            if shape_arg.target != exir_ops.backend.tosa.CONST_SHAPE.default:
+                return None
+            shape_arg = shape_arg.args[0]
+
+        if not isinstance(shape_arg, Sequence) or isinstance(shape_arg, str):
+            return None
+        if not all(isinstance(dim, int) for dim in shape_arg):
+            return None
+        return list(shape_arg)
+
+    @staticmethod
+    def _create_const_shape(graph: torch.fx.Graph, shape: list[int]) -> Node:
+        node = graph.create_node(
+            "call_function",
+            exir_ops.backend.tosa.CONST_SHAPE.default,
+            args=(shape,),
+        )
+        node.meta = {
+            "val": shape,
+            TosaSpecialDtype.meta_key(): TosaSpecialDtype.SHAPE,
+        }
+        return node
+
+    def _try_fuse_slice(self, node: Node) -> bool:
+        if not self._is_slice(node):
+            return False
+
+        input_node = node.args[_SLICE_INPUT]
+        if not isinstance(input_node, Node) or not self._is_slice(input_node):
+            return False
+
+        first_start = self._constant_shape(input_node.args[_SLICE_START])
+        second_start = self._constant_shape(node.args[_SLICE_START])
+        second_size = self._constant_shape(node.args[_SLICE_SIZE])
+        if first_start is None or second_start is None or second_size is None:
+            return False
+        if not (len(first_start) == len(second_start) == len(second_size)):
+            return False
+
+        fused_start = [outer + inner for outer, inner in zip(first_start, second_start)]
+        graph = node.graph
+
+        with graph.inserting_before(node):
+            fused_start_node = self._create_const_shape(graph, fused_start)
+            fused_size_node = self._create_const_shape(graph, second_size)
+            fused_slice = create_node(
+                graph,
+                exir_ops.backend.tosa.SLICE.default,
+                args=(input_node.args[_SLICE_INPUT], fused_start_node, fused_size_node),
+                from_node=node,
+                inherit_qparams=True,
+            )
+
+        node.replace_all_uses_with(fused_slice)
+        return True

backends/arm/test/ops/test_slice.py

@@ -94,7 +94,7 @@ def test_slice_tensor_tosa_FP_fp8(test_data):
         exir_op,
         tosa_extensions=[tosa_extension],
     )
-    pipeline.count_tosa_ops({"SLICE": 3})
+    pipeline.count_tosa_ops({"SLICE": 1})
     pipeline.run()
 
 

backends/arm/test/passes/test_fuse_consecutive_slices_pass.py

@@ -0,0 +1,110 @@
+# 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 Tuple
+
+import executorch.backends.arm.tosa.dialect  # noqa: F401
+import torch
+from executorch.backends.arm._passes import FuseConsecutiveSlicesPass, RewriteSlicePass
+from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
+from executorch.backends.arm.tosa.mapping import TosaSpecialDtype
+from executorch.backends.arm.tosa.specification import (
+    TosaLoweringContext,
+    TosaSpecification,
+)
+from executorch.exir.dialects._ops import ops as exir_ops
+
+input_t = Tuple[torch.Tensor]
+TOSA_FP_SPEC = TosaSpecification.create_from_string("TOSA-1.0+FP")
+
+
+class ConsecutiveSlice(torch.nn.Module):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        y = torch.ops.aten.slice.Tensor(x, 1, 1, 7, 1)
+        return torch.ops.aten.slice.Tensor(y, 2, 2, 5, 1)
+
+
+def _const_shape(graph: torch.fx.Graph, shape: list[int]) -> torch.fx.Node:
+    node = graph.create_node(
+        "call_function",
+        exir_ops.backend.tosa.CONST_SHAPE.default,
+        args=(shape,),
+    )
+    node.meta = {
+        "val": shape,
+        TosaSpecialDtype.meta_key(): TosaSpecialDtype.SHAPE,
+    }
+    return node
+
+
+def _slice(
+    graph: torch.fx.Graph,
+    input_node: torch.fx.Node,
+    start: list[int],
+    size: list[int],
+    output_shape: tuple[int, ...],
+) -> torch.fx.Node:
+    slice_node = graph.create_node(
+        "call_function",
+        exir_ops.backend.tosa.SLICE.default,
+        args=(input_node, _const_shape(graph, start), _const_shape(graph, size)),
+    )
+    slice_node.meta["val"] = torch.ones(output_shape)
+    return slice_node
+
+
+def _slice_nodes(graph_module: torch.fx.GraphModule) -> list[torch.fx.Node]:
+    return [
+        node
+        for node in graph_module.graph.nodes
+        if node.op == "call_function"
+        and node.target == exir_ops.backend.tosa.SLICE.default
+    ]
+
+
+def _shape_arg_value(node: torch.fx.Node, arg_index: int) -> list[int]:
+    shape_node = node.args[arg_index]
+    assert isinstance(shape_node, torch.fx.Node)
+    assert shape_node.target == exir_ops.backend.tosa.CONST_SHAPE.default
+    return list(shape_node.args[0])  # type: ignore[arg-type]
+
+
+def test_fuse_consecutive_slices_combines_static_starts() -> None:
+    graph = torch.fx.Graph()
+    x = graph.placeholder("x")
+    x.meta["val"] = torch.ones((1, 8, 10))
+    first_slice = _slice(graph, x, [0, 1, 0], [1, 6, 10], (1, 6, 10))
+    second_slice = _slice(graph, first_slice, [0, 2, 3], [1, 2, 4], (1, 2, 4))
+    graph.output(second_slice)
+    graph_module = torch.fx.GraphModule({}, graph)
+
+    with TosaLoweringContext(TOSA_FP_SPEC):
+        result = FuseConsecutiveSlicesPass().call(graph_module)
+
+    slice_nodes = _slice_nodes(result.graph_module)
+
+    assert result.modified
+    assert len(slice_nodes) == 1
+    assert slice_nodes[0].args[0] is x
+    assert _shape_arg_value(slice_nodes[0], 1) == [0, 3, 3]
+    assert _shape_arg_value(slice_nodes[0], 2) == [1, 2, 4]
+
+
+def test_fuse_consecutive_slices_tosa_FP_pipeline() -> None:
+    pipeline = PassPipeline[input_t](
+        ConsecutiveSlice(),
+        (torch.randn(1, 8, 10),),
+        ops_before_pass={
+            "executorch_exir_dialects_edge__ops_aten_slice_copy_Tensor": 2
+        },
+        ops_after_pass={"backend__ops_tosa_SLICE_default": 1},
+        ops_not_after_pass=[
+            "executorch_exir_dialects_edge__ops_aten_slice_copy_Tensor"
+        ],
+        pass_list=[RewriteSlicePass, FuseConsecutiveSlicesPass],
+    )
+    pipeline.pop_stage(-1)
+
+    pipeline.run()

backends/arm/test/passes/test_fuse_duplicate_users_pass.py

@@ -5,10 +5,16 @@
 
 from typing import Dict, Tuple
 
+import executorch.backends.arm.tosa.dialect  # noqa: F401
 import torch
 from executorch.backends.arm._passes import FuseDuplicateUsersPass
 from executorch.backends.arm.test import common
 from executorch.backends.arm.test.tester.test_pipeline import PassPipeline
+from executorch.backends.arm.tosa.specification import (
+    TosaLoweringContext,
+    TosaSpecification,
+)
+from executorch.exir.dialects._ops import ops as exir_ops
 from torch.fx import Graph, GraphModule
 
 input_t = Tuple[torch.Tensor]  # Input x
@@ -61,6 +67,36 @@ def _set_val(node, val):
     return node
 
 
+def _rescale_nodes(graph_module):
+    return [
+        node
+        for node in graph_module.graph.nodes
+        if node.op == "call_function"
+        and node.target == exir_ops.backend.tosa.RESCALE.default
+    ]
+
+
+def _graph_with_duplicate_rescale_users() -> GraphModule:
+    graph = Graph()
+    x = _set_val(graph.placeholder("x"), torch.ones(1, dtype=torch.int8))
+    rescale_args = (x, torch.int32, [1.0], 16, 0)
+    first_rescale = _set_val(
+        graph.call_function(exir_ops.backend.tosa.RESCALE.default, rescale_args),
+        torch.ones(1, dtype=torch.int32),
+    )
+    second_rescale = _set_val(
+        graph.call_function(exir_ops.backend.tosa.RESCALE.default, rescale_args),
+        torch.ones(1, dtype=torch.int32),
+    )
+    output = graph.output((first_rescale, second_rescale))
+    output.meta["val"] = (
+        torch.ones(1, dtype=torch.int32),
+        torch.ones(1, dtype=torch.int32),
+    )
+    graph.lint()
+    return GraphModule(torch.nn.Module(), graph)
+
+
 def _graph_with_users_not_in_node_order() -> GraphModule:
     graph = Graph()
     x = _set_val(graph.placeholder("x"), torch.ones(1))
@@ -116,3 +152,18 @@ def test_fuse_duplicate_users_preserves_graph_order_for_representative():
     result.graph_module.graph.lint()
     assert result.modified
     assert len(_add_node_names(result.graph_module)) == 1
+
+
+def test_fuse_duplicate_users_removes_identical_rescale_users():
+    graph_module = _graph_with_duplicate_rescale_users()
+
+    with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+INT")):
+        result = FuseDuplicateUsersPass()(graph_module)
+
+    rescale_nodes = _rescale_nodes(result.graph_module)
+
+    result.graph_module.graph.lint()
+    assert result.modified
+    assert len(rescale_nodes) == 1
+    output_node = result.graph_module.graph.output_node()
+    assert output_node.args[0] == (rescale_nodes[0], rescale_nodes[0])