Skip to content

Commit 2889983

Browse files
mcremon-metafacebook-github-bot
authored andcommitted
Introduce missing APIs to lower ExportedProgram objects directly (#8909)
Summary: Pull Request resolved: #8909 Currently, there is an imbalance in compilation APIs where the `export_to_executorch` function is used on a module, although that module is already from an exported program. This diff introduces APIs to further lower ExportedPrograms directly, without going back to module. Reviewed By: zonglinpeng Differential Revision: D70529392
1 parent 2aa7748 commit 2889983

File tree

1 file changed

+19
-9
lines changed

1 file changed

+19
-9
lines changed

backends/cadence/aot/compiler.py

Lines changed: 19 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -178,18 +178,11 @@ def export_program(
178178
return expo_program
179179

180180

181-
# Export the model and lower it to an EdgeProgramManager (in edge IR).
182-
def export_to_edge(
183-
model: torch.nn.Module,
184-
inputs: tuple[object, ...],
181+
def lower_ep_to_edge(
182+
expo_program: ExportedProgram,
185183
dump_graphs: bool = False,
186184
constant_methods: Optional[dict[str, object]] = None,
187185
) -> EdgeProgramManager:
188-
assert isinstance(model, torch.nn.Module), "model should be an nn.Module"
189-
190-
# Export the model into an ExportedProgram.
191-
expo_program = export_program(model, inputs)
192-
193186
# Call to_edge to convert the graph to edge IR.
194187
# Note: dim_order is skipped (https://github.com/pytorch/executorch/issues/3704)
195188
edge_prog_manager = to_edge(
@@ -215,6 +208,23 @@ def export_to_edge(
215208
logging.info(
216209
edge_prog_manager.exported_program().graph_module.graph.print_tabular()
217210
)
211+
return edge_prog_manager
212+
213+
214+
# Export the model and lower it to an EdgeProgramManager (in edge IR).
215+
def export_to_edge(
216+
model: torch.nn.Module,
217+
inputs: tuple[object, ...],
218+
dump_graphs: bool = False,
219+
constant_methods: Optional[dict[str, object]] = None,
220+
) -> EdgeProgramManager:
221+
assert isinstance(model, torch.nn.Module), "model should be an nn.Module"
222+
223+
# Export the model into an ExportedProgram.
224+
expo_program = export_program(model, inputs)
225+
226+
# Lower the model to edge IR.
227+
edge_prog_manager = lower_ep_to_edge(expo_program, dump_graphs, constant_methods)
218228

219229
return edge_prog_manager
220230

0 commit comments

Comments
 (0)