aot_neutron_compile.py

# Copyright 2024-2026 NXP
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.

# Example script to compile the model for the NXP Neutron NPU

import argparse
import io
import logging
import os
from collections import defaultdict

import executorch.extension.pybindings.portable_lib
import executorch.kernels.quantized  # noqa F401
import torch
from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec
from executorch.backends.nxp.edge_passes.neutron_edge_pass_manager import (
    NeutronEdgePassManager,
)
from executorch.backends.nxp.edge_passes.remove_additional_quantize_dequantize_nodes_pass import (
    RemoveAdditionalQDQClustersPass,
)
from executorch.backends.nxp.edge_passes.remove_io_quant_ops_pass import (
    RemoveIOQuantOpsPass,
)
from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner
from executorch.backends.nxp.nxp_backend import (
    core_aten_ops_exception_list,
    default_preserve_ops,
    generate_neutron_compile_spec,
)
from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer
from executorch.backends.nxp.quantizer.utils import calibrate_and_quantize
from executorch.devtools.visualization.visualization_utils import (
    visualize_with_clusters,
)
from executorch.examples.models import MODEL_NAME_TO_MODEL
from executorch.examples.models.model_factory import EagerModelFactory

from executorch.examples.nxp.experimental.cifar_net.cifar_net import (
    CifarNet,
    train_cifarnet_model,
    verify_cifarnet_model,
)
from executorch.examples.nxp.models.mlperf_tiny.anomaly_detection.mlperf_tiny_anomaly_detection import (
    MLPerfTinyAnomalyDetection,
)
from executorch.examples.nxp.models.mlperf_tiny.image_classification.mlperf_tiny_image_classification import (
    MLPerfTinyImageClassification,
)
from executorch.examples.nxp.models.mlperf_tiny.keyword_spotting.mlperf_tiny_keyword_spotting import (
    MLPerfTinyKeywordSpotting,
)
from executorch.examples.nxp.models.mlperf_tiny.visual_wake_words.mlperf_tiny_visual_wake_words import (
    MLPerfTinyVisualWakeWords,
)
from executorch.examples.nxp.models.mobilenet_v2 import MobileNetV2
from executorch.exir import (
    EdgeCompileConfig,
    ExecutorchBackendConfig,
    to_edge_transform_and_lower,
)
from executorch.extension.export_util import save_pte_program
from torch.export import export
from torchao.quantization.pt2e import (
    move_exported_model_to_eval,
    move_exported_model_to_train,
)
from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_qat_pt2e


MODELS = {
    "cifar10": CifarNet,
    "mobilenetv2": MobileNetV2,
    "mlperf_tiny_anomaly_detection": MLPerfTinyAnomalyDetection,
    "mlperf_tiny_image_classification": MLPerfTinyImageClassification,
    "mlperf_tiny_keyword_spotting": MLPerfTinyKeywordSpotting,
    "mlperf_tiny_visual_wake_words": MLPerfTinyVisualWakeWords,
}

FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s"
logging.basicConfig(level=logging.INFO, format=FORMAT)


def _print_ops_in_edge_program(edge_program):
    """Find all ops used in the `edge_program` and print them out along with their occurrence counts."""

    ops_and_counts = defaultdict(
        lambda: 0
    )  # Mapping ops to the numer of times they are used.
    for node in edge_program.graph.nodes:
        if "call" not in node.op:
            continue  # `placeholder` or `output`. (not an operator)

        if hasattr(node.target, "_schema"):
            # Regular op.
            # noinspection PyProtectedMember
            op = node.target._schema.schema.name
        else:
            # Builtin function.
            op = str(node.target)

        ops_and_counts[op] += 1

    # Sort the ops based on how many times they are used in the model.
    ops_and_counts = sorted(ops_and_counts.items(), key=lambda x: x[1], reverse=True)

    # Print the ops and use counts.
    for op, count in ops_and_counts:
        print(f"{op: <50} {count}x")


