Qualcomm AI Engine Direct - Support LPAI Backend

Summary:

- Add LPAI infra
  - Because of the accuracy isue with quantize and dequantize
    operations in LPAI, Q and DQ will fallback to CPU for the LPAI
backend.
  - Support runtime option for `fps`, `ftrt_ratio`, `clent_perf_type`,
    `affinity`, and `core_selection`. Refer to [QNN
doc](https://docs.qualcomm.com/doc/80-63442-10/topic/lpai_backend.html#full-json-scheme)
to get more details for these options.
- Support LPAI in `test_qnn_delegate.py`
- Refactor documentation for LPAI and GPU
- Add a script to sign the skel library

Author: shewu-quic

backends/qualcomm/README.md

@@ -16,24 +16,7 @@ Please check `generate_qnn_executorch_compiler_spec()` in
 [utils.py](utils/utils.py) for supported SoC and inference type.
 
 ### Supported Chipset
-- Snapdragon 8 Gen 1
-- Snapdragon 8 Gen 1+
-- Snapdragon 8 Gen 2
-- Snapdragon 8 Gen 3
-- Snapdragon 8 Elite
-- Snapdragon 8 Elite Gen 5
-- SA8295
-- SA8255
-- SA8797 (also used by SA8397)
-- SSG2115P
-- SSG2125P
-- SXR1230P
-- SXR2230P
-- SXR2330P
-- QCS9100
-- SAR2230P
-- SW6100
-- SM8845
+The supported Qualcomm SoCs are detailed in `QcomChipset`. Please navigate to [qc_schema.py](serialization/qc_schema.py).
 
 ### Adding more supported Chipset
 Currently, users cannot add additional chipset models because the chipset ID is not accessible to community users. If you have specific chipset models you wish to add, please contact one of the authors in the `Code Reviews` section at the bottom of this page.
@@ -42,6 +25,15 @@ Currently, users cannot add additional chipset models because the chipset ID is
 - Quantized
 - FP16
 
+### Supported Backend Type
+- QNN GPU
+- QNN HTP
+- QNN LPAI
+  - Currently, only LPAI Arch v6 is supported, which requires QNN SDK version 2.39 or higher.
+  - Please check `_soc_info_table` in [qc_schema.py](serialization/qc_schema.py) to find the supported chipsets.
+  - The skel library needs to be signed for LPAI to work. Please confirm if you are able to sign it. You can use [signed_library.sh](scripts/signed_library.sh) to sign the skel library.
+  
+
 ## Directory Structure
 
 ```

backends/qualcomm/_passes/fold_qdq.py

@@ -5,22 +5,42 @@
 # LICENSE file in the root directory of this source tree.
 import torch
 from executorch.backends.qualcomm.builders.node_visitor import dq_ops, q_ops
-from executorch.backends.qualcomm.builders.utils import is_parameter
-from executorch.backends.qualcomm.utils.constants import QCOM_BYPASS_NODE
+from executorch.backends.qualcomm.builders.utils import (
+    is_graph_input,
+    is_graph_output,
+    is_parameter,
+)
+from executorch.backends.qualcomm.serialization.qc_schema import (
+    QnnExecuTorchBackendType,
+)
+from executorch.backends.qualcomm.utils.constants import (
+    QCOM_BYPASS_NODE,
+    QCOM_FALLBACK_NODE,
+    QCOM_QUANT_ATTRS,
+    QCOM_QUANTIZED_IO,
+)
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
 from executorch.exir.passes import dead_code_elimination_pass
 
+from .utils import get_quant_attrs
+
 
 class FoldQDQ(ExportPass):
     """
     Erase QDQ pattern.
     """
 
-    def __init__(self, edge_program: torch.export.ExportedProgram, force_fold=False):
+    def __init__(
+        self,
+        edge_program: torch.export.ExportedProgram,
+        force_fold=False,
+        backend_type: QnnExecuTorchBackendType = QnnExecuTorchBackendType.kHtpBackend,
+    ):
         super(FoldQDQ, self).__init__()
         self.edge_program = edge_program
         self.force_fold = force_fold
+        self.backend_type = backend_type
 
     def _annotate_bypass(self, node):
         node.meta[QCOM_BYPASS_NODE] = True
@@ -35,20 +55,26 @@ def _fold_dq(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
             if n.target not in dq_ops:
                 continue
 
-            # skip parameters & buffers
-            if not self.force_fold and is_parameter(n.args[0], self.edge_program):
-                self._annotate_bypass(n)
-            else:
-                for user_n in user_list:
-                    user_n.replace_input_with(n, n.args[0])
-                graph_module.graph.erase_node(n)
+            if not self.force_fold and (
+                QCOM_BYPASS_NODE in n.meta or QCOM_FALLBACK_NODE in n.meta
+            ):
+                continue
+
+            for user_n in user_list:
+                user_n.replace_input_with(n, n.args[0])
+            graph_module.graph.erase_node(n)
 
     def _fold_q(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
         # remove q
         for n in graph_module.graph.nodes:
             if n.target not in q_ops:
                 continue
 
+            if not self.force_fold and (
+                QCOM_BYPASS_NODE in n.meta or QCOM_FALLBACK_NODE in n.meta
+            ):
+                continue
+
             to_be_removed = [n]
             source_n = n.args[0]
 
@@ -72,7 +98,54 @@ def _fold_q(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
             for n in to_be_removed:
                 graph_module.graph.erase_node(n)
 
+    def _preserve_qdq(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
+        for n in graph_module.graph.nodes:
+            # skip parameters & buffers
+            if n.target in dq_ops and is_parameter(n.args[0], self.edge_program):
+                self._annotate_bypass(n)
+                continue
+
+            # TODO: In LPAI backend v6, there is an accuracy drop for the quantize and dequantize operations.
+            # To address this, keep the quantize/dequantize operations at the model's input and output.
+            # For example, input -> q_1 (Fallback) -> dq_1 (Bypass) -> graph -> q_2 (Bypass) -> dq_2 (Fallback) -> output
+            # Here, q_1 and dq_2 will fallback to CPU, while q_2 and dq_1 will be bypassed in qnn_partition and folded in qnn_preprocess.
+            if self.backend_type == QnnExecuTorchBackendType.kLpaiBackend:
+                if (
+                    is_graph_input(n, self.edge_program)
+                    # For tagged quantized I/O, we should not fallback quantize operation.
+                    and QCOM_QUANTIZED_IO not in n.meta
+                ):
+                    user_list = list(n.users.keys())
+                    if len(user_list) > 0:
+                        q_node = user_list[0]
+                        q_node.meta[QCOM_FALLBACK_NODE] = True
+                        # Annotate the q_node since it will serve as the input for the first node during operator validation
+                        q_node.meta[QCOM_QUANT_ATTRS] = get_quant_attrs(
+                            self.edge_program, q_node
+                        )
+                        q_node.meta[QCOM_QUANTIZED_IO] = q_node.args[-1]
+                        dq_node = list(q_node.users.keys())[0]
+                        # Bypass dequantize op for graph validation by torch
+                        dq_node.meta[QCOM_BYPASS_NODE] = True
+                        # Make sure that the quantize operator isn't inserted for input in insert_io_qdq.py
+                        n.meta[QCOM_QUANTIZED_IO] = q_node.args[-1]
+                elif (
+                    is_graph_output(n)
+                    and n.target in dq_ops
+                    # For tagged quantized I/O, we should not fallback dequantize operation.
+                    and QCOM_QUANTIZED_IO not in n.args[0].args[0].meta
+                ):
+                    n.meta[QCOM_FALLBACK_NODE] = True
+                    q_node = n.args[0]
+                    # Bypass quantize op for graph validation by torch
+                    q_node.meta[QCOM_BYPASS_NODE] = True
+                    op_node = q_node.args[0]
+                    # Make sure that the dequantize operator isn't inserted for output in insert_io_qdq.py
+                    op_node.meta[QCOM_QUANTIZED_IO] = q_node.args[-1]
+
     def call(self, graph_module: torch.fx.GraphModule):
+        if not self.force_fold:
+            self._preserve_qdq(graph_module)
         self._fold_dq(graph_module)
         self._fold_q(graph_module)
         graph_module.recompile()

backends/qualcomm/_passes/qnn_pass_manager.py

@@ -59,6 +59,9 @@
 from executorch.backends.qualcomm._passes.utils import (
     get_passes_dependency_for_capture_program,
 )
+from executorch.backends.qualcomm.serialization.qc_schema import (
+    QnnExecuTorchBackendType,
+)
 from executorch.backends.qualcomm.utils.constants import (
     QCOM_PASS_ACTIVATE_KEY,
     QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY,
@@ -144,6 +147,7 @@ def get_to_edge_transform_passes(
         exported_program: ExportedProgram,
         passes_job: OrderedDict = None,
         dep_table: Dict = None,
+        backend_type: QnnExecuTorchBackendType = QnnExecuTorchBackendType.kHtpBackend,
     ):
         # TODO: remove this workaround when target could be correctly detected
         from executorch.backends.qualcomm.builders import node_visitor
@@ -176,6 +180,8 @@ def get_to_edge_transform_passes(
             kwargs = passes_job[p][QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY]
             if "edge_program" in kwargs:
                 kwargs["edge_program"] = exported_program
+            if "backend_type" in kwargs:
+                kwargs["backend_type"] = backend_type
             self.add_pass(p(**kwargs))
         assert isinstance(
             self.passes[-1], ResolveDebugHandle

backends/qualcomm/_passes/utils.py

@@ -64,7 +64,6 @@ def get_passes_dependency_for_capture_program():
         AnnotateQuantAttrs,
         AnnotateStack,
         AnnotateUnbind,
-        CanonicalizeConv,
         ConvertBmmToMatmul,
         DecomposeAny,
         DecomposeColIm,
@@ -103,7 +102,6 @@ def get_passes_dependency_for_capture_program():
         I64toI32: [RemoveRedundancy],
         LayoutTransform: [
             AnnotateQuantAttrs,
-            CanonicalizeConv,
             ExpandBroadcastTensorShape,
             FixedLinearKeepDim,
         ],

backends/qualcomm/partition/qnn_partitioner.py

@@ -12,10 +12,16 @@
 from executorch.backends.qualcomm.builders import node_visitor_manager
 from executorch.backends.qualcomm.builders.qnn_constants import OpContextLoader
 from executorch.backends.qualcomm.qnn_preprocess import QnnBackend
+from executorch.backends.qualcomm.serialization.qc_schema import (
+    QnnExecuTorchBackendType,
+)
 from executorch.backends.qualcomm.serialization.qc_schema_serialize import (
     flatbuffer_to_option,
 )
-from executorch.backends.qualcomm.utils.constants import QCOM_BYPASS_NODE
+from executorch.backends.qualcomm.utils.constants import (
+    QCOM_BYPASS_NODE,
+    QCOM_FALLBACK_NODE,
+)
 
 from executorch.backends.qualcomm.utils.qnn_manager_lifecycle import (
     get_current_qnn_manager,
@@ -78,6 +84,9 @@ def is_node_supported(self, _, node: torch.fx.Node) -> bool:
             )
             return False
 
+        if node.meta.get(QCOM_FALLBACK_NODE, False):
+            return False
+
         if (
             node.target in allow_list_operator
             # bypass if custom op appears
@@ -138,6 +147,9 @@ def __init__(
             skip_mutable_buffer (bool, optional): If True, mutable buffers are not delegated to QNN.
         """
         self.compiler_specs_snapshot = copy.deepcopy(compiler_specs)
+        self.backend = flatbuffer_to_option(
+            generate_qnn_executorch_option(self.compiler_specs_snapshot)
+        ).backend_options.backend_type
 
         self.delegation_spec = DelegationSpec(
             QnnBackend.__name__, self.compiler_specs_snapshot
@@ -188,14 +200,34 @@ def tag_nodes(
                 # since they will all be removed in following stage
                 node.meta["delegation_tag"] = delegation_tag
 
+    @staticmethod
+    def check_partitions(
+        backend: QnnExecuTorchBackendType, partitions: List[Any]
+    ) -> bool:
+        pl = len(partitions)
+        if backend == QnnExecuTorchBackendType.kLpaiBackend:
+            # By default, there are two partitions that are always created in LPAI backend:
+            # one for the quantized and one for the dequantized.
+            bypass_nodes = [
+                node
+                for partition in partitions
+                for node in partition.nodes
+                if node.meta.get(QCOM_BYPASS_NODE, False)
+            ]
+            pl -= len(bypass_nodes)
+        if pl == 0:
+            logging.warning("Nothing can be partitioned!")
+        else:
+            logging.info(f"Found {pl} subgraphs to be partitioned.")
+        return pl != 0
+
     # override
     def partition(self, edge_program: torch.export.ExportedProgram) -> PartitionResult:
         # Generate partitions by QNN op_support checker
         partitions = self.generate_partitions(edge_program)
         del self.op_support_checker
-
         # If partitions are found, handle tagging of nodes, constant data, and mutated buffers for delegation
-        if len(partitions) != 0:
+        if self.check_partitions(self.backend, partitions):
             self.tag_nodes(partitions, edge_program)
             tag_constant_data(edge_program)
             if not self.skip_mutable_buffer: