Dispatch thrust versions and upgrade rmm. (#7254)

Co-authored-by: Philip Hyunsu Cho <chohyu01@cs.washington.edu>

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 submodule.")
+endif (PLUGIN_RMM AND (NOT BUILD_WITH_CUDA_CUB))
 
 #-- Sanitizer
 if (USE_SANITIZER)

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
       """

python-package/xgboost/data.py

@@ -435,7 +435,11 @@ def _cudf_array_interfaces(data) -> Tuple[list, bytes]:
     interface is finished.
 
     """
+    try:
+        from cudf.api.types import is_categorical_dtype
+    except ImportError:
         from cudf.utils.dtypes import is_categorical_dtype
+
     cat_codes = []
     interfaces = []
     if _is_cudf_ser(data):
@@ -461,6 +465,9 @@ def _transform_cudf_df(
     feature_types: Optional[List[str]],
     enable_categorical: bool,
 ):
+    try:
+        from cudf.api.types import is_categorical_dtype
+    except ImportError:
         from cudf.utils.dtypes import is_categorical_dtype
 
     if feature_names is None:

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>>;
 
 /**

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

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