Introduce --interface-api={c,packed} parameter

apps/microtvm/zephyr/aot_demo/src/main.c

@@ -32,6 +32,7 @@
 
 #include "input_data.h"
 #include "output_data.h"
+#include "tvmgen_default.h"
 #include "zephyr_uart.h"
 
 #ifdef CONFIG_ARCH_POSIX
@@ -194,18 +195,18 @@ void main(void) {
  }
  TVMLogf("Zephyr AOT Runtime\n");
 
-void* inputs[1] = {
- input_data,
+struct tvmgen_default_inputs inputs = {
+.input_1 = input_data,
  };
-void* outputs[1] = {
- output_data,
+struct tvmgen_default_outputs outputs = {
+.output = output_data,
  };
 
  StackMemoryManager_Init(&app_workspace, g_aot_memory, WORKSPACE_SIZE);
 
  double elapsed_time = 0;
  TVMPlatformTimerStart();
- int ret_val = tvm_runtime_run(&tvmgen_default_network, inputs, outputs);
+ int ret_val = tvmgen_default_run(&inputs, &outputs);
  TVMPlatformTimerStop(&elapsed_time);
 
  if (ret_val != 0) {

include/tvm/runtime/module.h

@@ -230,8 +230,10 @@ constexpr const char* tvm_module_main = "__tvm_main__";
 constexpr const char* tvm_param_prefix = "__tvm_param__";
 /*! \brief A PackedFunc that looks up linked parameters by storage_id. */
 constexpr const char* tvm_lookup_linked_param = "_lookup_linked_param";
-/*! \brief The main AOT executor function */
+/*! \brief The main AOT executor function generated from TIR */
 constexpr const char* tvm_run_func_suffix = "run_model";
+/*! \brief Model entrypoint generated as an interface to the AOT function outside of TIR */
+constexpr const char* tvm_entrypoint_suffix = "run";
 } // namespace symbol
 
 // implementations of inline functions.

python/tvm/micro/interface_api.py

@@ -0,0 +1,79 @@
+# 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.
+
+"""Defines functions for generating a C interface header"""
+
+import os
+
+from tvm.relay.backend.utils import mangle_module_name
+
+
+def _emit_brief(header_file, module_name, description):
+ header_file.write("/*!\n")
+ header_file.write(f' * \\brief {description} for TVM module "{module_name}" \n')
+ header_file.write(" */\n")
+
+
+def generate_c_interface_header(module_name, inputs, outputs, output_path):
+ """Generates a C interface header for a given modules inputs and outputs
+
+ Parameters
+ ----------
+ module_name : str
+ Name of the module to be used in defining structs and naming the header
+ inputs : list[str]
+ List of module input names to be placed in generated structs
+ outputs : list[str]
+ List of module output names to be placed in generated structs
+ output_path : str
+ Path to the output folder to generate the header into
+ """
+
+ mangled_name = mangle_module_name(module_name)
+ metadata_header = os.path.join(output_path, f"{mangled_name}.h")
+ with open(metadata_header, "w") as header_file:
+ header_file.write(
+ "#include <stdint.h>\n"
+ f"#ifndef {mangled_name.upper()}_H_\n"
+ f"#define {mangled_name.upper()}_H_\n"
+ )
+
+ _emit_brief(header_file, module_name, "Input tensor pointers")
+ header_file.write(f"struct {mangled_name}_inputs {{\n")
+ for input_name in inputs:
+ header_file.write(f" void* {input_name};\n")
+ header_file.write("};\n\n")
+
+ _emit_brief(header_file, module_name, "Output tensor pointers")
+ header_file.write(f"struct {mangled_name}_outputs {{\n")
+ for output_name in outputs:
+ header_file.write(f" void* {output_name};\n")
+ header_file.write("};\n\n")
+
+ header_file.write(
+ "/*!\n"
+ f' * \\brief entrypoint function for TVM module "{module_name}"\n'
+ " * \\param inputs Input tensors for the module \n"
+ " * \\param outputs Output tensors for the module \n"
+ " */\n"
+ f"int32_t {mangled_name}_run(\n"
+ f" struct {mangled_name}_inputs* inputs,\n"
+ f" struct {mangled_name}_outputs* outputs\n"
+ ");\n"
+ )
+
+ header_file.write(f"#endif // {mangled_name.upper()}_H_\n")

python/tvm/micro/model_library_format.py

@@ -25,7 +25,9 @@
 import tarfile
 import typing
 
+from tvm.ir.type import TupleType
 from .._ffi import get_global_func
+from .interface_api import generate_c_interface_header
 from ..contrib import utils
 from ..driver import build_module
 from ..runtime import ndarray as _nd
@@ -55,7 +57,6 @@ def _populate_codegen_dir(mod, codegen_dir: str, module_name: str = None):
 
  """
  dso_modules = mod._collect_dso_modules()
- dso_module_handles = [m.handle.value for m in dso_modules]
  non_dso_modules = mod._collect_from_import_tree(lambda m: m not in dso_modules)
  if non_dso_modules:
  raise UnsupportedInModelLibraryFormatError(
@@ -213,6 +214,39 @@ def _build_function_memory_map(function_metadata):
  return ret
 
 
+def _get_main_relay_func(mod: executor_factory.ExecutorFactoryModule):
+ main_func = mod.function_metadata[MAIN_FUNC_NAME_STR]
+ target = list(main_func.relay_primfuncs.keys())[0]
+ return main_func.relay_primfuncs[target]
+
+
+def _convert_tuple_to_outputs(ret_type, offset=0):
+ outputs = []
+ added_fields = len(ret_type.fields)
+ for output_index in range(added_fields):
+ next_output = offset + len(outputs)
+ if isinstance(ret_type.fields[output_index], TupleType):
+ outputs.extend(_convert_tuple_to_outputs(ret_type.fields[output_index], next_output))
+ else:
+ outputs.append(f"output{next_output}")
+ return outputs
+
+
+def _get_inputs_and_outputs_from_module(mod):
+ main_func = _get_main_relay_func(mod)
+ inputs = [argument.name_hint for argument in main_func.params]
+
+ outputs = ["output"]
+ if isinstance(main_func.ret_type, TupleType):
+ outputs = _convert_tuple_to_outputs(main_func.ret_type)
+
+ return inputs, outputs
+
+
+def _should_generate_interface_header(mod):
+ return any(target.attrs.get("interface-api") == "c" for target in mod.target.values())
+
+
 def _make_tar(source_dir, tar_file_path):
  """Build a tar file from source_dir."""
  with tarfile.open(tar_file_path, "w") as tar_f:
@@ -260,6 +294,12 @@ def _export_graph_model_library_format(
  codegen_dir.mkdir()
  _populate_codegen_dir(mod.lib, codegen_dir, mod.libmod_name)
 
+ if _should_generate_interface_header(mod):
+ 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)
+
  parameters_dir = tempdir / "parameters"
  parameters_dir.mkdir()
  param_filename = parameters_dir / f"{mod.libmod_name}.params"

src/relay/backend/aot_executor_codegen.cc

@@ -650,7 +650,7 @@ class AOTExecutorCodegen : public ExprVisitor {
  /*! \brief mod */
  runtime::Module* mod_;
  /*! \brief list of input expressions (i.e., variable passed by the user) */
- std::vector<Expr> input_vars_;
+ std::vector<Var> input_vars_;
  /*! \brief input and output variables belonging to the main function signature */
  Array<tir::Var> main_signature_;
  /*! \brief target device */
@@ -782,8 +782,12 @@ class AOTExecutorCodegen : public ExprVisitor {
  ret.lowered_funcs.Set(target_host_str, mod_run);
  }
  ret.function_metadata = std::move(function_metadata_);
- ret.metadata = runtime::Metadata(input_vars_.size(), return_sid_.size(),
- runtime::kTvmExecutorAot, mod_name);
+
+ std::vector<String> input_var_names(input_vars_.size());
+ std::transform(input_vars_.begin(), input_vars_.end(), input_var_names.begin(),
+ [](Var input_var) -> String { return input_var->name_hint(); });
+ ret.metadata =
+ runtime::Metadata(input_var_names, return_sid_.size(), runtime::kTvmExecutorAot, mod_name);
  return ret;
  }
 };

src/runtime/meta_data.h

@@ -54,8 +54,8 @@ inline String get_name_mangled(const String& module_name, const String& name) {
  */
 class MetadataNode : public Object {
  public:
- /*! \brief number of inputs of the main function */
-int num_inputs = 1;
+ /*! \brief input information for the main function */
+Array<String> inputs;
  /*! \brief number of outputs of the main function */
  int num_outputs = 1;
  /*! \brief the executor to be used to run the model */
@@ -73,9 +73,9 @@ class MetadataNode : public Object {
  */
 class Metadata : public ObjectRef {
  public:
- TVM_DLL Metadata(int num_inputs, int num_outputs, String executor, String mod_name) {
+ TVM_DLL Metadata(Array<String> inputs, int num_outputs, String executor, String mod_name) {
  auto n = make_object<MetadataNode>();
- n->num_inputs = num_inputs;
+ n->inputs = inputs;
  n->num_outputs = num_outputs;
  n->executor = executor;
  n->mod_name = mod_name;

src/target/source/source_module.cc

@@ -192,25 +192,26 @@ class CSourceCrtMetadataModuleNode : public runtime::ModuleNode {
  << "}\n";
  }
 
- void GenerateEntrypointForUnpackedAPI(const std::string& run_func) {
+ void GenerateEntrypointForUnpackedAPI(const std::string& entrypoint_name,
+ const std::string& run_func) {
  code_ << "TVM_DLL int32_t " << run_func << "(";
- int total_args = (metadata_->num_inputs + metadata_->num_outputs);
- for (int i = 0; i < total_args; ++i) {
- code_ << "arg" << i;
+unsigned int total_args = (metadata_->inputs.size() + metadata_->num_outputs);
+ for (unsigned int i = 0; i < total_args; ++i) {
+ code_ << "void* arg" << i;
  if (i + 1 != total_args) {
  code_ << ",";
  }
  }
  code_ << ");\n";
- code_ << "static int32_t " << ::tvm::runtime::symbol::tvm_module_main;
+ code_ << "int32_t " << entrypoint_name;
  code_ << "(void* args, void* type_code, int num_args, void* out_value, void* "
  "out_type_code, void* resource_handle) {\n";
  code_ << "return " << run_func << "(";
- for (int i = 0; i < metadata_->num_inputs; ++i) {
+ for (unsigned int i = 0; i < metadata_->inputs.size(); ++i) {
  code_ << "((DLTensor*)(((TVMValue*)args)[" << i << "].v_handle))[0].data,";
  }
  for (int i = 0; i < metadata_->num_outputs; ++i) {
- int j = metadata_->num_inputs + i;
+ int j = metadata_->inputs.size() + i;
  code_ << "((DLTensor*)(((TVMValue*)args)[" << j << "].v_handle))[0].data";
  if (i + 1 != metadata_->num_outputs) {
  code_ << ",";
@@ -220,37 +221,83 @@ class CSourceCrtMetadataModuleNode : public runtime::ModuleNode {
  code_ << "}\n";
  }
 
- void GenerateEntrypointForPackedAPI(const std::string& run_func) {
+ void GenerateEntrypointForPackedAPI(const std::string& entrypoint_name,
+ const std::string& run_func) {
  code_ << "TVM_DLL int32_t " << run_func;
  code_ << "(void* args, void* type_code, int num_args, void* out_value, void* "
  "out_type_code, void* resource_handle);\n";
- code_ << "static int32_t " << ::tvm::runtime::symbol::tvm_module_main;
+ code_ << "int32_t " << entrypoint_name;
  code_ << "(void* args, void* type_code, int num_args, void* out_value, void* "
  "out_type_code, void* resource_handle) {\n";
  code_ << "return " << run_func;
  code_ << "(args, type_code, num_args, out_value, out_type_code, resource_handle);\n";
  code_ << "}\n";
  }
 
+ void GenerateCInterfaceEntrypoint(const std::string& entrypoint_name, const std::string& run_func,
+ const std::string& mod_name) {
+ code_ << "#include <" << mod_name << ".h>\n";
+ code_ << "TVM_DLL int32_t " << run_func << "(";
+ unsigned int total_args = (metadata_->inputs.size() + metadata_->num_outputs);
+ for (unsigned int i = 0; i < total_args; ++i) {
+ code_ << "void* arg" << i;
+ if (i + 1 != total_args) {
+ code_ << ",";
+ }
+ }
+ code_ << ");\n";
+ code_ << "int32_t " << entrypoint_name << "(";
+ code_ << "struct " << runtime::get_name_mangled(mod_name, "inputs") << "* inputs,"
+ << "struct " << runtime::get_name_mangled(mod_name, "outputs") << "* outputs"
+ << ") {";
+ code_ << "return " << run_func << "(";
+ for (const auto& input : metadata_->inputs) {
+ code_ << "inputs->" << input << ",";
+ }
+ if (metadata_->num_outputs == 1) {
+ code_ << "outputs->output";
+ } else {
+ for (int i = 0; i < metadata_->num_outputs; ++i) {
+ code_ << "outputs->output" << i;
+ if (i + 1 != metadata_->num_outputs) {
+ code_ << ",";
+ }
+ }
+ }
+ code_ << ");\n";
+ code_ << "}\n";
+ }
+
  void GenerateAOTDescriptor() {
- const std::string run_func = ::tvm::runtime::symbol::tvm_run_func_suffix;
- const std::string run_func_mangled = runtime::get_name_mangled(metadata_->mod_name, run_func);
+ const std::string run_func_suffix = ::tvm::runtime::symbol::tvm_run_func_suffix;
+ const std::string tvm_entrypoint_suffix = ::tvm::runtime::symbol::tvm_entrypoint_suffix;
+ const std::string run_func_mangled =
+ runtime::get_name_mangled(metadata_->mod_name, run_func_suffix);
+ const std::string entrypoint_mangled =
+ runtime::get_name_mangled(metadata_->mod_name, tvm_entrypoint_suffix);
  const std::string network_mangled = runtime::get_name_mangled(metadata_->mod_name, "network");
- code_ << "#include \"tvm/runtime/crt/internal/aot_executor/aot_executor.h\"\n";
+ auto unpacked_api = target_->GetAttr<Bool>("unpacked-api").value_or(Bool(false));
+ auto interface_api = target_->GetAttr<String>("interface-api").value_or(String("packed"));
+
  code_ << "#include \"tvm/runtime/c_runtime_api.h\"\n";
  code_ << "#ifdef __cplusplus\n";
- code_ << "extern \"C\"\n";
+ code_ << "extern \"C\" {\n";
  code_ << "#endif\n";
- if (target_->GetAttr<Bool>("unpacked-api").value_or(Bool(false))) {
- GenerateEntrypointForUnpackedAPI(run_func_mangled);
+
+ if (unpacked_api) {
+ if (interface_api == "c") {
+ GenerateCInterfaceEntrypoint(entrypoint_mangled, run_func_mangled, metadata_->mod_name);
+ } else {
+ GenerateEntrypointForUnpackedAPI(entrypoint_mangled, run_func_mangled);
+ }
  } else {
- GenerateEntrypointForPackedAPI(run_func_mangled);
+ ICHECK_EQ(interface_api, "packed") << "Packed interface required for packed operators";
+ GenerateEntrypointForPackedAPI(entrypoint_mangled, run_func_mangled);
  }
- code_ << "const tvm_model_t " << network_mangled << " = {\n"
- << " .run_func = &" << ::tvm::runtime::symbol::tvm_module_main << ",\n"
- << " .num_input_tensors = " << metadata_->num_inputs << ",\n"
- << " .num_output_tensors = " << metadata_->num_outputs << ", \n"
- << "};\n";
+
+ code_ << "#ifdef __cplusplus\n";
+ code_ << "}\n";
+ code_ << "#endif\n";
  }
 
  void CreateSource() {

src/target/target_kind.cc

@@ -299,6 +299,7 @@ TVM_REGISTER_TARGET_KIND("llvm", kDLCPU)
  .add_attr_option<String>("runtime")
  .add_attr_option<Bool>("link-params", Bool(false))
  .add_attr_option<Bool>("unpacked-api")
+ .add_attr_option<String>("interface-api")
  .set_default_keys({"cpu"});
 
 TVM_REGISTER_TARGET_KIND("c", kDLCPU)
@@ -310,6 +311,7 @@ TVM_REGISTER_TARGET_KIND("c", kDLCPU)
  .add_attr_option<String>("executor")
  .add_attr_option<Integer>("workspace-byte-alignment")
  .add_attr_option<Bool>("unpacked-api")
+ .add_attr_option<String>("interface-api")
  .set_default_keys({"cpu"});
 
 TVM_REGISTER_TARGET_KIND("cuda", kDLCUDA)