[AOT][DeviceAPI] Connecting devices structure to relevant operators

This patch adds support for passing the device context via the unpacked API in AOT, generating an additional struct if necessary:

```c
/*!
 * \brief Device context pointers for TVM module "default"
 */
struct tvmgen_default_devices {
  void* npu;
};
```

Which is then added as an argument to the entry function:
```c
/*!
 * \brief entrypoint function for TVM module "default"
 * \param inputs Input tensors for the module
 * \param outputs Output tensors for the module
 * \param devices Device context pointers for the module
 */
int32_t tvmgen_default_run(
  struct tvmgen_default_inputs* inputs,
  struct tvmgen_default_outputs* outputs,
  struct tvmgen_default_devices* devices
);
```

I've temporarily added the collection of external code generators to the TE compiler pending proper annotation of the eventual functions.

Co-authored-by: Grant Watson <grant.watson@arm.com>

apps/microtvm/ethosu/Makefile

@@ -35,7 +35,7 @@ RANLIB = arm-none-eabi-ranlib
 PKG_CFLAGS = ${PKG_COMPILE_OPTS} \
 	-I${STANDALONE_CRT_PATH}/include \
 	-I${STANDALONE_CRT_PATH}/src/runtime/crt/include \
-	-Iinclude \
+	-I${PWD}/include \
 	-I${CORSTONE_300_PATH} \
 	-I${ETHOSU_PATH}/core_driver/include \
 	-I${CMSIS_PATH}/Device/ARM/${ARM_CPU}/Include/ \
@@ -95,7 +95,7 @@ ${BUILD_DIR}/ethosu_core_driver/libethosu_core_driver.a:
 	$(QUIET)cd $(abspath $(BUILD_DIR)/ethosu_core_driver) && $(MAKE)
 
 # Build demo application
-$(BUILD_DIR)/demo: src/demo.c $(BUILD_DIR)/stack_allocator.o $(BUILD_DIR)/crt_backend_api.o ${BUILD_DIR}/libcodegen.a ${BUILD_DIR}/libcmsis_startup.a ${BUILD_DIR}/ethosu_core_driver/libethosu_core_driver.a ${BUILD_DIR}/libuart.a
+$(BUILD_DIR)/demo: src/demo.c src/tvm_ethosu_runtime.c $(BUILD_DIR)/stack_allocator.o $(BUILD_DIR)/crt_backend_api.o ${BUILD_DIR}/libcodegen.a ${BUILD_DIR}/libcmsis_startup.a ${BUILD_DIR}/ethosu_core_driver/libethosu_core_driver.a ${BUILD_DIR}/libuart.a
 	$(QUIET)mkdir -p $(@D)
 	$(QUIET)$(CC) $(PKG_CFLAGS) -o $@ $^ $(PKG_LDFLAGS)
 

apps/microtvm/ethosu/include/tvm_ethosu_runtime.h

@@ -0,0 +1,30 @@
+/*
+ * 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.
+ */
+
+#ifndef TVM_RUNTIME_CONTRIB_ETHOSU_ETHOSU_RUNTIME_H_
+#define TVM_RUNTIME_CONTRIB_ETHOSU_ETHOSU_RUNTIME_H_
+
+#include <ethosu_driver.h>
+#include <stddef.h>
+#include <stdint.h>
+
+int32_t TVMEthosULaunch(struct ethosu_driver* resource_handle, void* cms_data, size_t cms_data_size,
+                        uint64_t* base_addrs, size_t* base_addrs_size, int num_tensors);
+
+#endif  // TVM_RUNTIME_CONTRIB_ETHOSU_ETHOSU_RUNTIME_H_

apps/microtvm/ethosu/src/demo.c

