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
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,
    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.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

from .experimental.cifar_net.cifar_net import (
    CifarNet,
    train_cifarnet_model,
    verify_cifarnet_model,
)
from .models.mobilenet_v2 import MobilenetV2

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_and_inputs_from_name(model_name: str, use_random_dataset: bool):
    """Given the name of an example pytorch model, return it, 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():
        if use_random_dataset:
            if model_name != "mobilenetv2":
                raise NotImplementedError(
                    f"Random dataset for model {model_name} is not implemented."
                )
            m = models[model_name](use_random_dataset=use_random_dataset)
        else:
            m = models[model_name]()

        model = m.get_eager_model()
        example_inputs = m.get_example_inputs()
        calibration_inputs = m.get_calibration_inputs(64)
    # 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]
        )
    else:
        raise RuntimeError(
            f"Model '{model_name}' is not a valid name. Use --help for a list of available models."
        )

    return model, example_inputs, calibration_inputs


models = {
    "cifar10": CifarNet,
    "mobilenetv2": MobilenetV2,
}


if __name__ == "__main__":  # noqa C901
    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(
        "-s",
        "--so_library",
        required=False,
        default=None,
        help="Path to custome 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 conversion to Neutron microcode by Neutron Converter, 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(
        "--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.",
    )

    args = 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 = get_model_and_inputs_from_name(
        args.model_name, args.use_random_dataset
    )
    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:
            match args.model_name:
                case "cifar10":
                    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)
                case _:
                    raise ValueError(
                        f"QAT training is not supported for model '{args.model_name}'"
                    )
        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,
    )
    partitioners = (
        [
            NeutronPartitioner(
                compile_spec,
                neutron_target_spec,
                post_quantization_state_dict=module.state_dict(),
            )
        ]
        if args.delegate
        else []
    )

    edge_program_manager = to_edge_transform_and_lower(
        export(module, example_inputs, strict=True),
        transform_passes=NeutronEdgePassManager(),
        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)
        )
    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 ""
    )
    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)