[Build][Tests] Enable FP16 for GPU builds in CI (dmlc#4030)

* Enable FP16 for GPU builds in CI * Limit default GPU archs to pascal and above * Disable FP16 dispatching for cuda architectures less than 60 * Fix linting * Fix typos
linlinling · May 26, 2022 · 7a065a9 · 7a065a9
1 parent d1124b7
commit 7a065a9
Show file tree

Hide file tree

Showing 3 changed files with 22 additions and 2 deletions.
diff --git a/src/array/cuda/atomic.cuh b/src/array/cuda/atomic.cuh
@@ -234,14 +234,17 @@ __device__ __forceinline__ double AtomicAdd<double>(double* addr, double val) {
 
 #ifdef USE_FP16
 #if defined(CUDART_VERSION) && CUDART_VERSION >= 10000
+// half make sure we have half support
+#if __CUDA_ARCH__ >= 600
 template <>
 __device__ __forceinline__ half AtomicAdd<half>(half* addr, half val) {
 #if __CUDA_ARCH__ >= 700
   return atomicAdd(addr, val);
 #else
   return *addr + val;
-#endif  // __CUDA_ARCH__
+#endif  // __CUDA_ARCH__ >= 700
 }
+#endif  // __CUDA_ARCH__ >= 600
 #endif  // defined(CUDART_VERSION) && CUDART_VERSION >= 10000
 #endif  // USE_FP16
 

diff --git a/src/array/cuda/utils.h b/src/array/cuda/utils.h
@@ -21,6 +21,7 @@ namespace cuda {
 #define CUDA_MAX_NUM_THREADS 1024
 
 #ifdef USE_FP16
+#if __CUDA_ARCH__ >= 600
 #define SWITCH_BITS(bits, DType, ...)                           \
   do {                                                          \
     if ((bits) == 16) {                                         \
@@ -36,6 +37,22 @@ namespace cuda {
       LOG(FATAL) << "Data type not recognized with bits " << bits; \
     }                                                           \
   } while (0)
+#else
+#define SWITCH_BITS(bits, DType, ...)                           \
+  do {                                                          \
+    if ((bits) == 16) {                                         \
+      LOG(FATAL) << "FP16 only supported on CUDA architectures >= 60"; \
+    } else if ((bits) == 32) {                                  \
+      typedef float DType;                                      \
+      { __VA_ARGS__ }                                           \
+    } else if ((bits) == 64) {                                  \
+      typedef double DType;                                     \
+      { __VA_ARGS__ }                                           \
+    } else {                                                    \
+      LOG(FATAL) << "Data type not recognized with bits " << bits; \
+    }                                                           \
+  } while (0)
+#endif  // __CUDA_ARCH__ >= 600
 #else  // USE_FP16
 #define SWITCH_BITS(bits, DType, ...)                           \
   do {                                                          \

diff --git a/tests/scripts/build_dgl.sh b/tests/scripts/build_dgl.sh
@@ -20,7 +20,7 @@ if [[ $arch == *"x86"* ]]; then
 fi
 
 if [ "$1" == "gpu" ]; then
-    CMAKE_VARS="-DUSE_CUDA=ON -DUSE_NCCL=ON $CMAKE_VARS"
+    CMAKE_VARS="-DUSE_CUDA=ON -DUSE_NCCL=ON -DUSE_FP16=ON $CMAKE_VARS"
 fi
 
 if [ -d build ]; then