Qualcomm AI Engine Direct - Enable custom operator

backends/qualcomm/CMakeLists.txt

@@ -131,6 +131,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)
@@ -157,13 +158,13 @@ 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 qcir_utils)
 target_link_libraries(
   qnn_device PRIVATE qnn_executorch_logging qnn_implementation qnn_logger
 )
 target_link_libraries(
-  qnn_backend_cache PRIVATE qnn_sys_implementation qcir_utils
+  qnn_backend_cache PRIVATE qnn_sys_implementation qcir_utils qnn_schema
 )
 target_link_libraries(
   qnn_context PRIVATE qnn_implementation qnn_logger qnn_backend qnn_device

backends/qualcomm/builders/README.md

@@ -2,14 +2,18 @@
 Thank you for contributing to Qualcomm AI Engine Direct delegate for ExecuTorch. Reading and following these guidelines will help you quickly get the essentials of implementing operator builder to unblock yourself and land pull requests more efficiently.
 
 ## Sections
-* [References](#references)
-* [Getting Started](#getting-started)
-    * [Identify Unsupported Operator](#identify-unsupported-operator)
-    * [Check Operator Spec](#check-operator-spec)
-    * [Implementation](#implementation)
-    * [Quantizer Annotation](#quantizer-annotation)
-* [Issues](#issues)
-* [Pull Requests](#pull-requests)
+- [Contribution for More Operators](#contribution-for-more-operators)
+  - [Sections](#sections)
+  - [References](#references)
+    - [Qualcomm AI Engine Direct](#qualcomm-ai-engine-direct)
+    - [PyTorch](#pytorch)
+  - [Getting Started](#getting-started)
+    - [Identify Unsupported Operator](#identify-unsupported-operator)
+    - [Check Operator Spec](#check-operator-spec)
+    - [Implementation](#implementation)
+    - [Quantizer Annotation](#quantizer-annotation)
+  - [Issues](#issues)
+  - [Pull Requests](#pull-requests)
 
 ## References
 ### Qualcomm AI Engine Direct
@@ -175,7 +179,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

@@ -58,7 +58,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 = {
@@ -382,51 +384,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_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, QNN_TENSOR_TYPE_MAP, register_node_visitor
+from .node_visitor import NodeVisitor, QNN_TENSOR_TYPE_MAP
+from .node_visitor_manager import register_node_visitor
 from .qnn_constants import OpArgmin, OpCast, 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

@@ -13,7 +13,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

@@ -10,7 +10,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