Fix CUDA version detection in CUB

This fixes the problem with CUB using deprecated shfl/vote instructions when CUB is compiled with clang (e.g. some TensorFlow builds).
NVIDIA · Sep 23, 2019 · fd6e7a6 · fd6e7a6
1 parent 135da60
commit fd6e7a6
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/cub/util_arch.cuh b/cub/util_arch.cuh
@@ -44,7 +44,8 @@ namespace cub {
 
 #ifndef DOXYGEN_SHOULD_SKIP_THIS    // Do not document
 
-#if (__CUDACC_VER_MAJOR__ >= 9) && !defined(CUB_USE_COOPERATIVE_GROUPS)
+#if !defined(CUB_USE_COOPERATIVE_GROUPS) && \
+    (__CUDACC_VER_MAJOR__ >= 9 || CUDA_VERSION >= 9000)
     #define CUB_USE_COOPERATIVE_GROUPS
 #endif
 

diff --git a/cub/util_type.cuh b/cub/util_type.cuh
@@ -37,7 +37,7 @@
 #include <limits>
 #include <cfloat>
 
-#if (__CUDACC_VER_MAJOR__ >= 9)
+#if (__CUDACC_VER_MAJOR__ >= 9 || CUDA_VERSION >= 9000)
     #include <cuda_fp16.h>
 #endif
 
@@ -1063,7 +1063,7 @@ struct FpLimits<double>
 };
 
 
-#if (__CUDACC_VER_MAJOR__ >= 9)
+#if (__CUDACC_VER_MAJOR__ >= 9 || CUDA_VERSION >= 9000)
 template <>
 struct FpLimits<__half>
 {