Skip to content
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 29 additions & 28 deletions examples/nxp/aot_neutron_compile.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,22 +18,20 @@
from executorch.backends.nxp.backend.ir.edge_passes.remove_io_quant_ops_pass import (
RemoveIOQuantOpsPass,
)
from executorch.backends.nxp.edge_passes.neutron_edge_pass_manager import (
NeutronEdgePassManager,
)
from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner
from executorch.backends.nxp.nxp_backend import generate_neutron_compile_spec
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,
to_edge_transform_and_lower,
)
from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig
from executorch.extension.export_util import save_pte_program
from torch.export import export
from executorch.extension.export_util.utils import export_to_edge
from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e

from .experimental.cifar_net.cifar_net import CifarNet, test_cifarnet_model

from .models.mobilenet_v2 import MobilenetV2

FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
Expand Down Expand Up @@ -228,7 +226,7 @@ def _get_batch_size(data):

module = exported_program.module()

# 4. Quantize if required
# 3. Quantize if required
if args.quantize:
if calibration_inputs is None:
logging.warning(
Expand All @@ -254,39 +252,42 @@ def _get_batch_size(data):
quantized_str = "quantized " if args.quantize else ""
print(f"\nAccuracy of the {quantized_str}`{args.model_name}`: {accuracy}\n")

# 5. Export to edge program
partitioner_list = []
if args.delegate is True:
partitioner_list = [
# 4. Export to edge program
edge_compile_config = EdgeCompileConfig()
edge_program_manager = export_to_edge(
module,
example_inputs,
edge_compile_config=edge_compile_config,
)

logging.debug(f"Exported graph:\n{edge_program_manager.exported_program().graph}")

edge_program_manager = NeutronEdgePassManager()(edge_program_manager)

if args.remove_quant_io_ops:
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This step is done before delegation, but in executorch_pipeline.py it is done after delegation. In our local repo it is placed before delegation. Why are there differences?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it should be placed before the delegation. @MartinPavella, can you confirm it since you are the author of the updated executorch_pipeline.py?

edge_program_manager = edge_program_manager.transform(
[RemoveIOQuantOpsPass(edge_program_manager=edge_program_manager)]
)

# 5. Delegate to Neutron
if args.delegate:
logging.info("Executing Neutron Partitioner and Delegate")
edge_program_manager = edge_program_manager.to_backend(
NeutronPartitioner(
generate_neutron_compile_spec(
args.target,
args.neutron_converter_flavor,
operators_not_to_delegate=args.operators_not_to_delegate,
)
)
]

edge_program = to_edge_transform_and_lower(
export(module, example_inputs, strict=True),
partitioner=partitioner_list,
compile_config=EdgeCompileConfig(
_check_ir_validity=False,
),
)
logging.debug(f"Exported graph:\n{edge_program.exported_program().graph}")

if args.remove_quant_io_ops:
edge_program = edge_program.transform(
[RemoveIOQuantOpsPass(edge_program_manager=edge_program)]
)
logging.debug(
f"Exported graph (RemoveIOQuantOpsPass):\n{edge_program.exported_program().graph}"
f"Lowered graph:\n{edge_program_manager.exported_program().graph}"
)

# 6. Export to ExecuTorch program
try:
exec_prog = edge_program.to_executorch(
exec_prog = edge_program_manager.to_executorch(
config=ExecutorchBackendConfig(extract_delegate_segments=False)
)
except RuntimeError as e:
Expand Down
Loading