Skip to content

Commit c48a7b2

Browse files
Arm backend: Remove statically proven TOSA no-ops (pytorch#21141)
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 | 70 | 65 | 5 | | Qwen3-VL vision attention | 124 | 113 | 11 | | Phi-3 attention | 116 | 114 | 2 | Add positive, negative, symbolic-shape, graph-output, and follow-on cleanup coverage. Change-Id: I915a870d1bea4d27a89293c730e3f9f5f157cd3e Signed-off-by: Yufeng Shi <yufeng.shi@arm.com>
1 parent c967b70 commit c48a7b2

3 files changed

Lines changed: 497 additions & 5 deletions

File tree

backends/arm/_passes/arm_pass_manager.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
AccumulateIndexPutPass,
1616
BroadcastArgsPass,
1717
CanonicalizeGatherPass,
18+
CanonicalizeViewCopyPermutePass,
1819
CastInt64BuffersToInt32Pass,
1920
CastToInt32Pass,
2021
ComputeConstantOpsAOTPass,
@@ -652,6 +653,10 @@ def _tosa_pipeline(
652653
DecomposePermuteForU55Pass(),
653654
RewriteSlicePass(),
654655
FuseConsecutiveSlicesPass(),
656+
# Remove rewritten PAD/SLICE no-ops before cleaning up any
657+
# permute pairs they expose.
658+
RemoveNoopPass(),
659+
CanonicalizeViewCopyPermutePass(),
655660
InsertConstShapesPass(),
656661
InsertDataLayoutCastsPass(),
657662
]

backends/arm/_passes/remove_noop_pass.py

Lines changed: 73 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,13 @@
66

77

88
import logging
9-
from typing import Set, Type
9+
from typing import Any, Set, Type
1010

1111
from executorch.backends.arm._passes import ArmOpTargetedPass
1212

1313
from executorch.exir.dialects._ops import ops as exir_ops
14-
from executorch.exir.pass_base import ExportPass
14+
from executorch.exir.pass_base import ExportPass, PassResult, ProxyValue
15+
from torch.fx import GraphModule
1516

1617
logger = logging.getLogger(__name__)
1718

@@ -20,24 +21,91 @@ class RemoveNoopPass(ArmOpTargetedPass):
2021
"""Remove no-ops from graph_module."""
2122

2223
_passes_required_after: Set[Type[ExportPass]] = set()
24+
_single_input_concat_ops = (
25+
exir_ops.edge.aten.cat.default,
26+
exir_ops.edge.aten.concatenate.default,
27+
)
2328
target_ops = (
2429
exir_ops.edge.dim_order_ops._clone_dim_order.default,
2530
exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
2631
exir_ops.edge.aten.alias_copy.default,
2732
exir_ops.edge.aten.copy.default,
2833
exir_ops.edge.aten.detach_copy.default,
34+
*_single_input_concat_ops,
35+
exir_ops.backend.tosa.PAD.default,
36+
exir_ops.backend.tosa.SLICE.default,
2937
)
3038

31-
def call_operator(self, op, args, kwargs, meta):
39+
@staticmethod
40+
def _get_static_shape(value: Any) -> tuple[int, ...] | None:
41+
if isinstance(value, ProxyValue):
42+
# Shape calculations may carry concrete example values without being
43+
# compile-time constants. Only accept explicit CONST_SHAPE nodes.
44+
if value.node.target != exir_ops.backend.tosa.CONST_SHAPE.default:
45+
return None
46+
value = value.data
47+
48+
# Reject unresolved symbolic dimensions rather than relying on hints.
49+
if not isinstance(value, (list, tuple)) or not all(
50+
type(dim) is int for dim in value
51+
):
52+
return None
53+
return tuple(value)
54+
55+
def call(self, graph_module: GraphModule) -> PassResult:
56+
result = super().call(graph_module)
57+
# Removing a no-op can leave its shape operands without users.
58+
removed_dead_code = result.graph_module.graph.eliminate_dead_code()
59+
if removed_dead_code:
60+
result.graph_module.graph.lint()
61+
result.graph_module.recompile()
62+
return PassResult(result.graph_module, result.modified or removed_dead_code)
63+
64+
def call_operator(self, op, args, kwargs, meta, updated=False):
3265
if op not in self.target_ops:
33-
return super().call_operator(op, args, kwargs, meta)
66+
return super().call_operator(op, args, kwargs, meta, updated)
67+
68+
if op in self._single_input_concat_ops:
69+
inputs = args[0]
70+
# Concatenating one tensor returns that tensor unchanged.
71+
if isinstance(inputs, (list, tuple)) and len(inputs) == 1:
72+
return inputs[0]
73+
return super().call_operator(op, args, kwargs, meta, updated)
74+
75+
if op == exir_ops.backend.tosa.PAD.default:
76+
padding = self._get_static_shape(args[1])
77+
# PAD is an identity when every before/after padding value is zero.
78+
if (
79+
padding is not None
80+
and len(padding) == 2 * len(args[0].data.shape)
81+
and all(value == 0 for value in padding)
82+
):
83+
return args[0]
84+
return super().call_operator(op, args, kwargs, meta, updated)
85+
86+
if op == exir_ops.backend.tosa.SLICE.default:
87+
starts = self._get_static_shape(args[1])
88+
sizes = self._get_static_shape(args[2])
89+
input_shape = self._get_static_shape(args[0].data.shape)
90+
# Starting at zero and selecting the full input shape is an identity.
91+
if (
92+
starts is not None
93+
and sizes is not None
94+
and input_shape is not None
95+
and starts == (0,) * len(input_shape)
96+
and sizes == input_shape
97+
):
98+
return args[0]
99+
return super().call_operator(op, args, kwargs, meta, updated)
34100

35101
input_dtype = args[0].data.dtype
36102
output_dtype = kwargs.get("dtype", input_dtype)
37103

104+
# A clone-like operation that changes dtype is not a no-op.
38105
if input_dtype != output_dtype:
39-
return super().call_operator(op, args, kwargs, meta)
106+
return super().call_operator(op, args, kwargs, meta, updated)
40107

108+
# copy writes its source (args[1]); the remaining targets return args[0].
41109
if op == exir_ops.edge.aten.copy.default:
42110
return args[1]
43111
return args[0]

0 commit comments

Comments
 (0)