Fixing tvmgpu issue & not restoring tvmop checks (#18818)

* fix the error message of reshape() * Fixing issue #16655 reshape() error message * test pr * fixing #17840 * fixing issue #17840 * Update compile.py * Update ndarray.py * Update c_api.cc * Update op_module.cc * Update op_module.h * Update op_module.h * Update op_module.h * fixing tvmgpu issue & not restoring tvmop checks Co-authored-by: Ubuntu <ubuntu@ip-172-31-34-249.ap-northeast-1.compute.internal> Co-authored-by: Ubuntu <ubuntu@ip-172-31-37-194.ap-northeast-1.compute.internal> Co-authored-by: jinboci <cijinbo@outlook.com>
apache · Aug 22, 2020 · a847df7 · a847df7
1 parent 4a7282f
commit a847df7
Show file tree

Hide file tree

Showing 4 changed files with 19 additions and 0 deletions.
diff --git a/contrib/tvmop/compile.py b/contrib/tvmop/compile.py
@@ -152,6 +152,8 @@ def get_cuda_arch(arch):
     # we create libtvmop.o first, which gives us chance to link tvm_runtime together with the libtvmop
     # to allow mxnet find external helper functions in libtvm_runtime
     func_binary.save(arguments.target_path + "/libtvmop.o")
+    if len(func_binary.imported_modules):
+        func_binary.imported_modules[0].save(arguments.target_path + "/libtvmop.cubin")
     ld_path = arguments.target_path if arguments.ld_path is None else arguments.ld_path
     create_shared(arguments.target_path + "/libtvmop.so",
                   arguments.target_path + "/libtvmop.o",

diff --git a/src/c_api/c_api.cc b/src/c_api/c_api.cc
@@ -1583,6 +1583,15 @@ int MXGetVersion(int *out) {
 int MXLoadTVMOp(const char *libpath) {
   API_BEGIN();
   tvm::runtime::TVMOpModule::Get()->Load(libpath);
+  tvm::runtime::TVMOpModule *global_module =  tvm::runtime::TVMOpModule::Get();
+  global_module->Load(libpath);
+#if MXNET_USE_CUDA
+  std::string libpathstr(libpath);
+  std::string cubinpath = libpathstr.substr(0, libpathstr.size() - 11) + "libtvmop.cubin";
+  tvm::runtime::TVMOpModule cubin_module;
+  cubin_module.Load(cubinpath);
+  global_module->Import(cubin_module);
+#endif
   API_END();
 }
 

diff --git a/src/operator/tvmop/op_module.cc b/src/operator/tvmop/op_module.cc
@@ -46,6 +46,12 @@ void TVMOpModule::Load(const std::string &filepath) {
   *module_ptr_ = module;
 }
 
+void TVMOpModule::Import(const TVMOpModule& module) {
+  CHECK(module_ptr_ != nullptr) << "module_ptr_ is not initialized.";
+  std::lock_guard<std::mutex> lock(mutex_);
+  module_ptr_->Import(*(module.module_ptr_));
+}
+
 PackedFunc GetFunction(const std::shared_ptr<Module> &module,
                        const std::string &op_name,
                        const std::vector<mxnet::TBlob> &args) {

diff --git a/src/operator/tvmop/op_module.h b/src/operator/tvmop/op_module.h
@@ -44,6 +44,8 @@ class TVMOpModule {
   // Load TVM operators binary
   void Load(const std::string& filepath);
 
+  void Import(const TVMOpModule& module);
+
   void Call(const std::string& func_name,
             const mxnet::OpContext& ctx,
             const std::vector<mxnet::TBlob>& args) const;