Update to CCCL 2.2.0.

ci/build_cpp.sh

@@ -9,10 +9,14 @@ export CMAKE_GENERATOR=Ninja
 
 rapids-print-env
 
+LIBRMM_CHANNEL=$(rapids-get-pr-conda-artifact rmm 1404 cpp)
+
 version=$(rapids-generate-version)
 
 rapids-logger "Begin cpp build"
 
-RAPIDS_PACKAGE_VERSION=${version} rapids-conda-retry mambabuild conda/recipes/libraft
+RAPIDS_PACKAGE_VERSION=${version} rapids-conda-retry mambabuild \
+    --channel "${LIBRMM_CHANNEL}" \
+    conda/recipes/libraft
 
 rapids-upload-conda-to-s3 cpp

ci/build_python.sh

@@ -12,6 +12,8 @@ rapids-print-env
 rapids-logger "Begin py build"
 
 CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
+LIBRMM_CHANNEL=$(rapids-get-pr-conda-artifact rmm 1404 cpp)
+RMM_CHANNEL=$(rapids-get-pr-conda-artifact rmm 1404 python)
 
 version=$(rapids-generate-version)
 git_commit=$(git rev-parse HEAD)
@@ -29,18 +31,24 @@ done
 rapids-conda-retry mambabuild \
   --no-test \
   --channel "${CPP_CHANNEL}" \
+  --channel "${LIBRMM_CHANNEL}" \
+  --channel "${RMM_CHANNEL}" \
   conda/recipes/pylibraft
 
 rapids-conda-retry mambabuild \
   --no-test \
   --channel "${CPP_CHANNEL}" \
+  --channel "${LIBRMM_CHANNEL}" \
+  --channel "${RMM_CHANNEL}" \
   --channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
   conda/recipes/raft-dask
 
 # Build ann-bench for each cuda and python version
 rapids-conda-retry mambabuild \
   --no-test \
   --channel "${CPP_CHANNEL}" \
+  --channel "${LIBRMM_CHANNEL}" \
+  --channel "${RMM_CHANNEL}" \
   --channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
   conda/recipes/raft-ann-bench
 
@@ -51,6 +59,8 @@ if [[ ${RAPIDS_CUDA_MAJOR} == "11" ]]; then
   rapids-conda-retry mambabuild \
   --no-test \
   --channel "${CPP_CHANNEL}" \
+  --channel "${LIBRMM_CHANNEL}" \
+  --channel "${RMM_CHANNEL}" \
   --channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
   conda/recipes/raft-ann-bench-cpu
 fi

ci/test_cpp.sh

@@ -19,13 +19,15 @@ conda activate test
 set -u
 
 CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
+LIBRMM_CHANNEL=$(rapids-get-pr-conda-artifact rmm 1404 cpp)
 RAPIDS_TESTS_DIR=${RAPIDS_TESTS_DIR:-"${PWD}/test-results"}/
 mkdir -p "${RAPIDS_TESTS_DIR}"
 
 rapids-print-env
 
 rapids-mamba-retry install \
   --channel "${CPP_CHANNEL}" \
+  --channel "${LIBRMM_CHANNEL}" \
   libraft-headers libraft libraft-tests
 
 rapids-logger "Check GPU usage"

ci/test_python.sh

@@ -21,6 +21,8 @@ set -u
 rapids-logger "Downloading artifacts from previous jobs"
 CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
 PYTHON_CHANNEL=$(rapids-download-conda-from-s3 python)
+LIBRMM_CHANNEL=$(rapids-get-pr-conda-artifact rmm 1404 cpp)
+RMM_CHANNEL=$(rapids-get-pr-conda-artifact rmm 1404 python)
 
 RAPIDS_TESTS_DIR=${RAPIDS_TESTS_DIR:-"${PWD}/test-results"}
 RAPIDS_COVERAGE_DIR=${RAPIDS_COVERAGE_DIR:-"${PWD}/coverage-results"}
@@ -31,6 +33,8 @@ rapids-print-env
 rapids-mamba-retry install \
   --channel "${CPP_CHANNEL}" \
   --channel "${PYTHON_CHANNEL}" \
