Introduce --interface-api={c,packed} parameter (#8280)

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

This introduces structures generated to provide a documented and stable user
friendly interface to a TVM generated model, as can be seen in the AOT
demo application:
```
struct tvmgen_default_inputs inputs = {
  .input_1 = input_data,
};
struct tvmgen_default_outputs outputs = {
  .output = output_data,
};
int ret_val = tvmgen_default_run(&inputs, &outputs, NULL, NULL);
```

To facilitate this, some other changes are included:
* Removed dependency on `aot_executor.{c,h}` in tests, pending the
discussion in the interface RFC as to whether we keep them.
* Moved creation of test DLTensor's into the AOT test utils, in future this
can be replaced by loading via the Python API or otherwise
* Introduce `parametrize_aot_options` which can be used to test
permutations of AOT which work together - for now this filters C
interface and packed operators
* Updated demo application to generate the header for demonstration
purposes, we should consider porting the demo application to Model
Library Format and using the toolchain in the Zephyr App via CMake
instead?

This patch builds upon the improvements @giuseros made to AOT testing
and name mangling from #8014

* Tweak metadata variable description and MLF target loop

* Remove direct usage of `relay::Var` in meta_data.h

This looks like the only place that could be causing the Windows CI failures, so trying removing the additional header in meta_data.h

* Linting fix

* Post-rebase files fixing

These tests were somehow transmuted in transit, I've updated them to the
most recent variant of the test helpers.

* Strip back interface API to just inputs and outputs

This removes any speculative structures from the generated code and cleans up some of the documentation.

* Add header guards and tweak documentation

Author: Mousius

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)