Move examples export to two stage export (#238)

Summary:
Pull Request resolved: #238

Following the alignment, this diff moves examples export to two stage export
api.

Will follow up with changes to quantization examples.

However, before landing, must land the torch nightly update to executorch

Reviewed By: guangy10

Differential Revision: D49025972

fbshipit-source-id: 0b2c71dd7d562a714ff723b87b660d564a96cc37

Author: kimishpatel

examples/backend/xnnpack_examples.py

@@ -9,11 +9,12 @@
 import argparse
 import logging
 
+import torch._export as export
+
 from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
     XnnpackFloatingPointPartitioner,
     XnnpackQuantizedPartitioner,
 )
-
 from executorch.exir import CaptureConfig, EdgeCompileConfig
 from executorch.exir.backend.backend_api import to_backend
 from executorch.exir.backend.canonical_partitioners.duplicate_dequant_node_pass import (
@@ -77,6 +78,8 @@
     )
 
     model = model.eval()
+    # pre-autograd export. eventually this will become torch.export
+    model = export.capture_pre_autograd_graph(model, example_inputs)
 
     partitioner = XnnpackFloatingPointPartitioner
     if args.quantize:
@@ -85,10 +88,11 @@
         # TODO(T161849167): Partitioner will eventually be a single partitioner for both fp32 and quantized models
         partitioner = XnnpackQuantizedPartitioner
 