@@ -45,7 +45,12 @@ int main(int argc, char** argv) {
   struct tvmgen_default_inputs inputs = {
       .input = input,
   };
-  tvmgen_default_run(&inputs, &outputs);
+  struct ethosu_driver* driver = ethosu_reserve_driver();
+  struct tvmgen_default_devices devices = {
+      .ethos_u = driver,
+  };
+  tvmgen_default_run(&inputs, &outputs, &devices);
+  ethosu_release_driver(driver);
 
   // Calculate index of max value
   uint8_t max_value = 0;

apps/microtvm/ethosu/src/tvm_ethosu_runtime.c

@@ -0,0 +1,34 @@
+/*
+ * 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.
+ */
+
+#include "tvm_ethosu_runtime.h"
+
+#include <ethosu_driver.h>
+
+int32_t TVMEthosULaunch(struct ethosu_driver* driver, void* cms_data, size_t cms_data_size,
+                        uint64_t* base_addrs, size_t* base_addrs_size, int num_tensors) {
+  int32_t result =
+      ethosu_invoke(driver, cms_data, cms_data_size, base_addrs, base_addrs_size, num_tensors);
+
+  // Map errors in invoke to TVM errors
+  if (result != 0) {
+    return -1;
+  }
+  return 0;
+}

python/tvm/driver/tvmc/common.py

@@ -27,6 +27,7 @@
 
 import tvm
 
+from tvm.driver import tvmc
 from tvm import relay
 from tvm import transform
 from tvm._ffi import registry
@@ -206,6 +207,7 @@ def parse_target(target):
         a key-value for all options passed via CLI; 'raw',
         containing the plain string for this codegen
     """
+    codegen_names = tvmc.composite_target.get_codegen_names()
     codegens = []
 
     tvm_target_kinds = tvm.target.Target.list_kinds()
@@ -232,7 +234,7 @@ def parse_target(target):
     for codegen_def in split_codegens:
         # the first is expected to be the name
         name = codegen_def[0]
-        is_tvm_target = name in tvm_target_kinds
+        is_tvm_target = name in tvm_target_kinds and name not in codegen_names
         raw_target = " ".join(codegen_def)
         all_opts = codegen_def[1:] if len(codegen_def) > 1 else []
         opts = {}

python/tvm/driver/tvmc/composite_target.py

@@ -65,7 +65,7 @@
         "pass_pipeline": partition_for_ethosn78,
     },
     "ethos-u": {
-        "config_key": "relay.ext.ethosu.options",
+        "config_key": "relay.ext.ethos-u.options",
         "pass_pipeline": partition_for_ethosu,
     },
     "bnns": {

python/tvm/driver/tvmc/target.py

@@ -18,6 +18,7 @@
 This file contains functions for processing target inputs for the TVMC CLI
 """
 
+from tvm.driver import tvmc
 from tvm.target import Target
 
 # We can't tell the type inside an Array but all current options are strings so
@@ -27,6 +28,11 @@
 INTERNAL_TO_HELP = {"runtime.String": " string", "IntImm": "", "Array": " options"}
 
 
+def _valid_target_kinds():
+    codegen_names = tvmc.composite_target.get_codegen_names()
+    return filter(lambda target: target not in codegen_names, Target.list_kinds())
+
+
 def _generate_target_kind_args(parser, kind):
     target_group = parser.add_argument_group(f"target {kind.name}")
     for target_option, target_type in kind.options.items():
@@ -45,8 +51,7 @@ def generate_target_args(parser):
         help="compilation target as plain string, inline JSON or path to a JSON file",
         required=True,
     )
-    target_kinds = Target.list_kinds()
-    for target_kind in target_kinds:
+    for target_kind in _valid_target_kinds():
         target = Target(target_kind)
         _generate_target_kind_args(parser, target.kind)
 
@@ -55,7 +60,7 @@ def _reconstruct_target_kind_args(args, kind):
     kind_options = {}
     for target_option, target_type in kind.options.items():
         if target_type in INTERNAL_TO_NATIVE_TYPE:
-            var_name = f"target_{kind.name}_{target_option.replace('-', '_')}"
+            var_name = f"target_{kind.name.replace('-', '_')}_{target_option.replace('-', '_')}"
             option_value = getattr(args, var_name)
             if option_value is not None:
                 kind_options[target_option] = getattr(args, var_name)
@@ -64,9 +69,8 @@ def _reconstruct_target_kind_args(args, kind):
 
 def reconstruct_target_args(args):
     """Reconstructs the target options from the arguments"""
-    target_kinds = Target.list_kinds()
     reconstructed = {}
-    for target_kind in target_kinds:
+    for target_kind in _valid_target_kinds():
         target = Target(target_kind)
         kind_options = _reconstruct_target_kind_args(args, target.kind)
         if kind_options:

python/tvm/micro/model_library_format.py

@@ -44,13 +44,13 @@ class UnsupportedInModelLibraryFormatError(Exception):
     """Raised when export_model_library_format does not support the given Module tree."""
 
 
-def generate_c_interface_header(module_name, inputs, outputs, include_path):
+def generate_c_interface_header(module_name, inputs, outputs, devices, include_path):
     """Generate C Interface header to be included in MLF"""
     mangled_name = to_c_variable_style(prefix_generated_name(module_name))
     metadata_header = os.path.join(include_path, f"{mangled_name}.h")
 
     interface_c_create = tvm._ffi.get_global_func("runtime.InterfaceCCreate")
