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 · Nov 5, 2020 · daaa127 · daaa127
1 parent f53dbc7
commit daaa127
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 11 deletions.
diff --git a/cub/util_arch.cuh b/cub/util_arch.cuh
@@ -45,9 +45,10 @@ namespace cub {
 
 #ifndef DOXYGEN_SHOULD_SKIP_THIS    // Do not document
 
-#if ((__CUDACC_VER_MAJOR__ >= 9) || defined(__NVCOMPILER_CUDA__)) && \
-        !defined(CUB_USE_COOPERATIVE_GROUPS)
-    #define CUB_USE_COOPERATIVE_GROUPS
+#if ((__CUDACC_VER_MAJOR__ >= 9) || defined(__NVCOMPILER_CUDA__) ||            \
+     CUDA_VERSION >= 9000) &&                                                  \
+  !defined(CUB_USE_COOPERATIVE_GROUPS)
+#define CUB_USE_COOPERATIVE_GROUPS
 #endif
 
 /// In device code, CUB_PTX_ARCH expands to the PTX version for which we are

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>
 {
@@ -1143,7 +1143,7 @@ template <> struct NumericTraits<unsigned long long> :  BaseTraits<UNSIGNED_INTE
 
 template <> struct NumericTraits<float> :               BaseTraits<FLOATING_POINT, true, false, unsigned int, float> {};
 template <> struct NumericTraits<double> :              BaseTraits<FLOATING_POINT, true, false, unsigned long long, double> {};
-#if (__CUDACC_VER_MAJOR__ >= 9)
+#if (__CUDACC_VER_MAJOR__ >= 9 || CUDA_VERSION >= 9000)
     template <> struct NumericTraits<__half> :          BaseTraits<FLOATING_POINT, true, false, unsigned short, __half> {};
 #endif
 

diff --git a/test/test_device_radix_sort.cu b/test/test_device_radix_sort.cu
@@ -37,7 +37,7 @@
 #include <algorithm>
 #include <typeinfo>
 
-#if (__CUDACC_VER_MAJOR__ >= 9)
+#if (__CUDACC_VER_MAJOR__ >= 9 || CUDA_VERSION >= 9000)
     #include <cuda_fp16.h>
 #endif
 
@@ -733,7 +733,7 @@ void Test(
     ValueT      *h_reference_values)
 {
     // Key alias type
-#if (__CUDACC_VER_MAJOR__ >= 9)
+#if (__CUDACC_VER_MAJOR__ >= 9 || CUDA_VERSION >= 9000)
     typedef typename If<Equals<KeyT, half_t>::VALUE, __half, KeyT>::Type KeyAliasT;
 #else
     typedef KeyT KeyAliasT;
@@ -1240,7 +1240,7 @@ int main(int argc, char** argv)
 
     printf("\n-------------------------------\n");
 
-#if (__CUDACC_VER_MAJOR__ >= 9)
+#if (__CUDACC_VER_MAJOR__ >= 9 || CUDA_VERSION >= 9000)
     Test<CUB,           half_t,             NullType, IS_DESCENDING>(num_items, 1, RANDOM, entropy_reduction, 0, bits);
 #endif
     Test<CUB,           float,              NullType, IS_DESCENDING>(num_items, 1, RANDOM, entropy_reduction, 0, bits);
@@ -1299,7 +1299,7 @@ int main(int argc, char** argv)
         TestGen<long long>            (num_items, num_segments);
         TestGen<unsigned long long>   (num_items, num_segments);
 
-#if (__CUDACC_VER_MAJOR__ >= 9)
+#if (__CUDACC_VER_MAJOR__ >= 9 || CUDA_VERSION >= 9000)
         TestGen<half_t>                (num_items, num_segments);
 #endif
         TestGen<float>                (num_items, num_segments);
@@ -1313,4 +1313,3 @@ int main(int argc, char** argv)
 
     return 0;
 }
-