def _get_model_info_from_name(
    model_name: str,
    dataset_path: str | None,
    use_random_dataset: bool,
    num_samples: int | None,
):
    """Given the name of an example pytorch model and args, return the model, its class instance (can be None), example inputs and calibration inputs (can be None).

    Raises RuntimeError if there is no example model corresponding to the given name.
    """

    calibration_inputs = None
    # Case 1: Model is defined in this file
    if model_name in MODELS.keys():
        model_cls = MODELS[model_name]
        # TODO: establish a common interface or a factory for all models
        if model_cls is CifarNet:
            if use_random_dataset:
                raise NotImplementedError(
                    f"Random dataset for model {model_name} is not implemented."
                )
            model_cls_inst = model_cls()

        elif model_cls in (
            MLPerfTinyImageClassification,
            MLPerfTinyKeywordSpotting,
            MLPerfTinyVisualWakeWords,
            MLPerfTinyAnomalyDetection,
            MobileNetV2,
        ):
            model_cls_inst = model_cls(
                dataset_path=dataset_path,
                use_random_dataset=use_random_dataset,
                num_samples=num_samples,
                balanced_dataset=False,
            )

        else:
            model_cls_inst = model_cls(use_random_dataset=use_random_dataset)

        model = model_cls_inst.get_eager_model()
        example_inputs = model_cls_inst.get_example_inputs()
        calibration_inputs = model_cls_inst.get_calibration_inputs(64)

        return model, example_inputs, calibration_inputs, model_cls_inst
    # Case 2: Model is defined in executorch/examples/models/
    elif model_name in MODEL_NAME_TO_MODEL.keys():
        logging.warning(
            "Using a model from examples/models not all of these are currently supported"
        )
        model, example_inputs, _, _ = EagerModelFactory.create_model(
            *MODEL_NAME_TO_MODEL[model_name]
        )

        return model, example_inputs, calibration_inputs, None
    else:
        raise RuntimeError(
            f"Model '{model_name}' is not a valid name. Use --help for a list of available models."
        )


def _get_arg_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-m",
        "--model_name",
        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,
        default="imxrt700",
        help="Platform for running the delegated model",
    )
    parser.add_argument(
        "-q",
        "--quantize",
        action="store_true",
        required=False,
        default=False,
        help="Produce a quantized model",
    )
    parser.add_argument(
        "--use_qat",
        action="store_true",
        required=False,
        default=False,
        help="Use QAT mode for quantization (performs two QAT training epochs)",
    )
    parser.add_argument(
        "--use_profiling",
        action="store_true",
        required=False,
        default=False,
        help="Enable profiling for eIQ Neutron NPU delegated model",
    )
    parser.add_argument(
        "-s",
        "--so_library",
        required=False,
        default=None,
        help="Path to custom kernel library",
    )
    parser.add_argument(
        "--debug", action="store_true", help="Set the logging level to debug."
    )
    parser.add_argument(
        "-t",
        "--test",
        action="store_true",
        required=False,
        default=False,
        help="Test the selected model and print the accuracy between 0 and 1.",
    )
    parser.add_argument(
        "-r",
        "--remove-quant-io-ops",
        action="store_true",
        required=False,
        default=False,
        help="Remove I/O De/Quantize nodes. Model will start to accept quantized "
        "inputs and produce quantized outputs.",
    )
    parser.add_argument(
        "--operators_not_to_delegate",
        required=False,
        default=[],
        type=str,
        nargs="*",
        help="List of operators not to delegate. E.g., --operators_not_to_delegate aten::convolution aten::mm",
    )
    parser.add_argument(
        "--visualize",
        choices=["show", "store"],
        help="Visualize the lowered program. `show` launches a browser tab with the visualization. `store` stores the "
        "visualization in a json file for later inspection. See `docs/source/visualize-with-clusters.md` for details.",
    )
    parser.add_argument(
        "--use_channels_last_dim_order",
        required=False,
        default=False,
        action="store_true",
        help="The model (including the Neutron backend) will use the channels last dim order, which can result in "
        "faster inference. The inputs must also be provided in the channels last dim order.",
    )
    parser.add_argument(
        "--dump_kernel_selection_code",
        required=False,
        default=False,
        action="store_true",
        help="During compilation to Neutron microcode by Neutron Compiler, a kernel selection file will be dumped in "
        "the working directory. This file can be used for reduction of Neutron Firmware size in the built app."
        "See `docs/source/backends/nxp/nxp-kernel-selection.md` for details.",
    )
    parser.add_argument(
        "--use_random_dataset",
        required=False,
        default=False,
        action="store_true",
        help="The calibration and testing datasets will be generated randomly instead of being downloaded.",
    )
    parser.add_argument(
        "--num_random_samples",
        required=False,
        default=None,
        type=int,
        help="Number of random samples to generate, required when `--use_random_dataset` flag is set.",
    )
    parser.add_argument(
        "-dst",
        "--dataset_path",
        required=False,
        default=None,
        help="Path to custom dataset archive (.pt, .xz) to use for calibration.",
    )
    parser.add_argument(
        "--fetch_constants_to_sram",
        required=False,
        default=False,
        action="store_true",
        help="This feature allows running models which do not fit into SRAM by offloading them to an external memory.",
    )

    return parser


