66
77
88import logging
9- from typing import Set , Type
9+ from typing import Any , Set , Type
1010
1111from executorch .backends .arm ._passes import ArmOpTargetedPass
1212
1313from 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
1617logger = 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