# 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, AnnotateUnbind 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) file_extension = ".bin" if len(binary) == 0: binary = processed_bytes file_extension = ".dlc" with open( f"{ctx_path}/{execution_plan.name}_{i}{file_extension}", "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: skip_decompositions = [ skip_decomp_op for skip_decomp_op in skip_decompositions if skip_decomp_op not in AnnotateStack.decomp_ops + AnnotateUnbind.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" qnn_executorch_options.saver = True qnn_executorch_options.saver_output_dir = "saver_output" 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