From 64d045f44fd86a266b4a4d227971af3c1b80c995 Mon Sep 17 00:00:00 2001 From: Oscar Andersson Date: Wed, 2 Sep 2026 11:17:38 +0200 Subject: [PATCH 1/2] Arm backend: Rewrite cat->slice patterns Add RewriteCatSlicePass to rewrite static unit-stride slices along a concat dimension in terms of their overlapping original inputs. Only rewrite complete groups of slice users when the replacement does not increase the operation count. Preserve qparams for cat's tensor-list input and void materializing no-op slices for complete inputs. This avoids materializing an intermediate concat when only part of its result is needed. Change-Id: I93b2f987948ff715063e3dcfdf6d8efd501434f9 Signed-off-by: Oscar Andersson --- backends/arm/_passes/__init__.py | 1 + backends/arm/_passes/arm_pass_manager.py | 2 + .../arm/_passes/rewrite_cat_slice_pass.py | 304 ++++++++++++++++++ .../passes/test_rewrite_cat_slice_pass.py | 207 ++++++++++++ 4 files changed, 514 insertions(+) create mode 100644 backends/arm/_passes/rewrite_cat_slice_pass.py create mode 100644 backends/arm/test/passes/test_rewrite_cat_slice_pass.py diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index 032da72113f..d25887903b7 100644 --- a/backends/arm/_passes/__init__.py +++ b/backends/arm/_passes/__init__.py @@ -205,6 +205,7 @@ from .rewrite_bool_to_fp32_cast_via_int8_pass import ( # noqa RewriteBoolToFp32CastViaInt8Pass, ) +from .rewrite_cat_slice_pass import RewriteCatSlicePass # noqa from .rewrite_conv_pass import RewriteConvPass # noqa from .rewrite_high_rank_singleton_permute_pass import ( # noqa RewriteHighRankSingletonPermutePass, diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index 36665a2619d..cf2c6f5b776 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -161,6 +161,7 @@ RewriteAvgPool2dPass, RewriteBoolBitwiseToLogicalPass, RewriteBoolToFp32CastViaInt8Pass, + RewriteCatSlicePass, RewriteConvPass, RewriteHighRankSingletonPermutePass, RewriteIndexPutPass, @@ -708,6 +709,7 @@ def _tosa_pipeline( MoveDataMovementOpsToSmallerDtypePass(), MatchArgRanksPass(exported_program), RewriteHighRankSingletonPermutePass(), + RewriteCatSlicePass(), FuseConsecutiveConcatsPass(), DecomposePermuteForU55Pass(), RewriteSlicePass(), diff --git a/backends/arm/_passes/rewrite_cat_slice_pass.py b/backends/arm/_passes/rewrite_cat_slice_pass.py new file mode 100644 index 00000000000..2edf8f6ca67 --- /dev/null +++ b/backends/arm/_passes/rewrite_cat_slice_pass.py @@ -0,0 +1,304 @@ +# 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, cast, Set, Type + +import torch +from executorch.backends.arm._passes import ArmPass +from executorch.backends.arm._passes.arm_pass_utils import create_node +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult +from torch.fx import GraphModule, Node + + +_CAT = exir_ops.edge.aten.cat.default +_SLICE = exir_ops.edge.aten.slice_copy.Tensor + + +class RewriteCatSlicePass(ArmPass): + """Replace concat slices with concats of their overlapping inputs. + + Rewrites static unit-stride slices along a concat dimension without + materializing the source concat. Applies only when every source user is a + compatible slice and the replacements do not increase the operation count. + + """ + + _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 _try_rewrite_cat_slices(node): + modified = True + + if modified: + graph.eliminate_dead_code() + graph.lint() + graph_module.recompile() + + return PassResult(graph_module, modified) + + +def _is_cat(node: Node) -> bool: + return node.op == "call_function" and node.target == _CAT + + +def _is_slice(node: Node) -> bool: + return node.op == "call_function" and node.target == _SLICE + + +def _cat_inputs(node: Node) -> list[Node] | None: + if not node.args or not isinstance(node.args[0], (list, tuple)): + return None + inputs = list(node.args[0]) + return ( + inputs if all(isinstance(input_node, Node) for input_node in inputs) else None + ) + + +def _cat_dim(node: Node) -> int | None: + if len(node.args) > 1 and isinstance(node.args[1], int): + return node.args[1] + dim = node.kwargs.get("dim") + return dim if isinstance(dim, int) else None + + +def _node_shape(node: Node) -> tuple[int, ...] | None: + shape = getattr(node.meta.get("val"), "shape", None) + if shape is None or not all(isinstance(size, int) for size in shape): + return None + return tuple(shape) + + +def _slice_args(node: Node, rank: int) -> tuple[int, int, int] | None: + if len(node.args) < 4: + return None + dim_arg, start_arg, end_arg = node.args[1:4] + step_arg = node.args[4] if len(node.args) > 4 else 1 + if not all( + isinstance(value, int) for value in (dim_arg, start_arg, end_arg, step_arg) + ): + return None + dim = cast(int, dim_arg) + start = cast(int, start_arg) + end = cast(int, end_arg) + step = cast(int, step_arg) + dim = dim + rank if dim < 0 else dim + return (dim, start, end) if 0 <= dim < rank and step == 1 else None + + +def _slice_range(start: int, end: int, dim_size: int) -> tuple[int, int] | None: + start = max(0, min(dim_size, start + dim_size if start < 0 else start)) + end = max(0, min(dim_size, end + dim_size if end < 0 else end)) + return (start, end) if start < end else None + + +def _overlapping_inputs( + inputs: list[Node], dim: int, start: int, end: int, rank: int +) -> list[tuple[Node, int, int, int, int]] | None: + overlaps: list[tuple[Node, int, int, int, int]] = [] + offset = 0 + for input_index, input_node in enumerate(inputs): + input_shape = _node_shape(input_node) + if input_shape is None or len(input_shape) != rank: + return None + input_size = input_shape[dim] + input_end = offset + input_size + overlap_start = max(start, offset) + overlap_end = min(end, input_end) + if overlap_start < overlap_end: + overlaps.append( + ( + input_node, + input_index, + overlap_start - offset, + overlap_end - offset, + input_size, + ) + ) + offset = input_end + return overlaps + + +def _create_slice( + graph: torch.fx.Graph, + input_node: Node, + dim: int, + start: int, + end: int, + from_node: Node, + input_qparams: Any | None, +) -> Node: + with graph.inserting_before(from_node): + slice_node = create_node( + graph, + _SLICE, + args=(input_node, dim, start, end, 1), + from_node=from_node, + inherit_qparams=True, + ) + val = input_node.meta.get("val") + if val is not None and hasattr(val, "new_empty") and hasattr(val, "shape"): + shape = list(val.shape) + shape[dim] = end - start + slice_node.meta["val"] = val.new_empty(tuple(shape)) + if input_qparams is not None: + slice_node.meta["input_qparams"] = {0: input_qparams} + return slice_node + + +def _replacement_inputs( + node: Node, + overlaps: list[tuple[Node, int, int, int, int]], + dim: int, + input_qparams: Any | None, +) -> list[Node]: + inputs: list[Node] = [] + for input_node, _, start, end, input_size in overlaps: + if ( + start == 0 + and end == input_size + and (len(overlaps) > 1 or input_qparams is None) + ): + inputs.append(input_node) + else: + inputs.append( + _create_slice( + node.graph, + input_node, + dim, + start, + end, + node, + input_qparams, + ) + ) + return inputs + + +def _replacement_op_count( + overlaps: list[tuple[Node, int, int, int, int]], input_qparams: Any | None +) -> int: + slice_count = sum( + not ( + start == 0 + and end == input_size + and (len(overlaps) > 1 or input_qparams is None) + ) + for _, _, start, end, input_size in overlaps + ) + return slice_count + (len(overlaps) > 1) + + +def _source_replacement_op_count( + source: Node, + inputs: list[Node], + dim: int, + rank: int, + dim_size: int, + input_qparams: Any | None, +) -> int | None: + op_count = 0 + for user in source.users: + slice_args = _slice_args(user, rank) + if slice_args is None or slice_args[0] != dim: + return None + slice_range = _slice_range(slice_args[1], slice_args[2], dim_size) + if slice_range is None: + return None + overlaps = _overlapping_inputs(inputs, dim, *slice_range, rank) + if not overlaps: + return None + op_count += _replacement_op_count(overlaps, input_qparams) + return op_count + + +def _can_fuse_slice(node: Node, source: Node) -> bool: + if not _is_slice(node) or not node.args or node.args[0] is not source: + return False + + inputs = _cat_inputs(source) + dim = _cat_dim(source) + source_shape = _node_shape(source) + if inputs is None or dim is None or source_shape is None: + return False + rank = len(source_shape) + dim = dim + rank if dim < 0 else dim + slice_args = _slice_args(node, rank) + if not 0 <= dim < rank or slice_args is None or slice_args[0] != dim: + return False + slice_range = _slice_range(slice_args[1], slice_args[2], source_shape[dim]) + if slice_range is None: + return False + + return bool(_overlapping_inputs(inputs, dim, *slice_range, rank)) + + +def _try_rewrite_cat_slices(source: Node) -> bool: + if not _is_cat(source) or not source.users: + return False + + users = list(source.users) + if not all(_can_fuse_slice(user, source) for user in users): + return False + + inputs = _cat_inputs(source) + dim = _cat_dim(source) + source_shape = _node_shape(source) + if inputs is None or dim is None or source_shape is None: + return False + rank = len(source_shape) + dim = dim + rank if dim < 0 else dim + if not 0 <= dim < rank: + return False + + source_input_qparams = source.meta.get("input_qparams") + source_input_qparams = ( + source_input_qparams if isinstance(source_input_qparams, dict) else None + ) + input_qparams = ( + source_input_qparams.get(0) if source_input_qparams is not None else None + ) + replacement_op_count = _source_replacement_op_count( + source, + inputs, + dim, + rank, + source_shape[dim], + input_qparams, + ) + if replacement_op_count is None or replacement_op_count > 1 + len(users): + return False + + for node in users: + slice_args = _slice_args(node, rank) + if slice_args is None: + return False + slice_range = _slice_range(slice_args[1], slice_args[2], source_shape[dim]) + if slice_range is None: + return False + overlaps = _overlapping_inputs(inputs, dim, *slice_range, rank) + if not overlaps: + return False + replacement_inputs = _replacement_inputs(node, overlaps, dim, input_qparams) + if len(replacement_inputs) == 1: + replacement = replacement_inputs[0] + else: + with node.graph.inserting_before(node): + replacement = create_node( + node.graph, + _CAT, + args=(replacement_inputs, dim), + from_node=node, + inherit_qparams=True, + ) + if input_qparams is not None: + replacement.meta["input_qparams"] = {0: input_qparams} + node.replace_all_uses_with(replacement) + + return True diff --git a/backends/arm/test/passes/test_rewrite_cat_slice_pass.py b/backends/arm/test/passes/test_rewrite_cat_slice_pass.py new file mode 100644 index 00000000000..3fe4556b15f --- /dev/null +++ b/backends/arm/test/passes/test_rewrite_cat_slice_pass.py @@ -0,0 +1,207 @@ +# 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 cast + +import torch +from executorch.backends.arm._passes import ( + FuseConsecutiveConcatsPass, + RewriteCatSlicePass, +) +from executorch.backends.test.graph_builder import GraphBuilder +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import NodeMetadata +from torch.fx import GraphModule, Node + + +_CAT = exir_ops.edge.aten.cat.default +_SLICE = exir_ops.edge.aten.slice_copy.Tensor + + +def _concat_slice_graph( + input_channels: list[int], + slice_ranges: list[tuple[int, int]], + *, + inserted_channels: int | None = None, + input_qparams: dict[int, str] | None = None, + include_step: bool = True, +) -> tuple[GraphModule, list[Node], Node]: + builder = GraphBuilder() + inputs = [ + builder.placeholder(f"input_{index}", torch.randn(1, channels, 8, 8)) + for index, channels in enumerate(input_channels) + ] + base_concat = builder.call_operator(_CAT, (inputs, 1)) + slices = [ + builder.call_operator( + _SLICE, + ( + (base_concat, 1, start, end, 1) + if include_step + else (base_concat, 1, start, end) + ), + ) + for start, end in slice_ranges + ] + if inserted_channels is None: + output = slices[0] + else: + inserted = builder.placeholder( + "inserted", torch.randn(1, inserted_channels, 8, 8) + ) + output = builder.call_operator( + _CAT, + ([slices[0], inserted, *slices[1:]], 1), + meta=NodeMetadata({"input_qparams": input_qparams or {}}), + ) + builder.output([output]) + return builder.get_graph_module(), [input.node for input in inputs], output.node + + +def _run(graph_module: GraphModule) -> bool: + slice_result = RewriteCatSlicePass().call(graph_module) + concat_result = FuseConsecutiveConcatsPass().call(graph_module) + return slice_result.modified or concat_result.modified + + +def _call_nodes(graph_module: GraphModule) -> list[Node]: + return [node for node in graph_module.graph.nodes if node.op == "call_function"] + + +def test_rewrite_cat_slice_inserts_on_boundary() -> None: + graph_module, inputs, output = _concat_slice_graph( + [3, 4, 2], + [(0, 7), (7, 9)], + inserted_channels=1, + input_qparams={0: "left", 1: "inserted", 2: "right"}, + ) + + assert _run(graph_module) + assert [node.target for node in _call_nodes(graph_module)] == [_CAT] + assert list(cast(list[Node], output.args[0])) == [ + inputs[0], + inputs[1], + next(node for node in graph_module.graph.nodes if node.name == "inserted"), + inputs[2], + ] + assert output.meta["input_qparams"] == { + 0: "left", + 1: "left", + 2: "inserted", + 3: "right", + } + + +def test_rewrite_cat_slice_rejects_multi_user_op_growth() -> None: + graph_module, inputs, output = _concat_slice_graph( + [3, 4, 2], [(0, 5), (5, 9)], inserted_channels=1 + ) + + assert not _run(graph_module) + assert output in _call_nodes(graph_module) + + +def test_rewrite_cat_slice_rewrites_partial_coverage() -> None: + graph_module, inputs, output = _concat_slice_graph( + [3, 4], [(0, 3)], inserted_channels=1 + ) + + assert _run(graph_module) + assert [node.target for node in _call_nodes(graph_module)] == [_CAT] + assert list(cast(list[Node], output.args[0])) == [ + inputs[0], + next(node for node in graph_module.graph.nodes if node.name == "inserted"), + ] + + +def test_rewrite_cat_slice_rejects_noncontiguous_multi_user_op_growth() -> None: + graph_module, inputs, output = _concat_slice_graph( + [3, 4, 2], [(0, 5), (6, 9)], inserted_channels=1 + ) + + assert not _run(graph_module) + assert output in _call_nodes(graph_module) + + +def test_rewrite_cat_slice_rejects_empty_slice() -> None: + graph_module, _, output = _concat_slice_graph([0, 0], [(0, 0)]) + + assert not _run(graph_module) + assert output in _call_nodes(graph_module) + + +def test_rewrite_cat_slice_rejects_shared_concat() -> None: + builder = GraphBuilder() + a = builder.placeholder("a", torch.randn(1, 3, 8, 8)) + b = builder.placeholder("b", torch.randn(1, 4, 8, 8)) + source = builder.call_operator(_CAT, ([a, b], 1)) + sliced = builder.call_operator(_SLICE, (source, 1, 0, 3, 1)) + builder.output([source, sliced]) + graph_module = builder.get_graph_module() + + assert not RewriteCatSlicePass().call(graph_module).modified + + +def test_rewrite_cat_slice_rejects_three_piece_single_user_rewrite() -> None: + graph_module, _, output = _concat_slice_graph([3, 3, 3], [(1, 8)]) + + assert not RewriteCatSlicePass().call(graph_module).modified + assert output in _call_nodes(graph_module) + + +def test_rewrite_cat_slice_rewrites_quantized_full_coverage() -> None: + builder = GraphBuilder() + a = builder.placeholder("a", torch.randn(1, 3, 8, 8)) + b = builder.placeholder("b", torch.randn(1, 4, 8, 8)) + source = builder.call_operator(_CAT, ([a, b], 1)) + source.node.meta["input_qparams"] = {0: "cat_qparams"} + sliced = builder.call_operator(_SLICE, (source, 1, 0, 7, 1)) + builder.output([sliced]) + graph_module = builder.get_graph_module() + + assert RewriteCatSlicePass().call(graph_module).modified + assert [node.target for node in _call_nodes(graph_module)] == [_CAT] + + +def test_rewrite_cat_slice_accepts_implicit_step() -> None: + graph_module, inputs, output = _concat_slice_graph( + [3, 4], + [(0, 3), (3, 9223372036854775807)], + inserted_channels=1, + include_step=False, + ) + + assert _run(graph_module) + assert [node.target for node in _call_nodes(graph_module)] == [_CAT] + assert list(cast(list[Node], output.args[0])) == [ + inputs[0], + next(node for node in graph_module.graph.nodes if node.name == "inserted"), + inputs[1], + ] + + +def test_rewrite_cat_slice_preserves_source_input_qparams() -> None: + builder = GraphBuilder() + a = builder.placeholder("a", torch.randn(1, 3, 8, 8)) + b = builder.placeholder("b", torch.randn(1, 4, 8, 8)) + source = builder.call_operator(_CAT, ([a, b], 1)) + source.node.meta["input_qparams"] = {0: "cat_qparams"} + sliced = builder.call_operator(_SLICE, (source, 1, 0, 5, 1)) + sliced.node.meta["input_qparams"] = {0: "source_qparams"} + sliced.node.meta["output_qparams"] = {0: "slice_qparams"} + builder.output([sliced]) + graph_module = builder.get_graph_module() + + assert RewriteCatSlicePass().call(graph_module).modified + + slice_nodes = [node for node in _call_nodes(graph_module) if node.target == _SLICE] + fused_cat = next(node for node in _call_nodes(graph_module) if node.target == _CAT) + assert [node.args for node in slice_nodes] == [(b.node, 1, 0, 2, 1)] + assert [node.meta["input_qparams"] for node in slice_nodes] == [{0: "cat_qparams"}] + assert all( + node.meta["output_qparams"] == {0: "slice_qparams"} for node in slice_nodes + ) + assert fused_cat.meta["input_qparams"] == {0: "cat_qparams"} + assert fused_cat.meta["output_qparams"] == {0: "slice_qparams"} From 5dd9b9f94e990df963529e77d8c019eaeaf25ad2 Mon Sep 17 00:00:00 2001 From: Oscar Andersson Date: Wed, 16 Sep 2026 09:29:59 +0200 Subject: [PATCH 2/2] Fix failing test Signed-off-by: Oscar Andersson Change-Id: Iefce2ef195b47b3b7aeb7c01e87f2b8789d8c943 --- .../test/passes/test_rewrite_cat_slice_pass.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/backends/arm/test/passes/test_rewrite_cat_slice_pass.py b/backends/arm/test/passes/test_rewrite_cat_slice_pass.py index 3fe4556b15f..bf0b3073521 100644 --- a/backends/arm/test/passes/test_rewrite_cat_slice_pass.py +++ b/backends/arm/test/passes/test_rewrite_cat_slice_pass.py @@ -79,18 +79,23 @@ def test_rewrite_cat_slice_inserts_on_boundary() -> None: ) assert _run(graph_module) - assert [node.target for node in _call_nodes(graph_module)] == [_CAT] - assert list(cast(list[Node], output.args[0])) == [ + call_nodes = _call_nodes(graph_module) + assert [node.target for node in call_nodes] == [_CAT, _CAT] + output_inputs = cast(list[Node], output.args[0]) + nested_concat = output_inputs[0] + assert list(cast(list[Node], nested_concat.args[0])) == [ inputs[0], inputs[1], + ] + assert output_inputs == [ + nested_concat, next(node for node in graph_module.graph.nodes if node.name == "inserted"), inputs[2], ] assert output.meta["input_qparams"] == { 0: "left", - 1: "left", - 2: "inserted", - 3: "right", + 1: "inserted", + 2: "right", }