From 9dd9e0070c376795a608ce7901ff6510a49d35ba Mon Sep 17 00:00:00 2001
From: fis <jm.yuan@outlook.com>
Date: Fri, 24 Sep 2021 14:41:52 +0800
Subject: [PATCH] Dispatch thrust versions and upgrade rmm.

---
 CMakeLists.txt                |  3 +++
 Jenkinsfile                   |  2 +-
 src/common/device_helpers.cuh | 10 +++++++---
 tests/ci_build/Dockerfile.gpu |  2 +-
 tests/ci_build/Dockerfile.rmm |  2 +-
 5 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index bd21e9b67fb2..b3c2395df22b 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -112,6 +112,9 @@ endif (ENABLE_ALL_WARNINGS)
 if (BUILD_STATIC_LIB AND (R_LIB OR JVM_BINDINGS))
   message(SEND_ERROR "Cannot build a static library libxgboost.a when R or JVM packages are enabled.")
 endif (BUILD_STATIC_LIB AND (R_LIB OR JVM_BINDINGS))
+if (PLUGIN_RMM AND (NOT BUILD_WITH_CUDA_CUB))
+  message(SEND_ERROR "Cannot build with RMM using cub cubmodule.")
+endif (PLUGIN_RMM AND (NOT BUILD_WITH_CUDA_CUB))
 
 #-- Sanitizer
 if (USE_SANITIZER)
diff --git a/Jenkinsfile b/Jenkinsfile
index a038d27264dd..7ed17a966577 100644
--- a/Jenkinsfile
+++ b/Jenkinsfile
@@ -249,7 +249,7 @@ def BuildCUDA(args) {
       docker_args = "--build-arg CUDA_VERSION_ARG=${args.cuda_version}"
       sh """
       rm -rf build/
-      ${dockerRun} ${container_type} ${docker_binary} ${docker_args} tests/ci_build/build_via_cmake.sh --conda-env=gpu_test -DUSE_CUDA=ON -DUSE_NCCL=ON -DPLUGIN_RMM=ON ${arch_flag}
+      ${dockerRun} ${container_type} ${docker_binary} ${docker_args} tests/ci_build/build_via_cmake.sh --conda-env=gpu_test -DUSE_CUDA=ON -DUSE_NCCL=ON -DPLUGIN_RMM=ON -DBUILD_WITH_CUDA_CUB=ON ${arch_flag}
       ${dockerRun} ${container_type} ${docker_binary} ${docker_args} bash -c "cd python-package && rm -rf dist/* && python setup.py bdist_wheel --universal"
       ${dockerRun} ${container_type} ${docker_binary} ${docker_args} python tests/ci_build/rename_whl.py python-package/dist/*.whl ${commit_id} manylinux2014_x86_64
       """
diff --git a/src/common/device_helpers.cuh b/src/common/device_helpers.cuh
index c2c7ff07b463..61e0fd553659 100644
--- a/src/common/device_helpers.cuh
+++ b/src/common/device_helpers.cuh
@@ -708,6 +708,10 @@ constexpr std::pair<int, int> CUDAVersion() {
 #endif  // defined(__CUDACC_VER_MAJOR__)
 }
 
+constexpr std::pair<int32_t, int32_t> ThrustVersion() {
+  return std::make_pair(THRUST_MAJOR_VERSION, THRUST_MINOR_VERSION);
+}
+
 namespace detail {
 template <typename T>
 using TypedDiscardCTK114 = thrust::discard_iterator<T>;
@@ -721,9 +725,9 @@ class TypedDiscard : public thrust::discard_iterator<T> {
 
 template <typename T>
 using TypedDiscard =
-    std::conditional_t<((CUDAVersion().first == 11 &&
-                         CUDAVersion().second >= 4) ||
-                        CUDAVersion().first > 11),
+    std::conditional_t<((ThrustVersion().first == 1 &&
+                         ThrustVersion().second >= 12) ||
+                        ThrustVersion().first > 1),
                        detail::TypedDiscardCTK114<T>, detail::TypedDiscard<T>>;
 
 /**
diff --git a/tests/ci_build/Dockerfile.gpu b/tests/ci_build/Dockerfile.gpu
index b86e2133705b..72f7b7853c9c 100644
--- a/tests/ci_build/Dockerfile.gpu
+++ b/tests/ci_build/Dockerfile.gpu
@@ -19,7 +19,7 @@ ENV PATH=/opt/python/bin:$PATH
 # Create new Conda environment with cuDF, Dask, and cuPy
 RUN \
     conda create -n gpu_test -c rapidsai-nightly -c rapidsai -c nvidia -c conda-forge -c defaults \
-        python=3.7 cudf=21.08* rmm=21.08* cudatoolkit=$CUDA_VERSION_ARG dask dask-cuda dask-cudf cupy=9.1* \
+        python=3.8 cudf=21.10* rmm=21.10* cudatoolkit=$CUDA_VERSION_ARG dask dask-cuda=21.10* dask-cudf=21.10* cupy=9.1* \
         numpy pytest scipy scikit-learn pandas matplotlib wheel python-kubernetes urllib3 graphviz hypothesis
 
 ENV GOSU_VERSION 1.10
diff --git a/tests/ci_build/Dockerfile.rmm b/tests/ci_build/Dockerfile.rmm
index 65ec7266fbfa..3d4906d4a25a 100644
--- a/tests/ci_build/Dockerfile.rmm
+++ b/tests/ci_build/Dockerfile.rmm
@@ -29,7 +29,7 @@ ENV PATH=/opt/python/bin:$PATH
 # Create new Conda environment with RMM
 RUN \
     conda create -n gpu_test -c rapidsai-nightly -c rapidsai -c nvidia -c conda-forge -c defaults \
-        python=3.7 rmm=21.08* cudatoolkit=$CUDA_VERSION_ARG
+        python=3.8 rmm=21.10* cudatoolkit=$CUDA_VERSION_ARG
 
 ENV GOSU_VERSION 1.10