[REFACTOR][RUNTIME] Add LibraryModule that merges systemlib and dso.

Historically we have two variations of modules(DSOModule and SystemLibModule) that both exposes module via symbols. This PR creates a common implementation for both, and introduce a Library base class that allows us to have different implementations of GetSymbol. It paves ways for future library related module enhancements.
apache · Dec 9, 2019 · c70442d · c70442d
1 parent 5fe5cee
commit c70442d
Show file tree

Hide file tree

Showing 16 changed files with 333 additions and 292 deletions.
diff --git a/apps/android_deploy/app/src/main/jni/tvm_runtime.h b/apps/android_deploy/app/src/main/jni/tvm_runtime.h
@@ -27,12 +27,12 @@
 #include "../src/runtime/c_runtime_api.cc"
 #include "../src/runtime/cpu_device_api.cc"
 #include "../src/runtime/workspace_pool.cc"
-#include "../src/runtime/module_util.cc"
-#include "../src/runtime/system_lib_module.cc"
+#include "../src/runtime/library_module.cc"
+#include "../src/runtime/system_library.cc"
 #include "../src/runtime/module.cc"
 #include "../src/runtime/registry.cc"
 #include "../src/runtime/file_util.cc"
-#include "../src/runtime/dso_module.cc"
+#include "../src/runtime/dso_library.cc"
 #include "../src/runtime/thread_pool.cc"
 #include "../src/runtime/object.cc"
 #include "../src/runtime/threading_backend.cc"

diff --git a/apps/android_rpc/app/src/main/jni/tvm_runtime.h b/apps/android_rpc/app/src/main/jni/tvm_runtime.h
@@ -39,12 +39,12 @@
 #include "../src/runtime/c_runtime_api.cc"
 #include "../src/runtime/cpu_device_api.cc"
 #include "../src/runtime/workspace_pool.cc"
-#include "../src/runtime/module_util.cc"
-#include "../src/runtime/system_lib_module.cc"
+#include "../src/runtime/library_module.cc"
+#include "../src/runtime/system_library.cc"
 #include "../src/runtime/module.cc"
 #include "../src/runtime/registry.cc"
 #include "../src/runtime/file_util.cc"
-#include "../src/runtime/dso_module.cc"
+#include "../src/runtime/dso_library.cc"
 #include "../src/runtime/rpc/rpc_session.cc"
 #include "../src/runtime/rpc/rpc_event_impl.cc"
 #include "../src/runtime/rpc/rpc_server_env.cc"

diff --git a/apps/howto_deploy/run_example.sh b/apps/howto_deploy/run_example.sh
@@ -6,9 +6,9 @@
 # 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
@@ -29,4 +29,4 @@ echo "Run the cpp deployment with all in normal library..."
 lib/cpp_deploy_normal
 
 echo "Run the python deployment with all in normal library..."
-python python_deploy.py
+python3 python_deploy.py
diff --git a/apps/howto_deploy/tvm_runtime_pack.cc b/apps/howto_deploy/tvm_runtime_pack.cc
@@ -40,7 +40,7 @@
 #include "../../src/runtime/c_runtime_api.cc"
 #include "../../src/runtime/cpu_device_api.cc"
 #include "../../src/runtime/workspace_pool.cc"
-#include "../../src/runtime/module_util.cc"
+#include "../../src/runtime/library_module.cc"
 #include "../../src/runtime/module.cc"
 #include "../../src/runtime/registry.cc"
 #include "../../src/runtime/file_util.cc"
@@ -55,8 +55,8 @@
 // Likely we only need to enable one of the following
 // If you use Module::Load, use dso_module
 // For system packed library, use system_lib_module
-#include "../../src/runtime/dso_module.cc"
-#include "../../src/runtime/system_lib_module.cc"
+#include "../../src/runtime/dso_library.cc"
+#include "../../src/runtime/system_library.cc"
 
 // Graph runtime
 #include "../../src/runtime/graph/graph_runtime.cc"

