CUDA 11.0

dmlc · Feb 13, 2022 · 9b3cd3e · 9b3cd3e
1 parent 2951520
commit 9b3cd3e
Showing 1 changed file with 11 additions and 4 deletions.
diff --git a/src/common/device_helpers.cuh b/src/common/device_helpers.cuh
@@ -1558,17 +1558,24 @@ class CUDAEvent {
   cudaEvent_t Get() const { return event_; }
 };
 
-#ifndef cudaEventWaitDefault
-#define cudaEventWaitDefault 0x00
-#endif
-
 class CUDAStreamView {
   cudaStream_t stream_{nullptr};
 
  public:
   explicit CUDAStreamView(cudaStream_t s) : stream_{s} {}
   void Wait(CUDAEvent const &e) {
+    // #if CUDA
+#if defined(__CUDACC_VER_MAJOR__)
+#if __CUDACC_VER_MAJOR__ == 11 && __CUDACC_VER_MINOR__ == 0
+    // CUDA == 11.0
+    dh::safe_cuda(cudaStreamWaitEvent(stream_, e.Get(), 0));
+#else
+    // CUDA > 11.0
     dh::safe_cuda(cudaStreamWaitEvent(stream_, e.Get(), cudaEventWaitDefault));
+#endif  // __CUDACC_VER_MAJOR__ == 11 && __CUDACC_VER_MINOR__ == 0:
+#else   // clang
+  dh::safe_cuda(cudaStreamWaitEvent(stream_, e.Get(), cudaEventWaitDefault));
+#endif  //  defined(__CUDACC_VER_MAJOR__)
   }
   operator cudaStream_t() const {  // NOLINT
     return stream_;