Skip to content

Commit 377f506

Browse files
eigen-kfacebook-github-bot
authored andcommitted
Make core aten exception ops modifiable at runtime.
Summary: Right now we run the ir verifier with a hardcoded set of whitelisted ops. This diff modifies the API to pass the user-defined ops to extend the exception list. See the attached task for more details. Differential Revision: D75997936
1 parent c2aa614 commit 377f506

File tree

1 file changed

+4
-2
lines changed

1 file changed

+4
-2
lines changed

backends/cadence/aot/compiler.py

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -217,6 +217,7 @@ def _lower_ep_to_edge(
217217
expo_program: ExportedProgram,
218218
dump_graphs: bool = False,
219219
constant_methods: Optional[dict[str, object]] = None,
220+
core_aten_exceptions: Optional[list[torch._ops.OpOverload]] = None,
220221
) -> EdgeProgramManager:
221222
"""
222223
Lower an ExportedProgram to an EdgeProgramManager (in edge IR).
@@ -237,7 +238,7 @@ def _lower_ep_to_edge(
237238
torch.ops.aten.unfold.default,
238239
torch.ops.aten.angle.default,
239240
torch.ops.aten.rms_norm.default,
240-
],
241+
] + (core_aten_exceptions or []),
241242
),
242243
constant_methods=constant_methods,
243244
preserve_ops=(torch.ops.aten.rms_norm.default,),
@@ -264,7 +265,8 @@ def export_to_edge(
264265
expo_program = trace(model, inputs)
265266

266267
# Lower the model to edge IR.
267-
edge_prog_manager = _lower_ep_to_edge(expo_program, dump_graphs, constant_methods)
268+
edge_prog_manager = _lower_ep_to_edge(expo_program, dump_graphs, constant_methods,
269+
core_aten_exceptions = [torch.ops.aten.linear.default])
268270

269271
return edge_prog_manager
270272

0 commit comments

Comments
 (0)