Skip to content

Commit c8e543d

Browse files
mcremon-metafacebook-github-bot
authored andcommitted
Introduce missing APIs to lower ExportedProgram objects directly
Summary: 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. Differential Revision: D70529392
1 parent 24671a9 commit c8e543d

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(
@@ -213,6 +206,23 @@ def export_to_edge(
213206
logging.info(
214207
edge_prog_manager.exported_program().graph_module.graph.print_tabular()
215208
)
209+
return edge_prog_manager
210+
211+
212+
# Export the model and lower it to an EdgeProgramManager (in edge IR).
213+
def export_to_edge(
214+
model: torch.nn.Module,
215+
inputs: tuple[object, ...],
216+
dump_graphs: bool = False,
217+
constant_methods: Optional[dict[str, object]] = None,
218+
) -> EdgeProgramManager:
219+
assert isinstance(model, torch.nn.Module), "model should be an nn.Module"
220+
221+
# Export the model into an ExportedProgram.
222+
expo_program = export_program(model, inputs)
223+
224+
# Lower the model to edge IR.
225+
edge_prog_manager = lower_ep_to_edge(expo_program, dump_graphs, constant_methods)
216226

217227
return edge_prog_manager
218228

0 commit comments

Comments
 (0)