utils.py

# Copyright (c) Qualcomm Innovation Center, Inc.
# All rights reserved
#
# This source code is licensed under the BSD-style license found in the
# LICENSE file in the root directory of this source tree.
import operator
import re
import time
import warnings
from collections import OrderedDict
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import executorch.backends.qualcomm.python.PyQnnManagerAdaptor as PyQnnManagerAdaptor

import executorch.exir as exir

import torch

from executorch.backends.qualcomm._passes import AnnotateStack
from executorch.backends.qualcomm._passes.qnn_pass_manager import QnnPassManager

from executorch.backends.qualcomm.builders.node_visitor import (
    QNN_QUANT_TYPE_MAP,
    QNN_TENSOR_TYPE_MAP,
)
from executorch.backends.qualcomm.builders.qnn_constants import OpContextLoader
from executorch.backends.qualcomm.partition.qnn_partitioner import (
    generate_qnn_executorch_option,
    get_skip_decomp_table,
    QnnPartitioner,
)
from executorch.backends.qualcomm.serialization.qc_schema import (
    _soc_info_table,
    HtpArch,
    QcomChipset,
    QnnExecuTorchBackendOptions,
    QnnExecuTorchBackendType,
    QnnExecuTorchHtpBackendOptions,
    QnnExecuTorchHtpPerformanceMode,
    QnnExecuTorchHtpPrecision,
    QnnExecuTorchLogLevel,
    QnnExecuTorchOptions,
    QnnExecuTorchProfileLevel,
)
from executorch.backends.qualcomm.serialization.qc_schema_serialize import (
    flatbuffer_to_option,
    option_to_flatbuffer,
)
from executorch.backends.qualcomm.utils.constants import (
    QCOM_QNN_COMPILE_SPEC,
    QCOM_QUANTIZED_IO,
)

from executorch.exir import (
    EdgeCompileConfig,
    ExecutorchProgramManager,
    ExirExportedProgram,
    to_edge,
)
from executorch.exir.backend.compile_spec_schema import CompileSpec
from executorch.exir.capture import ExecutorchBackendConfig
from executorch.exir.lowered_backend_module import LoweredBackendModule
from executorch.exir.program._program import (
    EdgeProgramManager,
    to_edge_transform_and_lower,
)
from torch._decomp import core_aten_decompositions, remove_decompositions
from torch.export.exported_program import ExportedProgram
from torch.fx import passes
from torch.fx.passes.operator_support import OperatorSupportBase
from torch.library import Library


class _AnnotationSkipper(OperatorSupportBase):
    """
    Class used to partition out unwanted graph nodes.
    e.g. - nodes are prevented from quantization annotation
         - nodes have been grouped together as a submodule

    Attributes
    ----------
    fp_node_id_set : set
        a set contains nodes' name to be left in fp precision
    fp_node_op_set : set
        a set contains nodes' target (aten dialect) to be left in fp precision
    skip_annotated_submodule : bool
        flag to skip annotated submodule or not

    Methods
    -------
    should_delegate(n: torch.fx.Node)
        identify the residual nodes haven't be lowered with fixed-precision
    should_skip(n: torch.fx.Node)
        identify the nodes should be kept out with fixed-precision or not
    is_node_supported(_, node: torch.fx.Node)
        overridden method for graph partitioning
    """

    def __init__(
        self,
        fp_node_id_set: set = None,
        fp_node_op_set: set = None,
        skip_annotated_submodule: bool = False,
    ):
        self.fp_node_id_set = fp_node_id_set
        self.fp_node_op_set = fp_node_op_set
        self.skip_annotated_submodule = skip_annotated_submodule

    def should_delegate(self, n: torch.fx.Node):
        return n.op == "call_function" and n.target != operator.getitem

    def should_skip(self, n: torch.fx.Node):
        return n.name in self.fp_node_id_set or n.target in self.fp_node_op_set

    def is_node_supported(self, _, node: torch.fx.Node) -> bool:
        if self.skip_annotated_submodule:
            if node.op == "get_attr":
                return all(self.should_delegate(user) for user in node.users)
            return self.should_delegate(node)

        if any(
            [
                node.op in ("placeholder", "output"),
                self.should_skip(node),
                # check if parameters belong to fallbacked operator
                (
                    node.op == "get_attr"
                    and all(self.should_skip(user) for user in node.users)
                ),
            ]
        ):
            print(f"[QNN Quantizer Annotation]: {node.name} | Skipped")
            return False

        return True


def qnn_capture_config():
    return exir.CaptureConfig(enable_aot=True)


def qnn_edge_config() -> exir.EdgeCompileConfig:
    return exir.EdgeCompileConfig(
        _check_ir_validity=False,
    )


def convert_linear_to_conv2d(module: torch.nn.Module):
    class Conv2D(torch.nn.Module):
        def __init__(self, weight, bias=None):
            super().__init__()
            use_bias = bias is not None
            self.conv = torch.nn.Conv2d(
                in_channels=weight.shape[0],
                out_channels=weight.shape[1],
                kernel_size=1,
                padding=0,
                bias=use_bias,
            )
            self.conv.weight = torch.nn.Parameter(weight.reshape(*weight.shape, 1, 1))
            if use_bias:
                self.conv.bias = torch.nn.Parameter(bias)

        def forward(self, x):
            rank = x.dim()
            x = x.unsqueeze(-1) if rank == 3 else x.reshape(1, *x.shape, 1)
            x = torch.transpose(x, 1, 2)
            res = self.conv(x)
            res = torch.transpose(res, 1, 2)
            res = res.squeeze(-1) if rank == 3 else res.reshape(*res.shape[1:3])
            return res

    def replace_linear(module: torch.nn.Module):
        attr_strs = dir(module)
        if isinstance(module, torch.nn.ModuleList):
            attr_strs += [str(i) for i in range(len(module))]

        for attr_str in attr_strs:
            target_attr = getattr(module, attr_str)
            if isinstance(target_attr, torch.nn.Linear):
                setattr(module, attr_str, Conv2D(target_attr.weight, target_attr.bias))

        for _, sub_module in module.named_children():
            sub_module = replace_linear(sub_module)
        return module

    return replace_linear(module)