+    capture_config = CaptureConfig(enable_aot=True)
+
     edge = export_to_edge(
         model,
         example_inputs,
-        capture_config=CaptureConfig(enable_aot=True),
         edge_compile_config=EdgeCompileConfig(
             # TODO(T162080278): Duplicated Dequant nodes will be in quantizer spec
             _check_ir_validity=False

examples/export/export_and_delegate.py

@@ -10,6 +10,7 @@
 import logging
 
 import torch
+import torch._export as export
 from executorch.exir.backend.backend_api import to_backend
 from executorch.exir.backend.test.backend_with_compiler_demo import (
     BackendWithCompilerDemo,
@@ -18,7 +19,7 @@
 
 from ..models import MODEL_NAME_TO_MODEL
 
-from .utils import export_to_edge
+from ..utils import export_to_edge
 
 
 FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
@@ -59,6 +60,8 @@ 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()
+    # pre-autograd export. eventually this will become torch.export
+    m = export.capture_pre_autograd_graph(m, m_inputs)
     edge = export_to_edge(m, m_inputs)
     logging.info(f"Exported graph:\n{edge.exported_program.graph}")
 
@@ -78,7 +81,11 @@ 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 = export_to_edge(CompositeModule(), m_inputs)
+    m = CompositeModule()
+    m = m.eval()
+    # pre-autograd export. eventually this will become torch.export
+    m = export.capture_pre_autograd_graph(m, m_inputs)
+    composited_edge = export_to_edge(m, m_inputs)
 
     # The graph module is still runnerable
     composited_edge.exported_program.graph_module(*m_inputs)
@@ -125,7 +132,10 @@ def get_example_inputs(self):
             return (torch.randn(2, 2), torch.randn(2, 2), torch.randn(2, 2))
 
     m = Model()
-    edge = export_to_edge(m, m.get_example_inputs())
+    m_inputs = m.get_example_inputs()
+    # pre-autograd export. eventually this will become torch.export
+    m = export.capture_pre_autograd_graph(m, m_inputs)
+    edge = export_to_edge(m, m_inputs)
     logging.info(f"Exported graph:\n{edge.exported_program.graph}")
 
     # Lower to backend_with_compiler_demo
@@ -159,6 +169,8 @@ 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()
+    # pre-autograd export. eventually this will become torch.export
+    m = export.capture_pre_autograd_graph(m, m_inputs)
     edge = export_to_edge(m, m_inputs)
     logging.info(f"Exported graph:\n{edge.exported_program.graph}")
 

examples/export/test/TARGETS

@@ -10,7 +10,6 @@ python_unittest(
         "//caffe2:torch",
         "//executorch/examples/export:utils",
         "//executorch/examples/models:models",
-        "//executorch/exir:lib",
         "//executorch/extension/pybindings:portable_lib",  # @manual
     ],
 )

examples/export/test/test_export.py

@@ -9,6 +9,7 @@
 from typing import Any, Callable
 
 import torch
+import torch._export as export
 
 from executorch.examples.export.utils import export_to_edge
 from executorch.examples.models import MODEL_NAME_TO_MODEL
@@ -32,9 +33,10 @@ def _assert_eager_lowered_same_result(
         takes the eager mode output and ET output, and returns True if they
         match.
         """
-        import executorch.exir as exir
 
-        edge_model = export_to_edge(eager_model, example_inputs)
+        eager_model = eager_model.eval()
+        model = export.capture_pre_autograd_graph(eager_model, example_inputs)
+        edge_model = export_to_edge(model, example_inputs)
 
         executorch_prog = edge_model.to_executorch()
 

examples/export/utils.py

@@ -6,8 +6,15 @@
 
 import logging
 
+from typing import Tuple
+
 import executorch.exir as exir
 
+import torch
+import torch._export as export
+from executorch.exir.program import ExirExportedProgram
+from executorch.exir.tracer import Value
+
 
 _CAPTURE_CONFIG = exir.CaptureConfig(enable_aot=True)
 
@@ -17,26 +24,55 @@
 )
 
 
-def export_to_edge(
-    model,
-    example_inputs,
+def _to_core_aten(
+    model: torch.fx.GraphModule,
+    example_inputs: Tuple[Value, ...],
     capture_config=_CAPTURE_CONFIG,
+) -> ExirExportedProgram:
+    # post autograd export. eventually this will become .to_core_aten
+    if not isinstance(model, torch.fx.GraphModule):
+        raise ValueError(
+            f"Expected passed in model to be an instance of fx.GraphModule, got {type(model)}"
+        )
+    core_aten_exir_ep = exir.capture(model, example_inputs, capture_config)
+    logging.info(f"Core ATen graph:\n{core_aten_exir_ep.exported_program.graph}")
+    return core_aten_exir_ep
+
+
+def _core_aten_to_edge(
+    core_aten_exir_ep: ExirExportedProgram,
     edge_compile_config=_EDGE_COMPILE_CONFIG,
-):
-    m = model.eval()
-    edge = exir.capture(m, example_inputs, capture_config).to_edge(edge_compile_config)
+) -> ExirExportedProgram:
+    edge = core_aten_exir_ep.to_edge(edge_compile_config)
     logging.info(f"Exported graph:\n{edge.exported_program.graph}")
     return edge
 
 
+def export_to_edge(
+    model: torch.fx.GraphModule,
+    example_inputs: Tuple[Value, ...],
+    capture_config=_CAPTURE_CONFIG,
+    edge_compile_config=_EDGE_COMPILE_CONFIG,
+) -> ExirExportedProgram:
+    core_aten_exir_ep = _to_core_aten(model, example_inputs, capture_config)
+    return _core_aten_to_edge(core_aten_exir_ep, edge_compile_config)
+
+
 def export_to_exec_prog(
     model,
     example_inputs,
     capture_config=_CAPTURE_CONFIG,
     edge_compile_config=_EDGE_COMPILE_CONFIG,
     backend_config=None,
 ):
-    edge_m = export_to_edge(model, example_inputs, capture_config, edge_compile_config)
+    m = model.eval()
+    # pre-autograd export. eventually this will become torch.export
+    m = export.capture_pre_autograd_graph(m, example_inputs)
+
+    core_aten_exir_ep = _to_core_aten(m, example_inputs)
+
+    edge_m = _core_aten_to_edge(core_aten_exir_ep, edge_compile_config)
+
     exec_prog = edge_m.to_executorch(backend_config)
     return exec_prog
 

examples/quantization/example.py

@@ -28,7 +28,7 @@
     XNNPACKQuantizer,
 )
 
-from ..export.export_example import export_to_exec_prog, save_pte_program
+from ..export.utils import export_to_edge, save_pte_program
 from ..models import MODEL_NAME_TO_MODEL
 from ..models.model_factory import EagerModelFactory
 from ..recipes.xnnpack_optimization import MODEL_NAME_TO_OPTIONS
@@ -154,23 +154,26 @@ def verify_xnnpack_quantizer_matching_fx_quant_model(model_name, model, example_
         end = time.perf_counter()
         # logging.info(f"Verify time: {end - start}s")
 
+    model = model.eval()
+    # pre-autograd export. eventually this will become torch.export
+    model = export.capture_pre_autograd_graph(model, example_inputs)
     start = time.perf_counter()
     quantized_model = quantize(model, example_inputs)
     end = time.perf_counter()
     # logging.info(f"Quantize time: {end - start}s")
 
     # TODO[T163161310]: takes a long time to export to exec prog and save inception_v4 quantized model
     if args.model_name != "ic4":
+
         start = time.perf_counter()
-        prog = export_to_exec_prog(
-            quantized_model,
-            copy.deepcopy(example_inputs),
-            edge_compile_config=EdgeCompileConfig(_check_ir_validity=False),
+        edge_compile_config = EdgeCompileConfig(_check_ir_validity=False)
+        edge_m = export_to_edge(
+            quantized_model, example_inputs, edge_compile_config=edge_compile_config
         )
         end = time.perf_counter()
-        # logging.info(f"export_to_exec_prog time: {end - start}s")
+
         start = time.perf_counter()
+        prog = edge_m.to_executorch(None)
         save_pte_program(prog.buffer, f"{args.model_name}_quantized")
         end = time.perf_counter()
-        # logging.info(f"save_pte_program time: {end - start}s")
     logging.info("finished")

examples/quantization/utils.py

@@ -7,7 +7,6 @@
 import copy
 import logging
 
-import torch._export as export
 from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e
 from torch.ao.quantization.quantizer.xnnpack_quantizer import (
     get_symmetric_quantization_config,
@@ -17,14 +16,12 @@
 
 def quantize(model, example_inputs):
     """This is the official recommended flow for quantization in pytorch 2.0 export"""
-    m = model.eval()
-    m = export.capture_pre_autograd_graph(m, copy.deepcopy(example_inputs))
-    logging.info(f"Original model: {m}")
+    logging.info(f"Original model: {model}")
     quantizer = XNNPACKQuantizer()
     # if we set is_per_channel to True, we also need to add out_variant of quantize_per_channel/dequantize_per_channel
     operator_config = get_symmetric_quantization_config(is_per_channel=False)
     quantizer.set_global(operator_config)
-    m = prepare_pt2e(m, quantizer)
+    m = prepare_pt2e(model, quantizer)
     # calibration
     m(*example_inputs)
     m = convert_pt2e(m)