[BYOC][NNAPI] This PR intorduce NNAPI to TVM

This PR introduces a new BYOC backend for Android Neural Networks API (NNAPI),
enabling execution of neural networks on custom accelerators. This feature adds
a new codegen and runtime for NNAPI, supporting operations such as element-wise
ops, nn.dense, and nn.conv2d for CNN model with static shape.

Co-authored-by: Ming-Long Huang <[email protected]>
Co-authored-by: HMZ <[email protected]>

Author: 3 people

CMakeLists.txt

@@ -125,6 +125,8 @@ tvm_option(USE_ARM_COMPUTE_LIB "Build with Arm Compute Library" OFF)
 tvm_option(USE_ARM_COMPUTE_LIB_GRAPH_EXECUTOR "Build with Arm Compute Library graph executor" OFF)
 tvm_option(USE_TENSORRT_CODEGEN "Build with TensorRT Codegen support" OFF)
 tvm_option(USE_TENSORRT_RUNTIME "Build with TensorRT runtime" OFF)
+tvm_option(USE_NNAPI_CODEGEN "Build with NNAPI Codegen support" OFF)
+tvm_option(USE_NNAPI_RUNTIME "Build with NNAPI runtime" OFF)
 tvm_option(USE_RUST_EXT "Build with Rust based compiler extensions, STATIC, DYNAMIC, or OFF" OFF)
 tvm_option(USE_VITIS_AI "Build with VITIS-AI Codegen support" OFF)
 tvm_option(SUMMARIZE "Print CMake option summary after configuring" OFF)
@@ -602,6 +604,7 @@ include(cmake/modules/contrib/BNNS.cmake)
 include(cmake/modules/contrib/ONNX.cmake)
 include(cmake/modules/contrib/ArmComputeLib.cmake)
 include(cmake/modules/contrib/TensorRT.cmake)
+include(cmake/modules/contrib/NNAPI.cmake)
 include(cmake/modules/contrib/VitisAI.cmake)
 include(cmake/modules/contrib/Verilator.cmake)
 include(cmake/modules/contrib/UMA.cmake)

cmake/modules/LibInfo.cmake

@@ -144,6 +144,8 @@ function(add_lib_info src_file)
     TVM_INFO_USE_MSC="${USE_MSC}"
     TVM_INFO_USE_CCACHE="${USE_CCACHE}"
     TVM_INFO_USE_NVSHMEM="${USE_NVSHMEM}"
+    TVM_INFO_USE_NNAPI_CODEGEN="${USE_NNAPI_CODEGEN}"
+    TVM_INFO_USE_NNAPI_RUNTIME="${USE_NNAPI_RUNTIME}"
     TVM_INFO_BACKTRACE_ON_SEGFAULT="${BACKTRACE_ON_SEGFAULT}"
   )
 

cmake/modules/contrib/NNAPI.cmake

