frontend.py Source File

frontend.py Source File#

Mobilint SDK qb Compiler: frontend.py Source File
Mobilint SDK qb Compiler v1.0
MCS002-KR
 
import pathlib
from pathlib import Path
from typing import Any, Optional
 
import torch
 
from qbcompiler.compiler.compiler import Compiler
from qbcompiler.compiler.compiler import _quantization_task
from qbcompiler.configs import *
from qbcompiler.version import __version__
 
 
 
class Model_Dict(Compiler):
    """
    @brief Wrapper around the Mobilint compiler to support compilation and inference workflows.
 
    @details Capabilities include:
    - Compilation of models into MXQ artifacts runnable on Mobilint NPUs.
    - Inference using the full-precision high-level compiled model on CPU or GPU.
    - Inference with the quantized model on CPU or GPU.
    """
 
    def __init__(
        self,
        model,
        backend="onnx",
        device="cpu",
        feed_dict=None,
        dynamic_axes=None,
        in_dformats=None,
        yolo_decode_include=False,
        use_custom_mask=False,
        exclude_first_subgraph=False,
        **kwargs,
    ):
        """
        @brief Initialize the Mobilint compiler wrapper.
 
        @param model string or model instance. Model path or in-memory model to compile.
        @param backend string. Framework identifier for the model (for example "onnx"). Defaults to "onnx".
        @param device string. Target device for inference ("cpu" or "gpu"). Defaults to "cpu".
        @param feed_dict dict. Example inputs for shape resolution.
        @param dynamic_axes dict. Marks model axes as dynamic.
        @param in_dformats dict. Describes input data formats.
        @param yolo_decode_include bool. Runs YOLO decode on NPU when @c True.
        @param use_custom_mask bool. For non-LLM models that include a mask, use a custom mask provided as a separate input instead of the causal mask.
        @param exclude_first_subgraph bool. Applies only when CPU offloading is enabled: exclude the first subgraph from the final graph if it is unsupported.
        @param kwargs dict. Additional compiler arguments.
        """
        super().__init__(
            model=model,
            backend=backend,
            device=device,
            feed_dict=feed_dict,
            in_dformats=in_dformats,
            dynamic_axes=dynamic_axes,
            yolo_decode_include=yolo_decode_include,
            use_custom_mask=use_custom_mask,
            exclude_first_subgraph=exclude_first_subgraph,
            **kwargs,
        )
 
    def compile(
        self,
        calib_data_path: str = None,
        save_path: str = None,
        save_subgraph_type: int = 0,
        output_subgraph_path="",
        model_nickname: str = None,
        save_sample=False,
        use_random_calib=False,
        inference_scheme="single",
        cpu_offload=False,
        optimize_option=1,
        sample_dtype: str = "float",
        buffer_mode=1,
        compile_config=None,  # CompileConfig object
        **kwargs,
    ):
        """
        @brief Compile the wrapped model into an MXQ artifact.
 
        @param calib_data_path string. Calibration dataset directory or meta file (.txt/.json) listing NumPy samples.
        @param save_path string. Output MXQ filename.
        @param save_subgraph_type int. Controls optional MBLT exports: 0 disables; 1 saves graph structure; 2 saves
        structure and weights; 3 splits structure into multiple subgraphs; 4 splits structure and weights.
        @param output_subgraph_path string. Destination for exported .mblt when @p save_subgraph_type is 1–4.
        @param model_nickname string. Friendly name stored in qbcompiler for faster recompile.
        @param save_sample bool. Persists inputs/outputs for debugging in ./sampleInOut.
        @param use_random_calib bool. Uses randomly generated calibration samples when @c True.
        @param inference_scheme string. NPU inference configuration ("single", "multi", "global", etc.).
        @param cpu_offload bool. Enables CPU offload for unsupported groups.
        @param optimize_option int. Compiler optimization selector.
        @param sample_dtype string. Data type used for saved sample outputs.
        @param buffer_mode int. Buffer serialization mode (0 naive, 1 mmap).
        @param compile_config CompileConfig. Compile configuration object.
        @param kwargs dict. Additional arguments forwarded to the compiler.
 
        """
        super().compile(
            calib_data_path=calib_data_path,
            save_path=save_path,
            save_subgraph_type=save_subgraph_type,
            output_subgraph_path=output_subgraph_path,
            model_nickname=model_nickname,
            save_sample=save_sample,
            inference_scheme=inference_scheme,
            use_random_calib=use_random_calib,
            cpu_offload=cpu_offload,
            optimize_option=optimize_option,
            sample_dtype=sample_dtype,
            buffer_mode=buffer_mode,
            compile_config=compile_config,
            **kwargs,
        )
 
 
