| /*! |
| * Copyright (c) 2017 by Contributors |
| * \file dso_dll_module.cc |
| * \brief Module to load from dynamic shared library. |
| */ |
| #include <tvm/runtime/module.h> |
| #include <tvm/runtime/registry.h> |
| #include <tvm/runtime/packed_func.h> |
| #include "module_util.h" |
| |
| #if defined(_WIN32) |
| #include <windows.h> |
| #else |
| #include <dlfcn.h> |
| #endif |
| |
| namespace tvm { |
| namespace runtime { |
| |
| // Module to load from dynamic shared libary. |
| // This is the default module TVM used for host-side AOT |
| class DSOModuleNode final : public ModuleNode { |
| public: |
| ~DSOModuleNode() { |
| if (lib_handle_) Unload(); |
| } |
| |
| const char* type_key() const final { |
| return "dso"; |
| } |
| |
| PackedFunc GetFunction( |
| const std::string& name, |
| const std::shared_ptr<ModuleNode>& 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_); |
| } |
| } |
| |
| private: |
| // Platform dependent handling. |
| #if defined(_WIN32) |
| // library handle |
| HMODULE lib_handle_{nullptr}; |
| // Load the library |
| void Load(const std::string& name) { |
| // use wstring version that is needed by LLVM. |
| std::wstring wname(name.begin(), name.end()); |
| lib_handle_ = LoadLibraryW(wname.c_str()); |
| 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_); |
| } |
| #else |
| // Library handle |
| void* lib_handle_{nullptr}; |
| // load the library |
| void Load(const std::string& name) { |
| lib_handle_ = dlopen(name.c_str(), RTLD_LAZY | RTLD_LOCAL); |
| CHECK(lib_handle_ != nullptr) |
| << "Failed to load dynamic shared library " << name |
| << " " << dlerror(); |
| } |
| void* GetSymbol(const char* name) { |
| return dlsym(lib_handle_, name); |
| } |
| void Unload() { |
| dlclose(lib_handle_); |
| } |
| #endif |
| }; |
| |
| TVM_REGISTER_GLOBAL("module.loadfile_so") |
| .set_body([](TVMArgs args, TVMRetValue* rv) { |
| std::shared_ptr<DSOModuleNode> n = std::make_shared<DSOModuleNode>(); |
| n->Init(args[0]); |
| *rv = runtime::Module(n); |
| }); |
| } // namespace runtime |
| } // namespace tvm |