@@ -0,0 +1,39 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# NNAPI Codegen
+if(USE_NNAPI_CODEGEN)
+    message(STATUS "Build with NNAPI codegen")
+
+    tvm_file_glob(GLOB COMPILER_NNAPI_SRCS src/relax/backend/contrib/nnapi/*.cc)
+    tvm_file_glob(GLOB RUNTIME_NNAPI_SRCS src/runtime/contrib/nnapi/*.cc)
+    list(APPEND COMPILER_SRCS ${COMPILER_NNAPI_SRCS})
+    if(NOT USE_NNAPI_RUNTIME)
+        list(APPEND COMPILER_SRCS ${RUNTIME_NNAPI_SRCS})
+    endif()
+endif()
+
+# NNAPI Runtime
+if(USE_NNAPI_RUNTIME)
+    message(STATUS "Build with NNAPI runtime")
+
+    tvm_file_glob(GLOB RUNTIME_NNAPI_SRCS src/runtime/contrib/nnapi/*.cc)
+    list(APPEND RUNTIME_SRCS ${RUNTIME_NNAPI_SRCS})
+    list(APPEND TVM_RUNTIME_LINKER_LIBS neuralnetworks log)
+
+    add_definitions(-DTVM_GRAPH_EXECUTOR_NNAPI)
+endif()

python/tvm/relax/backend/contrib/nnapi.py

@@ -0,0 +1,322 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""Pattern table for NNAPI backend"""
+from typing import (
+    Mapping,
+    Optional,
+    Tuple,
+    List,
+)
+from tvm.ir import IRModule
+from tvm.relax.transform import FuseOpsByPattern, MergeCompositeFunctions
+from tvm.relax.dpl.pattern import (
+    DFPattern,
+    wildcard,
+    is_op,
+)
+
+from ..pattern_registry import get_patterns_with_prefix, register_patterns
+
+
+def elementwise_binary_patterns() -> List[Tuple[str, DFPattern, Mapping[str, DFPattern]]]:
+    """
+    Returns a list of tuples representing elementwise binary operation patterns mapped
+    between NNAPI and Relax frameworks.
+    """
+    def _elementwise_binary_pattern(
+        pattern_name: str,
+        op_name: str,
+    ) -> Tuple[str, DFPattern, Mapping[str, DFPattern]]:
+        input0 = wildcard()
+        input1 = wildcard()
+
+        pattern = is_op(op_name)(input0, input1)
+
+        return (pattern_name, pattern, {})
+
+    return [
+        _elementwise_binary_pattern("nnapi.add", "relax.add"),
+        _elementwise_binary_pattern("nnapi.mul", "relax.multiply"),
+        _elementwise_binary_pattern("nnapi.div", "relax.divide"),
+        _elementwise_binary_pattern("nnapi.sub", "relax.subtract"),
+        _elementwise_binary_pattern("nnapi.pow", "relax.power"),
+        _elementwise_binary_pattern("nnapi.equal", "relax.equal"),
+        _elementwise_binary_pattern("nnapi.greater", "relax.greater"),
+        _elementwise_binary_pattern("nnapi.greater_equal", "relax.greater_equal"),
+        _elementwise_binary_pattern("nnapi.less", "relax.less"),
+        _elementwise_binary_pattern("nnapi.less_equal", "relax.less_equal"),
+        _elementwise_binary_pattern("nnapi.not_equal", "relax.not_equal"),
+        _elementwise_binary_pattern("nnapi.maximum", "relax.maximum"),
+        _elementwise_binary_pattern("nnapi.minimum", "relax.minimum"),
+    ]
+
+
+def unary_patterns() -> List[Tuple[str, DFPattern, Mapping[str, DFPattern]]]:
+    """
+    Returns a list of tuples representing unary operation patterns mapped
+    between NNAPI and Relax frameworks.
+    """
+    def _unary_pattern(
+        pattern_name: str, op_name: str
+    ) -> Tuple[str, DFPattern, Mapping[str, DFPattern]]:
+        input0 = wildcard()
+        pattern = is_op(op_name)(input0)
+        return (pattern_name, pattern, {})
+
+    return [
+        _unary_pattern("nnapi.floor", "relax.floor"),
+        _unary_pattern("nnapi.relu", "relax.nn.relu"),
+        _unary_pattern("nnapi.logistic", "relax.sigmoid"),
+        _unary_pattern("nnapi.softmax", "relax.nn.softmax"),
+        _unary_pattern("nnapi.tanh", "relax.tanh"),
+        _unary_pattern("nnapi.abs", "relax.abs"),
+        _unary_pattern("nnapi.exp", "relax.exp"),
+        _unary_pattern("nnapi.log", "relax.log"),
+        _unary_pattern("nnapi.neg", "relax.negative"),
+        _unary_pattern("nnapi.cast", "relax.astype"),
+        _unary_pattern("nnapi.sqrt", "relax.sqrt"),
+        _unary_pattern("nnapi.rsqrt", "relax.rsqrt"),
+    ]
+
+
+def matmul_pattern() -> Tuple[str, DFPattern, Mapping[str, DFPattern]]:
+    """
+    Returns a tuple representing matmul operation patterns mapped
+    between NNAPI and Relax frameworks.
+    """
+    input0 = wildcard()
+    input1 = wildcard()
+    pattern = is_op("relax.matmul")(input0, input1)
+    return ("nnapi.batch_matmul", pattern, {})
+
+
+def permute_dims_pattern() -> Tuple[str, DFPattern, Mapping[str, DFPattern]]:
+    """
+    Returns a tuple representing permute operation patterns mapped
+    between NNAPI and Relax frameworks.
+    """
+    input0 = wildcard()
+    pattern = is_op("relax.permute_dims")(input0)
+    return ("nnapi.transpose", pattern, {})
+
+
+def astype_pattern() -> Tuple[str, DFPattern, Mapping[str, DFPattern]]:
+    """
+    Returns a tuple representing astype operation patterns mapped
+    between NNAPI and Relax frameworks.
+    """
+    input0 = wildcard().has_dtype("float16") | wildcard().has_dtype("float32")
+    pattern = is_op("relax.astype")(input0).has_dtype("float16") | is_op("relax.astype")(
+        input0
+    ).has_dtype("float32")
+
+    return ("nnapi.cast", pattern, {})
+
+
+def mean_pattern() -> Tuple[str, DFPattern, Mapping[str, DFPattern]]:
+    """
+    Returns a tuple representing mean operation patterns mapped
+    between NNAPI and Relax frameworks.
+    """
+    input0 = wildcard()
+    pattern = is_op("relax.mean")(input0)
+
+    return ("nnapi.mean", pattern, {})
+
+
+def conv2d_pattern() -> Tuple[str, DFPattern, Mapping[str, DFPattern]]:
+    """
+    Returns a tuple representing conv2d operation patterns mapped
+    between NNAPI and Relax frameworks.
+    """
+    input0 = wildcard()
+    input1 = wildcard()
+    input2 = wildcard()
+    conv = is_op("relax.nn.conv2d")(input0, input1)
+    pattern = is_op("relax.add")(conv, input2)
+    return ("nnapi.conv2d", pattern, {})
+
+
+def max_pool2d_pattern() -> Tuple[str, DFPattern, Mapping[str, DFPattern]]:
+    """
+    Returns a tuple representing max_pool2d operation patterns mapped
+    between NNAPI and Relax frameworks.
+    """
+    input0 = wildcard()
+    pattern = is_op("relax.nn.max_pool2d")(input0)
+    return ("nnapi.max_pool_2d", pattern, {})
+
+
+register_patterns(
+    [
+        *elementwise_binary_patterns(),
+        *unary_patterns(),
+        matmul_pattern(),
+        permute_dims_pattern(),
+        astype_pattern(),
+        mean_pattern(),
+        conv2d_pattern(),
+        max_pool2d_pattern(),
+    ]
+)
+
+
+def min_feature_level(pattern_name: str) -> int:
+    """
+    Returns the minimum feature level required to support a given NNAPI operation pattern.
+
+    Args:
+        pattern_name (str): The name of the NNAPI operation pattern
+        (e.g., "nnapi.add", "nnapi.conv2d").
+
+    Returns:
+        int: The minimum feature level for the specified pattern, or 1 if the pattern is not found.
+    """
+
+    levels = {
+        "nnapi.add": 1,
+        "nnapi.average_pool_2d": 1,
+        "nnapi.concatenation": 1,
+        "nnapi.conv2d": 1,
+        "nnapi.depthwise_conv_2d": 1,
+        "nnapi.depth_to_space": 1,
+        "nnapi.dequantize": 1,
+        "nnapi.embedding_lookup": 1,
+        "nnapi.floor": 1,
+        "nnapi.fully_connected": 1,
+        "nnapi.hashtable_lookup": 1,
+        "nnapi.l2_normalization": 1,
+        "nnapi.l2_pool_2d": 1,
+        "nnapi.local_response_normalization": 1,
+        "nnapi.logistic": 1,
+        "nnapi.lsh_projection": 1,
+        "nnapi.lstm": 1,
+        "nnapi.max_pool_2d": 1,
+        "nnapi.mul": 1,
+        "nnapi.relu": 1,
+        "nnapi.relu1": 1,
+        "nnapi.relu6": 1,
+        "nnapi.reshape": 1,
+        "nnapi.resize_bilinear": 1,
+        "nnapi.rnn": 1,
+        "nnapi.softmax": 1,
+        "nnapi.space_to_depth": 1,
+        "nnapi.svdf": 1,
+        "nnapi.tanh": 1,
+        "nnapi.batch_to_space_nd": 2,
+        "nnapi.div": 2,
+        "nnapi.mean": 2,
+        "nnapi.pad": 2,
+        "nnapi.space_to_batch_nd": 2,
+        "nnapi.squeeze": 2,
+        "nnapi.strided_slice": 2,
+        "nnapi.sub": 2,
+        "nnapi.transpose": 2,
+        "nnapi.abs": 3,
+        "nnapi.argmax": 3,
+        "nnapi.argmin": 3,
+        "nnapi.axis_aligned_bbox_transform": 3,
+        "nnapi.bidirectional_sequence_lstm": 3,
+        "nnapi.bidirectional_sequence_rnn": 3,
+        "nnapi.box_with_nms_limit": 3,
+        "nnapi.cast": 3,
+        "nnapi.channel_shuffle": 3,
+        "nnapi.detection_postprocessing": 3,
+        "nnapi.equal": 3,
+        "nnapi.exp": 3,
+        "nnapi.expand_dims": 3,
+        "nnapi.gather": 3,
+        "nnapi.generate_proposals": 3,
+        "nnapi.greater": 3,
+        "nnapi.greater_equal": 3,
+        "nnapi.grouped_conv_2d": 3,
+        "nnapi.heatmap_max_keypoint": 3,
+        "nnapi.instance_normalization": 3,
+        "nnapi.less": 3,
+        "nnapi.less_equal": 3,
+        "nnapi.log": 3,
+        "nnapi.logical_and": 3,
+        "nnapi.logical_not": 3,
+        "nnapi.logical_or": 3,
+        "nnapi.log_softmax": 3,
+        "nnapi.maximum": 3,
+        "nnapi.minimum": 3,
+        "nnapi.neg": 3,
+        "nnapi.not_equal": 3,
+        "nnapi.pad_v2": 3,
+        "nnapi.pow": 3,
+        "nnapi.prelu": 3,
+        "nnapi.quantize": 3,
+        "nnapi.quantized_16bit_lstm": 3,
+        "nnapi.random_multinomial": 3,
+        "nnapi.reduce_all": 3,
+        "nnapi.reduce_any": 3,
+        "nnapi.reduce_max": 3,
+        "nnapi.reduce_min": 3,
+        "nnapi.reduce_prod": 3,
+        "nnapi.reduce_sum": 3,
+        "nnapi.roi_align": 3,
+        "nnapi.roi_pooling": 3,
+        "nnapi.rsqrt": 3,
+        "nnapi.select": 3,
+        "nnapi.sin": 3,
+        "nnapi.slice": 3,
+        "nnapi.split": 3,
+        "nnapi.sqrt": 3,
+        "nnapi.tile": 3,
+        "nnapi.topk_v2": 3,
+        "nnapi.transpose_conv_2d": 3,
+        "nnapi.unidirectional_sequence_lstm": 3,
+        "nnapi.unidirectional_sequence_rnn": 3,
+        "nnapi.resize_nearest_neighbor": 3,
+        "nnapi.quantized_lstm": 4,
+        "nnapi.if": 4,
+        "nnapi.while": 4,
+        "nnapi.elu": 4,
+        "nnapi.hard_swish": 4,
+        "nnapi.fill": 4,
+        "nnapi.rank": 4,
+        "nnapi.batch_matmul": 6,
+        "nnapi.pack": 6,
+        "nnapi.mirror_pad": 7,
+        "nnapi.reverse": 7,
+    }
+    return levels[pattern_name]
+
+
+def partition_for_nnapi(mod: IRModule, feature_level: Optional[int] = None) -> IRModule:
+    """Partition the graph greedily offloading supported operators to NNAPI.
+
+    Parameters
+    ----------
+    mod : tvm.ir.IRModule
+        The module to run passes on.
+    feature_level : Optional[int]
+        The maximum NNAPI feature level.
+
+    Returns
+    -------
+    mod : tvm.ir.IRModule
+        Annotated and partitioned module.
+    """
+    patterns = get_patterns_with_prefix("nnapi")
+    if feature_level is not None:
+        patterns = [pat for pat in patterns if feature_level >= min_feature_level(pat.name)]
+    mod = FuseOpsByPattern(patterns, bind_constants=False, annotate_codegen=False)(mod)
+    mod = MergeCompositeFunctions()(mod)
+    return mod