def mxq_compile(
    model,
    calib_data_path: str = "",
    save_subgraph_type: int = 0,
    output_subgraph_path="",
    model_nickname: str = None,
    save_path: str = None,
    backend="onnx",
    device="cpu",
    feed_dict=None,
    dynamic_axes=None,
    in_dformats=None,
    save_sample=False,
    inference_scheme="single",
    yolo_decode_include=False,
    use_random_calib=False,
    cpu_offload=False,
    optimize_option=1,
    sample_dtype="float",
    buffer_mode=1,
    input_shape_dict=dict(),
    use_custom_mask=False,
    exclude_first_subgraph=False,
    compile_config: Optional[CompileConfig] = None,
    resource_management_config: Optional[ResourceManagementConfig] = None,
    quantization_config: Optional[QuantizationConfig] = None,
    advanced_quantization_config: Optional[AdvancedQuantizationConfig] = None,
    llm_config: Optional[LlmConfig] = None,
    input_process_config: Optional[InputProcessConfig] = None,
    **kwargs,
):
    """
    @brief Compile a model into a Mobilint eXeCUtable (MXQ) package for execution on Mobilint NPUs.
 
    @param model string or model instance. Model path. When using @c backend="onnx", this should be the path to an ONNX
    model file. When @c backend is "torchscript", provide a TorchScript module; when "torch", provide a standard
    PyTorch model. For @c backend="tf", pass the directory that contains the TensorFlow SavedModel graph and assets. For
    @c backend="tflite", pass the TF Lite model path.
    @param calib_data_path string. Path to the calibration dataset. Accepts either a text/json file that lists NumPy
    files or a directory that contains the pre-processed NumPy files.
    @param save_subgraph_type int. Controls optional MBLT subgraph exports: 0 disables exports; 1 saves only the graph
    structure; 2 saves graph structure plus weights; 3 saves the graph structure split into multiple subgraphs; 4 saves
    both structure and weights split into multiple subgraphs.
    @param output_subgraph_path string. Destination path for the exported .mblt file when @p save_subgraph_type is 1, 2,
    3, or 4. The resulting file can be used for visualization.
    @param model_nickname string. Friendly name stored in qbcompiler. Defaults to the model basename (for example, the file
    "/workspace/onnx/resnet50.onnx" becomes "resnet50"). Falls back to "temporary" when it cannot be inferred.
    @param save_path string. Output MXQ filename. Defaults to "{model_nickname}.mxq" when omitted.
    @param backend string. Framework used to generate the Mobilint IR. Must be one of "onnx", "tf", "tflite",
    "torchscript", or "torch". Defaults to "onnx".
    @param device string. Compilation and inference device: "cpu" or "gpu". Defaults to "cpu".
    @param feed_dict dict. Example input tensors for shape inference and inference validation.
    @param dynamic_axes dict. Marks model axes as dynamic. Keys are input names and values map dimension indices to
    aliases (for example {"input": {2: "seq_len"}}).
    @param in_dformats dict. Describes input tensor data formats.
    @param save_sample bool. Stores input/output tensors in ./sampleInOut when @c True to aid debugging. Defaults to
    @c False.
    @param inference_scheme string. NPU inference scheme. One of "single", "multi", "global", "global4", or
    "global8".
    @param yolo_decode_include bool. Determines whether YOLO decode runs on NPU.
    @param use_random_calib bool. Generates random calibration data to validate model compilability. Defaults to @c
    False.
    @param cpu_offload bool. Enables CPU offloading during NPU inference. Defaults to @c False.
    @param optimize_option int. Compiler optimization strategy selector.
    @param sample_dtype string. Data type for saved sample inference outputs: "float" or "int8".
    @param buffer_mode int. Buffer serialization mode: 0 uses a naïve buffer, 1 uses an mmap-backed buffer. Defaults to
    1.
    @param input_shape_dict:multi-shape compilation specification for STT/TTS models only (e.g., Conformer-CTC, MeloTTS).
    The dynamic axis is defined with respect to each input tensor’s shape. During compilation, the same axis is varied across all model inputs
    using the provided values. Currently, only a single entry key ("multi_shape0") is supported. Example: {"multi_shape0": {"axis": 2, "values": [100, 200, 300]}}
    @param use_custom_mask bool. For non-LLM models that include a mask, use a custom mask provided as a separate input instead of the causal mask.
    For LLMs that incorporate RoPE, setting LLMConfig.dynamic_rope to True not only enables the use of dynamic RoPE provided via an additional input, but also activates the custom mask.
    @param exclude_first_subgraph bool. Applies only when CPU offloading is enabled: exclude the first subgraph from the final graph if it is unsupported.
    @param compile_config string. Path to JSON configuration file containing all quantization-related settings
    (resourceManagement, quantization, advancedQuantization, llmConfig). If provided, this takes precedence over
    individual sub-config parameters.
    @param resource_management_config ResourceManagementConfig. Resource management configuration object.
    @param quantization_config QuantizationConfig. Quantization configuration object.
    @param advanced_quantization_config AdvancedQuantizationConfig. Advanced quantization configuration object.
    @param llm_config LlmConfig. LLM configuration object.
    @param input_process_config InputProcessConfig. Unified input processing configuration containing
        uint8 input handling, image channels, and preprocessing pipeline settings.
    @param kwargs dict. Additional compiler arguments propagated to the backend implementation. Moreover, If you want to override only
    a subset of configuration values, you can pass them to `mxq_compile` as keyword arguments (`kwargs`), and they will be applied.
    Note that `AdvancedQuantizationConfig` and its sub-configurations are not updated automatically from `kwargs`. In contrast, `ResourceManagementConfig`,
    `QuantizationConfig`, and `LlmConfig` support automatic updates via each configuration class’s `from_kwargs` method. See Example 3 below for a usage example.
    @return None.
 
    @par Using configuration
    There are two ways to configure quantization settings:
 
    1. Load all settings from a JSON config file:
    @code{.py}
    from qbcompiler import mxq_compile
 
    mxq_compile(
        model="path/to/model.onnx",
        calib_data_path="path/to/calib",
        compile_config="path/to/config.json",
        device="gpu",
    )
    @endcode
 
    2. Pass individual sub-config objects:
    @code{.py}
    from qbcompiler import mxq_compile
    from qbcompiler.configs import (
        ResourceManagementConfig,
        QuantizationConfig,
        AdvancedQuantizationConfig,
        LlmConfig,
    )
 
    resource_mgmt = ResourceManagementConfig(weight_dtype="float32")
    quant_cfg = QuantizationConfig(...)
    adv_quant_cfg = AdvancedQuantizationConfig(...)
    llm_cfg = LlmConfig(...)
 
    mxq_compile(
        model="path/to/model.onnx",
        calib_data_path="path/to/calib",
        resource_management_config=resource_mgmt,
        quantization_config=quant_cfg,
        advanced_quantization_config=adv_quant_cfg,
        llm_config=llm_cfg,
        device="gpu",
    )
    @endcode
 
    3. Automatically applied partial configuration overrides:
    @code{.py}
    from qbcompiler import mxq_compile
 
    mxq_compile(
        model="path/to/model.onnx",
        calib_data_path="path/to/calib",
        quantization_method=1, # per channel quantization
        quantization_mode=1, # max percentile quantization
        percentile=0.999, # percentile value for max percentile quantization
        quantization_output=0, # per layer quantization for the output layer
        device="gpu",
    )
    @endcode
    Please refer to the mxq_compile function and the quantization configuration section for the meaning of the quantization-related numeric values.
 
 
    @par Compiling models with custom inputs (ONNX/Torch/TensorFlow)
    Provide NumPy inputs whenever the model omits shape information so that qbcompiler can infer unknown dimensions and data
    formats.
 
    @code{.py}
    from qbcompiler import mxq_compile
    import numpy as np
 
    example_input = {
        "input_node_name_1": np.random.randn(1, 3, 224, 224).astype(np.float32),
        "input_node_name_2": np.random.randn(1, 8, 224, 224).astype(np.float32),
        "input_node_name_3": np.random.randn(1, 4, 56, 56).astype(np.float32),
    }
 
    in_dformats = {
        "input_node_name_1": "NCHW",
        "input_node_name_2": "NCHW",
        "input_node_name_3": "NCHW",
    }
 
    onnx_model_path = "path/to/your/model.onnx"
    mxq_compile(
        model=onnx_model_path,
        feed_dict=example_input,
        in_dformats=in_dformats,
        backend="onnx",
        compile_config="path/to/config.json",
    )
    @endcode
    """
    if model_nickname and isinstance(model_nickname, str):
        model_name = model_nickname
    elif isinstance(model, str):  # in the case of model path
        model_name = Path(model).stem
    else:  # in the case of model instance
        if save_path is None:
            raise ValueError("Please set the save_path with .mxq extension")
        model_name = "model"
    if save_path is None:
        save_path = f"{model_name}.mxq"
 
    # Handle config loading
    if compile_config is not None:
        # Accept both CompileConfig object and JSON file path
        if isinstance(compile_config, CompileConfig):
            compile_cfg = compile_config
        elif isinstance(compile_config, str):
            compile_cfg = CompileConfig.from_json(
                compile_config,
                calib_data_path=calib_data_path,
                use_random_calib=use_random_calib,
                save_path=save_path,
                save_msgpack_name=None,
                model_nickname=model_name,
                model_path=None,
                inference_scheme=inference_scheme,
                save_sample=save_sample,
                cpu_offload=cpu_offload,
                optimize_option=optimize_option,
                sample_dtype=sample_dtype,
                buffer_mode=buffer_mode,
                version=__version__,
                input_shape_dict=input_shape_dict,
            )
        else:
            raise ValueError(
                "compile_config must be either a CompileConfig object or a path to JSON configuration file"
            )
    else:
        pairs = [
            (
                "resource_management_config",
                resource_management_config,
                ResourceManagementConfig,
            ),
            ("quantization_config", quantization_config, QuantizationConfig),
            (
                "advanced_quantization_config",
                advanced_quantization_config,
                AdvancedQuantizationConfig,
            ),
            ("llm_config", llm_config, LlmConfig),
            ("input_process_config", input_process_config, InputProcessConfig),
        ]
        for name, obj, expected_cls in pairs:
            if obj is None:
                continue
            if not isinstance(obj, expected_cls):
                raise TypeError(
                    f"Invalid type for '{name}': expected {expected_cls.__name__} or None, "
                    f"but got {type(obj).__name__}."
                )
 
        compile_cfg = get_compile_config(
            calib_data_path=calib_data_path,
            use_random_calib=use_random_calib,
            save_path=save_path,
            save_msgpack_name=None,
            model_nickname=model_name,
            model_path=None,
            inference_scheme=inference_scheme,
            save_sample=save_sample,
            cpu_offload=cpu_offload,
            optimize_option=optimize_option,
            sample_dtype=sample_dtype,
            buffer_mode=buffer_mode,
            version=__version__,
            input_shape_dict=input_shape_dict,
            resource_management_config=resource_management_config,
            quantization_config=quantization_config,
            llm_config=llm_config,
            advanced_quantization_config=advanced_quantization_config,  # AdvancedQuantizationConfig config don't use "from_kwargs" because of shared key names, such as "apply".
            input_process_config=input_process_config,
            **kwargs,
        )
 
    if isinstance(model, str) and model.endswith(".mblt") and os.path.isfile(model):
        _quantization_task(model, device, compile_config=compile_cfg)
    else:
        model_dict = Model_Dict(
            model=model,
            backend=backend,
            device=device,
            feed_dict=feed_dict,
            in_dformats=in_dformats,
            dynamic_axes=dynamic_axes,
            yolo_decode_include=yolo_decode_include,
            use_custom_mask=use_custom_mask,
            exclude_first_subgraph=exclude_first_subgraph,
            **kwargs,
        )
        model_dict.compile(
            model_nickname=model_name,
            calib_data_path=calib_data_path,
            save_path=save_path,
            save_subgraph_type=save_subgraph_type,
            output_subgraph_path=output_subgraph_path,
            save_sample=save_sample,
            inference_scheme=inference_scheme,
            use_random_calib=use_random_calib,
            cpu_offload=cpu_offload,
            optimize_option=optimize_option,
            sample_dtype=sample_dtype,
            buffer_mode=buffer_mode,
            compile_config=compile_cfg,
        )
        torch.cuda.empty_cache()
        del model_dict
 
 
