From 6b311013718a7a798ec7cc23c11425d15da9fa4b Mon Sep 17 00:00:00 2001 From: Yufeng Shi Date: Fri, 17 Jul 2026 16:41:08 +0100 Subject: [PATCH] Arm backend: Remove statically proven TOSA no-ops Remove single-input concatenations, zero-width padding, and full-range slices when their identity behavior is statically proven. Reject unresolved symbolic shapes and eliminate constants left without users. Run cleanup before transpose fusion and before duplicate/output cleanup so newly exposed patterns are simplified safely. TOSA-FP operator counts: | Model | Before | After | Reduction | |----------------------------|-------:|------:|----------:| | Gemma AltUp | 5 | 2 | 3 | | Gemma conformer light-conv | 72 | 67 | 5 | | Qwen3-VL vision attention | 126 | 112 | 14 | | Phi-3 attention | 114 | 112 | 2 | Add positive, negative, symbolic-shape, graph-output, and follow-on cleanup coverage. Change-Id: I915a870d1bea4d27a89293c730e3f9f5f157cd3e Signed-off-by: Yufeng Shi --- backends/arm/_passes/arm_pass_manager.py | 5 + backends/arm/_passes/remove_noop_pass.py | 78 +++- .../passes/test_remove_data_layout_noops.py | 419 ++++++++++++++++++ 3 files changed, 497 insertions(+), 5 deletions(-) create mode 100644 backends/arm/test/passes/test_remove_data_layout_noops.py diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index 522fcc9c8d7..947125fbd52 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -15,6 +15,7 @@ AccumulateIndexPutPass, BroadcastArgsPass, CanonicalizeGatherPass, + CanonicalizeViewCopyPermutePass, CastInt64BuffersToInt32Pass, CastToInt32Pass, ComputeConstantOpsAOTPass, @@ -652,6 +653,10 @@ def _tosa_pipeline( DecomposePermuteForU55Pass(), RewriteSlicePass(), FuseConsecutiveSlicesPass(), + # Remove rewritten PAD/SLICE no-ops before cleaning up any + # permute pairs they expose. + RemoveNoopPass(), + CanonicalizeViewCopyPermutePass(), InsertConstShapesPass(), InsertDataLayoutCastsPass(), ] diff --git a/backends/arm/_passes/remove_noop_pass.py b/backends/arm/_passes/remove_noop_pass.py index 5fafc848003..2885dc8e36a 100644 --- a/backends/arm/_passes/remove_noop_pass.py +++ b/backends/arm/_passes/remove_noop_pass.py @@ -6,12 +6,13 @@ import logging -from typing import Set, Type +from typing import Any, Set, Type from executorch.backends.arm._passes import ArmOpTargetedPass from executorch.exir.dialects._ops import ops as exir_ops -from executorch.exir.pass_base import ExportPass +from executorch.exir.pass_base import ExportPass, PassResult, ProxyValue +from torch.fx import GraphModule logger = logging.getLogger(__name__) @@ -20,24 +21,91 @@ class RemoveNoopPass(ArmOpTargetedPass): """Remove no-ops from graph_module.""" _passes_required_after: Set[Type[ExportPass]] = set() + _single_input_concat_ops = ( + exir_ops.edge.aten.cat.default, + exir_ops.edge.aten.concatenate.default, + ) target_ops = ( exir_ops.edge.dim_order_ops._clone_dim_order.default, exir_ops.edge.dim_order_ops._to_dim_order_copy.default, exir_ops.edge.aten.alias_copy.default, exir_ops.edge.aten.copy.default, exir_ops.edge.aten.detach_copy.default, + *_single_input_concat_ops, + exir_ops.backend.tosa.PAD.default, + exir_ops.backend.tosa.SLICE.default, ) - def call_operator(self, op, args, kwargs, meta): + @staticmethod + def _get_static_shape(value: Any) -> tuple[int, ...] | None: + if isinstance(value, ProxyValue): + # Shape calculations may carry concrete example values without being + # compile-time constants. Only accept explicit CONST_SHAPE nodes. + if value.node.target != exir_ops.backend.tosa.CONST_SHAPE.default: + return None + value = value.data + + # Reject unresolved symbolic dimensions rather than relying on hints. + if not isinstance(value, (list, tuple)) or not all( + type(dim) is int for dim in value + ): + return None + return tuple(value) + + def call(self, graph_module: GraphModule) -> PassResult: + result = super().call(graph_module) + # Removing a no-op can leave its shape operands without users. + removed_dead_code = result.graph_module.graph.eliminate_dead_code() + if removed_dead_code: + result.graph_module.graph.lint() + result.graph_module.recompile() + return PassResult(result.graph_module, result.modified or removed_dead_code) + + def call_operator(self, op, args, kwargs, meta, updated=False): if op not in self.target_ops: - return super().call_operator(op, args, kwargs, meta) + return super().call_operator(op, args, kwargs, meta, updated) + + if op in self._single_input_concat_ops: + inputs = args[0] + # Concatenating one tensor returns that tensor unchanged. + if isinstance(inputs, (list, tuple)) and len(inputs) == 1: + return inputs[0] + return super().call_operator(op, args, kwargs, meta, updated) + + if op == exir_ops.backend.tosa.PAD.default: + padding = self._get_static_shape(args[1]) + # PAD is an identity when every before/after padding value is zero. + if ( + padding is not None + and len(padding) == 2 * len(args[0].data.shape) + and all(value == 0 for value in padding) + ): + return args[0] + return super().call_operator(op, args, kwargs, meta, updated) + + if op == exir_ops.backend.tosa.SLICE.default: + starts = self._get_static_shape(args[1]) + sizes = self._get_static_shape(args[2]) + input_shape = self._get_static_shape(args[0].data.shape) + # Starting at zero and selecting the full input shape is an identity. + if ( + starts is not None + and sizes is not None + and input_shape is not None + and starts == (0,) * len(input_shape) + and sizes == input_shape + ): + return args[0] + return super().call_operator(op, args, kwargs, meta, updated) input_dtype = args[0].data.dtype output_dtype = kwargs.get("dtype", input_dtype) + # A clone-like operation that changes dtype is not a no-op. if input_dtype != output_dtype: - return super().call_operator(op, args, kwargs, meta) + return super().call_operator(op, args, kwargs, meta, updated) + # copy writes its source (args[1]); the remaining targets return args[0]. if op == exir_ops.edge.aten.copy.default: return args[1] return args[0] diff --git a/backends/arm/test/passes/test_remove_data_layout_noops.py b/backends/arm/test/passes/test_remove_data_layout_noops.py new file mode 100644 index 00000000000..2719178337b --- /dev/null +++ b/backends/arm/test/passes/test_remove_data_layout_noops.py @@ -0,0 +1,419 @@ +# 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 executorch.backends.arm.tosa.dialect # noqa: F401 +import pytest +import torch +from executorch.backends.arm._passes import ( + CanonicalizeViewCopyPermutePass, + EnsureUniqueOutputNodesPass, + ExirToTosaPass, + FuseDuplicateUsersPass, + InsertRescalePass, + RemoveNoopPass, + RewriteSlicePass, +) +from executorch.backends.arm._passes.arm_pass_manager import ArmPassManager +from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec +from executorch.backends.arm.tosa.dialect.lib import TosaValueError +from executorch.backends.arm.tosa.specification import ( + TosaLoweringContext, + TosaSpecification, +) +from executorch.exir import EdgeCompileConfig, to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import export, ExportedProgram +from torch.fx import Graph, GraphModule, Node + + +def _call( + graph: Graph, + target, + args, + val, + kwargs=None, +) -> Node: + node = graph.call_function(target, args, kwargs or {}) + node.meta["val"] = val + return node + + +def _count_target(graph_module: GraphModule, target) -> int: + return sum(node.target == target for node in graph_module.graph.nodes) + + +def _run_remove_noop(graph_module: GraphModule) -> GraphModule: + spec = TosaSpecification.create_from_string("TOSA-1.0+FP") + with TosaLoweringContext(spec): + result = RemoveNoopPass()(graph_module) + assert result is not None + return result.graph_module + + +def _tosa_data_layout_graph( + target, + shape_values: tuple[list[int], ...], + output_shape: tuple[int, ...], +) -> GraphModule: + graph = Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.ones(2, 3) + shape_nodes = [ + _call( + graph, + exir_ops.backend.tosa.CONST_SHAPE.default, + (values,), + values, + ) + for values in shape_values + ] + kwargs = {"value": 0.0} if target == exir_ops.backend.tosa.PAD.default else {} + result = _call( + graph, + target, + (x, *shape_nodes), + torch.ones(output_shape), + kwargs, + ) + graph.output((result,)) + return GraphModule(torch.nn.Module(), graph) + + +@pytest.mark.parametrize( + "target", + [ + exir_ops.edge.aten.cat.default, + exir_ops.edge.aten.concatenate.default, + ], +) +def test_remove_single_input_concat(target): + graph = Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.ones(2, 3) + concat = _call(graph, target, ([x], 0), torch.ones(2, 3)) + graph.output((concat,)) + + result = _run_remove_noop(GraphModule(torch.nn.Module(), graph)) + + assert _count_target(result, target) == 0 + assert result.graph.output_node().args[0][0].op == "placeholder" + + +def test_keep_multi_input_concat(): + graph = Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.ones(2, 3) + concat = _call( + graph, + exir_ops.edge.aten.cat.default, + ([x, x], 0), + torch.ones(4, 3), + ) + graph.output((concat,)) + + result = _run_remove_noop(GraphModule(torch.nn.Module(), graph)) + + assert _count_target(result, exir_ops.edge.aten.cat.default) == 1 + + +def test_remove_full_slice_and_unused_shape_constants(): + graph_module = _tosa_data_layout_graph( + exir_ops.backend.tosa.SLICE.default, + ([0, 0], [2, 3]), + (2, 3), + ) + + result = _run_remove_noop(graph_module) + + assert _count_target(result, exir_ops.backend.tosa.SLICE.default) == 0 + assert _count_target(result, exir_ops.backend.tosa.CONST_SHAPE.default) == 0 + assert result.graph.output_node().args[0][0].op == "placeholder" + + +def test_keep_partial_slice(): + graph_module = _tosa_data_layout_graph( + exir_ops.backend.tosa.SLICE.default, + ([1, 0], [1, 3]), + (1, 3), + ) + + result = _run_remove_noop(graph_module) + + assert _count_target(result, exir_ops.backend.tosa.SLICE.default) == 1 + + +def test_keep_slice_with_nonconstant_shape_operands(): + graph = Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.ones(2, 3) + starts = graph.placeholder("starts") + starts.meta["val"] = [0, 0] + sizes = graph.placeholder("sizes") + sizes.meta["val"] = [2, 3] + sliced = _call( + graph, + exir_ops.backend.tosa.SLICE.default, + (x, starts, sizes), + torch.ones(2, 3), + ) + graph.output((sliced,)) + + result = _run_remove_noop(GraphModule(torch.nn.Module(), graph)) + + assert _count_target(result, exir_ops.backend.tosa.SLICE.default) == 1 + + +def test_keep_full_slice_with_symbolic_shape(): + class FullSlice(torch.nn.Module): + def forward(self, x): + return torch.ops.aten.slice.Tensor(x, 0, 0, torch.iinfo(torch.int64).max, 1) + + exported_program = export( + FullSlice(), + (torch.ones(4, 3),), + dynamic_shapes={ + "x": {0: torch.export.Dim("batch", min=2, max=8)}, + }, + strict=True, + ) + edge_program = to_edge( + exported_program, + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ).exported_program() + spec = TosaSpecification.create_from_string("TOSA-1.0+FP") + with TosaLoweringContext(spec): + rewritten = RewriteSlicePass()(edge_program.graph_module).graph_module + + slice_nodes = [ + node + for node in rewritten.graph.nodes + if node.target == exir_ops.backend.tosa.SLICE.default + ] + assert len(slice_nodes) == 1 + assert isinstance(slice_nodes[0].args[2][0], torch.SymInt) + + result = _run_remove_noop(rewritten) + + assert _count_target(result, exir_ops.backend.tosa.SLICE.default) == 1 + + +def test_remove_zero_pad_and_unused_shape_constant(): + graph_module = _tosa_data_layout_graph( + exir_ops.backend.tosa.PAD.default, + ([0, 0, 0, 0],), + (2, 3), + ) + + result = _run_remove_noop(graph_module) + + assert _count_target(result, exir_ops.backend.tosa.PAD.default) == 0 + assert _count_target(result, exir_ops.backend.tosa.CONST_SHAPE.default) == 0 + + +def test_keep_nonzero_pad(): + graph_module = _tosa_data_layout_graph( + exir_ops.backend.tosa.PAD.default, + ([0, 0, 1, 0],), + (2, 4), + ) + + result = _run_remove_noop(graph_module) + + assert _count_target(result, exir_ops.backend.tosa.PAD.default) == 1 + + +@pytest.mark.parametrize("padding", [[], [0], [0, 0]]) +def test_reject_malformed_zero_pad(padding: list[int]): + graph_module = _tosa_data_layout_graph( + exir_ops.backend.tosa.PAD.default, + (padding,), + (2, 3), + ) + + with pytest.raises(TosaValueError, match="Padding length"): + _run_remove_noop(graph_module) + + +def test_noop_removal_exposes_transpose_cleanup(): + graph = Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.ones(2, 3) + first_permute = _call( + graph, + exir_ops.edge.aten.permute_copy.default, + (x, [1, 0]), + torch.ones(3, 2), + ) + padding = _call( + graph, + exir_ops.backend.tosa.CONST_SHAPE.default, + ([0, 0, 0, 0],), + [0, 0, 0, 0], + ) + padded = _call( + graph, + exir_ops.backend.tosa.PAD.default, + (first_permute, padding), + torch.ones(3, 2), + {"value": 0.0}, + ) + second_permute = _call( + graph, + exir_ops.edge.aten.permute_copy.default, + (padded, [1, 0]), + torch.ones(2, 3), + ) + graph.output((second_permute,)) + + graph_module = _run_remove_noop(GraphModule(torch.nn.Module(), graph)) + result = CanonicalizeViewCopyPermutePass()(graph_module).graph_module + + assert _count_target(result, exir_ops.backend.tosa.PAD.default) == 0 + assert _count_target(result, exir_ops.edge.aten.permute_copy.default) == 0 + + +def test_noop_removal_exposes_duplicate_user_cleanup(): + graph = Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.ones(2, 3) + y = graph.placeholder("y") + y.meta["val"] = torch.ones(2, 3) + padding = _call( + graph, + exir_ops.backend.tosa.CONST_SHAPE.default, + ([0, 0, 0, 0],), + [0, 0, 0, 0], + ) + padded = _call( + graph, + exir_ops.backend.tosa.PAD.default, + (x, padding), + torch.ones(2, 3), + {"value": 0.0}, + ) + first_add = _call( + graph, + exir_ops.edge.aten.add.Tensor, + (padded, y), + torch.ones(2, 3), + ) + second_add = _call( + graph, + exir_ops.edge.aten.add.Tensor, + (x, y), + torch.ones(2, 3), + ) + graph.output((first_add, second_add)) + + graph_module = _run_remove_noop(GraphModule(torch.nn.Module(), graph)) + graph_module = FuseDuplicateUsersPass()(graph_module).graph_module + result = EnsureUniqueOutputNodesPass()(graph_module).graph_module + + assert _count_target(result, exir_ops.edge.aten.add.Tensor) == 1 + outputs = result.graph.output_node().args[0] + assert outputs[0] is not outputs[1] + + +class _PipelineInspectionPassManager(ArmPassManager): + def _transform( + self, exported_program: ExportedProgram, graph_module: GraphModule + ) -> GraphModule: + return graph_module + + +class _DataLayoutNoopPipelineModule(torch.nn.Module): + def forward(self, x, y): + selected = torch.select(x, 1, 0) + permuted = x.permute(0, 2, 1) + permuted = torch.nn.functional.pad(permuted, (0, 0, 0, 0, 0, 0)) + permuted = permuted.permute(0, 2, 1) + padded = torch.nn.functional.pad(x, (0, 0, 0, 0, 0, 0)) + return selected, permuted, padded + y, x + y + + +def test_data_layout_noop_cleanup_backend_pipeline(): + exported_program = export( + _DataLayoutNoopPipelineModule(), + (torch.ones(2, 1, 3), torch.ones(2, 1, 3)), + strict=True, + ) + edge_program = to_edge( + exported_program, + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ).exported_program() + + assert ( + _count_target(edge_program.graph_module, exir_ops.edge.aten.select_copy.int) + == 1 + ) + assert ( + _count_target( + edge_program.graph_module, + exir_ops.edge.aten.constant_pad_nd.default, + ) + == 2 + ) + assert ( + _count_target( + edge_program.graph_module, exir_ops.edge.aten.permute_copy.default + ) + == 2 + ) + assert _count_target(edge_program.graph_module, exir_ops.edge.aten.add.Tensor) == 2 + + graph_module = ArmPassManager( + TosaCompileSpec("TOSA-1.0+FP") + ).transform_to_backend_pipeline(edge_program, edge_program.graph_module) + + assert _count_target(graph_module, exir_ops.backend.tosa.PAD.default) == 0 + assert _count_target(graph_module, exir_ops.backend.tosa.SLICE.default) == 0 + assert _count_target(graph_module, exir_ops.edge.aten.permute_copy.default) == 0 + assert _count_target(graph_module, exir_ops.backend.tosa.TRANSPOSE.default) == 0 + + add_nodes = [ + node + for node in graph_module.graph.nodes + if node.target == exir_ops.backend.tosa.ADD.default + ] + identity_nodes = [ + node + for node in graph_module.graph.nodes + if node.target == exir_ops.backend.tosa.IDENTITY.default + ] + assert len(add_nodes) == 1 + assert len(identity_nodes) == 2 + assert all(node.args[0] is add_nodes[0] for node in identity_nodes) + + outputs = graph_module.graph.output_node().args[0] + assert len(outputs) == len(set(outputs)) + graph_module.graph.lint() + + +def test_data_layout_noop_cleanup_pipeline_order(): + exported_program = export(torch.nn.Identity(), (torch.ones(1),), strict=True) + pass_manager = _PipelineInspectionPassManager(TosaCompileSpec("TOSA-1.0+FP")) + + pass_manager.transform_to_backend_pipeline( + exported_program, exported_program.graph_module + ) + + pass_types = [type(pass_) for pass_ in pass_manager.passes] + exir_to_tosa_index = pass_types.index(ExirToTosaPass) + cleanup_indices = [ + index + for index, pass_type in enumerate(pass_types) + if pass_type is RemoveNoopPass + ] + pre_tosa_cleanup = max( + index for index in cleanup_indices if index < exir_to_tosa_index + ) + post_tosa_cleanup = min( + index for index in cleanup_indices if index > exir_to_tosa_index + ) + + assert pass_types[pre_tosa_cleanup + 1] is CanonicalizeViewCopyPermutePass + assert pass_types[post_tosa_cleanup + 1] is FuseDuplicateUsersPass + assert pass_types[post_tosa_cleanup + 2] is InsertRescalePass + assert pass_types[post_tosa_cleanup + 3] is EnsureUniqueOutputNodesPass