Qualcomm AI Engine Direct - Enable custom operator (#8726)

Summary:
- Support to register op package in QNN Backend
- Add example script to run torch custom op with QNN Op package
- Allow op package override torch built-in operator
- Add op package example
- Modify the flag of dlopen for QNN library
- Generate custom op based on the meta and _schema.arguments of
torch.fx.Node
- Add README for the custom op

Reproduce commands:
```bash
# Follow the README to Install qpm
# Follow the README to install hexagon-sdk and hexagon-tool
# install hexagon sdk 5.4.0 for SM8650
# qpm-cli --install hexagonsdk5.x --version 5.4.0.3 --path /path/to/Qualcomm/Hexagon_SDK/hexagon-sdk-5.4.0
# install hexagon sdk 6.0.0 for x86
# qpm-cli --install hexagonsdk6.x --version 6.0.0.2 --path /path/to/Qualcomm/Hexagon_SDK/hexagon-sdk-6.0.0
# install hexagon tool 8.8.02 for x86
# qpm-cli --extract hexagon8.8 --version 8.8.02.1 --path /path/to/Qualcomm/Hexagon_SDK/hexagon-sdk-6.0.0/tools/HEXAGON_Tools/8.8.02

export HEXAGON_SDK_ROOT=/path/to/hexagon-sdk-5.4.0
export ANDROID_NDK_ROOT=/path/to/android-ndk-r26c
# use clang-9.0.0
export X86_CXX=/path/to/clang++
# run custom op with example script
python3 examples/qualcomm/custom_op/custom_ops_1.py --build_folder build-android -s <device_serial> -H <host> -m SM8650 --op_package_dir examples/qualcomm/custom_op/example_op_package_htp/ExampleOpPackage --build_op_package
# run custom op with unit test
python3 backends/qualcomm/tests/test_qnn_delegate.py TestUtilScript.test_custom_op -b build-android -s <device_serial> -H <host> -m SM8650 --op_package_dir examples/qualcomm/custom_op/example_op_package_htp/ExampleOpPackage -r </path/to/executorch> -a </path/to/artifacts>
```

Author: shewu-quic

backends/qualcomm/CMakeLists.txt

@@ -130,6 +130,7 @@ add_library(qnn_implementation STATIC)
 add_library(qnn_logger STATIC)
 add_library(qnn_manager STATIC)
 add_library(qnn_mem_manager STATIC)
+add_library(qnn_op_package_manager STATIC)
 add_library(qnn_profiler STATIC)
 add_library(qnn_schema INTERFACE ${_qnn_schema__outputs})
 add_library(qnn_sys_function_interface INTERFACE)
@@ -152,7 +153,7 @@ target_link_libraries(
 target_link_libraries(qnn_executorch_logging PRIVATE qnn_schema)
 target_link_libraries(qnn_profiler PRIVATE qnn_executorch_logging)
 target_link_libraries(qnn_logger PRIVATE qnn_implementation ${android_log})
-target_link_libraries(qnn_backend PRIVATE qnn_implementation qnn_logger)
+target_link_libraries(qnn_backend PRIVATE qnn_implementation qnn_logger qnn_op_package_manager)
 target_link_libraries(qnn_custom_protocol PRIVATE qnn_logger)
 target_link_libraries(
   qnn_device PRIVATE qnn_executorch_logging qnn_implementation qnn_logger

backends/qualcomm/builders/README.md

@@ -176,7 +176,8 @@ import torch
 from executorch.backends.qualcomm.utils.constants import QCOM_DATA
 # op builder will inherit NodeVisitor and have its own implementation
 # register_node_visitor for book-keeping the dictionary of target name v.s. callback
-from .node_visitor import NodeVisitor, register_node_visitor
+from .node_visitor import NodeVisitor
+from .node_visitor_manager import register_node_visitor
 # the definitions required to build operator in QNN
 from .qnn_constants import OpLayerNorm, QNN_OP_PACKAGE_NAME_QTI_AISW
 # utility to get parameter value when creating tensor in QNN

backends/qualcomm/builders/node_visitor.py

@@ -63,7 +63,9 @@
     torch.int64: PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_INT_64,
     torch.uint8: PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_UINT_8,
     torch.uint16: PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_UINT_16,
+    torch.uint32: PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_UINT_32,
     float: PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_FLOAT_32,
+    int: PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_UINT_32,
 }
 
 PER_CHANNEL_ENCODING = {
@@ -470,51 +472,3 @@ def define_node(
     ) -> PyQnnWrapper.PyQnnOpWrapper:
         """Convert torch.fx.Node to OpWrapper"""
         raise NotImplementedError("NodeVisitor must be extended!")
-
-
-# This will hold mapping of all node names to the visitor class
-_node_visitor_dict = {}
-
-
-def register_node_visitor(visitor):
-    """Register node visitor into _node_visitor_dict"""
-    assert (
-        isinstance(visitor, type)
-        and issubclass(visitor, NodeVisitor)
-        and hasattr(visitor, "target")
-    ), f"Illformed NodeVisitor subclass, can't register!, got: {visitor}"
-    for target in visitor.target:
-        _node_visitor_dict[target] = visitor
-
-
-def generate_node_to_external_map(
-    edge_program: torch.export.ExportedProgram,
-) -> Dict[torch.fx.Node, int]:
-    node_to_external_map = {}
-    for node in edge_program.graph_module.graph.nodes:
-        # The order in which we visit the placeholder node is same as the *args
-        # order for the forward(*args) signature for this gm. Using the order of
-        # the nodes as external_id to extract the right arg from *args at runtime
-        if is_graph_input(node, edge_program):
-            node_to_external_map[node] = len(node_to_external_map)
-    for node in edge_program.graph_module.graph.nodes:
-        if is_graph_output(node):
-            node_to_external_map[node] = len(node_to_external_map)
-    return node_to_external_map
-
-
-def get_node_visitors(
-    edge_program: torch.export.ExportedProgram,
-    enable_tensor_dump=False,
-) -> Dict[str, NodeVisitor]:
-    """Create a new class instance at runtime, and put them in a dict"""
-    node_to_external_map = generate_node_to_external_map(edge_program)
-    node_visitors = {}
-    for target, visitor in _node_visitor_dict.items():
-        assert callable(
-            visitor
-        ), f"Expeting a callable class, but got {visitor} of type {type(visitor)}"
-        node_visitors[target] = visitor(
-            node_to_external_map, edge_program, enable_tensor_dump
-        )
-    return node_visitors

backends/qualcomm/builders/node_visitor_manager.py

@@ -0,0 +1,77 @@
+# 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.
+
+from typing import Dict, List
+
+import torch
+from executorch.backends.qualcomm.serialization.qc_schema import (
+    QnnExecuTorchOpPackageInfo,
+)
+
+from .node_visitor import NodeVisitor
+from .op_custom_op import CustomOp
+from .utils import is_graph_input, is_graph_output
+
+
+# This will hold mapping of all node names to the visitor class
+_node_visitor_dict = {}
+
+
+def register_node_visitor(visitor):
+    """Register node visitor into _node_visitor_dict"""
+    assert (
+        isinstance(visitor, type)
+        and issubclass(visitor, NodeVisitor)
+        and hasattr(visitor, "target")
+    ), f"Informed NodeVisitor subclass, can't register!, got: {visitor}"
+    for target in visitor.target:
+        _node_visitor_dict[target] = visitor
+
+
+def generate_node_to_external_map(
+    edge_program: torch.export.ExportedProgram,
+) -> Dict[torch.fx.Node, int]:
+    node_to_external_map = {}
+    for node in edge_program.graph_module.graph.nodes:
+        # The order in which we visit the placeholder node is same as the *args
+        # order for the forward(*args) signature for this gm. Using the order of
+        # the nodes as external_id to extract the right arg from *args at runtime
+        if is_graph_input(node, edge_program):
+            node_to_external_map[node] = len(node_to_external_map)
+    for node in edge_program.graph_module.graph.nodes:
+        if is_graph_output(node):
+            node_to_external_map[node] = len(node_to_external_map)
+    return node_to_external_map
+
+
+def get_node_visitors(
+    edge_program: torch.export.ExportedProgram,
+    enable_tensor_dump=False,
+    op_package_infos: List[QnnExecuTorchOpPackageInfo] = None,
+) -> Dict[str, NodeVisitor]:
+    """Create a new class instance at runtime, and put them in a dict"""
+    node_to_external_map = generate_node_to_external_map(edge_program)
+    node_visitors = {}
+    for target, visitor in _node_visitor_dict.items():
+        assert callable(
+            visitor
+        ), f"Expecting a callable class, but got {visitor} of type {type(visitor)}"
+        node_visitors[target] = visitor(
+            node_to_external_map, edge_program, enable_tensor_dump
+        )
+    if op_package_infos:
+        custom_ops = []
+        for op_package_info in op_package_infos:
+            if op_package_info.custom_op_name not in custom_ops:
+                custom_op_builder = CustomOp(
+                    op_package_info,
+                    node_to_external_map,
+                    edge_program,
+                    enable_tensor_dump,
+                )
+                node_visitors[op_package_info.custom_op_name] = custom_op_builder
+                custom_ops.append(op_package_info.custom_op_name)
+    return node_visitors

backends/qualcomm/builders/op_abs.py

@@ -9,7 +9,8 @@
 
 import torch
 
-from .node_visitor import NodeVisitor, register_node_visitor
+from .node_visitor import NodeVisitor
+from .node_visitor_manager import register_node_visitor
 from .qnn_constants import OpElementWiseAbs, QNN_OP_PACKAGE_NAME_QTI_AISW
 
 

backends/qualcomm/builders/op_adaptive_avg_pool2d.py

@@ -11,7 +11,8 @@
 
 import torch
 
-from .node_visitor import NodeVisitor, register_node_visitor
+from .node_visitor import NodeVisitor
+from .node_visitor_manager import register_node_visitor
 from .qnn_constants import OpPoolAvg2d, QNN_OP_PACKAGE_NAME_QTI_AISW
 
 

backends/qualcomm/builders/op_add.py

@@ -9,7 +9,8 @@
 
 import torch
 
-from .node_visitor import NodeVisitor, register_node_visitor
+from .node_visitor import NodeVisitor
+from .node_visitor_manager import register_node_visitor
 from .qnn_constants import OpElementWiseAdd, QNN_OP_PACKAGE_NAME_QTI_AISW
 
 

backends/qualcomm/builders/op_amax.py

@@ -12,7 +12,8 @@
 import torch
 from executorch.backends.qualcomm.utils.constants import QCOM_AXIS_ORDER, QCOM_DATA
 
-from .node_visitor import NodeVisitor, register_node_visitor
+from .node_visitor import NodeVisitor
+from .node_visitor_manager import register_node_visitor
 from .qnn_constants import OpReduceMax, QNN_OP_PACKAGE_NAME_QTI_AISW
 
 

backends/qualcomm/builders/op_and.py

@@ -9,7 +9,8 @@
 
 import torch
 
-from .node_visitor import NodeVisitor, register_node_visitor
+from .node_visitor import NodeVisitor
+from .node_visitor_manager import register_node_visitor
 from .qnn_constants import OpElementWiseAnd, QNN_OP_PACKAGE_NAME_QTI_AISW
 
 

backends/qualcomm/builders/op_arange.py

@@ -9,7 +9,8 @@
 
 import torch
 
-from .node_visitor import NodeVisitor, register_node_visitor
+from .node_visitor import NodeVisitor
+from .node_visitor_manager import register_node_visitor
 
 
 @register_node_visitor

backends/qualcomm/builders/op_argmin.py

@@ -10,7 +10,8 @@
 import torch
 from executorch.backends.qualcomm.utils.constants import QCOM_AXIS_ORDER, QCOM_DATA
 
-from .node_visitor import NodeVisitor, register_node_visitor
+from .node_visitor import NodeVisitor
+from .node_visitor_manager import register_node_visitor
 from .qnn_constants import OpArgmin, QNN_OP_PACKAGE_NAME_QTI_AISW
 
 

backends/qualcomm/builders/op_avg_pool2d.py

@@ -12,7 +12,8 @@
 import torch
 from executorch.backends.qualcomm.utils.constants import QCOM_DATA
 
-from .node_visitor import NodeVisitor, register_node_visitor
+from .node_visitor import NodeVisitor
+from .node_visitor_manager import register_node_visitor
 from .qnn_constants import OpPoolAvg2d, QNN_OP_PACKAGE_NAME_QTI_AISW
 
 

backends/qualcomm/builders/op_batch_norm.py

@@ -18,7 +18,8 @@
 )
 from executorch.exir.dialects._ops import ops as exir_ops
 
-from .node_visitor import NodeVisitor, register_node_visitor
+from .node_visitor import NodeVisitor
+from .node_visitor_manager import register_node_visitor
 from .qnn_constants import OpBatchnorm, QNN_OP_PACKAGE_NAME_QTI_AISW
 from .utils import get_parameter
 

backends/qualcomm/builders/op_bmm.py

@@ -9,7 +9,8 @@
 
 import torch
 
-from .node_visitor import NodeVisitor, register_node_visitor
+from .node_visitor import NodeVisitor
+from .node_visitor_manager import register_node_visitor
 from .qnn_constants import OpMatMul, QNN_OP_PACKAGE_NAME_QTI_AISW
 
 

backends/qualcomm/builders/op_cat.py

@@ -12,7 +12,8 @@
 import torch
 from executorch.backends.qualcomm.utils.constants import QCOM_AXIS_ORDER, QCOM_DATA
 
-from .node_visitor import NodeVisitor, register_node_visitor
+from .node_visitor import NodeVisitor
+from .node_visitor_manager import register_node_visitor
 from .qnn_constants import OpConcat, QNN_OP_PACKAGE_NAME_QTI_AISW
 
 

backends/qualcomm/builders/op_ceil.py

@@ -9,7 +9,8 @@
 
 import torch
 
-from .node_visitor import NodeVisitor, register_node_visitor
+from .node_visitor import NodeVisitor
+from .node_visitor_manager import register_node_visitor
 from .qnn_constants import OpElementWiseCeil, QNN_OP_PACKAGE_NAME_QTI_AISW
 
 

backends/qualcomm/builders/op_clamp.py

@@ -11,7 +11,8 @@
 import torch
 from executorch.backends.qualcomm.utils.constants import QCOM_DATA
 
-from .node_visitor import NodeVisitor, register_node_visitor
+from .node_visitor import NodeVisitor
+from .node_visitor_manager import register_node_visitor
 from .qnn_constants import OpReluMinMax, QNN_OP_PACKAGE_NAME_QTI_AISW
 
 

backends/qualcomm/builders/op_conv2d.py

@@ -12,7 +12,8 @@
 import torch
 from executorch.backends.qualcomm.utils.constants import QCOM_DATA
 
-from .node_visitor import NodeVisitor, register_node_visitor
+from .node_visitor import NodeVisitor
+from .node_visitor_manager import register_node_visitor
 from .qnn_constants import (
     OpConv2d,
     OpDepthWiseConv2d,

backends/qualcomm/builders/op_cos.py

@@ -9,7 +9,8 @@
 
 import torch
 
-from .node_visitor import NodeVisitor, register_node_visitor
+from .node_visitor import NodeVisitor
+from .node_visitor_manager import register_node_visitor
 from .qnn_constants import OpElementWiseCos, QNN_OP_PACKAGE_NAME_QTI_AISW
 
 

backends/qualcomm/builders/op_cum_sum.py

@@ -11,7 +11,8 @@
 import torch
 from executorch.backends.qualcomm.utils.constants import QCOM_AXIS_ORDER, QCOM_DATA
 
-from .node_visitor import NodeVisitor, register_node_visitor
+from .node_visitor import NodeVisitor
+from .node_visitor_manager import register_node_visitor
 from .qnn_constants import OpCumulativeSum, QNN_OP_PACKAGE_NAME_QTI_AISW