def dump_context_from_pte(pte_path):
    """
    Dump compiled binaries under the same directory of pte_path.
    For partitioned graph, there will be multiple files with names f"{graph_name}_{index}".
    Where 'graph_name' comes from the compiler_specs and 'index' represents the execution order.

    Args:
        pte_path (str): The path of generated pte.
    """
    import os

    from executorch.exir._serialize._program import deserialize_pte_binary

    with open(pte_path, "rb") as f:
        program_data = f.read()

    program = deserialize_pte_binary(program_data)

    ctx_path = os.path.dirname(pte_path)
    dummy_compiler_specs = generate_qnn_executorch_compiler_spec(
        soc_model=QcomChipset.SM8650,
        backend_options=generate_htp_compiler_spec(use_fp16=False),
    )
    qnn_mgr = PyQnnManagerAdaptor.QnnManager(
        generate_qnn_executorch_option(dummy_compiler_specs)
    )
    qnn_mgr.Init()
    for execution_plan in program.execution_plan:
        for i, delegate in enumerate(execution_plan.delegates):
            if delegate.id == "QnnBackend":
                processed_bytes = program.backend_delegate_data[
                    delegate.processed.index
                ].data
                binary = qnn_mgr.StripProtocol(processed_bytes)
                with open(f"{ctx_path}/{execution_plan.name}_{i}.bin", "wb") as f:
                    f.write(binary)


def update_spill_fill_size(
    exported_program: ExportedProgram | List[LoweredBackendModule],
):
    # check if user specifies to use multi_contexts
    # this is a generic approach in case there exists multiple backends
    def get_program_info(program):
        def process_exported_program(prog):
            max_sf_buf_size, module_map = 0, {}
            for _, m in prog.graph_module._modules.items():
                # currently only 1 compile spec is expected in each partition
                options = flatbuffer_to_option(m.compile_specs[0].value)
                if (
                    options.backend_options.backend_type
                    == QnnExecuTorchBackendType.kHtpBackend
                    and options.backend_options.htp_options.use_multi_contexts
                ):
                    qnn_mgr = PyQnnManagerAdaptor.QnnManager(
                        m.compile_specs[0].value, m.processed_bytes
                    )
                    assert qnn_mgr.Init().value == 0, "failed to load context binary"
                    max_sf_buf_size = max(
                        max_sf_buf_size, qnn_mgr.GetSpillFillBufferSize()
                    )
                    module_map[m] = options
                    qnn_mgr.Destroy()
            return max_sf_buf_size, module_map

        def process_lowered_module(module):
            qnn_mgr = PyQnnManagerAdaptor.QnnManager(
                module.compile_specs[0].value, module.processed_bytes
            )
            assert qnn_mgr.Init().value == 0, "failed to load context binary"
            spill_fill_size = qnn_mgr.GetSpillFillBufferSize()
            qnn_mgr.Destroy()
            return spill_fill_size, {
                module: flatbuffer_to_option(module.compile_specs[0].value)
            }

        dispatch = {
            ExportedProgram: process_exported_program,
            LoweredBackendModule: process_lowered_module,
        }
        return dispatch[type(program)](program)

    def update_program(max_sf_buf_size, module_map):
        def set_spec(module, options):
            spec = CompileSpec(QCOM_QNN_COMPILE_SPEC, option_to_flatbuffer(options))
            if isinstance(module, ExportedProgram):
                module.compile_specs[0] = spec
            else:
                module._compile_specs[0] = spec

        for module, options in module_map.items():
            options.backend_options.htp_options.max_sf_buf_size = max_sf_buf_size
            set_spec(module, options)

    max_sf_size, modules_map = 0, {}
    if isinstance(exported_program, list):
        for prog in exported_program:
            max_sf_buf_size, module_map = get_program_info(prog)
            max_sf_size = max(max_sf_size, max_sf_buf_size)
            modules_map.update(module_map)
    else:
        max_sf_size, module_map = get_program_info(exported_program)
    update_program(max_sf_size, module_map)

    return max_sf_size


def canonicalize_program(obj):
    update_spill_fill_size(obj)


def get_decomp_table(passes_job) -> Dict[torch._ops.OperatorBase, Callable]:
    source_decompositions = core_aten_decompositions()
    # The below super ops are supported by QNN
    skip_decompositions = get_skip_decomp_table()

    # If we want to annotate the decomposed ops, then we should decompose the operation.
    if passes_job and passes_job.get(AnnotateStack, False):
        skip_decompositions = [
            skip_decomp_op
            for skip_decomp_op in skip_decompositions
            if skip_decomp_op not in AnnotateStack.decomp_ops
        ]
    remove_decompositions(source_decompositions, skip_decompositions)

    return source_decompositions


