Add export_to_edge util

Summary:
Adding export_to_edge util to avoid possiblity of using different configs for
exir.capture and to_edge.

Reviewed By: kirklandsign

Differential Revision: D48435836

fbshipit-source-id: f7029d04bfb845cf1ec4ad919b44b7defdfc6c83

Author: kimishpatel

examples/export/export_and_delegate.py

@@ -18,7 +18,7 @@
 
 from ..models import MODEL_NAME_TO_MODEL
 
-from .utils import _CAPTURE_CONFIG, _EDGE_COMPILE_CONFIG
+from .utils import export_to_edge
 
 """
 BackendWithCompilerDemo is a test demo backend, only supports torch.mm and torch.add, here are some examples
@@ -52,7 +52,7 @@ def export_compsite_module_with_lower_graph():
     m, m_inputs = MODEL_NAME_TO_MODEL.get("add_mul")()
     m = m.eval()
     m_inputs = m.get_example_inputs()
-    edge = exir.capture(m, m_inputs, _CAPTURE_CONFIG).to_edge(_EDGE_COMPILE_CONFIG)
+    edge = export_to_edge(m, m_inputs)
     print("Exported graph:\n", edge.exported_program.graph)
 
     # Lower AddMulModule to the demo backend
@@ -71,11 +71,7 @@ def forward(self, *args):
             return torch.sub(self.lowered_graph(*args), torch.ones(1))
 
     # Get the graph for the composite module, which includes lowered graph
-    composited_edge = exir.capture(
-        CompositeModule(),
-        m_inputs,
-        _CAPTURE_CONFIG,
-    ).to_edge(_EDGE_COMPILE_CONFIG)
+    composited_edge = export_to_edge(CompositeModule(), m_inputs)
 
     # The graph module is still runnerable
     composited_edge.exported_program.graph_module(*m_inputs)
@@ -122,9 +118,7 @@ def get_example_inputs(self):
             return (torch.randn(2, 2), torch.randn(2, 2), torch.randn(2, 2))
 
     m = Model()
-    edge = exir.capture(m, m.get_example_inputs(), _CAPTURE_CONFIG).to_edge(
-        _EDGE_COMPILE_CONFIG
-    )
+    edge = export_to_edge(m, m.get_example_inputs())
     print("Exported graph:\n", edge.exported_program.graph)
 
     # Lower to backend_with_compiler_demo
@@ -158,7 +152,7 @@ def export_and_lower_the_whole_graph():
     m, m_inputs = MODEL_NAME_TO_MODEL.get("add_mul")()
     m = m.eval()
     m_inputs = m.get_example_inputs()
-    edge = exir.capture(m, m_inputs, _CAPTURE_CONFIG).to_edge(_EDGE_COMPILE_CONFIG)
+    edge = export_to_edge(m, m_inputs)
     print("Exported graph:\n", edge.exported_program.graph)
 
     # Lower AddMulModule to the demo backend

examples/export/export_example.py

@@ -12,16 +12,11 @@
 
 from ..models import MODEL_NAME_TO_MODEL
 
-from .utils import _CAPTURE_CONFIG, _EDGE_COMPILE_CONFIG
+from .utils import export_to_edge
 
 
 def export_to_pte(model_name, model, example_inputs):
-    m = model.eval()
-    edge = exir.capture(m, example_inputs, _CAPTURE_CONFIG).to_edge(
-        _EDGE_COMPILE_CONFIG
-    )
-    print("Exported graph:\n", edge.exported_program.graph)
-
+    edge = export_to_edge(model, example_inputs)
     exec_prog = edge.to_executorch()
 
     buffer = exec_prog.buffer

examples/export/test/test_export.py

@@ -10,7 +10,7 @@
 
 import torch
 
-from executorch.examples.export.utils import _CAPTURE_CONFIG, _EDGE_COMPILE_CONFIG
+from executorch.examples.export.utils import export_to_edge
 from executorch.examples.models import MODEL_NAME_TO_MODEL
 
 # pyre-ignore[21]: Could not find module `executorch.extension.pybindings.portable`.
@@ -34,9 +34,7 @@ def _assert_eager_lowered_same_result(
         """
         import executorch.exir as exir
 
-        edge_model = exir.capture(eager_model, example_inputs, _CAPTURE_CONFIG).to_edge(
-            _EDGE_COMPILE_CONFIG
-        )
+        edge_model = export_to_edge(eager_model, example_inputs)
 
         executorch_prog = edge_model.to_executorch()
         # pyre-ignore

examples/export/utils.py

@@ -15,3 +15,12 @@
 _EDGE_COMPILE_CONFIG = exir.EdgeCompileConfig(
     _check_ir_validity=False,
 )
+
+
+def export_to_edge(model, example_inputs):
+    m = model.eval()
+    edge = exir.capture(m, example_inputs, _CAPTURE_CONFIG).to_edge(
+        _EDGE_COMPILE_CONFIG
+    )
+    print("Exported graph:\n", edge.exported_program.graph)
+    return edge