Enable stream-per-thread

wweic · Feb 22, 2021 · 8db2517 · 8db2517
1 parent 84359a9
commit 8db2517
Show file tree

Hide file tree

Showing 6 changed files with 8 additions and 0 deletions.
diff --git a/src/contrib/tf_op/tvm_dso_op_kernels.cc b/src/contrib/tf_op/tvm_dso_op_kernels.cc
@@ -18,6 +18,7 @@
  */
 
 #ifdef TF_TVMDSOOP_ENABLE_GPU
+#define CUDA_API_PER_THREAD_DEFAULT_STREAM
 #include <cuda_runtime.h>
 #endif
 #include <dlpack/dlpack.h>

diff --git a/src/runtime/contrib/cublas/cublas_utils.h b/src/runtime/contrib/cublas/cublas_utils.h
@@ -24,6 +24,7 @@
 #ifndef TVM_RUNTIME_CONTRIB_CUBLAS_CUBLAS_UTILS_H_
 #define TVM_RUNTIME_CONTRIB_CUBLAS_CUBLAS_UTILS_H_
 
+#define CUDA_API_PER_THREAD_DEFAULT_STREAM
 #include <cublas_v2.h>
 #include <cuda_runtime.h>
 #include <cuda_runtime_api.h>

diff --git a/src/runtime/cuda/cuda_common.h b/src/runtime/cuda/cuda_common.h
@@ -24,6 +24,7 @@
 #ifndef TVM_RUNTIME_CUDA_CUDA_COMMON_H_
 #define TVM_RUNTIME_CUDA_CUDA_COMMON_H_
 
+#define CUDA_API_PER_THREAD_DEFAULT_STREAM
 #include <cuda_runtime.h>
 #include <tvm/runtime/packed_func.h>
 

diff --git a/src/runtime/cuda/cuda_device_api.cc b/src/runtime/cuda/cuda_device_api.cc
@@ -21,6 +21,7 @@
  * \file cuda_device_api.cc
  * \brief GPU specific API
  */
+#define CUDA_API_PER_THREAD_DEFAULT_STREAM
 #include <cuda.h>
 #include <cuda_runtime.h>
 #include <dmlc/thread_local.h>

diff --git a/src/runtime/cuda/cuda_module.cc b/src/runtime/cuda/cuda_module.cc
@@ -22,6 +22,7 @@
  */
 #include "cuda_module.h"
 
+#define CUDA_API_PER_THREAD_DEFAULT_STREAM
 #include <cuda.h>
 #include <cuda_runtime.h>
 #include <tvm/runtime/registry.h>

diff --git a/src/target/opt/build_cuda_on.cc b/src/target/opt/build_cuda_on.cc
@@ -26,6 +26,7 @@
 #if defined(__linux__)
 #include <sys/stat.h>
 #endif
+#define CUDA_API_PER_THREAD_DEFAULT_STREAM
 #include <cuda_runtime.h>
 #include <nvrtc.h>
 
@@ -143,7 +144,9 @@ runtime::Module BuildCUDA(IRModule mod, Target target) {
   }
   std::string fmt = "ptx";
   std::string ptx;
+  code = "#define CUDA_API_PER_THREAD_DEFAULT_STREAM \n\n" + code;
   if (const auto* f = Registry::Get("tvm_callback_cuda_compile")) {
+    std::cout << "code is \n" << code << std::endl;
     ptx = (*f)(code).operator std::string();
     // Dirty matching to check PTX vs cubin.
     // TODO(tqchen) more reliable checks