def mblt_compile(
    model: str | Any,
    mblt_save_path: str,
    backend="onnx",
    device="cpu",
    feed_dict=None,
    dynamic_axes=None,
    in_dformats=None,
    yolo_decode_include=False,
    cpu_offload=False,
    use_custom_mask=False,
    exclude_first_subgraph=False,
    **kwargs,
):
    """
    @brief Export a model to the Mobilint .mblt format without producing an MXQ package.
 
    @param model string or model instance. Source model or path to compile.
    @param mblt_save_path string. Output path for the .mblt artifact.
    @param backend string. Framework identifier ("onnx", "tf", "tflite", "torchscript", or "torch"). Defaults to
    "onnx".
    @param device string. Compilation device ("cpu" or "gpu"). Defaults to "cpu".
    @param feed_dict dict. Example inputs used for shape inference.
    @param dynamic_axes dict. Declares dynamic axes per input name.
    @param in_dformats dict. Input dataformat metadata.
    @param yolo_decode_include bool. Runs YOLO decode on NPU when True.
    @param cpu_offload bool. Enables CPU offloading for unsupported groups.
    @param use_custom_mask bool. For non-LLM models that include a mask, use a custom mask provided as a separate input instead of the causal mask.
    For LLMs that incorporate RoPE, setting LLMConfig.dynamic_rope to True not only enables the use of dynamic RoPE provided via an additional input, but also activates the custom mask.
    @param exclude_first_subgraph bool. Applies only when CPU offloading is enabled: exclude the first subgraph from the final graph if it is unsupported.
    @param kwargs dict. Additional arguments forwarded to the compiler.
    @return None.
    """
    model_dict = Model_Dict(
        model=model,
        backend=backend,
        device=device,
        feed_dict=feed_dict,
        dynamic_axes=dynamic_axes,
        in_dformats=in_dformats,
        yolo_decode_include=yolo_decode_include,
        use_custom_mask=use_custom_mask,
        exclude_first_subgraph=exclude_first_subgraph,
        **kwargs,
    )
 
    model_dict.mblt_compile(
        model,
        mblt_save_path=mblt_save_path,
        backend=backend,
        cpu_offload=cpu_offload,
        **kwargs,
    )
 
 
 
 
