From a847df78f92d1e6824ab1682dbdc6d1269780a7b Mon Sep 17 00:00:00 2001
From: Jinbo Ci <32102965+jinboci@users.noreply.github.com>
Date: Sat, 22 Aug 2020 08:46:04 +0800
Subject: [PATCH] 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>
---
 contrib/tvmop/compile.py        | 2 ++
 src/c_api/c_api.cc              | 9 +++++++++
 src/operator/tvmop/op_module.cc | 6 ++++++
 src/operator/tvmop/op_module.h  | 2 ++
 4 files changed, 19 insertions(+)

diff --git a/contrib/tvmop/compile.py b/contrib/tvmop/compile.py
index 6341e70fe766..f15e5a727428 100644
--- 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
index faa030dcb459..4614b4e542b2 100644
--- 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
index c75e5a990086..d833ae0585c2 100644
--- 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
index 269a0aa50c11..8a25b3b9951e 100644
--- 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;