Commit 559036a
Skip redundant run_decompositions when no ops match decomp table (pytorch#18496)
Summary:
Adds an early-exit check to _gen_edge_manager_for_partitioners: before
calling program.run_decompositions(table), scan the graph for ops that
appear in the decomposition table. If none are found, skip the call
entirely.
Each run_decompositions call performs a full re-export of the program
via make_fx(), re-tracing every node through FakeTensor dispatch.
On the EDGE_DO_NOT_DECOMP path this function is called up to 3 times;
the early-exit eliminates at least one redundant call where the previous
pass already decomposed all matching ops.
The check recursively walks control flow submodules (cond/map/scan) to
avoid incorrectly skipping when decomposable ops are nested.
## Benchmark
Model: small CNN feature extractor (~50K params, 9 conv layers with
LayerNorm, targeting Ethos-U55 via the ARM/TOSA lowering pipeline).
Graph: ~1200 nodes.
lower() before: 82 s
lower() after: 71 s
Delta: -11 s (-13 %)
Differential Revision: D964899031 parent ae07d06 commit 559036a
1 file changed
Lines changed: 36 additions & 5 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
1161 | 1161 | | |
1162 | 1162 | | |
1163 | 1163 | | |
1164 | | - | |
| 1164 | + | |
| 1165 | + | |
| 1166 | + | |
| 1167 | + | |
| 1168 | + | |
| 1169 | + | |
| 1170 | + | |
| 1171 | + | |
| 1172 | + | |
| 1173 | + | |
| 1174 | + | |
| 1175 | + | |
| 1176 | + | |
| 1177 | + | |
| 1178 | + | |
| 1179 | + | |
| 1180 | + | |
| 1181 | + | |
| 1182 | + | |
| 1183 | + | |
| 1184 | + | |
| 1185 | + | |
| 1186 | + | |
| 1187 | + | |
| 1188 | + | |
| 1189 | + | |
| 1190 | + | |
1165 | 1191 | | |
1166 | 1192 | | |
1167 | 1193 | | |
| |||
1196 | 1222 | | |
1197 | 1223 | | |
1198 | 1224 | | |
1199 | | - | |
| 1225 | + | |
| 1226 | + | |
1200 | 1227 | | |
1201 | 1228 | | |
1202 | 1229 | | |
| |||
1216 | 1243 | | |
1217 | 1244 | | |
1218 | 1245 | | |
1219 | | - | |
| 1246 | + | |
| 1247 | + | |
1220 | 1248 | | |
1221 | 1249 | | |
1222 | 1250 | | |
| |||
1230 | 1258 | | |
1231 | 1259 | | |
1232 | 1260 | | |
1233 | | - | |
| 1261 | + | |
| 1262 | + | |
1234 | 1263 | | |
1235 | 1264 | | |
1236 | 1265 | | |
| |||
1243 | 1272 | | |
1244 | 1273 | | |
1245 | 1274 | | |
1246 | | - | |
| 1275 | + | |
| 1276 | + | |
| 1277 | + | |
1247 | 1278 | | |
1248 | 1279 | | |
1249 | 1280 | | |
| |||
0 commit comments