NXP backend: Replace use of export_to_edge with to_edge_transform_and_lower

Author: StrycekSimon

examples/nxp/aot_neutron_compile.py

@@ -23,9 +23,13 @@
 from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer
 from executorch.examples.models import MODEL_NAME_TO_MODEL
 from executorch.examples.models.model_factory import EagerModelFactory
-from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig
+from executorch.exir import (
+    EdgeCompileConfig,
+    ExecutorchBackendConfig,
+    to_edge_transform_and_lower,
+)
 from executorch.extension.export_util import save_pte_program
-from executorch.extension.export_util.utils import export_to_edge
+from torch.export import export
 from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
 
 from .experimental.cifar_net.cifar_net import CifarNet, test_cifarnet_model
@@ -141,14 +145,6 @@ def _get_batch_size(data):
         required=True,
         help=f"Provide model name. Valid ones: {set(models.keys())}",
     )
-    parser.add_argument(
-        "-d",
-        "--delegate",
-        action="store_true",
-        required=False,
-        default=False,
-        help="Flag for producing eIQ NeutronBackend delegated model",
-    )
     parser.add_argument(
         "--target",
         required=False,
@@ -249,37 +245,28 @@ def _get_batch_size(data):
         quantized_str = "quantized " if args.quantize else ""
         print(f"\nAccuracy of the {quantized_str}`{args.model_name}`: {accuracy}\n")
 
-    # 4. Export to edge program
-    edge_compile_config = EdgeCompileConfig()
-    edge_program_manager = export_to_edge(
-        module,
-        example_inputs,
-        edge_compile_config=edge_compile_config,
+    # 4. Transform and lower
+
+    compile_spec = generate_neutron_compile_spec(
+        args.target,
+        operators_not_to_delegate=args.operators_not_to_delegate,
+        neutron_converter_flavor=args.neutron_converter_flavor,
     )
+    partitioner = NeutronPartitioner(compile_spec)
 
-    logging.debug(f"Exported graph:\n{edge_program_manager.exported_program().graph}")
+    edge_program_manager = to_edge_transform_and_lower(
+        export(module, example_inputs, strict=True),
+        partitioner=[partitioner],
+        compile_config=EdgeCompileConfig(),
+    )
 
     edge_program_manager = NeutronEdgePassManager(
         remove_io_quant_ops=args.remove_quant_io_ops
     )(edge_program_manager)
 
-    # 5. Delegate to Neutron
-    if args.delegate:
-        logging.info("Executing Neutron Partitioner and Delegate")
-
-        compile_spec = generate_neutron_compile_spec(
-            args.target,
-            operators_not_to_delegate=args.operators_not_to_delegate,
-            neutron_converter_flavor=args.neutron_converter_flavor,
-        )
-        partitioner = NeutronPartitioner(compile_spec)
-
-        edge_program_manager = edge_program_manager.to_backend(partitioner)
-        logging.debug(
-            f"Lowered graph:\n{edge_program_manager.exported_program().graph}"
-        )
+    logging.debug(f"Lowered graph:\n{edge_program_manager.exported_program().graph}")
 
-    # 6. Export to ExecuTorch program
+    # 5. Export to ExecuTorch program
     try:
         exec_prog = edge_program_manager.to_executorch(
             config=ExecutorchBackendConfig(extract_delegate_segments=False)
@@ -301,8 +288,6 @@ def executorch_program_to_str(ep, verbose=False):
 
     logging.debug(f"Executorch program:\n{executorch_program_to_str(exec_prog)}")
 
-    # 7. Serialize to *.pte
-    model_name = f"{args.model_name}" + (
-        "_nxp_delegate" if args.delegate is True else ""
-    )
+    # 6. Serialize to *.pte
+    model_name = f"{args.model_name}"
     save_pte_program(exec_prog, model_name)

examples/nxp/run_aot_example.sh

@@ -13,6 +13,6 @@ cd $EXECUTORCH_DIR
 
 # Run the AoT example
 python -m examples.nxp.aot_neutron_compile --quantize \
-    --delegate --neutron_converter_flavor SDK_25_06 -m ${MODEL}
+    --neutron_converter_flavor SDK_25_06 -m ${MODEL}
 # verify file exists
 test -f ${MODEL}_nxp_delegate.pte