def run_mblt_model(
    path_to_mblt: str,
    save_all_outputs: bool = True,
    inference_output_path: Optional[str] = None,
    target_subgraph: Optional[int] = None,
):
    from qbcompiler.model_dict.common.model_dict import ValueDict
    from qbcompiler.model_dict.inference.model import ModelInference
    from qbcompiler.model_dict.serialize import load_mblt_model
 
    assert path_to_mblt.endswith(".mblt")
    if not save_all_outputs:
        raise NotImplementedError(
            f"save_all_outputs={save_all_outputs} is not supported yet."
        )
    if target_subgraph:
        raise NotImplementedError(
            f"target_subgraph={target_subgraph} is not supported yet."
        )
    import numpy
 
    if not inference_output_path:
        inference_output_path = str(pathlib.Path(path_to_mblt).with_suffix(".infer"))
 
    vd = ValueDict()
    md, wd = load_mblt_model(path_to_mblt, load_weight=True)
    for sg in md.subgraphs:
        for iidx in sg.inputs:
            inact = sg.activations[iidx]
            if inact.name in set(md.inputs):
                vd[inact.name] = numpy.random.random(tuple(map(int, inact.src_shape)))
 
    model_inference = ModelInference(
        model_dict=md,
        weight_dict=wd,
        value_dict=vd,
        save_all_outputs=save_all_outputs,
        inference_all_outputs_write_path=inference_output_path,
    )
 
    model_inference.run_inference(use_value_dict_data_for_next_input=False)