add cuda 10.1 support (#4468)

dmlc · May 14, 2019 · df2cdac · df2cdac
1 parent c6f2a7e
commit df2cdac
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 9 deletions.
diff --git a/Jenkinsfile b/Jenkinsfile
@@ -58,6 +58,7 @@ pipeline {
             'build-gpu-cuda8.0': { BuildCUDA(cuda_version: '8.0') },
             'build-gpu-cuda9.0': { BuildCUDA(cuda_version: '9.0') },
             'build-gpu-cuda10.0': { BuildCUDA(cuda_version: '10.0') },
+            'build-gpu-cuda10.1': { BuildCUDA(cuda_version: '10.1') },
             'build-jvm-packages': { BuildJVMPackages(spark_version: '2.4.3') },
             'build-jvm-doc': { BuildJVMDoc() }
           ])
@@ -74,9 +75,10 @@ pipeline {
             'test-python-gpu-cuda8.0': { TestPythonGPU(cuda_version: '8.0') },
             'test-python-gpu-cuda9.0': { TestPythonGPU(cuda_version: '9.0') },
             'test-python-gpu-cuda10.0': { TestPythonGPU(cuda_version: '10.0') },
-            'test-python-mgpu-cuda10.0': { TestPythonGPU(cuda_version: '10.0', multi_gpu: true) },
-            'test-cpp-gpu': { TestCppGPU(cuda_version: '10.0') },
-            'test-cpp-mgpu': { TestCppGPU(cuda_version: '10.0', multi_gpu: true) },
+            'test-python-gpu-cuda10.1': { TestPythonGPU(cuda_version: '10.1') },
+            'test-python-mgpu-cuda10.1': { TestPythonGPU(cuda_version: '10.1', multi_gpu: true) },
+            'test-cpp-gpu': { TestCppGPU(cuda_version: '10.1') },
+            'test-cpp-mgpu': { TestCppGPU(cuda_version: '10.1', multi_gpu: true) },
             'test-jvm-jdk8': { CrossTestJVMwithJDK(jdk_version: '8', spark_version: '2.4.3') },
             'test-jvm-jdk11': { CrossTestJVMwithJDK(jdk_version: '11') },
             'test-jvm-jdk12': { CrossTestJVMwithJDK(jdk_version: '12') },

diff --git a/src/common/device_helpers.cuh b/src/common/device_helpers.cuh
@@ -27,10 +27,6 @@
 #include "../common/io.h"
 #endif
 
-#if __CUDACC_VER_MAJOR__ == 10 && __CUDACC_VER_MINOR__ == 1
-#error "CUDA 10.1 is not supported, see #4264."
-#endif
-
 namespace dh {
 
 #define HOST_DEV_INLINE XGBOOST_DEVICE __forceinline__

diff --git a/tests/ci_build/Dockerfile.gpu_build b/tests/ci_build/Dockerfile.gpu_build
@@ -21,7 +21,7 @@ RUN \
 # NCCL2 (License: https://docs.nvidia.com/deeplearning/sdk/nccl-sla/index.html)
 RUN \
     export CUDA_SHORT=`echo $CUDA_VERSION | egrep -o '[0-9]+\.[0-9]'` && \
-    if [ "${CUDA_SHORT}" != "10.0" ]; then \
+    if [ "${CUDA_SHORT}" != "10.0" ] && [ "${CUDA_SHORT}" != "10.1" ]; then \
     wget https://developer.download.nvidia.com/compute/redist/nccl/v2.2/nccl_2.2.13-1%2Bcuda${CUDA_SHORT}_x86_64.txz && \
     tar xf "nccl_2.2.13-1+cuda${CUDA_SHORT}_x86_64.txz" && \
     cp nccl_2.2.13-1+cuda${CUDA_SHORT}_x86_64/include/nccl.h /usr/include && \
@@ -31,7 +31,7 @@ RUN \
     wget https://developer.download.nvidia.com/compute/machine-learning/repos/rhel7/x86_64/nvidia-machine-learning-repo-rhel7-1.0.0-1.x86_64.rpm && \
     rpm -i nvidia-machine-learning-repo-rhel7-1.0.0-1.x86_64.rpm && \
     yum -y update && \
-    yum install -y libnccl-2.4.2-1+cuda10.0 libnccl-devel-2.4.2-1+cuda10.0 libnccl-static-2.4.2-1+cuda10.0 && \
+    yum install -y libnccl-2.4.2-1+cuda${CUDA_SHORT} libnccl-devel-2.4.2-1+cuda${CUDA_SHORT} libnccl-static-2.4.2-1+cuda${CUDA_SHORT} && \
     rm -f nvidia-machine-learning-repo-rhel7-1.0.0-1.x86_64.rpm; fi
 
 ENV PATH=/opt/python/bin:$PATH