[BYOC][ETHOSN] Introduce the Ethos-N BYOC integration

CMakeLists.txt

@@ -8,6 +8,7 @@ include(cmake/util/FindOpenCL.cmake)
 include(cmake/util/FindVulkan.cmake)
 include(cmake/util/FindLLVM.cmake)
 include(cmake/util/FindROCM.cmake)
+include(cmake/util/FindEthosN.cmake)
 
 if(EXISTS ${CMAKE_CURRENT_BINARY_DIR}/config.cmake)
   include(${CMAKE_CURRENT_BINARY_DIR}/config.cmake)
@@ -44,6 +45,7 @@ tvm_option(INSTALL_DEV "Install compiler infrastructure" OFF)
 tvm_option(HIDE_PRIVATE_SYMBOLS "Compile with -fvisibility=hidden." OFF)
 tvm_option(USE_TF_TVMDSOOP "Build with TensorFlow TVMDSOOp" OFF)
 tvm_option(USE_FALLBACK_STL_MAP "Use TVM's POD compatible Map" OFF)
+tvm_option(USE_ETHOSN "Build with Arm Ethos-N" OFF)
 
 # 3rdparty libraries
 tvm_option(DLPACK_PATH "Path to DLPACK" "3rdparty/dlpack/include")
@@ -307,6 +309,7 @@ include(cmake/modules/Metal.cmake)
 include(cmake/modules/ROCM.cmake)
 include(cmake/modules/LLVM.cmake)
 include(cmake/modules/Micro.cmake)
+include(cmake/modules/contrib/EthosN.cmake)
 include(cmake/modules/contrib/BLAS.cmake)
 include(cmake/modules/contrib/CODEGENC.cmake)
 include(cmake/modules/contrib/DNNL.cmake)

cmake/config.cmake

@@ -210,6 +210,16 @@ set(USE_DNNL_CODEGEN OFF)
 set(USE_ARM_COMPUTE_LIB OFF)
 set(USE_ARM_COMPUTE_LIB_GRAPH_RUNTIME OFF)
 
+# Whether to build with Arm Ethos-N support
+# Possible values:
+# - OFF: disable Arm Ethos-N support
+# - path/to/arm-ethos-N-stack: use a specific version of the
+#   Ethos-N driver stack
+set(USE_ETHOSN OFF)
+# If USE_ETHOSN is enabled, use ETHOSN_HW (ON) if Ethos-N hardware is available on this machine
+# otherwise use ETHOSN_HW (OFF) to use the software test infrastructure
+set(USE_ETHOSN_HW OFF)
+
 # Build ANTLR parser for Relay text format
 # Possible values:
 # - ON: enable ANTLR by searching default locations (cmake find_program for antlr4 and /usr/local for jar)

cmake/modules/contrib/EthosN.cmake