-    interface_c_module = interface_c_create(module_name, inputs, outputs)
+    interface_c_module = interface_c_create(module_name, inputs, outputs, devices)
 
     with open(metadata_header, "w") as header_file:
         header_file.write(interface_c_module.get_source())
@@ -318,7 +318,8 @@ def _export_graph_model_library_format(
         include_path = codegen_dir / "host" / "include"
         include_path.mkdir()
         inputs, outputs = _get_inputs_and_outputs_from_module(mod)
-        generate_c_interface_header(mod.libmod_name, inputs, outputs, include_path)
+        devices = mod.get_devices()
+        generate_c_interface_header(mod.libmod_name, inputs, outputs, devices, include_path)
 
     parameters_dir = tempdir / "parameters"
     parameters_dir.mkdir()

python/tvm/relay/backend/contrib/ethosu/_ffi_api.py

@@ -17,4 +17,4 @@
 """FFI APIs for relay transformation passes."""
 import tvm._ffi  # type: ignore
 
-tvm._ffi._init_api("relay.ext.ethosu", __name__)
+tvm._ffi._init_api("relay.ext.ethos-u", __name__)

python/tvm/relay/backend/contrib/ethosu/codegen.py

@@ -24,7 +24,7 @@
 from tvm.relay.backend.contrib.ethosu import util
 
 
-@tvm._ffi.register_func("relay.ext.ethosu")
+@tvm._ffi.register_func("relay.ext.ethos-u")
 def ethosu_compiler(external_function):
     """The entry-point to a compile a external relay function of
     NPU compatible operators to generated command stream.
@@ -38,11 +38,11 @@ def ethosu_compiler(external_function):
     input_size = util.calculate_size_bytes(external_function.params[0])
     output_size = util.calculate_size_bytes(external_function.body)
     cmms, encoded_constants, scratch_size = _compile(external_function)
-    ethosu_runtime = tvm._ffi.get_global_func("runtime.module.ethosu.create")
+    ethosu_runtime = tvm._ffi.get_global_func("runtime.module.ethos-u.create")
     return ethosu_runtime(func_name, cmms, encoded_constants, scratch_size, input_size, output_size)
 
 
-@tvm._ffi.register_func("relay.ext.ethosu.constant_updater")
+@tvm._ffi.register_func("relay.ext.ethos-u.constant_updater")
 def constant_updater(expr, symbol):  # pylint: disable=unused-argument
     """
     The constant updater process happen after lowering in the core compiler.

python/tvm/relay/backend/contrib/ethosu/legalize.py

@@ -126,7 +126,7 @@ class EthosUConv2DRewriter(DFPatternCallback):
 
     def __init__(self):
         super().__init__(require_type=True)
-        self.pattern = (wildcard().has_attr({"Composite": "ethosu.qnn_conv2d"}))(wildcard())
+        self.pattern = (wildcard().has_attr({"Composite": "ethos-u.qnn_conv2d"}))(wildcard())
 
     def callback(
         self, pre: tvm.relay.Expr, post: tvm.relay.Expr, node_map: tvm.ir.container.Map

python/tvm/relay/backend/contrib/ethosu/util.py

@@ -139,7 +139,7 @@ def round_up(a: int, b: int) -> int:
 
 def get_accelerator_config():
     """Get the variant of the accelerator to compile for"""
-    compiler_attrs = tvm.get_global_func("relay.ext.ethosu.get_compiler_attrs")()
+    compiler_attrs = tvm.get_global_func("relay.ext.ethos-u.get_compiler_attrs")()
     return compiler_attrs.accelerator_config
 
 

python/tvm/relay/backend/contrib/ethosu/vela_api.py

@@ -381,7 +381,7 @@ def get_accelerator_config() -> vapi.NpuAccelerator:
         "ethos-u55-64": vapi.NpuAccelerator.Ethos_U55_64,
         "ethos-u55-32": vapi.NpuAccelerator.Ethos_U55_32,
     }
-    compiler_attrs = tvm.get_global_func("relay.ext.ethosu.get_compiler_attrs")()
+    compiler_attrs = tvm.get_global_func("relay.ext.ethos-u.get_compiler_attrs")()
     accel_config_str = compiler_attrs.accelerator_config
     assert accel_config_str in npu_accel_str_map.keys(), f"{accel_config_str} is not supported"
     return npu_accel_str_map[accel_config_str]

python/tvm/relay/backend/executor_factory.py

@@ -83,16 +83,22 @@ class AOTExecutorFactoryModule(ExecutorFactoryModule):
         The parameters of module
     function_metadata : Map of String to FunctionInfo
         This holds a map function names to their information
+    devices : List[str]
+        List of devices used in the module
     """
 
-    def __init__(self, ir_mod, target, libmod, libmod_name, params, function_metadata):
+    def __init__(self, ir_mod, target, libmod, libmod_name, params, function_metadata, devices):
         self.ir_mod = ir_mod
         self.target = target
         self.lib = libmod
         self.libmod_name = libmod_name
         self.params = params
         self.iter_cnt = 0
         self.function_metadata = function_metadata
+        self.devices = devices
+
+    def get_devices(self):
+        return self.devices
 
     def get_params(self):
         return self.params
@@ -148,6 +154,9 @@ def __init__(
     def export_library(self, file_name, fcompile=None, addons=None, **kwargs):
         return self.module.export_library(file_name, fcompile, addons, **kwargs)
 
+    def get_devices(self):
+        return []
+
     def get_params(self):
         return self.params
 

python/tvm/relay/build_module.py

@@ -101,6 +101,7 @@ def __init__(self):
         self._set_params_func = self.mod["set_params"]
         self._get_params_func = self.mod["get_params"]
         self._get_function_metadata = self.mod["get_function_metadata"]
+        self._get_devices = self.mod["get_devices"]
 
     def build(
         self, mod, target=None, target_host=None, params=None, executor="graph", mod_name=None
@@ -231,6 +232,10 @@ def get_function_metadata(self):
         each PrimFunc"""
         return self._get_function_metadata()
 
+    def get_devices(self):
+        """Returns a list of devices configured in this module"""
+        return self._get_devices()
+
     def get_params(self):
         """Return the updated weights."""
         params = self._get_params_func()
@@ -358,14 +363,21 @@ def build(ir_mod, target=None, target_host=None, params=None, mod_name="default"
             mod=ir_mod, target=target, params=params, executor=executor, mod_name=mod_name
         )
         func_metadata = bld_mod.get_function_metadata()
+        devices = bld_mod.get_devices()
 
         if executor == "aot":
             executor_factory = _executor_factory.AOTExecutorFactoryModule(
-                ir_mod, target, runtime_mod, mod_name, params, func_metadata
+                ir_mod, target, runtime_mod, mod_name, params, func_metadata, devices
             )
         elif executor == "graph":
             executor_factory = _executor_factory.GraphExecutorFactoryModule(
-                ir_mod, target, executor_config, runtime_mod, mod_name, params, func_metadata
+                ir_mod,
+                target,
+                executor_config,
+                runtime_mod,
+                mod_name,
+                params,
+                func_metadata,
             )
         else:
             assert False, "Executor " + executor + " not supported"

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

@@ -176,7 +176,7 @@ class QnnConv2DParams:
     and extract quantization information of all the associated tensors.
     """
 
-    composite_name = "ethosu.qnn_conv2d"
+    composite_name = "ethos-u.qnn_conv2d"
     # The NPU only supports padding upto the numbers as follows
     padding_bounds = [31, 31, 32, 32]
     activation_map = {"clip": "CLIP"}
@@ -265,7 +265,7 @@ class QnnDepthwiseConv2DParams(QnnConv2DParams):
     and extract the parameter information.
     """
 
-    composite_name = "ethosu.depthwise_conv2d"
+    composite_name = "ethos-u.depthwise_conv2d"
     # The hardware only supports padding upto the numbers as follows
     padding_bounds = [31, 31, 32, 32]
 
@@ -331,7 +331,7 @@ def qnn_depthwise_conv2d_pattern() -> tvm.relay.dataflow_pattern.DFPattern:
     return clip_or_req
 
 
-@register_pattern_table("ethosu")
+@register_pattern_table("ethos-u")
 def pattern_table() -> List[Tuple[str, tvm.relay.dataflow_pattern.DFPattern, Callable]]:
     return [
         (
@@ -371,10 +371,10 @@ def partition_for_ethosu(
     if params:
         mod["main"] = bind_params_by_name(mod["main"], params)
 
-    pattern = relay.op.contrib.get_pattern_table("ethosu")
+    pattern = relay.op.contrib.get_pattern_table("ethos-u")
     mod = relay.transform.InferType()(mod)
     mod = relay.transform.MergeComposite(pattern)(mod)
-    mod = relay.transform.AnnotateTarget("ethosu")(mod)
+    mod = relay.transform.AnnotateTarget("ethos-u")(mod)
     mod = relay.transform.MergeCompilerRegions()(mod)
     mod = relay.transform.InferType()(mod)
     mod = relay.transform.PartitionGraph()(mod)