if __name__ == "__main__":  # noqa C901
    args = _get_arg_parser().parse_args()

    if args.debug:
        logging.basicConfig(level=logging.DEBUG, format=FORMAT, force=True)

    neutron_target_spec = NeutronTargetSpec(target=args.target)

    # 1. pick model from one of the supported lists
    model, example_inputs, calibration_inputs, model_cls_inst = (
        _get_model_info_from_name(
            args.model_name,
            args.dataset_path,
            args.use_random_dataset,
            args.num_random_samples,
        )
    )
    model = model.eval()

    if args.use_channels_last_dim_order:
        # Turn the model to channels last.
        model.to(memory_format=torch.channels_last)

        # The dim order of the example inputs will define the dim order of the intermediate tensors in the model.
        example_inputs = tuple(
            i.to(memory_format=torch.channels_last) for i in example_inputs
        )

    else:
        # Notify the user of this option.
        print(
            "HINT: Converting your model to channels last may significantly improve inference speed. You can use the "
            "flag `--use_channels_last_dim_order`. See `docs/source/backends/nxp/nxp-dim-order.md` for more information."
        )

    # 2. Export the model to ATEN
    exported_program = torch.export.export(model, example_inputs, strict=True)

    module = exported_program.module()

    # 3. Quantize if required
    if args.quantize:
        quantizer = NeutronQuantizer(neutron_target_spec, is_qat=args.use_qat)
        if args.use_qat:
            if not isinstance(
                model_cls_inst,
                (
                    CifarNet,
                    MLPerfTinyImageClassification,
                    MLPerfTinyKeywordSpotting,
                    MLPerfTinyVisualWakeWords,
                    MLPerfTinyAnomalyDetection,
                    MobileNetV2,
                ),
            ):
                raise ValueError(
                    f"QAT training is not supported for model '{args.model_name}'"
                )

            if isinstance(model_cls_inst, CifarNet):
                print("Starting two epochs of QAT training with CifarNet model...")
                module = prepare_qat_pt2e(module, quantizer)
                module = move_exported_model_to_train(module)
                module = train_cifarnet_model(module, num_epochs=2)
                module = move_exported_model_to_eval(module)
                module = convert_pt2e(module)

            else:
                print(
                    f"Starting two epochs of QAT training with {args.model_name} model..."
                )
                module = prepare_qat_pt2e(module, quantizer)
                module = move_exported_model_to_train(module)
                module = model_cls_inst.train_model_fn(
                    module,
                    num_epochs=2,
                    channels_last=args.use_channels_last_dim_order,
                )
                module = move_exported_model_to_eval(module)
                module = convert_pt2e(module)

        else:
            if calibration_inputs is None:
                logging.warning(
                    "No calibration inputs available, using the example inputs instead"
                )
                calibration_inputs = example_inputs
            module = calibrate_and_quantize(module, calibration_inputs, quantizer)

    if args.so_library is not None:
        logging.debug(f"Loading libraries: {args.so_library}")
        torch.ops.load_library(args.so_library)

    if args.test:
        match args.model_name:
            case "cifar10":
                accuracy = verify_cifarnet_model(module)

            case _:
                raise NotImplementedError(
                    f"Testing of model `{args.model_name}` is not yet supported."
                )

        quantized_str = "quantized " if args.quantize else ""
        print(f"\nAccuracy of the {quantized_str}`{args.model_name}`: {accuracy}\n")

    # 4. Transform and lower

    compile_spec = generate_neutron_compile_spec(
        args.target,
        operators_not_to_delegate=args.operators_not_to_delegate,
        fetch_constants_to_sram=args.fetch_constants_to_sram,
        dump_kernel_selection_code=args.dump_kernel_selection_code,
        use_profiling=args.use_profiling,
    )
    partitioners = (
        [
            NeutronPartitioner(
                compile_spec,
                neutron_target_spec,
                post_quantization_state_dict=module.state_dict(),
                preserve_ops=default_preserve_ops,
            )
        ]
        if args.delegate
        else []
    )

    edge_program_manager = to_edge_transform_and_lower(
        export(module, example_inputs, strict=True),
        transform_passes=NeutronEdgePassManager(),
        generate_etrecord=args.use_profiling,
        partitioner=partitioners,
        compile_config=EdgeCompileConfig(
            _core_aten_ops_exception_list=core_aten_ops_exception_list,
        ),
    )

    if args.remove_quant_io_ops:
        edge_program_manager = edge_program_manager.transform(
            [RemoveIOQuantOpsPass(edge_program_manager=edge_program_manager)]
        )

    edge_program_manager = edge_program_manager.transform(
        NeutronEdgePassManager([RemoveAdditionalQDQClustersPass()])
    )

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

    # 5. Export to ExecuTorch program
    try:
        exec_prog = edge_program_manager.to_executorch(
            config=ExecutorchBackendConfig(extract_delegate_segments=False)
        )

        # Generate ETRecord if profiling flag is set
        if args.use_profiling:
            etrecord_path = os.path.join("etrecord", f"{args.model_name}_etrecord.bin")
            # Create directory if it doesn't exist
            os.makedirs(os.path.dirname(etrecord_path), exist_ok=True)
            # Save ETRecord
            exec_prog.get_etrecord().save(etrecord_path)
            # Notify the user about profiling enablement and ETRecord generation.
            logging.info(
                "The model was converted with profiling enabled. The time spent generating the profiling dump is traced as the "
                "final delegate operation and can be ignored, as no dump is produced for non‑profilable models."
            )
            logging.info(f"The ETRecord for the model was saved to {etrecord_path}.")

    except RuntimeError as e:
        if "Missing out variants" in str(e.args[0]):
            raise RuntimeError(
                e.args[0]
                + ".\nThis likely due to an external so library not being loaded. Supply a path to it with the "
                "--so_library flag."
            ).with_traceback(e.__traceback__) from None
        else:
            raise e

    def executorch_program_to_str(ep, verbose=False):
        f = io.StringIO()
        ep.dump_executorch_program(out=f, verbose=verbose)
        return f.getvalue()

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

    # 6. Serialize to *.pte
    model_name = (
        f"{args.model_name}"
        + ("_nxp_delegate" if args.delegate is True else "")
        + ("_profile" if args.use_profiling is True else "")
    )
    save_pte_program(exec_prog, model_name)

    # 7. Optionally visualize the model.
    if args.visualize == "show":
        visualize_with_clusters(exec_prog.exported_program())
    elif args.visualize == "store":
        file_name = f"{args.model_name}-visualization.json"
        logging.info(
            f"Saved the graph visualization in `{file_name}`. It can be opened using the ModelExplorer."
        )
        visualize_with_clusters(exec_prog.exported_program(), file_name)