Skip to content

Commit 5da974a

Browse files
authored
Introduce missing APIs to lower ExportedProgram objects directly
Differential Revision: D70529392 Pull Request resolved: #8909
1 parent 2aa7748 commit 5da974a

File tree

1 file changed

+22
-9
lines changed

1 file changed

+22
-9
lines changed

backends/cadence/aot/compiler.py

+22-9
Original file line numberDiff line numberDiff line change
@@ -178,18 +178,14 @@ 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-
186+
"""
187+
Lower an ExportedProgram to an EdgeProgramManager (in edge IR).
188+
"""
193189
# Call to_edge to convert the graph to edge IR.
194190
# Note: dim_order is skipped (https://github.com/pytorch/executorch/issues/3704)
195191
edge_prog_manager = to_edge(
@@ -215,6 +211,23 @@ def export_to_edge(
215211
logging.info(
216212
edge_prog_manager.exported_program().graph_module.graph.print_tabular()
217213
)
214+
return edge_prog_manager
215+
216+
217+
# Export the model and lower it to an EdgeProgramManager (in edge IR).
218+
def export_to_edge(
219+
model: torch.nn.Module,
220+
inputs: tuple[object, ...],
221+
dump_graphs: bool = False,
222+
constant_methods: Optional[dict[str, object]] = None,
223+
) -> EdgeProgramManager:
224+
assert isinstance(model, torch.nn.Module), "model should be an nn.Module"
225+
226+
# Export the model into an ExportedProgram.
227+
expo_program = export_program(model, inputs)
228+
229+
# Lower the model to edge IR.
230+
edge_prog_manager = lower_ep_to_edge(expo_program, dump_graphs, constant_methods)
218231

219232
return edge_prog_manager
220233

0 commit comments

Comments
 (0)