+  --channel "${LIBRMM_CHANNEL}" \
+  --channel "${RMM_CHANNEL}" \
   libraft libraft-headers pylibraft raft-dask
 
 rapids-logger "Check GPU usage"

cpp/CMakeLists.txt

@@ -177,8 +177,8 @@ include(cmake/modules/ConfigureCUDA.cmake)
 rapids_cpm_init()
 
 if(NOT BUILD_CPU_ONLY)
-  # thrust before rmm/cuco so we get the right version of thrust/cub
-  include(cmake/thirdparty/get_thrust.cmake)
+  # CCCL before rmm/cuco so we get the right version of CCCL
+  include(cmake/thirdparty/get_cccl.cmake)
   include(cmake/thirdparty/get_rmm.cmake)
   include(cmake/thirdparty/get_cutlass.cmake)
 
@@ -206,7 +206,7 @@ target_include_directories(
 
 if(NOT BUILD_CPU_ONLY)
   # Keep RAFT as lightweight as possible. Only CUDA libs and rmm should be used in global target.
-  target_link_libraries(raft INTERFACE rmm::rmm cuco::cuco nvidia::cutlass::cutlass raft::Thrust)
+  target_link_libraries(raft INTERFACE rmm::rmm cuco::cuco nvidia::cutlass::cutlass CCCL::Thrust)
 endif()
 
 target_compile_features(raft INTERFACE cxx_std_17 $<BUILD_INTERFACE:cuda_std_17>)
@@ -628,16 +628,6 @@ Imported Targets:
 
 set(code_string ${nvtx_export_string})
 
-string(
-  APPEND
-  code_string
-  [=[
-if(NOT TARGET raft::Thrust)
-  thrust_create_target(raft::Thrust FROM_OPTIONS)
-endif()
-]=]
-)
-
 string(
   APPEND
   code_string

cpp/cmake/thirdparty/get_thrust.cmake → cpp/cmake/thirdparty/get_cccl.cmake

@@ -12,13 +12,10 @@
 # the License.
 # =============================================================================
 
-# Use CPM to find or clone thrust
-function(find_and_configure_thrust)
-        include(${rapids-cmake-dir}/cpm/thrust.cmake)
-
-        rapids_cpm_thrust( NAMESPACE raft
-                           BUILD_EXPORT_SET raft-exports
-                           INSTALL_EXPORT_SET raft-exports)
+# Use CPM to find or clone CCCL
+function(find_and_configure_cccl)
+        include(${rapids-cmake-dir}/cpm/cccl.cmake)
+        rapids_cpm_cccl(BUILD_EXPORT_SET raft-exports INSTALL_EXPORT_SET raft-exports)
 endfunction()
 
-find_and_configure_thrust()
+find_and_configure_cccl()

cpp/template/cmake/thirdparty/fetch_rapids.cmake

@@ -14,6 +14,8 @@
 # Use this variable to update RAPIDS and RAFT versions
 set(RAPIDS_VERSION "24.02")
 
+set(rapids-cmake-repo bdice/rapids-cmake)
+set(rapids-cmake-branch cuco-cccl-2.2.0)
 if(NOT EXISTS ${CMAKE_CURRENT_BINARY_DIR}/RAFT_RAPIDS.cmake)
     file(DOWNLOAD https://raw.githubusercontent.com/rapidsai/rapids-cmake/branch-${RAPIDS_VERSION}/RAPIDS.cmake
             ${CMAKE_CURRENT_BINARY_DIR}/RAFT_RAPIDS.cmake)

fetch_rapids.cmake

@@ -11,6 +11,8 @@
 # or implied. See the License for the specific language governing permissions and limitations under
 # the License.
 # =============================================================================
+set(rapids-cmake-repo bdice/rapids-cmake)
+set(rapids-cmake-branch cuco-cccl-2.2.0)
 if(NOT EXISTS ${CMAKE_CURRENT_BINARY_DIR}/RAFT_RAPIDS.cmake)
   file(DOWNLOAD https://raw.githubusercontent.com/rapidsai/rapids-cmake/branch-24.02/RAPIDS.cmake
        ${CMAKE_CURRENT_BINARY_DIR}/RAFT_RAPIDS.cmake