Skip to content

Commit f5f213a

Browse files
eigen-kfacebook-github-bot
authored andcommitted
Make core aten exception ops modifiable at runtime. (#11497)
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 a7a9130 commit f5f213a

File tree

1 file changed

+9
-2
lines changed

1 file changed

+9
-2
lines changed

backends/cadence/aot/compiler.py

Lines changed: 9 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,8 @@ 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+
]
242+
+ (core_aten_exceptions or []),
241243
),
242244
constant_methods=constant_methods,
243245
preserve_ops=(torch.ops.aten.rms_norm.default,),
@@ -264,7 +266,12 @@ def export_to_edge(
264266
expo_program = trace(model, inputs)
265267

266268
# Lower the model to edge IR.
267-
edge_prog_manager = _lower_ep_to_edge(expo_program, dump_graphs, constant_methods)
269+
edge_prog_manager = _lower_ep_to_edge(
270+
expo_program,
271+
dump_graphs,
272+
constant_methods,
273+
core_aten_exceptions=[torch.ops.aten.linear.default],
274+
)
268275

269276
return edge_prog_manager
270277

0 commit comments

Comments
 (0)