def to_edge_transform_and_lower_to_qnn(
    module: Union[torch.nn.Module, torch.fx.GraphModule],
    inputs: Tuple[torch.Tensor],
    compiler_specs: List[CompileSpec],
    constant_methods: Optional[Dict[str, Any]] = None,
    dynamic_shapes: Optional[Dict] = None,
    dep_table: Optional[Dict] = None,
    passes_job: Optional[OrderedDict] = None,
    skip_node_id_set: Optional[set] = None,
    skip_node_op_set: Optional[set] = None,
) -> EdgeProgramManager:
    """
    Transforms and lowers a given PyTorch module to QNN backend.

    Args:
        module (Union[torch.nn.Module, torch.fx.GraphModule]): The PyTorch module or fx.GraphModule to be transformed.
        inputs (Tuple[torch.Tensor]): The input tensors for the module.
        compiler_specs (List[CompileSpec]): Compiler specs for Qualcomm AI Engine Direct.
        constant_methods (Optional[Dict[str, Any]]): An optional dictionary of method name to the constant value
            returned by that method in eager mode. Often used to store config information on
            Edge models.
        dynamic_shapes (Optional[Dict]): Information about dynamic shapes.
        dep_table (Optional[Dict]): Dependency table for the transformation passes.
        passes_job (Optional[OrderedDict]): Ordered dictionary of transformation passes.
        skip_node_id_set (Optional[set]): Set of node IDs to skip during partitioning.
        skip_node_op_set (Optional[set]): Set of node operations to skip during partitioning.

    Returns:
        EdgeProgramManager: The manager for the edge program after transformation and lowering.
    """
    ep = torch.export.export(module, inputs, dynamic_shapes=dynamic_shapes, strict=True)
    # This transformation is primarily intended for the LiftConstantScalarOperands pass
    # to avoid creating temporary tensors in the operation builder.
    # However, this pass will create a get_attr node, which should be converted
    # into a lifted tensor constant by the lift_constant_tensor_pass.
    # If placed in the to_edge_transform_passes, it will be executed
    # after the lift_constant_tensor_pass, causing the operation builder
    # to fail to correctly retrieve the parameter by the get_parameter.
    ep = QnnPassManager().transform_for_export_pipeline(ep)
    transform_passes = QnnPassManager().get_to_edge_transform_passes(
        ep, passes_job=passes_job, dep_table=dep_table
    )
    qnn_partitioner = QnnPartitioner(
        compiler_specs,
        skip_node_id_set=skip_node_id_set,
        skip_node_op_set=skip_node_op_set,
    )
    edge_program_manager = to_edge_transform_and_lower(
        ep,
        transform_passes=transform_passes,
        partitioner=[qnn_partitioner],
        constant_methods=constant_methods,
        compile_config=qnn_edge_config(),
    )
    return edge_program_manager


def capture_program(
    module: Union[torch.nn.Module, torch.fx.GraphModule],
    inputs: Tuple[torch.Tensor],
    passes_job: OrderedDict = None,
    dynamic_shapes: Dict = None,
) -> exir.ExirExportedProgram:
    """
    TODO: Deprecated capture_program with to_edge_transform_and_lower_to_qnn

    Captures and transforms a PyTorch module into an Exir exported program.

    Args:
        module (Union[torch.nn.Module, torch.fx.GraphModule]): The PyTorch module or fx.GraphModule to be captured.
        inputs (Tuple[torch.Tensor]): The input tensors for the module.
        passes_job (OrderedDict, optional): Ordered dictionary of transformation passes.
        dynamic_shapes (Dict, optional): Information about dynamic shapes.

    Returns:
        exir.ExirExportedProgram: The transformed Exir exported program ready for lowering to QNN backend.
    """
    warnings.warn(
        "capture_program is deprecated. Use to_edge_transform_and_lower_to_qnn instead.",
        DeprecationWarning,
        stacklevel=1,
    )
    ep = torch.export.export(module, inputs, dynamic_shapes=dynamic_shapes, strict=True)
    ep = QnnPassManager().transform_for_export_pipeline(ep)
    # TODO: Handle stack op. If we want to run annotate_decomposed pass for stack op,
    # we need to make stack op decompose, which means we need to find a method to
    # remove it from skip_decomp table
    decomposed_ep = ep.run_decompositions(get_decomp_table(passes_job))
    core_ep = ExirExportedProgram(decomposed_ep, False)
    edge_ep = core_ep.to_edge(qnn_edge_config())
    transform_passes = QnnPassManager().get_to_edge_transform_passes(
        edge_ep.exported_program, passes_job=passes_job
    )
    edge_ep.transform(*transform_passes)
    return edge_ep


def _partition_graph_into_submodules(gm, subgm_tag, subgm_cb, ptn):
    from torch.fx.passes.utils.fuser_utils import (
        erase_nodes,
        fuse_as_graphmodule,
        insert_subgm,
        legalize_graph,
        topo_sort,
    )

    partitions = ptn.propose_partitions()
    # insert meta for each partition group
    for i, partition in enumerate(partitions):
        for node in partition.nodes:
            node.meta[subgm_tag] = i

    for i in range(len(partitions)):
        # find nodes with same group id in current graph
        node_list = [
            node for node in gm.graph.nodes if node.meta.get(subgm_tag, "") == i
        ]
        # fuse group nodes into submodule
        sorted_nodes = topo_sort(node_list)
        submodule_name = f"{subgm_tag}_{i}"
        subgm, orig_inputs, orig_outputs = fuse_as_graphmodule(
            gm, sorted_nodes, submodule_name
        )
        # insert submodule & trim group nodes
        gm = insert_subgm(
            gm,
            subgm_cb(subgm, submodule_name),
            orig_inputs,
            orig_outputs,
        )
        erase_nodes(gm, sorted_nodes)
        legalize_graph(gm)

    gm.recompile()
    return gm


def _canonicalize_graph_with_lowered_module(gm, subgm_tag, compiler_specs):
    # return lowered program for user to debug
    edge_prog_mgrs = []
    # partition each submodule which went through convert_pt2e
    for node in gm.graph.nodes:
        if node.op == "call_module" and subgm_tag in node.name:
            # obtain sample inputs through meta
            subgm_input = [
                torch.ones(arg.meta["val"].shape, dtype=arg.meta["val"].dtype)
                for arg in node.args
            ]
            # start lowering with given partitioner
            edge_prog_mgrs.append(
                to_edge_transform_and_lower_to_qnn(
                    gm.get_submodule(node.name), tuple(subgm_input), compiler_specs
                )
            )
            # replace submodule with lowered module
            gm.set_submodule(
                node.name,
                edge_prog_mgrs[-1].exported_program().graph_module,
            )
            # if node has multiple outputs, getitems will be default generated
            if all(n.target != operator.getitem for n in node.users):
                with gm.graph.inserting_after(node):
                    getitem_node = gm.graph.call_function(
                        operator.getitem,
                        (node, 0),
                    )
                    getitem_node.meta = node.meta
                    node.replace_all_uses_with(
                        replace_with=getitem_node,
                        delete_user_cb=lambda user: user.target != operator.getitem,
                    )

    gm.recompile()
    return gm, edge_prog_mgrs


