Qualcomm AI Engine Direct - Support LPAI Backend (#18151)

Author: shewu-quic

.ci/scripts/build-qnn-sdk.sh

@@ -40,6 +40,8 @@ set_up_aot() {
       -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON \
       -DEXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP=ON \
       -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \
+      -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \
+      -DEXECUTORCH_BUILD_KERNELS_QUANTIZED_AOT=ON \
       -DEXECUTORCH_ENABLE_EVENT_TRACER=ON \
       -DPYTHON_EXECUTABLE=python3
   cmake --build $PWD --target "PyQnnManagerAdaptor" -j$(nproc)

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 [sign_library.sh](scripts/sign_library.sh) to sign the skel library.
+  
+
 ## Directory Structure
 
 ```
@@ -59,8 +51,8 @@ backends/qualcomm
 |   |       # models on a device.
 |   └── backends # Backends supported by QNN.
 |       └── gpu / htp
-|           ├── aarch64 # Configuration required to run on device. (Device Part).
-|           └── x86_64 # Configuration required to compile graph on host. (AoT Part).
+|           ├── target # Configuration required to run on device. (Device Part).
+|           └── host # Configuration required to compile graph on host. (AoT Part).
 ├── scripts # Misc supporting scripts, not related to core functionality.
 ├── serialization # Contains files related to serializing QNN compiler options and SoC information
 ├── tests # Unit tests and model tests go here.

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)
         dead_code_elimination_pass(graph_module)

backends/qualcomm/_passes/qnn_pass_manager.py

@@ -63,6 +63,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,
@@ -151,6 +154,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
@@ -183,6 +187,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,
         DecomposeAcos,
         DecomposeAny,
@@ -109,7 +108,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,28 @@ 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:
+            assert (
+                pl == 1
+            ), "LPAI backend only supports fully delegation due to the accuracy issue of Q/DQ in the LPAI backend."
+        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: