Skip to content

Commit aaf0a4c

Browse files
authored
Define edge exception and preserve as global (#12300)
Summary: There is a list of ops provided as arguments to the to_edge_with_preserved_ops and EdgeCompileConfig. This PR pulls these lists out of the function and hoists it above the function as a global variable, allowing it to be imported into other code to align compiler behavior between different callsites. Differential Revision: D78008871
1 parent a8d7298 commit aaf0a4c

File tree

1 file changed

+17
-11
lines changed

1 file changed

+17
-11
lines changed

backends/cadence/aot/compiler.py

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -210,6 +210,21 @@ def quantize_pt2(
210210
return program
211211

212212

213+
TO_EDGE_OP_EXCEPTION_LIST: list[torch._ops.OpOverload] = [
214+
torch.ops.aten._linalg_det.default,
215+
torch.ops.aten._linalg_svd.default,
216+
torch.ops.aten._native_batch_norm_legit_functional.default,
217+
torch.ops.aten.linear.default,
218+
torch.ops.aten.linalg_vector_norm.default,
219+
torch.ops.aten.unfold.default,
220+
torch.ops.aten.angle.default,
221+
torch.ops.aten.rms_norm.default,
222+
]
223+
TO_EDGE_PRESERVE_OPS: tuple[torch._ops.OpOverload, ...] = (
224+
torch.ops.aten.rms_norm.default,
225+
)
226+
227+
213228
def _lower_ep_to_edge(
214229
expo_program: ExportedProgram,
215230
dump_graphs: bool = False,
@@ -226,20 +241,11 @@ def _lower_ep_to_edge(
226241
compile_config=EdgeCompileConfig(
227242
_skip_dim_order=True,
228243
# Allow specific non-core aten ops in the IR.
229-
_core_aten_ops_exception_list=[
230-
torch.ops.aten._linalg_det.default,
231-
torch.ops.aten._linalg_svd.default,
232-
torch.ops.aten._native_batch_norm_legit_functional.default,
233-
torch.ops.aten.linear.default,
234-
torch.ops.aten.linalg_vector_norm.default,
235-
torch.ops.aten.unfold.default,
236-
torch.ops.aten.angle.default,
237-
torch.ops.aten.rms_norm.default,
238-
]
244+
_core_aten_ops_exception_list=TO_EDGE_OP_EXCEPTION_LIST
239245
+ (core_aten_exceptions or []),
240246
),
241247
constant_methods=constant_methods,
242-
preserve_ops=(torch.ops.aten.rms_norm.default,),
248+
preserve_ops=TO_EDGE_PRESERVE_OPS,
243249
)
244250

245251
if dump_graphs:

0 commit comments

Comments
 (0)