def skip_annotation(
    nn_module: torch.nn.Module,
    quantizer,
    compiler_specs,
    sample_input: Tuple[torch.Tensor, ...],
    calibration_cb: Callable[[torch.fx.GraphModule], None],
    fp_node_id_set: set = None,
    fp_node_op_set: set = None,
    fallback_to_cpu: bool = True,
):
    r"""
    Exclude speific operators from quantizer annotation.
    Skipped operators will defaultly stay in CPU, set 'fallback_to_cpu'
    to False for trying to delegate them with FP16 precision.

    e.g.: consider following graph:
    bias_1 weight_1 input_1   bias_2 weight_2 input_2
      | (placeholder) |         | (placeholder) |
       \      |      /           \      |      /
        \     |     /             \     |     /
         \    |    /               \    |    /
           conv2d_1                 conv2d_2
           (torch.ops.aten.conv2d.default)
               \                       /
                \                     /
                 \_______     _______/
                         add_1
             (torch.ops.aten.add.default)
                           |
                         output

    If user wants to skip convolution op by names with
    'skip_node_id_set' = {"conv2d_1"}
    "bias_1 / weight_1 / input_1 / input_2 / conv2d_1"
    will be partitioned out and not annotated / lowered with QNN.

    [Generated graph]
    bias_1 weight_1 input_1   input_2
      | (placeholder) |          |
       \      |      /           |
        \     |     /            |
         \    |    /             |
           conv2d_1              |
              \                 /
               \               /
                \             /
               lowered_module_1
            (QNN fixed precision)
                      |
                    output

    If user wants to skip convolution op by target with
    'skip_node_op_set' = {torch.ops.aten.conv2d.default}
    "bias_1 / weight_1 / input_1 / conv2d_1,
     bias_2 / weight_2 / input_2 / conv2d_2"
    will be partitioned out and not annotated / lowered with QNN.

    [Generated graph]
    bias_1 weight_1 input_1   bias_2 weight_2 input_2
      | (placeholder) |         | (placeholder) |
       \      |      /           \      |      /
        \     |     /             \     |     /
         \    |    /               \    |    /
           conv2d_1                 conv2d_2
           (torch.ops.aten.conv2d.default)
               \                       /
                \                     /
                 \__               __/
                    lowered_module_1
                 (QNN fixed precision)
                           |
                         output

    If user wants to delegate the skipped conv2d from above graph
    with 'fallback_to_cpu' = False:

    [Generated graph]
       input_1         input_2
    (placeholder)   (placeholder)
          |               |
          \               /
          lowered_module_2
         (QNN fp16 precision)
                  |
                  |
          lowered_module_1
         (QNN fixed precision)
                  |
                output

    Args:
        nn_module (torch.nn.Module): The module to be lowered.
        quantizer (QnnQuantizer): Instance of QnnQuantizer.
        compiler_specs (List[CompileSpec]): Compiler specs for Qualcomm AI Engine Direct.
        sample_input ((torch.Tensor, ...)): Sample input tensors for graph exporting.
        calibration_cb (callable): Callback function for user-defined calibration.
        fp_node_id_set ({str, ...}): Set of operator names to be left in fp precision.
        fp_node_op_set ({torch.ops.aten.xxx, ...}): Set of operator targets to be left in fp precision.
        fallback_to_cpu (bool): Whether to lower skipped nodes to fp16 or not.

    Returns:
        exported_programs: List of programs lowered to QnnBackend (quantized graphs only).
    """
    from executorch.backends.qualcomm.serialization.qc_schema import (
        QnnExecuTorchHtpPrecision,
    )
    from executorch.backends.qualcomm.serialization.qc_schema_serialize import (
        flatbuffer_to_option,
    )
    from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e
    from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner

    def prepare_subgm(subgm, subgm_name):
        # prepare current submodule for quantization annotation
        subgm_prepared = prepare_pt2e(subgm, quantizer)
        # overwrite this attribute or name will be set to "GraphModule"
        # we could not identify each submodule if action is not performed
        subgm_prepared.__class__.__name__ = subgm_name
        return subgm_prepared

    fp_node_id_set = fp_node_id_set if fp_node_id_set is not None else set()
    fp_node_op_set = fp_node_op_set if fp_node_op_set is not None else set()
    graph_module = torch.export.export(nn_module, sample_input, strict=True).module()
    # define node support type
    capability_partitioner = CapabilityBasedPartitioner(
        graph_module,
        _AnnotationSkipper(fp_node_id_set, fp_node_op_set),
        allows_single_node_partition=True,
    )
    subgm_tag = "annotated_group"
    graph_module = _partition_graph_into_submodules(
        gm=graph_module,
        subgm_tag=subgm_tag,
        subgm_cb=prepare_subgm,
        ptn=capability_partitioner,
    )
    # perform calibration
    calibration_cb(graph_module)
    # convert sub modules which went through prepare_pt2e
    for node in graph_module.graph.nodes:
        if node.op == "call_module":
            graph_module.set_submodule(
                node.name, convert_pt2e(graph_module.get_submodule(node.name))
            )
    # canonicalize graph for lowering again
    graph_module, edge_prog_mgrs = _canonicalize_graph_with_lowered_module(
        gm=graph_module,
        subgm_tag=subgm_tag,
        compiler_specs=compiler_specs,
    )

    if not fallback_to_cpu:
        try:
            # change HTP compiler spec for hardware to enable fp16
            qnn_option = generate_qnn_executorch_option(compiler_specs)
            compile_option = flatbuffer_to_option(qnn_option)
            htp_options = compile_option.backend_options.htp_options
            htp_options.precision = QnnExecuTorchHtpPrecision.kHtpFp16
            compiler_specs[0].value = option_to_flatbuffer(compile_option)
        except:
            print(
                "Failed to change HTP compiler spec with 'use_fp16' as True,"
                " skipped operators will fallback to cpu,"
            )
            return graph_module, edge_prog_mgrs

        # try lowering skipped operator into fp16
        capability_partitioner = CapabilityBasedPartitioner(
            graph_module,
            _AnnotationSkipper(skip_annotated_submodule=True),
            allows_single_node_partition=True,
        )
        subgm_tag = "skipped_group"
        graph_module = _partition_graph_into_submodules(
            gm=graph_module,
            subgm_tag=subgm_tag,
            subgm_cb=lambda subgm, _: subgm,
            ptn=capability_partitioner,
        )
        graph_module, edge_prog_mgrs_fp = _canonicalize_graph_with_lowered_module(
            gm=graph_module,
            subgm_tag=subgm_tag,
            compiler_specs=compiler_specs,
        )
        edge_prog_mgrs.extend(edge_prog_mgrs_fp)

    return graph_module, edge_prog_mgrs