diff --git a/apps/ios_rpc/tvmrpc/TVMRuntime.mm b/apps/ios_rpc/tvmrpc/TVMRuntime.mm
@@ -27,12 +27,12 @@
 #include "../../../src/runtime/workspace_pool.cc"
 #include "../../../src/runtime/thread_pool.cc"
 #include "../../../src/runtime/threading_backend.cc"
-#include "../../../src/runtime/module_util.cc"
-#include "../../../src/runtime/system_lib_module.cc"
+#include "../../../src/runtime/library_module.cc"
+#include "../../../src/runtime/system_library.cc"
 #include "../../../src/runtime/module.cc"
 #include "../../../src/runtime/registry.cc"
 #include "../../../src/runtime/file_util.cc"
-#include "../../../src/runtime/dso_module.cc"
+#include "../../../src/runtime/dso_library.cc"
 #include "../../../src/runtime/ndarray.cc"
 #include "../../../src/runtime/object.cc"
 

diff --git a/golang/src/tvm_runtime_pack.cc b/golang/src/tvm_runtime_pack.cc
@@ -24,7 +24,7 @@
 #include "src/runtime/c_runtime_api.cc"
 #include "src/runtime/cpu_device_api.cc"
 #include "src/runtime/workspace_pool.cc"
-#include "src/runtime/module_util.cc"
+#include "src/runtime/library_module.cc"
 #include "src/runtime/module.cc"
 #include "src/runtime/registry.cc"
 #include "src/runtime/file_util.cc"
@@ -39,8 +39,8 @@
 // Likely we only need to enable one of the following
 // If you use Module::Load, use dso_module
 // For system packed library, use system_lib_module
-#include "src/runtime/dso_module.cc"
-#include "src/runtime/system_lib_module.cc"
+#include "src/runtime/dso_library.cc"
+#include "src/runtime/system_library.cc"
 
 // Graph runtime
 #include "src/runtime/graph/graph_runtime.cc"