@@ -0,0 +1,57 @@
+# 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.
+
+# Arm Ethos-N rules
+
+if(NOT USE_ETHOSN STREQUAL "OFF")
+  find_ethosn(${USE_ETHOSN})
+
+  if(NOT ETHOSN_FOUND)
+    message(FATAL_ERROR "Cannot find Ethos-N, USE_ETHOSN=" ${USE_ETHOSN})
+
+  else()
+    include_directories(${ETHOSN_INCLUDE_DIRS})
+    add_definitions(${ETHOSN_DEFINITIONS})
+
+    message(STATUS "Build with Ethos-N ${ETHOSN_PACKAGE_VERSION}")
+
+    file(GLOB ETHOSN_RUNTIME_CONTRIB_SRC
+      CONFIGURE_DEPENDS src/runtime/contrib/ethosn/ethosn_runtime.cc
+      CONFIGURE_DEPENDS src/runtime/contrib/ethosn/ethosn_device.cc)
+    list(APPEND RUNTIME_SRCS ${ETHOSN_RUNTIME_CONTRIB_SRC})
+
+    file(GLOB COMPILER_ETHOSN_SRCS
+      CONFIGURE_DEPENDS src/relay/backend/contrib/ethosn/*)
+    list(APPEND COMPILER_SRCS ${COMPILER_ETHOSN_SRCS})
+
+    list(APPEND TVM_LINKER_LIBS ${ETHOSN_COMPILER_LIBRARY}
+      ${ETHOSN_RUNTIME_LIBRARY})
+    list(APPEND TVM_RUNTIME_LINKER_LIBS ${ETHOSN_COMPILER_LIBRARY}
+      ${ETHOSN_RUNTIME_LIBRARY})
+
+    if(NOT MSVC)
+      set_source_files_properties(${COMPILER_ETHOSN_SRCS}
+        PROPERTIES COMPILE_DEFINITIONS "DMLC_ENABLE_RTTI=0")
+      set_source_files_properties(${COMPILER_ETHOSN_SRCS}
+        PROPERTIES COMPILE_FLAGS "-fno-rtti")
+    endif()
+  endif(NOT ETHOSN_FOUND)
+else()
+  if(USE_ETHOSN_HW)
+    message(FATAL_ERROR "Cannot enable Ethos-N HW if USE_ETHOSN=OFF")
+  endif()
+endif(NOT USE_ETHOSN STREQUAL "OFF")

cmake/util/FindEthosN.cmake

@@ -0,0 +1,94 @@
+# 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.
+
+#######################################################
+# Find Arm Ethos-N libraries
+#
+# Usage:
+#   find_ethosn(${USE_ETHOSN})
+#
+# - When USE_ETHOSN=/path/to/ethos-sdk-path, use the path from USE_ETHOSN
+# - Else, when environment variable ETHOSN_STACK is set, use that path
+# - When USE_ETHOSN=ON, use auto search
+#
+# Provide variables:
+#
+# - ETHOSN_FOUND
+# - ETHOSN_PACKAGE_VERSION
+# - ETHOSN_DEFINITIONS
+# - ETHOSN_INCLUDE_DIRS
+# - ETHOSN_COMPILER_LIBRARY
+# - ETHOSN_RUNTIME_LIBRARY
+
+macro(find_ethosn use_ethosn)
+  set(__use_ethosn ${use_ethosn})
+  if(IS_DIRECTORY ${__use_ethosn})
+    set(__ethosn_stack ${__use_ethosn})
+    message(STATUS "Arm Ethos-N driver stack PATH=" ${__use_ethosn})
+  elseif(IS_DIRECTORY $ENV{ETHOSN_STACK})
+     set(__ethosn_stack $ENV{ETHOSN_STACK})
+    message(STATUS "Arm Ethos-N driver stack from env=" ${__use_ethosn})
+  else()
+     set(__ethosn_stack "")
+  endif()
+
+  if(__ethosn_stack)
+    set(ETHOSN_INCLUDE_DIRS "")
+    # Compile-time support
+    find_path(_SL_DIR NAMES Support.hpp
+      PATHS ${__ethosn_stack}/include/ethosn_support_library)
+    string(REGEX REPLACE "/ethosn_support_library" "" _SL_DIR2 ${_SL_DIR})
+    list(APPEND ETHOSN_INCLUDE_DIRS "${_SL_DIR2}")
+
+    find_library(ETHOSN_COMPILER_LIBRARY NAMES EthosNSupport
+      PATHS ${__ethosn_stack}/lib)
+    find_library(ETHOSN_COMPILER_LIBRARY NAMES EthosNSupport)
+
+    set(ETHOSN_PACKAGE_VERSION "0.1.1")
+
+    if(USE_ETHOSN_HW STREQUAL "ON")
+      # Runtime hardware support
+      find_path(_DL_DIR NAMES Network.hpp
+        PATHS ${__ethosn_stack}/include/ethosn_driver_library)
+      string(REGEX REPLACE "/ethosn_driver_library" "" _DL_DIR2 ${_DL_DIR})
+      list(APPEND ETHOSN_INCLUDE_DIRS "${_DL_DIR2}")
+
+      find_library(ETHOSN_RUNTIME_LIBRARY NAMES EthosNDriver
+        PATHS ${__ethosn_stack}/lib)
+      find_library(ETHOSN_RUNTIME_LIBRARY NAMES EthosNDriver)
+      set(ETHOSN_DEFINITIONS -DETHOSN_HW)
+    endif ()
+
+    if(ETHOSN_COMPILER_LIBRARY)
+      set(ETHOSN_FOUND TRUE)
+    endif()
+  endif(__ethosn_stack)
+
+  if(NOT ETHOSN_FOUND)
+    if(__use_ethosn STREQUAL "ON")
+      message(WARNING "No cmake find_package available for Arm Ethos-N")
+    endif()
+
+  # additional libraries
+  else()
+    message(STATUS "Found ETHOSN_DEFINITIONS=${ETHOSN_DEFINITIONS}")
+    message(STATUS "Found ETHOSN_INCLUDE_DIRS=${ETHOSN_INCLUDE_DIRS}")
+    message(STATUS "Found ETHOSN_COMPILER_LIBRARY=${ETHOSN_COMPILER_LIBRARY}")
+    message(STATUS "Found ETHOSN_RUNTIME_LIBRARY=${ETHOSN_RUNTIME_LIBRARY}")
+  endif(NOT ETHOSN_FOUND)
+
+endmacro(find_ethosn)

python/tvm/relay/op/contrib/__init__.py

@@ -21,3 +21,4 @@
 from .arm_compute_lib import *
 from .dnnl import *
 from .coreml import *
+from .ethosn import *

python/tvm/relay/op/contrib/_ethosn.py

@@ -0,0 +1,22 @@
+# 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.
+
+"""Expose 'is supported' functions to Python."""
+
+import tvm._ffi
+
+tvm._ffi._init_api("relay.ethos-n.support", __name__)

python/tvm/relay/op/contrib/ethosn.py

@@ -0,0 +1,89 @@
+# 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.
+# pylint: disable=invalid-name, unused-argument
+"""Arm(R) Ethos(TM) -N NPU supported operators."""
+from enum import Enum
+import tvm.ir
+from ... import qnn as _qnn
+from . import _ethosn as support
+
+
+class Available(Enum):
+    UNAVAILABLE = 0
+    SW_ONLY = 1
+    SW_AND_HW = 2
+
+    def __bool__(self):
+        return self != Available.UNAVAILABLE
+
+
+def ethosn_available():
+    """Return whether Ethos-N software and hardware support is available"""
+    if not tvm.get_global_func("relay.ethos-n.query", True):
+        print("skip because Ethos-N module is not available")
+        return Available.UNAVAILABLE
+    hw = tvm.get_global_func("relay.ethos-n.query")()
+    return Available.SW_AND_HW if hw else Available.SW_ONLY
+
+
+@tvm.ir.register_op_attr("qnn.concatenate", "target.ethos-n")
+def qnn_concatenate(attrs, args):
+    """Check if a concatenate is supported by Ethos-N."""
+    if not ethosn_available():
+        return False
+
+    conc = _qnn.op.concatenate(*args, **attrs)
+    if not support.concatenate(conc):
+        return False
+
+    # Support library has some unenforced restrictions on qnn params
+    min_range = 1e9
+    max_range = -1e9
+    qnn_params = []
+    for i in range(len(args[1].fields)):
+        scale = args[1].fields[i].data.asnumpy()
+        zero_point = args[2].fields[i].data.asnumpy()
+        min_range = min(-1 * zero_point * scale, min_range)
+        max_range = max((255 - zero_point) * scale, max_range)
+        qnn_params.append((scale, zero_point))
+
+    scale = (max_range - min_range) / 255
+    zero_point = int(-min_range/scale)
+    if (scale, zero_point) in qnn_params:
+        return True
+
+    return False
+
+
+@tvm.ir.register_op_attr("split", "target.ethos-n")
+def split(attrs, args):
+    """Check if a split is supported by Ethos-N."""
+    if not ethosn_available():
+        return False
+
+    if isinstance(attrs["indices_or_sections"], tvm.tir.IntImm):
+        sp = tvm.relay.split(*args,
+                             indices_or_sections=attrs["indices_or_sections"].value,
+                             axis=attrs["axis"])
+    else:
+        sp = tvm.relay.split(*args,
+                             indices_or_sections=attrs["indices_or_sections"],
+                             axis=attrs["axis"])
+    if not support.split(sp.astuple()):
+        return False
+
+    return True

src/relay/backend/contrib/ethosn/capabilities.h

@@ -0,0 +1,81 @@
+/*
+ * 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.
+ */
+
+/*!
+ * \file src/relay/backend/contrib/ethosn/capabilities.h
+ * \brief The Ethos-N processor series has four variants, the Ethos-N37, Ethos-N57, Ethos-N77
+ * and the Ethos-N78. This release of the integration supports the first three variants.
+ * Configuration information for each variant is stored as a blob in this file. These blobs
+ * are passed into the Ethos-N support library, which in turn uses them to optimize the
+ * generated command-stream appropriately for the specified variant.
+ */
+
+#ifndef TVM_RELAY_BACKEND_CONTRIB_ETHOSN_CAPABILITIES_H_
+#define TVM_RELAY_BACKEND_CONTRIB_ETHOSN_CAPABILITIES_H_
+
+#include <vector>
+
+namespace tvm {
+namespace relay {
+namespace contrib {
+namespace ethosn {
+
+/* Ethos-N variants (N77, N57 and N37)
+ * variant[0] - N77
+ * variant[1] - N57
+ * variant[2] - N37
+ */
+static std::vector<char> variants[3] = {
+    {
+        0x02, 0x00, 0x00, 0x00, 0x74, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x00, 0x10, 0x00,
+        0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00,
+        0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00,
+        0x01, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00,
+        0x00, 0x01, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x01, 0x00,
+        0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x01,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    },
+    {
+        0x02, 0x00, 0x00, 0x00, 0x74, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x00, 0x10, 0x00,
+        0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00,
+        0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00,
+        0x01, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00,
+        0x00, 0x01, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x01, 0x00,
+        0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x01,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    },
+    {
+        0x02, 0x00, 0x00, 0x00, 0x74, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+        0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x10, 0x00, 0x10, 0x00,
+        0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x01, 0x00, 0x00, 0x00, 0x00,
+        0x10, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00,
+        0x01, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x10, 0x00, 0x00,
+        0x00, 0x01, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x04, 0x00, 0x00, 0x00, 0x01, 0x00,
+        0x00, 0x00, 0x08, 0x00, 0x00, 0x00, 0x40, 0x00, 0x00, 0x00, 0x00, 0x02, 0x00, 0x00, 0x01,
+        0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
+    }};
+
+}  // namespace ethosn
+}  // namespace contrib
+}  // namespace relay
+}  // namespace tvm
+
+#endif  // TVM_RELAY_BACKEND_CONTRIB_ETHOSN_CAPABILITIES_H_