def from_context_binary(  # noqa: C901
    ctx_path: str | bytes,
    op_name: str,
    soc_model: QcomChipset = QcomChipset.SM8650,
    custom_info: Dict = None,
):
    from pathlib import Path

    def implement_op(custom_op, op_name, outputs):
        @torch.library.impl(
            custom_op, str(op_name), dispatch_key="CompositeExplicitAutograd"
        )
        def op_impl(inputs: List[torch.Tensor]):
            return tuple(
                torch.zeros(tuple(v.shape), device="meta", dtype=v.dtype)
                for v in outputs.values()
            )

    def build_graph(
        inputs,
        outputs,
        qnn_in_order: Optional[List[int]] = None,
        executorch_in_order: Optional[List[int]] = None,
        executorch_out_order: Optional[List[int]] = None,
    ):
        # custom op declaration
        inputs_str = "Tensor[] inputs"
        func_proto = f"{op_name}({inputs_str}) -> Any"
        custom_op = Library(OpContextLoader.namespace, "FRAGMENT")
        custom_op.define(func_proto)
        # custom op implementation
        implement_op(custom_op, op_name, outputs)

        # model architecture mimicking context binary
        class Model(torch.nn.Module):
            """
            The args of forward() can be thought of as what executorch is accepting as input.
            The getattr inside the forward() can be thought of as qnn context binary.
            When we first pass in the input, we need to use the executorch's(nn.module) input order.
            After we get into forward(), we then need to convert input order to qnn's input order.
            Same as return, when qnn returns the value, we need to reorder them back to executorh's output order.
            """

            def __init__(self, qnn_in_order, executorch_out_order):
                super().__init__()
                self.qnn_in_order = qnn_in_order
                self.executorch_out_order = executorch_out_order

            def forward(self, *inputs):  # executorch
                if self.qnn_in_order:
                    inputs = tuple(inputs[i] for i in self.qnn_in_order)
                ret = getattr(
                    getattr(torch.ops, OpContextLoader.namespace), op_name
                ).default(inputs)
                return (
                    [ret[idx] for idx in self.executorch_out_order]
                    if self.executorch_out_order
                    else ret
                )

        inputs = (
            tuple(tuple(inputs.values())[i] for i in executorch_in_order)
            if executorch_in_order
            else tuple(inputs.values())
        )

        model = Model(qnn_in_order, executorch_out_order)
        prog = torch.export.export(model, inputs, strict=True)
        # bookkeeping for variables' life cycle
        return {
            "custom_op": custom_op,
            "custom_module": model,
            "exported_program": prog,
        }

    def build_tensor(tensors, dtype_map):
        ret = OrderedDict()
        for t in tensors:
            dtype = t.GetDataType()
            dtype_torch = dtype_map.get(dtype, None)
            assert dtype_torch is not None, f"unknown qnn data type {dtype}"
            ret[t.GetName()] = torch.zeros(tuple(t.GetDims()), dtype=dtype_torch)

        return ret

    def preprocess_binary(ctx_bin, compiler_specs):
        qnn_mgr = PyQnnManagerAdaptor.QnnManager(
            generate_qnn_executorch_option(compiler_specs),
        )
        return bytes(qnn_mgr.MakeBinaryInfo(ctx_bin))

    # dummy compiler spec would be fine, since we're not compiling
    backend_options = generate_htp_compiler_spec(use_fp16=False)
    compiler_specs = generate_qnn_executorch_compiler_spec(
        soc_model=soc_model,
        backend_options=backend_options,
        is_from_context_binary=True,
    )

    ctx_bin = (
        ctx_path
        if not isinstance(ctx_path, str)
        else preprocess_binary(Path(f"{ctx_path}").read_bytes(), compiler_specs)
    )

    dtype_map = {}
    for type_map in (QNN_QUANT_TYPE_MAP, QNN_TENSOR_TYPE_MAP):
        for k, v in type_map.items():
            dtype_map.setdefault(v, k)

    qnn_in_order, executorch_in_order, executorch_out_order = None, None, None
    if custom_info is not None:
        # since some context binaries might fail to open on host
        # if they are compiled with special flags:
        # e.g. weight sharing
        # use custom information here instead
        inputs = build_tensor(custom_info["graph_inputs"], dtype_map)
        outputs = build_tensor(custom_info["graph_outputs"], dtype_map)
        qnn_in_order = custom_info.get("qnn_in_order", None)
        executorch_in_order = custom_info.get("executorch_in_order", None)
        executorch_out_order = custom_info.get("executorch_out_order", None)
        graph_name = custom_info["graph_name"]
    else:
        # get context-binary io tensor info through qnn manager
        qnn_mgr = PyQnnManagerAdaptor.QnnManager(
            generate_qnn_executorch_option(compiler_specs),
            ctx_bin,
        )
        assert qnn_mgr.Init().value == 0, "failed to load context binary"
        # assume we only have one graph in current context
        graph_name = qnn_mgr.GetGraphNames()[0]
        qnn_mgr.AllocateTensor(graph_name)
        inputs = build_tensor(qnn_mgr.GetGraphInputs(graph_name), dtype_map)
        outputs = build_tensor(qnn_mgr.GetGraphOutputs(graph_name), dtype_map)
        qnn_mgr.Destroy()
    # generate graph specific for loading context
    bundle_prog = build_graph(
        inputs, outputs, qnn_in_order, executorch_in_order, executorch_out_order
    )
    bundle_prog.update({"inputs": inputs, "outputs": outputs})

    # TODO: to_edge() decorator alters the function call behavior, which
    # requires "self" when calling. To work around this issue,
    # temporarily remove the first parameter name.
    edge_prog_mgr = to_edge(
        {graph_name: bundle_prog["exported_program"]},
        # do not alter name for custom op
        compile_config=EdgeCompileConfig(_use_edge_ops=False),
    )

    # update meta with context binary
    for n in edge_prog_mgr._edge_programs[graph_name].graph.nodes:
        if n.op == "call_function" and OpContextLoader.namespace in str(n.target):
            n.meta[OpContextLoader.meta_ctx_bin] = ctx_bin
            break

    bundle_prog["edge_program_manager"] = edge_prog_mgr.to_backend(
        QnnPartitioner(compiler_specs)
    )
    return bundle_prog


