Qualcomm AI Engine Direct - Optimize performance of pcq embedding

backends/qualcomm/builders/op_embedding.py

@@ -9,6 +9,9 @@
 
 import numpy as np
 import torch
+from executorch.backends.qualcomm.utils.check_qnn_version import (
+    is_qnn_sdk_version_less_than,
+)
 from executorch.backends.qualcomm.utils.constants import (
     QCOM_DATA,
     QCOM_DTYPE,
@@ -35,15 +38,19 @@ class Embedding(NodeVisitor):
     def __init__(self, *args) -> None:
         super().__init__(*args)
 
-    def define_node(
+    def _is_pcq_embedding(self, weight_node: torch.fx.Node) -> bool:
+        return (
+            QCOM_QUANT_ATTRS in weight_node.meta
+            and weight_node.meta[QCOM_QUANT_ATTRS][QCOM_ENCODING]
+            in PER_CHANNEL_ENCODING
+        )
+
+    def _define_weight_and_indices(
         self,
         node: torch.fx.Node,
         nodes_to_wrappers: Dict[torch.fx.Node, PyQnnManager.TensorWrapper],
-    ) -> PyQnnManager.PyQnnOpWrapper:
+    ):
         weight_node = self.get_node(node.args[0])
-        is_pcq_embedding = QCOM_QUANT_ATTRS in weight_node.meta and weight_node.meta[
-            QCOM_QUANT_ATTRS
-        ][QCOM_ENCODING] in (PER_CHANNEL_ENCODING)
         weight_tensor = get_parameter(weight_node, self.edge_program)
         weight_tensor_wrapper = self.define_tensor(
             weight_node,
@@ -62,31 +69,131 @@ def define_node(
             PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
         )
+        return weight_node, weight_tensor, weight_tensor_wrapper, indices_tensor_wrapper
+
+    def _act_dtype(self, node: torch.fx.Node):
+        act_quant_encoding, act_quant_configs = self.get_quant_encoding_conf(node, node)
+        act_dtype = (
+            torch.uint16
+            if act_quant_configs[QCOM_DTYPE] == torch.int32
+            else act_quant_configs[QCOM_DTYPE]
+        )
+        return act_quant_encoding, act_quant_configs, act_dtype
+
+    def _build_gather_op(
+        self, node_name, gather_input_tensors, gather_output_tensors
+    ) -> PyQnnManager.PyQnnOpWrapper:
+        gather_op = PyQnnManager.PyQnnOpWrapper(
+            node_name,
+            QNN_OP_PACKAGE_NAME_QTI_AISW,
+            OpGather.op_name,
+        )
+        gather_op.AddInputTensors(gather_input_tensors)
+        gather_op.AddOutputTensors(gather_output_tensors)
+        # For now, default axis is zero.
+        gather_op.AddScalarParam(
+            OpGather.param_axis,
+            PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_INT_32,
+            {QCOM_DATA: np.int32(0)},
+        )
+        return gather_op
+
+    # Note that this pattern is only supported with QNN 2.48+.
+    # The weight is converted to the activation quantization before gather so
+    # the backend can fuse the convert into gather for better performance.
+    def define_node_optimize(
+        self,
+        node: torch.fx.Node,
+        nodes_to_wrappers: Dict[torch.fx.Node, PyQnnManager.TensorWrapper],
+    ) -> PyQnnManager.PyQnnOpWrapper:
+        op_wrapper_list = []
+        (
+            weight_node,
+            weight_tensor,
+            weight_tensor_wrapper,
+            indices_tensor_wrapper,
+        ) = self._define_weight_and_indices(node, nodes_to_wrappers)
+
+        gather_input_tensors = []
+        if self._is_pcq_embedding(weight_node):
+            act_quant_encoding, act_quant_configs, act_dtype = self._act_dtype(node)
+            convert_tensor_wrapper = self.define_custom_tensor_wrapper(
+                node_name=node.name + "_convert",
+                tensor_type=PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
+                dtype=QNN_QUANT_TYPE_MAP[act_dtype],
+                quant_encoding=act_quant_encoding,
+                quant_configs=act_quant_configs,
+                dims=weight_tensor.size(),
+                tensor=weight_tensor,
+                is_fake_tensor=True,
+                nodes_to_wrappers=nodes_to_wrappers,
+            )
+            convert_op = PyQnnManager.PyQnnOpWrapper(
+                node.name + "_convert",
+                QNN_OP_PACKAGE_NAME_QTI_AISW,
+                OpConvert.op_name,
+            )
+            convert_op.AddInputTensors([weight_tensor_wrapper])
+            convert_op.AddOutputTensors([convert_tensor_wrapper])
+            op_wrapper_list.append(convert_op)
+            gather_input_tensors.append(convert_tensor_wrapper)
+        else:
+            gather_input_tensors.append(weight_tensor_wrapper)
+        gather_input_tensors.append(indices_tensor_wrapper)
+
+        output_tensor = self.get_tensor(node, node)
+        output_tensor_wrapper = self.define_tensor(
+            node,
+            node,
+            output_tensor,
+            PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
+            nodes_to_wrappers,
+            node_name=node.name,
+        )
+
+        op_wrapper_list.append(
+            self._build_gather_op(
+                node.name, gather_input_tensors, [output_tensor_wrapper]
+            )
+        )
+        return op_wrapper_list
+
+    def define_node_legacy(
+        self,
+        node: torch.fx.Node,
+        nodes_to_wrappers: Dict[torch.fx.Node, PyQnnManager.TensorWrapper],
+    ) -> PyQnnManager.PyQnnOpWrapper:
+        (
+            weight_node,
+            _,
+            weight_tensor_wrapper,
+            indices_tensor_wrapper,
+        ) = self._define_weight_and_indices(node, nodes_to_wrappers)
+        is_pcq_embedding = self._is_pcq_embedding(weight_node)
 
         gather_input_tensors = [weight_tensor_wrapper, indices_tensor_wrapper]
 
         output_tensor = self.get_tensor(node, node)
         node_name = node.name
         if is_pcq_embedding:
             node_quant_attrs = node.meta[QCOM_QUANT_ATTRS].copy()
-            intermediate_quant_attrs = node.meta[QCOM_QUANT_ATTRS].copy()
+            weight_quant_attrs = weight_node.meta[QCOM_QUANT_ATTRS]
             # Based on QNN HTP quantization constraints,
             # we should set the scale to max of scales and per-tensor quantization for embedding op
+            intermediate_quant_attrs = node_quant_attrs.copy()
             intermediate_quant_attrs[QCOM_SCALE] = (
-                weight_node.meta[QCOM_QUANT_ATTRS][QCOM_SCALES].max().item()
+                weight_quant_attrs[QCOM_SCALES].max().item()
             )
             intermediate_quant_attrs[QCOM_ZERO_POINT] = (
-                weight_node.meta[QCOM_QUANT_ATTRS][QCOM_ZERO_POINTS].max().item()
+                weight_quant_attrs[QCOM_ZERO_POINTS].max().item()
             )
-            intermediate_quant_attrs[QCOM_DTYPE] = weight_node.meta[QCOM_QUANT_ATTRS][
-                QCOM_DTYPE
+            intermediate_quant_attrs[QCOM_DTYPE] = weight_quant_attrs[QCOM_DTYPE]
+            intermediate_quant_attrs[QCOM_QUANT_MAX] = weight_quant_attrs[
+                QCOM_QUANT_MAX
+            ]
+            intermediate_quant_attrs[QCOM_QUANT_MIN] = weight_quant_attrs[
+                QCOM_QUANT_MIN
             ]
-            intermediate_quant_attrs[QCOM_QUANT_MAX] = weight_node.meta[
-                QCOM_QUANT_ATTRS
-            ][QCOM_QUANT_MAX]
-            intermediate_quant_attrs[QCOM_QUANT_MIN] = weight_node.meta[
-                QCOM_QUANT_ATTRS
-            ][QCOM_QUANT_MIN]
             node.meta[QCOM_QUANT_ATTRS] = intermediate_quant_attrs
             node_name += "_intermediate"
         output_tensor_wrapper = self.define_tensor(
@@ -99,33 +206,15 @@ def define_node(
         )
         gather_output_tensors = [output_tensor_wrapper]
 
-        gather_op = PyQnnManager.PyQnnOpWrapper(
-            node_name,
-            QNN_OP_PACKAGE_NAME_QTI_AISW,
-            OpGather.op_name,
-        )
-        gather_op.AddInputTensors(gather_input_tensors)
-        gather_op.AddOutputTensors(gather_output_tensors)
-
-        # For now, default axis is zero.
-        gather_op.AddScalarParam(
-            OpGather.param_axis,
-            PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_INT_32,
-            {QCOM_DATA: np.int32(0)},
-        )
-
-        op_wrapper_list = [gather_op]
+        op_wrapper_list = [
+            self._build_gather_op(
+                node_name, gather_input_tensors, gather_output_tensors
+            )
+        ]
 
         if is_pcq_embedding:
             node.meta[QCOM_QUANT_ATTRS] = node_quant_attrs
-            act_quant_encoding, act_quant_configs = self.get_quant_encoding_conf(
-                node, node
-            )
-            act_dtype = (
-                torch.uint16
-                if act_quant_configs[QCOM_DTYPE] == torch.int32
-                else act_quant_configs[QCOM_DTYPE]
-            )
+            act_quant_encoding, act_quant_configs, act_dtype = self._act_dtype(node)
             convert_tensor_wrapper = self.define_custom_tensor_wrapper(
                 node_name=node.name,
                 tensor_type=PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
@@ -147,3 +236,12 @@ def define_node(
             op_wrapper_list.append(convert_op)
 
         return op_wrapper_list
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        nodes_to_wrappers: Dict[torch.fx.Node, PyQnnManager.TensorWrapper],
+    ) -> PyQnnManager.PyQnnOpWrapper:
+        if is_qnn_sdk_version_less_than("2.48"):
+            return self.define_node_legacy(node, nodes_to_wrappers)
+        return self.define_node_optimize(node, nodes_to_wrappers)

backends/qualcomm/export_utils.py

@@ -37,6 +37,10 @@
     QnnExecuTorchLpaiTargetEnv,
     QnnExecuTorchOpPackageOptions,
 )
+from executorch.backends.qualcomm.utils.check_qnn_version import (
+    get_sdk_build_id,
+    is_qnn_sdk_version_less_than,
+)
 from executorch.backends.qualcomm.utils.constants import (
     HEXAGON_SDK_ROOT,
     HEXAGON_TOOLS_ROOT,
@@ -47,10 +51,8 @@
     generate_lpai_compiler_spec,
     generate_qnn_executorch_compiler_spec,
     get_qnn_context_binary_alignment,
-    get_sdk_build_id,
     get_soc_to_htp_arch_map,
     get_soc_to_lpai_hw_ver_map,
-    is_qnn_sdk_version_less_than,
     to_edge_transform_and_lower_to_qnn,
 )
 from executorch.exir.capture._config import ExecutorchBackendConfig

backends/qualcomm/tests/test_qnn_delegate.py

@@ -44,6 +44,10 @@
     TestQNN,
     validate_context_binary,
 )
+from executorch.backends.qualcomm.utils.check_qnn_version import (
+    is_qnn_sdk_version_greater_than,
+    is_qnn_sdk_version_less_than,
+)
 from executorch.backends.qualcomm.utils.constants import (
     QCOM_ANNOTATION,
     QCOM_MODULE,
@@ -61,8 +65,6 @@
     generate_htp_compiler_spec,
     generate_lpai_compiler_spec,
     generate_qnn_executorch_compiler_spec,
-    is_qnn_sdk_version_greater_than,
-    is_qnn_sdk_version_less_than,
     PyQnnManagerAdaptor,
     rewrite_prepared_observer,
     skip_annotation,
@@ -3929,18 +3931,20 @@ def test_qnn_backend_embedding(self):
                 )
                 self.lower_module_and_test_output(modules[i], sample_input)
 
-    # TODO: Once the accuracy issue is fixed, enable this test.
-    @unittest.skip("Bad accuracy for HTP")
+    @unittest.skipIf(is_qnn_sdk_version_less_than("2.48"), "UT pass after QNN 2.48")
     def test_qnn_backend_embedding_per_channel(self):
         module = Embedding()  # noqa: F405
         sample_input = (torch.Tensor([1, 2, 4, 5]).to(torch.int32),)
-        qdq_module = self.get_qdq_module(
-            module,
-            sample_input,
-            quant_dtype=QuantDtype.use_16a8w,
-            is_embedding_per_channel=True,
-        )
-        self.lower_module_and_test_output(qdq_module, sample_input)
+        quant_dtype = [QuantDtype.use_16a8w, QuantDtype.use_16a4w]
+        for i, qdtype in enumerate(quant_dtype):
+            with self.subTest(i=i):
+                qdq_module = self.get_qdq_module(
+                    module,
+                    sample_input,
+                    quant_dtype=qdtype,
+                    is_embedding_per_channel=True,
+                )
+                self.lower_module_and_test_output(qdq_module, sample_input)
 
     def test_qnn_backend_equal(self):
         test_comb = [

backends/qualcomm/utils/check_qnn_version.py

@@ -0,0 +1,69 @@
+# 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 os
+import platform
+import re
+
+import executorch.backends.qualcomm.python.PyQnnManagerAdaptor as PyQnnManagerAdaptor
+
+
+def get_qnn_lib_name(base: str) -> str:
+    """Returns the platform-specific shared library filename for a QNN library."""
+    if platform.system().lower() == "windows":
+        return f"{base}.dll"
+    return f"lib{base}.so"
+
+
+def _get_qnn_host_lib_dir_name() -> str:
+    """Returns the QNN SDK library subdirectory name for the current x86-64 host OS."""
+    if platform.system().lower() == "windows":
+        return "x86_64-windows-msvc"
+    return "x86_64-linux-clang"
+
+
+def get_sdk_build_id():
+    htp_library_path = os.path.join(
+        os.environ.get("QNN_SDK_ROOT", None),
+        "lib",
+        _get_qnn_host_lib_dir_name(),
+        get_qnn_lib_name("QnnHtp"),
+    )
+    # The GetQnnSdkBuildId API can be used without needing to create a backend first, so it works regardless of which backend is used.
+    sdk_build_id = PyQnnManagerAdaptor.GetQnnSdkBuildId(htp_library_path)
+    return sdk_build_id
+
+
+def is_qnn_sdk_version_less_than(target_version):
+    current_version = get_sdk_build_id()
+
+    match = re.search(r"v(\d+)\.(\d+)", current_version)
+    if match:
+        current_major, current_minor = map(int, match.groups()[:2])
+    else:
+        raise ValueError(
+            f"Failed to get current major and minor version from QNN SDK Build id {current_version}"
+        )
+
+    target_major, target_minor = map(int, target_version.split(".")[:2])
+
+    return current_major == target_major and current_minor < target_minor
+
+
+def is_qnn_sdk_version_greater_than(target_version):
+    current_version = get_sdk_build_id()
+
+    match = re.search(r"v(\d+)\.(\d+)", current_version)
+    if match:
+        current_major, current_minor = map(int, match.groups()[:2])
+    else:
+        raise ValueError(
+            f"Failed to get current major and minor version from QNN SDK Build id {current_version}"
+        )
+
+    target_major, target_minor = map(int, target_version.split(".")[:2])
+
+    return current_major == target_major and current_minor > target_minor