Update OSS flow to call the right APIs (#13967)

Summary:
Pull Request resolved: #13967

Reorganize the calls so that we don't run into issues calling `eval()`

Reviewed By: zonglinpeng

Differential Revision: D81698817

Author: mcremon-meta

backends/cadence/aot/compiler.py

@@ -388,9 +388,6 @@ def quantize_and_export_to_cadence(
     )
 
 
-# Export the model and lower it to an EdgeProgramManager (in edge IR), and
-# apply passes specific to Cadence DSP execution. Return both to print the
-# differences.
 def export_to_executorch_gen_etrecord(
     model: torch.nn.Module,
     inputs: tuple[object, ...],
@@ -402,7 +399,33 @@ def export_to_executorch_gen_etrecord(
     memory_config: Optional[MemoryConfig] = None,
     dump_graphs: bool = False,
 ) -> ExecutorchProgramManager:
-    edge_prog_manager = export_to_edge(model, inputs, dump_graphs)
+    ep = torch.export.export(model, inputs, strict=True)
+    return _lower_ep_to_cadence_gen_etrecord(
+        ep,
+        output_dir=output_dir,
+        opt_level=opt_level,
+        mem_algo=mem_algo,
+        alloc_graph_input=alloc_graph_input,
+        alloc_graph_output=alloc_graph_output,
+        memory_config=memory_config,
+        dump_graphs=dump_graphs,
+    )
+
+
+# Export the model and lower it to an EdgeProgramManager (in edge IR), and
+# apply passes specific to Cadence DSP execution. Return both to print the
+# differences.
+def _lower_ep_to_cadence_gen_etrecord(
+    ep: ExportedProgram,
+    output_dir: Optional[str] = None,
+    opt_level: int = 1,
+    mem_algo: int = 0,
+    alloc_graph_input: bool = True,
+    alloc_graph_output: bool = True,
+    memory_config: Optional[MemoryConfig] = None,
+    dump_graphs: bool = False,
+) -> ExecutorchProgramManager:
+    edge_prog_manager = _lower_ep_to_edge(ep, dump_graphs)
     cadence_prog_manager = apply_exir_ops_passes(opt_level, edge_prog_manager)
 
     # Print some information to terminal

backends/cadence/aot/export_example.py

@@ -10,13 +10,14 @@
 
 import logging
 import tempfile
+import torch
 
 from executorch.backends.cadence.aot.ops_registrations import *  # noqa
-from typing import Any, Tuple
+from typing import Any, Optional, Tuple
 
 from executorch.backends.cadence.aot.compiler import (
+    _lower_ep_to_cadence_gen_etrecord,
     convert_pt2,
-    export_to_executorch_gen_etrecord,
     fuse_pt2,
     prepare_pt2,
 )
@@ -38,13 +39,12 @@ def export_model(
     model: nn.Module,
     example_inputs: Tuple[Any, ...],
     file_name: str = "CadenceDemoModel",
-    run_and_compare: bool = True,
-    eps_error: float = 1e-1,
-    eps_warn: float = 1e-5,
+    working_dir: Optional[str] = None,
 ):
     # create work directory for outputs and model binary
-    working_dir = tempfile.mkdtemp(dir="/tmp")
-    logging.debug(f"Created work directory {working_dir}")
+    if working_dir is None:
+        working_dir = tempfile.mkdtemp(dir="/tmp")
+        logging.debug(f"Created work directory {working_dir}")
 
     # Instantiate the quantizer
     quantizer = CadenceDefaultQuantizer()
@@ -66,9 +66,11 @@ def export_model(
     # the one used in prepare_and_convert_pt2)
     quantized_model = fuse_pt2(converted_model, quantizer)
 
+    ep = torch.export.export(quantized_model, example_inputs, strict=True)
+
     # Get edge program after Cadence specific passes
-    exec_prog: ExecutorchProgramManager = export_to_executorch_gen_etrecord(
-        quantized_model, example_inputs, output_dir=working_dir
+    exec_prog: ExecutorchProgramManager = _lower_ep_to_cadence_gen_etrecord(
+        ep, output_dir=working_dir
     )
 
     logging.info("Final exported graph:\n")
@@ -92,13 +94,24 @@ def export_model(
         f"Executorch bundled program buffer saved to {file_name} is {len(buffer)} total bytes"
     )
 
-    # TODO: move to test infra
-    if run_and_compare:
-        runtime.run_and_compare(
-            executorch_prog=exec_prog,
-            inputs=example_inputs,
-            ref_outputs=ref_outputs,
-            working_dir=working_dir,
-            eps_error=eps_error,
-            eps_warn=eps_warn,
-        )
+
+def export_and_run_model(
+    model: nn.Module,
+    example_inputs: Tuple[Any, ...],
+    file_name: str = "CadenceDemoModel",
+    eps_error: float = 1e-1,
+    eps_warn: float = 1e-5,
+):
+    # create work directory for outputs and model binary
+    working_dir = tempfile.mkdtemp(dir="/tmp")
+    logging.debug(f"Created work directory {working_dir}")
+    exec_prog = export_model(model, example_inputs, file_name, working_dir)
+    ref_outputs = model(*example_inputs)
+    runtime.run_and_compare(
+        executorch_prog=exec_prog,
+        inputs=example_inputs,
+        ref_outputs=ref_outputs,
+        working_dir=working_dir,
+        eps_error=eps_error,
+        eps_warn=eps_warn,
+    )

examples/cadence/models/babyllama.py

@@ -12,7 +12,7 @@
 
 import torch
 
-from executorch.backends.cadence.aot.export_example import export_model
+from executorch.backends.cadence.aot.export_example import export_and_run_model
 
 from executorch.examples.models.llama.llama_transformer import ModelArgs, Transformer
 
@@ -35,7 +35,7 @@ def main() -> None:
     model = Transformer(args)
     example_inputs = (torch.randint(0, 10, [b, seq], dtype=torch.int64),)
 
-    export_model(model, example_inputs)
+    export_and_run_model(model, example_inputs)
 
 
 if __name__ == "__main__":

examples/cadence/models/mobilenet_v2.py

@@ -13,7 +13,7 @@
 from executorch.backends.cadence.aot.ops_registrations import *  # noqa
 
 
-from executorch.backends.cadence.aot.export_example import export_model
+from executorch.backends.cadence.aot.export_example import export_and_run_model
 from torchvision.models import mobilenet_v2, MobileNet_V2_Weights
 
 
@@ -27,4 +27,4 @@
     model.eval()
     example_inputs = (torch.randn(1, 3, 64, 64),)
 
-    export_model(model, example_inputs)
+    export_and_run_model(model, example_inputs)

examples/cadence/models/resnet18.py

@@ -13,7 +13,7 @@
 from executorch.backends.cadence.aot.ops_registrations import *  # noqa
 
 
-from executorch.backends.cadence.aot.export_example import export_model
+from executorch.backends.cadence.aot.export_example import export_and_run_model
 from torchvision.models import resnet18, ResNet18_Weights
 
 
@@ -27,4 +27,4 @@
     model.eval()
     example_inputs = (torch.randn(1, 3, 64, 64),)
 
-    export_model(model, example_inputs)
+    export_and_run_model(model, example_inputs)

examples/cadence/models/resnet50.py

@@ -13,7 +13,7 @@
 from executorch.backends.cadence.aot.ops_registrations import *  # noqa
 
 
-from executorch.backends.cadence.aot.export_example import export_model
+from executorch.backends.cadence.aot.export_example import export_and_run_model
 from torchvision.models import resnet50, ResNet50_Weights
 
 
@@ -27,4 +27,4 @@
     model.eval()
     example_inputs = (torch.randn(1, 3, 64, 64),)
 
-    export_model(model, example_inputs)
+    export_and_run_model(model, example_inputs)

examples/cadence/models/rnnt_encoder.py

@@ -14,7 +14,7 @@
 
 from typing import List, Optional, Tuple
 
-from executorch.backends.cadence.aot.export_example import export_model
+from executorch.backends.cadence.aot.export_example import export_and_run_model
 from torchaudio.prototype.models import ConvEmformer
 
 
@@ -157,4 +157,4 @@ def infer(
         transcriber_lengths,
     )
 
-    export_model(encoder, example_inputs)
+    export_and_run_model(encoder, example_inputs)

examples/cadence/models/rnnt_joiner.py

@@ -14,7 +14,7 @@
 
 from typing import Tuple
 
-from executorch.backends.cadence.aot.export_example import export_model
+from executorch.backends.cadence.aot.export_example import export_and_run_model
 
 
 FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
@@ -62,4 +62,4 @@ def forward(
         target_encodings,
     )
 
-    export_model(model, example_inputs)
+    export_and_run_model(model, example_inputs)

examples/cadence/models/rnnt_predictor.py

@@ -14,7 +14,7 @@
 
 from typing import Tuple
 
-from executorch.backends.cadence.aot.export_example import export_model
+from executorch.backends.cadence.aot.export_example import export_and_run_model
 
 
 FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
@@ -66,4 +66,4 @@ def forward(
         predictor_lengths,
     )
 
-    export_model(model, example_inputs)
+    export_and_run_model(model, example_inputs)

examples/cadence/models/vision_transformer.py

@@ -13,7 +13,7 @@
 import torch
 import torchvision
 
-from executorch.backends.cadence.aot.export_example import export_model
+from executorch.backends.cadence.aot.export_example import export_and_run_model
 
 
 FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
@@ -25,4 +25,4 @@
     model = torchvision.models.vit_b_16()
     example_inputs = (torch.randn(1, 3, 224, 224),)
 
-    export_model(model, example_inputs)
+    export_and_run_model(model, example_inputs)