def draw_graph(title, path, graph_module: torch.fx.GraphModule):
    graph = passes.graph_drawer.FxGraphDrawer(graph_module, title)
    with open(f"{path}/{title}.svg", "wb") as f:
        f.write(graph.get_dot_graph().create_svg())


def generate_multi_graph_program(
    compiler_specs: List[CompileSpec],
    processed_bytes: List[bytes],
    input_nodes_dict: List[torch.fx.Node] = None,
    output_nodes_dict: List[torch.fx.Node] = None,
    backend_config: ExecutorchBackendConfig = None,
    constant_methods: Optional[Dict[str, Any]] = None,
) -> ExecutorchProgramManager:
    # compile multiple graphs in qcir into single context binary
    (
        graph_inputs,
        graph_outputs,
        qnn_in_order,
        executorch_in_order,
        executorch_out_order,
    ) = ({}, {}, {}, {}, {})
    qnn_mgr = PyQnnManagerAdaptor.QnnManager(
        generate_qnn_executorch_option(compiler_specs), processed_bytes
    )
    assert qnn_mgr.Init().value == 0, "failed to load processed bytes"
    binary_info = bytes(qnn_mgr.Compile())
    assert len(binary_info) != 0, "failed to generate QNN context binary"
    graph_names = qnn_mgr.GetGraphNames()
    for graph_name in graph_names:
        graph_inputs[graph_name] = qnn_mgr.GetGraphInputs(graph_name)
        graph_outputs[graph_name] = qnn_mgr.GetGraphOutputs(graph_name)

    # We need to obtain the order of the IOs to correctly map QNN with nn.module
    for graph_name in graph_names:
        if input_nodes_dict:
            # input
            input_names = [node.name for node in input_nodes_dict[graph_name]]
            qnn_input_names = [
                wrapper.GetName() for wrapper in graph_inputs[graph_name]
            ]
            # The input of intermideate module including call_function node
            # could not be reorder by node name
            if len(input_names) == len(qnn_input_names):
                input_order_list = []
                for input_name in input_names:
                    # e.g., input_0_tokens_0
                    pattern = rf"^input_(\d+)_({input_name})_(\d+)$"
                    for j in range(len(qnn_input_names)):
                        if re.match(pattern, qnn_input_names[j]):
                            input_order_list.append(j)
                            break
                assert len(input_order_list) == len(
                    input_names
                ), "Order list length is different from names"
                executorch_in_order[graph_name] = input_order_list
                qnn_in_order[graph_name] = sorted(
                    range(len(input_order_list)), key=lambda k: input_order_list[k]
                )
        if output_nodes_dict:
            # output
            get_item_list = output_nodes_dict[graph_name][0].args[0]
            output_order_list = [item.args[1] for item in get_item_list]
            executorch_out_order[graph_name] = output_order_list

    qnn_mgr.Destroy()

    # build custom ops with different graph signatures
    compiler_options = flatbuffer_to_option(compiler_specs[0].value)
    bundle_progs = [
        from_context_binary(
            ctx_path=binary_info,
            op_name=f"loader_{graph_name}_{int(time.time())}",
            soc_model=compiler_options.soc_info.soc_model,
            custom_info={
                "graph_inputs": graph_inputs[graph_name],
                "graph_outputs": graph_outputs[graph_name],
                "graph_name": graph_name,
                "qnn_in_order": qnn_in_order.get(graph_name, None),
                "executorch_in_order": executorch_in_order.get(graph_name, None),
                "executorch_out_order": executorch_out_order.get(graph_name, None),
            },
        )
        for graph_name in graph_names
    ]
    # leverage ExecutorchProgramManager for generating pte with multi-methods
    edge_prog_mgr = to_edge(
        {
            graph_name: bundle_prog["exported_program"]
            for graph_name, bundle_prog in zip(graph_names, bundle_progs)
        },
        constant_methods=constant_methods,
        # do not alter name for custom op
        compile_config=EdgeCompileConfig(_use_edge_ops=False),
    )
    # restore meta losed in generating EdgeProgramManager
    for graph_name in graph_names:
        for n in edge_prog_mgr._edge_programs[graph_name].graph.nodes:
            if graph_name in n.name:
                n.meta[OpContextLoader.meta_ctx_bin] = binary_info
                break

    edge_prog_mgr = edge_prog_mgr.to_backend(QnnPartitioner(compiler_specs))
    exec_prog = edge_prog_mgr.to_executorch(
        config=backend_config or ExecutorchBackendConfig()
    )
    return exec_prog, bundle_progs


