[RUNTIME] Add System Lib (#227)

* [RUNTIME] Add System Lib

* lint

* lint

* fix compile

Author: tqchen

include/tvm/codegen.h

@@ -38,9 +38,10 @@ runtime::Module Build(const Array<LoweredFunc>& funcs,
  *  the dependency from the shared library.
  *
  * \param m The host module with the imports.
+ * \param system_lib Whether expose as system library.
  * \return cstr The C string representation of the file.
  */
-std::string PackImportsToC(const runtime::Module& m);
+std::string PackImportsToC(const runtime::Module& m, bool system_lib);
 }  // namespace codegen
 }  // namespace tvm
 

include/tvm/runtime/c_runtime_api.h

@@ -350,6 +350,17 @@ TVM_DLL int TVMFuncListGlobalNames(int *out_size,
 TVM_DLL int TVMBackendGetFuncFromEnv(void* mod_node,
                                      const char* func_name,
                                      TVMFunctionHandle *out);
+/*!
+ * \brief Backend function to register system-wide library symbol.
+ *
+ * \note This API is supposed to be used by backend,
+ *  it is not supposed to be used by user.
+ *
+ * \param name The name of the symbol
+ * \param ptr The symbol address.
+ * \return 0 when no error is thrown, -1 when failure happens
+ */
+TVM_DLL int TVMBackendRegisterSystemLibSymbol(const char* name, void* ptr);
 
 /*!
  * \brief Backend function for running parallel for loop.

include/tvm/runtime/registry.h

@@ -94,7 +94,7 @@ class Registry {
 #define TVM_STR_CONCAT(__x, __y) TVM_STR_CONCAT_(__x, __y)
 
 #define TVM_FUNC_REG_VAR_DEF                                            \
-  static TVM_ATTRIBUTE_UNUSED ::tvm::runtime::Registry& __make_ ## TVMOp
+  static TVM_ATTRIBUTE_UNUSED ::tvm::runtime::Registry& __mk_ ## TVM
 
 /*!
  * \brief Register a function globally.

python/tvm/build.py

@@ -223,23 +223,31 @@ def build(sch,
         The target and option of the compilation.
         When the target is llvm, you can set options like:
 
-          * **-mtriple=<target triple>** or **-target**
+          - **-mtriple=<target triple>** or **-target**
 
             Specify the target triple, which is useful for cross
             compilation.
 
-          * **-mcpu=<cpuname>**
+          - **-mcpu=<cpuname>**
 
             Specify a specific chip in the current architecture to
             generate code for. By default this is infered from the
             target triple and autodetected to the current architecture.
 
-          * **-mattr=a1,+a2,-a3,...**
+          - **-mattr=a1,+a2,-a3,...**
 
             Override or control specific attributes of the target,
             such as whether SIMD operations are enabled or not. The
             default set of attributes is set by the current CPU.
 
+          - **-system-lib**
+
+            Build TVM system library module. System lib is a global module that contains
+            self registered functions in program startup. User can get the module using
+            :any:`tvm.module.system_lib`.
+            It is useful in environments where dynamic loading api like dlopen is banned.
+            The system lib will be available as long as the result code is linked by the program.
+
     target_host : str, optional
         Host compilation target, if target is device.
         When TVM compiles device specific program such as CUDA,

python/tvm/contrib/cc_compiler.py

@@ -25,7 +25,7 @@ def create_shared(output,
         The compile string.
     """
     cmd = [cc]
-    cmd += ["-shared"]
+    cmd += ["-shared", "-fPIC"]
 
     if sys.platform == "darwin":
         cmd += ["-undefined", "dynamic_lookup"]

python/tvm/module.py

@@ -71,16 +71,26 @@ def export_library(self, file_name):
         file_name : str
             The name of the shared library.
         """
+        if self.type_key == "stacktvm":
+            raise ValueError("Module[%s]: export_library requires llvm module,"
+                             " did you build with LLVM enabled?" % self.type_key)
+
         if self.type_key != "llvm":
             raise ValueError("Module[%s]: Only llvm support export shared" % self.type_key)
         temp = _util.tempdir()
         path_obj = temp.relpath("lib.o")
         self.save(path_obj)
         files = [path_obj]
+        try:
+            self.get_function("__tvm_module_startup")
+            is_system_lib = True
+        except AttributeError:
+            is_system_lib = False
+
         if self.imported_modules:
             path_cc = temp.relpath("devc.cc")
             with open(path_cc, "w") as f:
-                f.write(_PackImportsToC(self))
+                f.write(_PackImportsToC(self, is_system_lib))
             files.append(path_cc)
         _cc.create_shared(file_name, files)
 
@@ -116,6 +126,27 @@ def time_evaluator(self, func_name, ctx, number):
             raise NameError("time_evaluate is only supported when RPC is enabled")
 
 
+def system_lib():
+    """Get system-wide library module singleton.
+
+    System lib is a global module that contains self register functions in startup.
+    Unlike normal dso modules which need to be loaded explicitly.
+    It is useful in environments where dynamic loading api like dlopen is banned.
+
+    To build system lib function, simply specify target option ```llvm --system-lib```
+    The system lib will be available as long as the result code is linked by the program.
+
+    The system lib is intended to be linked and loaded during the entire life-cyle of the program.
+    If you want dynamic loading features, use dso modules instead.
+
+    Returns
+    -------
+    module : Module
+        The system-wide library module.
+    """
+    return _GetSystemLib()
+
+
 def load(path, fmt=""):
     """Load module from file
 

src/api/api_codegen.cc

@@ -23,7 +23,7 @@ TVM_REGISTER_API("codegen._Build")
 
 TVM_REGISTER_API("module._PackImportsToC")
 .set_body([](TVMArgs args, TVMRetValue *ret) {
-    *ret = PackImportsToC(args[0]);
+    *ret = PackImportsToC(args[0], args[1]);
   });
 }  // namespace codegen
 }  // namespace tvm

src/codegen/codegen.cc

@@ -35,7 +35,7 @@ runtime::Module Build(const Array<LoweredFunc>& funcs,
   return m;
 }
 
-std::string PackImportsToC(const runtime::Module& mod) {
+std::string PackImportsToC(const runtime::Module& mod, bool system_lib) {
   std::string bin;
   dmlc::MemoryStringStream ms(&bin);
   dmlc::Stream* stream = &ms;
@@ -55,29 +55,39 @@ std::string PackImportsToC(const runtime::Module& mod) {
      << "extern \"C\" {\n"
      << "#endif\n";
   os << "extern const char " << runtime::symbol::tvm_dev_mblob << "[];\n";
-  os << "extern const unsigned long " << runtime::symbol::tvm_dev_mblob_nbytes << ";\n";
+  uint64_t nbytes = bin.length();
   os << "const char " << runtime::symbol::tvm_dev_mblob
-     << "[" << bin.length() << "] = {\n  ";
+     << "[" << bin.length() + sizeof(nbytes) << "] = {\n  ";
   os << std::hex;
   size_t nunit = 80 / 4;
-  for (size_t i = 0; i < bin.length(); ++i) {
+  for (size_t i = 0; i < sizeof(nbytes); ++i) {
     // sperators
     if (i != 0) {
-      if (i % nunit == 0) {
-        os << ",\n  ";
-      } else {
-        os << ",";
-      }
+      os << ",";
+    }
+    os << "0x" << ((nbytes >> (i * 8)) & 0xffUL);
+  }
+  for (size_t i = 0; i < bin.length(); ++i) {
+    // sperators
+    if ((i + sizeof(nbytes)) % nunit == 0) {
+      os << ",\n  ";
+    } else {
+      os << ",";
     }
     int c = bin[i];
     os << "0x" << (c & 0xff);
   }
-  os << "\n};\n"
-     << "const unsigned long " << runtime::symbol::tvm_dev_mblob_nbytes
-     << " = " << std::dec << bin.length() << "UL;\n"
-     << "#ifdef __cplusplus\n"
+  os << "\n};\n";
+  if (system_lib) {
+    os << "extern int TVMBackendRegisterSystemLibSymbol(const char*, void*);\n";
+    os << "static int " << runtime::symbol::tvm_dev_mblob << "_reg_ = "
+       << "TVMBackendRegisterSystemLibSymbol(\"" << runtime::symbol::tvm_dev_mblob << "\", (void*)"
+       << runtime::symbol::tvm_dev_mblob << ");\n";
+  }
+  os << "#ifdef __cplusplus\n"
      << "}\n"
      << "#endif\n";
+
   return os.str();
 }
 }  // namespace codegen

src/codegen/llvm/codegen_llvm.cc

@@ -28,14 +28,16 @@ std::unique_ptr<CodeGenLLVM> CodeGenLLVM::Create(llvm::TargetMachine *tm) {
 
 void CodeGenLLVM::Init(const std::string& module_name,
                        llvm::TargetMachine* tm,
-                       llvm::LLVMContext* ctx) {
+                       llvm::LLVMContext* ctx,
+                       bool system_lib) {
   InitializeLLVM();
   static_assert(sizeof(TVMValue) == sizeof(double), "invariant");
   // static_assert(alignof(TVMValue) == alignof(double), "invariant");
   // clear maps
   var_map_.clear();
   str_map_.clear();
   func_handle_map_.clear();
+  export_system_symbols_.clear();
   // initialize types.
   if (ctx_ != ctx) {
     t_void_ = llvm::Type::getVoidTy(*ctx);
@@ -96,6 +98,13 @@ void CodeGenLLVM::Init(const std::string& module_name,
           t_int64_, t_int64_, t_f_tvm_par_for_lambda_->getPointerTo(), t_void_p_}
         , false),
       llvm::Function::ExternalLinkage, "TVMBackendParallelFor", module_.get());
+  if (system_lib) {
+    f_tvm_register_system_symbol_ = llvm::Function::Create(
+        llvm::FunctionType::get(t_int_, {t_char_->getPointerTo(), t_void_p_}, false),
+        llvm::Function::ExternalLinkage, "TVMBackendRegisterSystemLibSymbol", module_.get());
+  } else {
+    f_tvm_register_system_symbol_ = nullptr;
+  }
   this->InitTarget(tm);
   // initialize builder
   builder_.reset(new IRBuilder(*ctx));
@@ -125,9 +134,15 @@ void CodeGenLLVM::InitTarget(llvm::TargetMachine* tm) {
 void CodeGenLLVM::InitGlobalContext() {
   gv_mod_ctx_ = new llvm::GlobalVariable(
       *module_, t_void_p_, false,
-      llvm::GlobalValue::LinkOnceAnyLinkage, 0, "__tvm_module_ctx");
+      llvm::GlobalValue::LinkOnceAnyLinkage, 0,
+      tvm::runtime::symbol::tvm_module_ctx);
   gv_mod_ctx_->setAlignment(data_layout_->getTypeAllocSize(t_void_p_));
   gv_mod_ctx_->setInitializer(llvm::Constant::getNullValue(t_void_p_));
+
+  if (f_tvm_register_system_symbol_ != nullptr) {
+    export_system_symbols_.emplace_back(
+        std::make_pair(tvm::runtime::symbol::tvm_module_ctx, gv_mod_ctx_));
+  }
 }
 
 void CodeGenLLVM::InitFuncState() {
@@ -171,6 +186,11 @@ void CodeGenLLVM::AddFunction(const LoweredFunc& f) {
   builder_->SetInsertPoint(block);
   this->VisitStmt(f->body);
   builder_->CreateRet(ConstInt32(0));
+
+  if (f_tvm_register_system_symbol_ != nullptr) {
+    export_system_symbols_.emplace_back(
+        std::make_pair(f->name, builder_->CreatePointerCast(function_, t_void_p_)));
+  }
 }
 
 void CodeGenLLVM::AddMainFunction(const std::string& entry_func_name) {
@@ -225,13 +245,35 @@ void CodeGenLLVM::Optimize() {
 }
 
 std::unique_ptr<llvm::Module> CodeGenLLVM::Finish() {
+  this->AddStartupFunction();
   this->Optimize();
   var_map_.clear();
   str_map_.clear();
   func_handle_map_.clear();
+  export_system_symbols_.clear();
   return std::move(module_);
 }
 
+void CodeGenLLVM::AddStartupFunction() {
+  if (export_system_symbols_.size() != 0) {
+    llvm::FunctionType* ftype = llvm::FunctionType::get(t_void_, {}, false);
+    function_ = llvm::Function::Create(
+        ftype,
+        llvm::Function::InternalLinkage,
+        "__tvm_module_startup", module_.get());
+    llvm::BasicBlock* startup_entry = llvm::BasicBlock::Create(*ctx_, "entry", function_);
+    builder_->SetInsertPoint(startup_entry);
+    for (const auto& kv : export_system_symbols_) {
+      llvm::Value* name = GetConstString(kv.first);
+      builder_->CreateCall(
+          f_tvm_register_system_symbol_, {
+            name, builder_->CreateBitCast(kv.second, t_void_p_)});
+    }
+    llvm::appendToGlobalCtors(*module_, function_, 65535);
+    builder_->CreateRet(nullptr);
+  }
+}
+
 llvm::Type* CodeGenLLVM::LLVMType(const Type& t) const {
   llvm::Type* ret = nullptr;
   if (t.is_uint() || t.is_int()) {

src/codegen/llvm/codegen_llvm.h

@@ -12,6 +12,7 @@
 #include <tvm/codegen.h>
 #include <tvm/arithmetic.h>
 #include <memory>
+#include <utility>
 #include <vector>
 #include <string>
 #include "./llvm_common.h"
@@ -39,10 +40,12 @@ class CodeGenLLVM :
    * \param module_name The name of the module.
    * \param tm Target machine model
    * \param ctx The context.
+   * \param system_lib Whether to insert system library registration.
    */
   void Init(const std::string& module_name,
             llvm::TargetMachine* tm,
-            llvm::LLVMContext* ctx);
+            llvm::LLVMContext* ctx,
+            bool system_lib);
   /*!
    * \brief Compile and add function f to the current module.
    * \param f The function to be added.
@@ -163,7 +166,7 @@ class CodeGenLLVM :
   llvm::LLVMContext* ctx_{nullptr};
   // helpful data types
   llvm::Type* t_void_{nullptr};
-  llvm::Type* t_void_p_{nullptr};
+  llvm::PointerType* t_void_p_{nullptr};
   llvm::Type* t_int_{nullptr};
   llvm::Type* t_char_{nullptr};
   llvm::Type* t_int8_{nullptr};
@@ -188,6 +191,7 @@ class CodeGenLLVM :
   llvm::Function* f_tvm_get_func_from_env_{nullptr};
   llvm::Function* f_tvm_api_set_last_error_{nullptr};
   llvm::Function* f_tvm_parallel_for_{nullptr};
+  llvm::Function* f_tvm_register_system_symbol_{nullptr};
   // The acting body
   llvm::BasicBlock* block_{nullptr};
   /*! \brief native vector bits of current targetx*/
@@ -227,6 +231,8 @@ class CodeGenLLVM :
   llvm::BasicBlock* CheckCallSuccess(llvm::Value* retcode);
   // Add a function to set global module context
   void InitGlobalContext();
+  // Add module startup function if needed.
+  void AddStartupFunction();
   // add alias information.
   void AddAliasInfo(llvm::Instruction* load, const Variable* buffer, Expr index, Type type);
   // The definition of local variable.
@@ -243,6 +249,8 @@ class CodeGenLLVM :
   llvm::GlobalVariable* gv_mod_ctx_{nullptr};
   // global to packed function handle
   std::unordered_map<std::string, llvm::GlobalVariable*> func_handle_map_;
+  // List of symbols to be exported to TVM system lib.
+  std::vector<std::pair<std::string, llvm::Value*> > export_system_symbols_;
 };
 }  // namespace codegen
 }  // namespace tvm