diff --git a/include/tvm/runtime/module.h b/include/tvm/runtime/module.h
@@ -201,6 +201,7 @@ class ModuleNode : public Object {
 
  protected:
   friend class Module;
+  friend class ModuleInternal;
   /*! \brief The modules this module depend on */
   std::vector<Module> imports_;
 

diff --git a/src/codegen/llvm/llvm_module.cc b/src/codegen/llvm/llvm_module.cc
@@ -28,7 +28,7 @@
 #include "llvm_common.h"
 #include "codegen_llvm.h"
 #include "../../runtime/file_util.h"
-#include "../../runtime/module_util.h"
+#include "../../runtime/library_module.h"
 
 namespace tvm {
 namespace codegen {
@@ -286,7 +286,7 @@ class LLVMModuleNode final : public runtime::ModuleNode {
       *ctx_addr = this;
     }
     runtime::InitContextFunctions([this](const char *name) {
-        return GetGlobalAddr(name);
+        return reinterpret_cast<void*>(GetGlobalAddr(name));
       });
   }
   // Get global address from execution engine.

diff --git a/src/runtime/dso_module.cc → src/runtime/dso_library.cc b/src/runtime/dso_module.cc → src/runtime/dso_library.cc
@@ -18,14 +18,14 @@
  */
 
 /*!
- * \file dso_module.cc
- * \brief Module to load from dynamic shared library.
+ * \file dso_libary.cc
+ * \brief Create library module to load from dynamic shared library.
  */
 #include <tvm/runtime/module.h>
 #include <tvm/runtime/memory.h>
 #include <tvm/runtime/registry.h>
 #include <tvm/runtime/packed_func.h>
-#include "module_util.h"
+#include "library_module.h"
 
 #if defined(_WIN32)
 #include <windows.h>
@@ -36,58 +36,32 @@
 namespace tvm {
 namespace runtime {
 
-// Module to load from dynamic shared libary.
+// Dynamic shared libary.
 // This is the default module TVM used for host-side AOT
-class DSOModuleNode final : public ModuleNode {
+class DSOLibrary final : public Library {
  public:
-  ~DSOModuleNode() {
+  ~DSOLibrary() {
     if (lib_handle_) Unload();
   }
-
-  const char* type_key() const final {
-    return "dso";
-  }
-
-  PackedFunc GetFunction(
-      const std::string& name,
-      const ObjectPtr<Object>& sptr_to_self) final {
-    BackendPackedCFunc faddr;
-    if (name == runtime::symbol::tvm_module_main) {
-      const char* entry_name = reinterpret_cast<const char*>(
-          GetSymbol(runtime::symbol::tvm_module_main));
-      CHECK(entry_name!= nullptr)
-          << "Symbol " << runtime::symbol::tvm_module_main << " is not presented";
-      faddr = reinterpret_cast<BackendPackedCFunc>(GetSymbol(entry_name));
-    } else {
-      faddr = reinterpret_cast<BackendPackedCFunc>(GetSymbol(name.c_str()));
-    }
-    if (faddr == nullptr) return PackedFunc();
-    return WrapPackedFunc(faddr, sptr_to_self);
-  }
-
   void Init(const std::string& name) {
     Load(name);
-    if (auto *ctx_addr =
-        reinterpret_cast<void**>(GetSymbol(runtime::symbol::tvm_module_ctx))) {
-      *ctx_addr = this;
-    }
-    InitContextFunctions([this](const char* fname) {
-        return GetSymbol(fname);
-      });
-    // Load the imported modules
-    const char* dev_mblob =
-        reinterpret_cast<const char*>(
-            GetSymbol(runtime::symbol::tvm_dev_mblob));
-    if (dev_mblob != nullptr) {
-      ImportModuleBlob(dev_mblob, &imports_);
-    }
+  }
+
+  void* GetSymbol(const char* name) final {
+    return GetSymbol_(name);
   }
 
  private:
   // Platform dependent handling.
 #if defined(_WIN32)
   // library handle
   HMODULE lib_handle_{nullptr};
+
+  void* GetSymbol_(const char* name) {
+    return reinterpret_cast<void*>(
+        GetProcAddress(lib_handle_, (LPCSTR)name)); // NOLINT(*)
+  }
+
   // Load the library
   void Load(const std::string& name) {
     // use wstring version that is needed by LLVM.
@@ -96,12 +70,10 @@ class DSOModuleNode final : public ModuleNode {
     CHECK(lib_handle_ != nullptr)
         << "Failed to load dynamic shared library " << name;
   }
-  void* GetSymbol(const char* name) {
-    return reinterpret_cast<void*>(
-        GetProcAddress(lib_handle_, (LPCSTR)name)); // NOLINT(*)
-  }
+
   void Unload() {
     FreeLibrary(lib_handle_);
+    lib_handle_ = nullptr;
   }
 #else
   // Library handle
@@ -113,20 +85,23 @@ class DSOModuleNode final : public ModuleNode {
         << "Failed to load dynamic shared library " << name
         << " " << dlerror();
   }
-  void* GetSymbol(const char* name) {
+
+  void* GetSymbol_(const char* name) {
     return dlsym(lib_handle_, name);
   }
+
   void Unload() {
     dlclose(lib_handle_);
+    lib_handle_ = nullptr;
   }
 #endif
 };
 
 TVM_REGISTER_GLOBAL("module.loadfile_so")
 .set_body([](TVMArgs args, TVMRetValue* rv) {
-    auto n = make_object<DSOModuleNode>();
+    auto n = make_object<DSOLibrary>();
     n->Init(args[0]);
-    *rv = runtime::Module(n);
+    *rv = CreateModuleFromLibrary(n);
   });
 }  // namespace runtime
 }  // namespace tvm