def generate_composite_llama_program(
    llama_model: torch.nn.Module,
    graph_names: List[str],
    sample_inputs_list: List[Tuple[Any]],
    lower_module_dict: Dict[str, List[LoweredBackendModule]],
    call_delegate_node_name_dict: Dict[str, List[str]],
    call_delegate_inputs_dict: Dict[str, List[Tuple[str, int | None]]],
    outputs_dict: Dict[str, List[Tuple[str, int]]],
    embedding_quantize: str,
    backend_config: ExecutorchBackendConfig = None,
    constant_methods: Optional[Dict[str, Any]] = None,
) -> ExecutorchProgramManager:
    class CompositeLlamaModule(torch.nn.Module):
        def __init__(
            self,
            llama_model,
            lower_module_list,
            call_delegate_node_name_list,
            call_delegate_inputs_list,
            outputs_list,
            embedding_quantize,
        ) -> None:
            super().__init__()
            self.llama_model = llama_model
            self.lower_module_list = lower_module_list
            self.call_delegate_node_name_list = call_delegate_node_name_list
            self.call_delegate_inputs_list = call_delegate_inputs_list
            self.outputs_list = outputs_list
            self.embedding_quantize = embedding_quantize

        def reorder(
            self,
            call_delegate_inputs: List[Tuple[str, int | None]],
            module_inputs: dict[str, torch.Tensor],
            all_ret: dict[str, torch.Tensor],
        ) -> Tuple[torch.Tensor]:
            ret = []
            for name, index in call_delegate_inputs:
                if index is not None:
                    # Get tensor from previous results
                    ret.append(all_ret[name][index])
                else:
                    # Get tensor from the inputs of module
                    ret.append(module_inputs[name])
            return tuple(ret)

        def forward(
            self,
            tokens: torch.Tensor,
            atten_mask: torch.Tensor,
            input_pos: Optional[torch.Tensor] = None,
            *args,
        ) -> Tuple[torch.Tensor]:
            all_ret = {}
            module_input_dict = {
                "tokens": tokens,
                "atten_mask": atten_mask,
                "input_pos": input_pos,
            }
            for num, arg in enumerate(args):
                module_input_dict[f"args_{num}"] = arg

            if self.embedding_quantize:
                hidden_states = self.llama_model.tok_embeddings(tokens)
                module_input_dict["quantized_decomposed_embedding_4bit_dtype"] = (
                    hidden_states
                )

            for lower_module, call_delegate_node_name, call_delegate_inputs in zip(
                self.lower_module_list,
                self.call_delegate_node_name_list,
                self.call_delegate_inputs_list,
            ):
                inp = self.reorder(call_delegate_inputs, module_input_dict, all_ret)
                ret = lower_module(*inp)
                all_ret[call_delegate_node_name] = ret
            llama_outputs = []
            for output_src_name, index in self.outputs_list:
                llama_outputs.append(all_ret[output_src_name][index])
            return tuple(llama_outputs)

    progs_dict = {}
    for graph_name, sample_inputs in zip(graph_names, sample_inputs_list):
        composite_llama_module = CompositeLlamaModule(
            llama_model,
            lower_module_dict[graph_name],
            call_delegate_node_name_dict[graph_name],
            call_delegate_inputs_dict[graph_name],
            outputs_dict[graph_name],
            embedding_quantize,
        )
        prog = torch.export.export(composite_llama_module, sample_inputs, strict=True)
        progs_dict[graph_name] = prog
    # leverage ExecutorchProgramManager for generating pte with multi-methods
    edge_prog_mgr = to_edge(
        progs_dict,
        constant_methods=constant_methods,
        # do not alter name for custom op
        compile_config=EdgeCompileConfig(_check_ir_validity=False, _use_edge_ops=False),
    )
    exec_prog = edge_prog_mgr.to_executorch(
        config=backend_config or ExecutorchBackendConfig()
    )
    return exec_prog


def generate_htp_compiler_spec(
    use_fp16: bool,
    use_dlbc: bool = False,
    use_multi_contexts: bool = False,
) -> QnnExecuTorchBackendOptions:
    """
    Helper function generating backend options for QNN HTP

    Args:
        use_fp16: If true, the model is compiled to QNN HTP fp16 runtime.
            Note that not all SoC support QNN HTP fp16. Only premium tier SoC
            like Snapdragon 8 Gen 1 or newer can support HTP fp16.
        use_dlbc: Deep Learning Bandwidth Compression allows inputs to be
            compressed, such that the processing bandwidth can be lowered.
        use_multi_contexts: When multiple contexts are generated inside the same
            pte, it is possible to reserve a single spill-fill allocation that
            could be re-used across all the splits.

    Returns:
        QnnExecuTorchHtpBackendOptions: backend options for QNN HTP.
    """
    htp_options = QnnExecuTorchHtpBackendOptions()
    htp_options.precision = (
        QnnExecuTorchHtpPrecision.kHtpFp16
        if use_fp16
        else QnnExecuTorchHtpPrecision.kHtpQuantized
    )
    # This actually is not an option which can affect the compiled blob.
    # But we don't have other place to pass this option at execution stage.
    # TODO: enable voting mechanism in runtime and make this as an option
    htp_options.performance_mode = QnnExecuTorchHtpPerformanceMode.kHtpBurst
    htp_options.use_multi_contexts = use_multi_contexts
    htp_options.use_dlbc = use_dlbc
    return QnnExecuTorchBackendOptions(
        backend_type=QnnExecuTorchBackendType.kHtpBackend,
        htp_options=htp_options,
    )


