Skip to content

Commit bdf54da

Browse files
committed
NXP backend: Replace use of export_to_edge with to_edge_transform_and_lower
1 parent 5867f97 commit bdf54da

File tree

1 file changed

+21
-26
lines changed

1 file changed

+21
-26
lines changed

examples/nxp/aot_neutron_compile.py

Lines changed: 21 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,13 @@
2323
from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer
2424
from executorch.examples.models import MODEL_NAME_TO_MODEL
2525
from executorch.examples.models.model_factory import EagerModelFactory
26-
from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig
26+
from executorch.exir import (
27+
EdgeCompileConfig,
28+
ExecutorchBackendConfig,
29+
to_edge_transform_and_lower,
30+
)
2731
from executorch.extension.export_util import save_pte_program
28-
from executorch.extension.export_util.utils import export_to_edge
32+
from torch.export import export
2933
from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
3034

3135
from .experimental.cifar_net.cifar_net import CifarNet, test_cifarnet_model
@@ -249,37 +253,28 @@ def _get_batch_size(data):
249253
quantized_str = "quantized " if args.quantize else ""
250254
print(f"\nAccuracy of the {quantized_str}`{args.model_name}`: {accuracy}\n")
251255

252-
# 4. Export to edge program
253-
edge_compile_config = EdgeCompileConfig()
254-
edge_program_manager = export_to_edge(
255-
module,
256-
example_inputs,
257-
edge_compile_config=edge_compile_config,
256+
# 4. Transform and lower
257+
258+
compile_spec = generate_neutron_compile_spec(
259+
args.target,
260+
operators_not_to_delegate=args.operators_not_to_delegate,
261+
neutron_converter_flavor=args.neutron_converter_flavor,
258262
)
263+
partitioners = [NeutronPartitioner(compile_spec)] if args.delegate else []
259264

260-
logging.debug(f"Exported graph:\n{edge_program_manager.exported_program().graph}")
265+
edge_program_manager = to_edge_transform_and_lower(
266+
export(module, example_inputs, strict=True),
267+
partitioner=partitioners,
268+
compile_config=EdgeCompileConfig(),
269+
)
261270

262271
edge_program_manager = NeutronEdgePassManager(
263272
remove_io_quant_ops=args.remove_quant_io_ops
264273
)(edge_program_manager)
265274

266-
# 5. Delegate to Neutron
267-
if args.delegate:
268-
logging.info("Executing Neutron Partitioner and Delegate")
269-
270-
compile_spec = generate_neutron_compile_spec(
271-
args.target,
272-
operators_not_to_delegate=args.operators_not_to_delegate,
273-
neutron_converter_flavor=args.neutron_converter_flavor,
274-
)
275-
partitioner = NeutronPartitioner(compile_spec)
276-
277-
edge_program_manager = edge_program_manager.to_backend(partitioner)
278-
logging.debug(
279-
f"Lowered graph:\n{edge_program_manager.exported_program().graph}"
280-
)
275+
logging.debug(f"Lowered graph:\n{edge_program_manager.exported_program().graph}")
281276

282-
# 6. Export to ExecuTorch program
277+
# 5. Export to ExecuTorch program
283278
try:
284279
exec_prog = edge_program_manager.to_executorch(
285280
config=ExecutorchBackendConfig(extract_delegate_segments=False)
@@ -301,7 +296,7 @@ def executorch_program_to_str(ep, verbose=False):
301296

302297
logging.debug(f"Executorch program:\n{executorch_program_to_str(exec_prog)}")
303298

304-
# 7. Serialize to *.pte
299+
# 6. Serialize to *.pte
305300
model_name = f"{args.model_name}" + (
306301
"_nxp_delegate" if args.delegate is True else ""
307302
)

0 commit comments

Comments
 (0)