Created CSourceMetaData module for model metadata

apps/bundle_deploy/build_model.py

@@ -51,7 +51,6 @@ def build_module(opts):
  build_dir = os.path.abspath(opts.out_dir)
  if not os.path.isdir(build_dir):
  os.makedirs(build_dir)
-
  lib.save(os.path.join(build_dir, file_format_str.format(name="model", ext="o")))
  with open(
  os.path.join(build_dir, file_format_str.format(name="graph", ext="json")), "w"
@@ -85,7 +84,6 @@ def build_test_module(opts):
  build_dir = os.path.abspath(opts.out_dir)
  if not os.path.isdir(build_dir):
  os.makedirs(build_dir)
-
  lib.save(os.path.join(build_dir, file_format_str.format(name="test_model", ext="o")))
  with open(
  os.path.join(build_dir, file_format_str.format(name="test_graph", ext="json")), "w"

python/tvm/driver/build_module.py

@@ -424,4 +424,16 @@ def build(inputs, args=None, target=None, target_host=None, name="default_functi
  for mdev in device_modules:
  if mdev:
  rt_mod_host.import_module(mdev)
+
+ if not isinstance(target_host, Target):
+ target_host = Target(target_host)
+ if (
+ "system-lib" in target_host.attrs
+ and target_host.attrs["system-lib"].value == 1
+ and target_host.kind.name == "c"
+ ):
+ create_csource_metadata_module = tvm._ffi.get_global_func(
+ "runtime.CreateCSourceMetadataModule"
+ )
+ return create_csource_metadata_module([rt_mod_host], target_host)
  return rt_mod_host

python/tvm/micro/build.py

@@ -95,6 +95,7 @@ def path(self):
 # void* arg0 = (((TVMValue*)args)[0].v_handle);
 # int32_t arg0_code = ((int32_t*)arg_type_ids)[(0)];
 _CRT_GENERATED_LIB_OPTIONS["cflags"].append("-Wno-unused-variable")
+_CRT_GENERATED_LIB_OPTIONS["ccflags"].append("-Wno-unused-variable")
 
 
 # Many TVM-intrinsic operators (i.e. expf, in particular)
@@ -159,9 +160,6 @@ def build_static_runtime(
  mod_build_dir = workspace.relpath(os.path.join("build", "module"))
  os.makedirs(mod_build_dir)
  mod_src_dir = workspace.relpath(os.path.join("src", "module"))
- os.makedirs(mod_src_dir)
- mod_src_path = os.path.join(mod_src_dir, "module.c")
- module.save(mod_src_path, "cc")
 
  libs = []
  for mod_or_src_dir in (extra_libs or []) + RUNTIME_LIB_SRC_DIRS:
@@ -181,7 +179,15 @@ def build_static_runtime(
 
  libs.append(compiler.library(lib_build_dir, lib_srcs, lib_opts))
 
- libs.append(compiler.library(mod_build_dir, [mod_src_path], generated_lib_opts))
+ mod_src_dir = workspace.relpath(os.path.join("src", "module"))
+ os.makedirs(mod_src_dir)
+ libs.append(
+ module.export_library(
+ mod_build_dir,
+ workspace_dir=mod_src_dir,
+ fcompile=lambda bdir, srcs, **kwargs: compiler.library(bdir, srcs, generated_lib_opts),
+ )
+ )
 
  runtime_build_dir = workspace.relpath(f"build/runtime")
  os.makedirs(runtime_build_dir)

python/tvm/runtime/module.py

@@ -15,7 +15,7 @@
 # specific language governing permissions and limitations
 # under the License.
 
-# pylint: disable=invalid-name, unused-import, import-outside-toplevel
+# pylint: disable=invalid-name, unused-import, import-outside-toplevel, inconsistent-return-statements
 """Runtime Module namespace."""
 import os
 import ctypes
@@ -252,7 +252,7 @@ def _collect_dso_modules(self):
  def _dso_exportable(self):
  return self.type_key == "llvm" or self.type_key == "c"
 
- def export_library(self, file_name, fcompile=None, addons=None, **kwargs):
+ def export_library(self, file_name, fcompile=None, addons=None, workspace_dir=None, **kwargs):
  """Export the module and its imported device code one library.
 
  This function only works on host llvm modules.
@@ -268,8 +268,19 @@ def export_library(self, file_name, fcompile=None, addons=None, **kwargs):
  If fcompile has attribute object_format, will compile host library
  to that format. Otherwise, will use default format "o".
 
+ workspace_dir : str, optional
+ the path to a directory used to create intermediary
+ artifacts for the process exporting of the library.
+ If this is not provided a temporary dir will be created.
+
  kwargs : dict, optional
  Additional arguments passed to fcompile
+
+ Returns
+ -------
+ result of fcompile() : unknown, optional
+ If the compilation function returns an artifact it would be returned via
+ export_library, if any.
  """
  # NOTE: this function depends on contrib library features
  # which are only available in when TVM function is available.
@@ -292,22 +303,28 @@ def export_library(self, file_name, fcompile=None, addons=None, **kwargs):
  return
 
  modules = self._collect_dso_modules()
- temp = _utils.tempdir()
+ if workspace_dir is None:
+ temp = _utils.tempdir()
+ workspace_dir = temp.temp_dir
  files = addons if addons else []
  is_system_lib = False
  has_c_module = False
  llvm_target_triple = None
  for index, module in enumerate(modules):
  if fcompile is not None and hasattr(fcompile, "object_format"):
- object_format = fcompile.object_format
+ if module.type_key == "c":
+ object_format = "c"
+ has_c_module = True
+ else:
+ object_format = fcompile.object_format
  else:
  if module.type_key == "llvm":
  object_format = "o"
  else:
  assert module.type_key == "c"
- object_format = "cc"
+ object_format = "c"
  has_c_module = True
- path_obj = temp.relpath("lib" + str(index) + "." + object_format)
+ path_obj = os.path.join(workspace_dir, f"lib{index}.{object_format}")
  module.save(path_obj)
  files.append(path_obj)
  is_system_lib = (
@@ -330,25 +347,28 @@ def export_library(self, file_name, fcompile=None, addons=None, **kwargs):
 
  if self.imported_modules:
  if enabled("llvm") and llvm_target_triple:
- path_obj = temp.relpath("devc." + object_format)
+ path_obj = os.path.join(workspace_dir, f"devc.{object_format}")
  m = _ffi_api.ModulePackImportsToLLVM(self, is_system_lib, llvm_target_triple)
  m.save(path_obj)
  files.append(path_obj)
  else:
- path_cc = temp.relpath("devc.cc")
+ path_cc = os.path.join(workspace_dir, "devc.c")
  with open(path_cc, "w") as f:
  f.write(_ffi_api.ModulePackImportsToC(self, is_system_lib))
  files.append(path_cc)
 
- if has_c_module:
+ # The imports could contain a c module but the object format could be tar
+ # Thus, it would not recognize the following include paths as options
+ # which are there assuming a c compiler is the fcompile.
+ if has_c_module and not file_name.endswith(".tar"):
  options = []
  if "options" in kwargs:
  opts = kwargs["options"]
  options = opts if isinstance(opts, (list, tuple)) else [opts]
  opts = options + ["-I" + path for path in find_include_path()]
  kwargs.update({"options": opts})
 
- fcompile(file_name, files, **kwargs)
+return fcompile(file_name, files, **kwargs)
 
 
 def system_lib():

src/relay/backend/build_module.cc

@@ -510,18 +510,14 @@ class RelayBuildModule : public runtime::ModuleNode {
  // If we cannot decide the target is LLVM, we create an empty CSourceModule.
  // The code content is initialized with ";" to prevent complaining
  // from CSourceModuleNode::SaveToFile.
- ret_.mod = tvm::codegen::CSourceModuleCreate(";", "");
+ ret_.mod = tvm::codegen::CSourceModuleCreate(";", "", Array<String>{});
  }
  } else {
  ret_.mod = tvm::build(lowered_funcs, target_host_);
  }
 
- Array<tvm::runtime::Module> ext_mods = graph_codegen_->GetExternalModules();
- // TODO(zhiics) We should be able to completely switch to MetadataModule no
- // matter whether there are external modules or not.
- if (!ext_mods.empty()) {
- ret_.mod = tvm::codegen::CreateMetadataModule(ret_.params, ret_.mod, ext_mods);
- }
+ auto ext_mods = graph_codegen_->GetExternalModules();
+ ret_.mod = tvm::codegen::CreateMetadataModule(ret_.params, ret_.mod, ext_mods, GetTargetHost());
  }
 
  private:

src/relay/backend/contrib/codegen_c/codegen.cc

@@ -215,20 +215,19 @@ class CodegenC : public MemoizedExprTranslator<std::vector<Output>>, public Code
 
 class CSourceCodegen : public CSourceModuleCodegenBase {
  public:
- std::pair<std::string, Array<String>> GenCFunc(const Function& func) {
+ std::tuple<Array<String>, String, String> GenCFunc(const Function& func) {
  ICHECK(func.defined()) << "Input error: expect a Relay function.";
-
- // Record the external symbol for runtime lookup.
- auto sid = GetExtSymbol(func);
-
- CodegenC builder(sid);
+ CodegenC builder(GetExtSymbol(func));
  auto out = builder.VisitExpr(func->body);
- code_stream_ << builder.JIT(out);
-
- return {sid, builder.const_vars_};
+ return std::make_tuple(builder.const_vars_, builder.ext_func_id_, builder.JIT(out));
  }
 
  runtime::Module CreateCSourceModule(const ObjectRef& ref) override {
+ ICHECK(ref->IsInstance<FunctionNode>());
+ auto res = GenCFunc(Downcast<Function>(ref));
+ Array<String> variables = std::get<0>(res);
+ String func_name = std::get<1>(res);
+
  // Create headers
  code_stream_ << "#include <cstring>\n";
  code_stream_ << "#include <vector>\n";
@@ -259,18 +258,13 @@ class CSourceCodegen : public CSourceModuleCodegenBase {
  )op_macro";
 
  code_stream_ << operator_macro << "\n\n";
-
- ICHECK(ref->IsInstance<FunctionNode>());
- auto res = GenCFunc(Downcast<Function>(ref));
+ code_stream_ << std::get<2>(res);
  std::string code = code_stream_.str();
 
- String sym = std::get<0>(res);
- Array<String> variables = std::get<1>(res);
-
  // Create a CSource module
  const auto* pf = runtime::Registry::Get("runtime.CSourceModuleCreate");
  ICHECK(pf != nullptr) << "Cannot find csource module to create the external runtime module";
- return (*pf)(code, "c", sym, variables);
+ return (*pf)(code, "c", Array<String>{func_name}, variables);
  }
 
  private:

src/relay/backend/contrib/dnnl/codegen.cc

@@ -413,7 +413,8 @@ class DNNLModuleCodegen : public CSourceModuleCodegenBase {
  // Create a CSource module
  const auto* pf = runtime::Registry::Get("runtime.CSourceModuleCreate");
  ICHECK(pf != nullptr) << "Cannot find csource module to create the external runtime module";
- return (*pf)(code, "c", sym, variables);
+ // TODO(@manupa-arm): pass the function names to enable system-lib creation
+ return (*pf)(code, "c", Array<String>{sym}, variables);
  }
 
  private:

src/relay/backend/vm/compiler.cc

@@ -1146,11 +1146,9 @@ void VMCompiler::Codegen() {
  } else {
  // There is no function handled by TVM. We create a virtual main module
  // to make sure a DSO module will be also available.
- exec_->lib = codegen::CSourceModuleCreate(";", "");
- }
- if (!ext_mods.empty()) {
- exec_->lib = codegen::CreateMetadataModule(params_, exec_->lib, ext_mods);
+ exec_->lib = codegen::CSourceModuleCreate(";", "", Array<String>{});
  }
+ exec_->lib = codegen::CreateMetadataModule(params_, exec_->lib, ext_mods, target_host_);
 }
 
 ExprDeviceMap VMCompiler::AnalyzeContext() const {

src/target/func_registry_generator.cc

@@ -29,7 +29,7 @@
 namespace tvm {
 namespace target {
 
-std::string GenerateFuncRegistryNames(const std::vector<std::string>& function_names) {
+std::string GenerateFuncRegistryNames(const Array<String>& function_names) {
  std::stringstream ss;
  ss << (unsigned char)(function_names.size());
  for (auto f : function_names) {

src/target/func_registry_generator.h

@@ -24,13 +24,18 @@
 #ifndef TVM_TARGET_FUNC_REGISTRY_GENERATOR_H_
 #define TVM_TARGET_FUNC_REGISTRY_GENERATOR_H_
 
+#include <tvm/runtime/container.h>
+
 #include <string>
 #include <vector>
 
+using tvm::runtime::Array;
+using tvm::runtime::String;
+
 namespace tvm {
 namespace target {
 
-std::string GenerateFuncRegistryNames(const std::vector<std::string>& function_names);
+std::string GenerateFuncRegistryNames(const Array<String>& function_names);
 
 } // namespace target
 } // namespace tvm

src/target/llvm/codegen_cpu.cc

@@ -794,10 +794,10 @@ llvm::Value* CodeGenCPU::RuntimeTVMParallelBarrier() {
 void CodeGenCPU::AddStartupFunction() {
  if (registry_functions_.size() != 0) {
  ICHECK(is_system_lib_) << "Loading of --system-lib modules is yet to be defined for C runtime";
-std::vector<std::string> symbols;
+Array<String> symbols;
  std::vector<llvm::Constant*> funcs;
  for (auto sym : registry_functions_) {
- symbols.emplace_back(sym.first);
+ symbols.push_back(sym.first);
  funcs.emplace_back(llvm::ConstantExpr::getBitCast(
  sym.second, ftype_tvm_backend_packed_c_func_->getPointerTo()));
  }

src/target/llvm/llvm_module.cc

@@ -60,6 +60,13 @@ class LLVMModuleNode final : public runtime::ModuleNode {
  if (name == "__tvm_is_system_module") {
  bool flag = (mptr_->getFunction("__tvm_module_startup") != nullptr);
  return PackedFunc([flag](TVMArgs args, TVMRetValue* rv) { *rv = flag; });
+ } else if (name == "get_func_names") {
+ return PackedFunc(
+ [sptr_to_self, this](TVMArgs args, TVMRetValue* rv) { *rv = this->function_names_; });
+ } else if (name == "get_symbol") {
+ return PackedFunc(nullptr);
+ } else if (name == "get_const_vars") {
+ return PackedFunc(nullptr);
  } else if (name == "_get_target_triple") {
  std::string target_triple = tm_->getTargetTriple().str();
  // getTargetTriple() doesn't include other flags besides the triple. Add back flags which are
@@ -218,9 +225,10 @@ class LLVMModuleNode final : public runtime::ModuleNode {
  ICHECK(kv.second->IsInstance<PrimFuncNode>())
  << "Can only lower IR Module with PrimFuncs, but got " << kv.second->GetTypeKey();
  auto f = Downcast<PrimFunc>(kv.second);
+ auto global_symbol = f->GetAttr<String>(tvm::attr::kGlobalSymbol);
+ ICHECK(global_symbol.defined());
+ function_names_.push_back(global_symbol.value());
  if (f->HasNonzeroAttr(tir::attr::kIsEntryFunc)) {
- auto global_symbol = f->GetAttr<String>(tvm::attr::kGlobalSymbol);
- ICHECK(global_symbol.defined());
  entry_func = global_symbol.value();
  }
  funcs.push_back(f);
@@ -377,6 +385,8 @@ class LLVMModuleNode final : public runtime::ModuleNode {
  std::unique_ptr<llvm::Module> module_;
  // the context.
  std::shared_ptr<llvm::LLVMContext> ctx_;
+ /* \brief names of the functions declared in this module */
+ Array<String> function_names_;
 };
 
 TVM_REGISTER_GLOBAL("target.build.llvm")