def generate_qnn_executorch_compiler_spec(
    soc_model: QcomChipset,
    backend_options: QnnExecuTorchBackendOptions,
    debug: bool = False,
    saver: bool = False,
    online_prepare: bool = False,
    dump_intermediate_outputs: bool = False,
    profile: bool = False,
    optrace: bool = False,
    shared_buffer: bool = False,
    is_from_context_binary: bool = False,
    multiple_graphs: bool = False,
    weight_sharing: bool = False,
    graph_name: str = "forward",
) -> List[CompileSpec]:
    """
    Helper function generating compiler specs for Qualcomm AI Engine Direct

    Args:
        soc_model: The SoC you plan to run the compiled model. Please check
            QcomChipset for supported SoC.
            SM8450 (Snapdragon 8 Gen 1)
            SM8475(Snapdragon 8 Gen 1+)
            SM8550(Snapdragon 8 Gen 2)
            SM8650(Snapdragon 8 Gen 3)
            SM8750(Snapdragon 8 Elite)
        backend_options: Options required by different backends.
        debug: Enable verbose logging. Disclaimer: this option must change in
            the near future.
        online_prepare: Compose QNN graph on device if set to True
        saver: Instead of compiling the model, run QNN Saver. Please check
            documents of Qualcomm AI Engine Direct SDK. This feature is usually
            for debugging purpose.
        dump_intermediate_outputs: If tensor dump is enabled, all intermediate tensors output will be dumped.
            This option exists for debugging accuracy issues
        profile: Enable profile the performance of per operator.
            Note that for now only support kProfileDetailed to
            profile the performance of each operator with cycle unit.
        shared_buffer: Enables usage of shared buffer between application
            and backend for graph I/O.
        is_from_context_binary: True if current graph comes from pre-built context binary.
        multiple_graphs: True if multiple methods are expected to have in single .pte file.
            Please see test cases for post-processing example.
        weight_sharing: Used with multiple_graphs, where model size will be reduced when operations have the same weights across multiple graphs.
        graph_name: Assign unique graph name if 'multiple_graphs' is used.

    Returns:
        List[CompileSpec]: Compiler specs for Qualcomm AI Engine Direct.

    Raises:
        ValueError: The value QcomChipset is currently not supported.
        ValueError: Confliction between compiler specs.
    """
    _supported_soc_models = {soc_model.value for soc_model in QcomChipset}
    if soc_model not in _supported_soc_models:
        raise ValueError(f"unknown SoC model for QNN: {soc_model}")

    if profile and dump_intermediate_outputs:
        warnings.warn(
            "It is not recommended to turn on both profiling and dump_intermediate_outputs the same time"
            ", because dump_intermediate_outputs will cause performance drop.",
            stacklevel=1,
        )

    if weight_sharing and not multiple_graphs:
        warnings.warn(
            "Weight sharing is intended for multiple graphs scenario, please ensure if there are multiple graphs",
            stacklevel=1,
        )

    qnn_executorch_options = QnnExecuTorchOptions(
        _soc_info_table[soc_model], backend_options
    )
    qnn_executorch_options.graph_name = graph_name
    qnn_executorch_options.log_level = (
        QnnExecuTorchLogLevel.kLogLevelDebug
        if debug
        else QnnExecuTorchLogLevel.kLogLevelWarn
    )

    qnn_executorch_options.dump_intermediate_outputs = dump_intermediate_outputs

    if saver:
        qnn_executorch_options.library_path = "libQnnSaver.so"

    if optrace:
        qnn_executorch_options.profile_level = QnnExecuTorchProfileLevel.kProfileOptrace
    elif profile:
        qnn_executorch_options.profile_level = (
            QnnExecuTorchProfileLevel.kProfileDetailed
        )
    else:
        qnn_executorch_options.profile_level = QnnExecuTorchProfileLevel.kProfileOff

    if (
        online_prepare
        and backend_options.backend_type == QnnExecuTorchBackendType.kHtpBackend
        and backend_options.htp_options.use_multi_contexts
    ):
        raise ValueError(
            "'use_multi_context' could not function in online prepare mode, "
            "please set 'online_prepare' to False"
        )

    qnn_executorch_options.shared_buffer = shared_buffer
    qnn_executorch_options.online_prepare = online_prepare
    qnn_executorch_options.is_from_context_binary = is_from_context_binary
    qnn_executorch_options.multiple_graphs = multiple_graphs

    if multiple_graphs:
        # enable weight sharing mechanism if multiple graphs appear
        if (
            backend_options.backend_type == QnnExecuTorchBackendType.kHtpBackend
            and weight_sharing
        ):
            backend_options.htp_options.use_weight_sharing = True

    return [
        CompileSpec(QCOM_QNN_COMPILE_SPEC, option_to_flatbuffer(qnn_executorch_options))
    ]


def get_soc_to_arch_map():
    return {
        "SA8295": HtpArch.V68,
        "SM8450": HtpArch.V69,
        "SM8475": HtpArch.V69,
        "SM8550": HtpArch.V73,
        "SM8650": HtpArch.V75,
        "SM8750": HtpArch.V79,
        "SSG2115P": HtpArch.V73,
        "SSG2125P": HtpArch.V73,
        "SXR1230P": HtpArch.V73,
        "SXR2230P": HtpArch.V69,
        "SXR2330P": HtpArch.V79,
    }


def get_soc_to_chipset_map():
    return {
        "SA8295": QcomChipset.SA8295,
        "SM8450": QcomChipset.SM8450,
        "SM8475": QcomChipset.SM8475,
        "SM8550": QcomChipset.SM8550,
        "SM8650": QcomChipset.SM8650,
        "SM8750": QcomChipset.SM8750,
        "SSG2115P": QcomChipset.SSG2115P,
        "SSG2125P": QcomChipset.SSG2125P,
        "SXR1230P": QcomChipset.SXR1230P,
        "SXR2230P": QcomChipset.SXR2230P,
        "SXR2330P": QcomChipset.SXR2330P,
    }


def tag_quant_io(gm: torch.fx.GraphModule, get_quant_io_dtype_fn: Callable):
    """
    Tag io nodes which get/output quantized tensor. No need to insert q/dq in qnn_preprocess
    """
    for node in gm.graph.nodes:
        if dtype := get_quant_io_dtype_fn(node):
            node.meta[QCOM_QUANTIZED_IO] = dtype