Move portable/util.py into extension/export_util

Pull Request resolved: #4141

`portable/util.py` is used by multiple example models including llm models. Since we are moving llm models' export libraries into `extension/llm/export`, it makes sense to move `portable/util.py`
into extension as well.

This change does that.

It moves the test into `examples/models` which seems to be a better place for these unit tests since they test the export flow for different models and also rely on some code under `examples/models` like
`EagerModelFactory`.

It also fixes the import hack inside `executorch/examples/models/llama2/builder.py`.

Differential Revision: [D59306348](https://our.internmc.facebook.com/intern/diff/D59306348/)
ghstack-source-id: 232553687

Author: helunwencser

examples/apple/mps/scripts/mps_example.py

@@ -23,6 +23,7 @@
 from executorch.exir.backend.backend_api import to_backend
 from executorch.exir.backend.backend_details import CompileSpec
 from executorch.exir.capture._config import ExecutorchBackendConfig
+from executorch.extension.export_util.utils import export_to_edge, save_pte_program
 from executorch.sdk import BundledProgram, generate_etrecord
 from executorch.sdk.bundled_program.config import MethodTestCase, MethodTestSuite
 from executorch.sdk.bundled_program.serialize import (
@@ -32,8 +33,6 @@
 from ....models import MODEL_NAME_TO_MODEL
 from ....models.model_factory import EagerModelFactory
 
-from ....portable.utils import export_to_edge, save_pte_program
-
 FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
 logging.basicConfig(level=logging.INFO, format=FORMAT)
 

examples/arm/aot_arm_compiler.py

@@ -18,13 +18,13 @@
     get_symmetric_quantization_config,
 )
 from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig
+from executorch.extension.export_util.utils import export_to_edge, save_pte_program
 
 # Quantize model if required using the standard export quantizaion flow.
 from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e
 
 from ..models import MODEL_NAME_TO_MODEL
 from ..models.model_factory import EagerModelFactory
-from ..portable.utils import export_to_edge, save_pte_program
 
 FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
 logging.basicConfig(level=logging.WARNING, format=FORMAT)

examples/models/llama2/TARGETS

@@ -86,9 +86,9 @@ runtime.python_library(
         "//executorch/examples/models:model_base",
         "//executorch/examples/models:models",
         "//executorch/examples/models/llama2/custom_ops:custom_ops_aot_py",
-        "//executorch/examples/portable:utils",
         "//executorch/exir:lib",
         "//executorch/extension/llm/export:export_lib",
+        "//executorch/extension/export_util:export_util",
         # one definition has to be included in the user of the libarary
         # depending on what library the client wants to use
         # "//executorch/extension/pybindings:aten_lib",

examples/models/llama2/builder.py

@@ -13,14 +13,6 @@
 from typing import Any, Callable, List, Optional
 
 import torch
-
-try:
-    from ...portable.utils import export_to_edge, save_pte_program
-except ImportError:
-    # Workaround to bypass the different paths between executorch pip package and directly python call
-    # TODO: remove this try catch workaround and have a standard wa to import portable.utils
-    # pyre-ignore: Undefined import [21]: Could not find a module corresponding to import `examples.portable.utils`.
-    from examples.portable.utils import export_to_edge, save_pte_program
 from executorch.backends.transforms.duplicate_dynamic_quant_chain import (
     DuplicateDynamicQuantChainPass,
 )
@@ -33,6 +25,8 @@
 from executorch.exir.passes import MemoryPlanningPass
 from executorch.exir.passes.quant_fusion_pass import QuantFusionPass
 from executorch.exir.passes.sym_shape_eval_pass import ConstraintBasedSymShapeEvalPass
+
+from executorch.extension.export_util.utils import export_to_edge, save_pte_program
 from torch._export import capture_pre_autograd_graph
 from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e
 from torch.ao.quantization.quantizer import Quantizer

examples/portable/test/__init__.py renamed to examples/models/test/__init__.py

@@ -10,12 +10,12 @@
 from executorch.examples.models import MODEL_NAME_TO_MODEL
 from executorch.examples.models.model_factory import EagerModelFactory
 
+from executorch.extension.export_util.utils import export_to_edge
+
 from executorch.extension.pybindings.portable_lib import (  # @manual
     _load_for_executorch_from_buffer,
 )
 
-from ..utils import export_to_edge
-
 
 class ExportTest(unittest.TestCase):
     def collect_executorch_and_eager_outputs(

examples/portable/test/test_export.py renamed to examples/models/test/test_export.py

@@ -12,10 +12,14 @@
 import torch
 
 from executorch.exir.capture import EdgeCompileConfig, ExecutorchBackendConfig
+from executorch.extension.export_util.utils import (
+    export_to_edge,
+    export_to_exec_prog,
+    save_pte_program,
+)
 
 from ...models import MODEL_NAME_TO_MODEL
 from ...models.model_factory import EagerModelFactory
-from ..utils import export_to_edge, export_to_exec_prog, save_pte_program
 
 
 FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"

examples/portable/scripts/export.py

@@ -17,9 +17,9 @@
 )
 from executorch.examples.models import MODEL_NAME_TO_MODEL
 from executorch.examples.models.model_factory import EagerModelFactory
-from executorch.examples.portable.utils import save_pte_program
 from executorch.exir.backend.backend_api import to_backend, validation_disabled
 from executorch.exir.capture._config import ExecutorchBackendConfig
+from executorch.extension.export_util.utils import save_pte_program
 from executorch.sdk import generate_etrecord
 
 from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e

examples/qualcomm/scripts/export_example.py

@@ -13,6 +13,7 @@
 import torch
 
 from executorch.exir import ExecutorchProgramManager
+from executorch.extension.export_util.utils import export_to_exec_prog
 from executorch.sdk import BundledProgram
 from executorch.sdk.bundled_program.config import (
     MethodInputType,
@@ -25,7 +26,6 @@
 
 from ...models import MODEL_NAME_TO_MODEL
 from ...models.model_factory import EagerModelFactory
-from ...portable.utils import export_to_exec_prog
 
 
 def save_bundled_program(

examples/sdk/scripts/export_bundled_program.py

@@ -13,11 +13,11 @@
 import torch
 from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner
 from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig
+from executorch.extension.export_util.utils import export_to_edge, save_pte_program
 from executorch.sdk import generate_etrecord
 
 from ..models import MODEL_NAME_TO_MODEL
 from ..models.model_factory import EagerModelFactory
-from ..portable.utils import export_to_edge, save_pte_program
 from . import MODEL_NAME_TO_OPTIONS
 from .quantization.utils import quantize
 

examples/xnnpack/aot_compiler.py

@@ -12,6 +12,7 @@
 import torch
 from executorch.exir import EdgeCompileConfig
 from executorch.exir.capture._config import ExecutorchBackendConfig
+from executorch.extension.export_util.utils import export_to_edge, save_pte_program
 from torch.ao.ns.fx.utils import compute_sqnr
 from torch.ao.quantization import (  # @manual
     default_per_channel_symmetric_qnnpack_qconfig,
@@ -30,7 +31,6 @@
 
 from ...models import MODEL_NAME_TO_MODEL
 from ...models.model_factory import EagerModelFactory
-from ...portable.utils import export_to_edge, save_pte_program
 
 from .. import MODEL_NAME_TO_OPTIONS
 from .utils import quantize

examples/xnnpack/quantization/example.py

@@ -28,7 +28,7 @@ def define_common_targets():
         deps = [
             ":models",
             "//executorch/backends/xnnpack/partition:xnnpack_partitioner",
-            "//executorch/examples/portable:utils",
+            "//executorch/extension/export_util:export_util",
             "//executorch/examples/xnnpack/quantization:quant_utils",
             "//executorch/exir:lib",
             "//executorch/exir/backend:backend_api",