fixing apache#17840

jinboci · Jun 9, 2020 · 087a100 · 087a100
1 parent 7627eaa
commit 087a100
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/contrib/tvmop/compile.py b/contrib/tvmop/compile.py
@@ -152,6 +152,12 @@ 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")
+    func_binary.imported_modules[0].save(arguments.target_path + "/libtvmop.cubin")
+    # dev_module = func_binary.imported_modules[0]
+    # print("-----GPU code-----")
+    # print(dev_module.get_source())
+    # print("-----GPU code END-----")
+
     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
@@ -1365,10 +1365,9 @@ int MXLoadTVMOp(const char *libpath) {
   API_BEGIN();
   tvm::runtime::TVMOpModule *libpath_module =  tvm::runtime::TVMOpModule::Get();
   libpath_module->Load(libpath);
-  tvm::runtime::TVMOpModule *cubin_module =  tvm::runtime::TVMOpModule::Get();
-  cubin_module->Load("/home/ubuntu/Documents/mxnet/build/libtvmop.cubin");
-  libpath_module->Import(*cubin_module);
-  std::cout << "_________ MXLoadTVMOp succeed __________" << std::endl;
+  tvm::runtime::TVMOpModule cubin_module;
+  cubin_module.Load("/home/ubuntu/Documents/mxnet/build/libtvmop.cubin");
+  libpath_module->Import(cubin_module);
   API_END();
 }