Add RMM as a dependency of the GPU code (#753)

Author: benfred

implicit/gpu/CMakeLists.txt

@@ -16,8 +16,8 @@ else()
     add_compile_options(-DCCCL_IGNORE_DEPRECATED_STREAM_REF_HEADER)
 
     # use rapids-cmake to install dependencies
-    set(rapids-cmake-version "25.02")
-    file(DOWNLOAD https://raw.githubusercontent.com/rapidsai/rapids-cmake/refs/heads/branch-25.02/RAPIDS.cmake
+    set(rapids-cmake-version "26.02")
+    file(DOWNLOAD https://raw.githubusercontent.com/rapidsai/rapids-cmake/refs/heads/release/26.02/RAPIDS.cmake
         ${CMAKE_BINARY_DIR}/RAPIDS.cmake)
     include(${CMAKE_BINARY_DIR}/RAPIDS.cmake)
     include(rapids-cmake)
@@ -37,11 +37,11 @@ else()
     include(${rapids-cmake-dir}/cpm/rmm.cmake)
     rapids_cpm_rmm()
 
-    rapids_cpm_find(raft 25.02
+    rapids_cpm_find(raft 26.02
         GLOBAL_TARGETS raft::raft
         CPM_ARGS
           GIT_REPOSITORY  https://github.com/rapidsai/raft.git
-          GIT_TAG         v25.02.00
+          GIT_TAG         v26.02.00
           SOURCE_SUBDIR   cpp
           OPTIONS
               "BUILD_TESTS OFF"

implicit/gpu/__init__.py

@@ -3,17 +3,28 @@
 import warnings
 
 HAS_CUDA = False
+HAS_RMM = False
+
 try:
+    # RMM is required to enable GPU support - install with 'pip install rmm-cu13'
+    # Note that we need to import rmm here before importing the _cuda.so extension
+    import rmm  # noqa
+
+    HAS_RMM = True
+
     from ._cuda import *  # noqa
 
     get_device_count()  # noqa pylint: disable=undefined-variable
     HAS_CUDA = True
 
 except RuntimeError as e:
+    import warnings
+
     warnings.warn(
         f"CUDA extension is built, but disabling GPU support because of '{e}'",
     )
 except ImportError as e:
-    warnings.warn(
-        f"Disabling GPU support because of '{e}'",
-    )
+    if HAS_RMM:
+        warnings.warn(
+            f"Disabling GPU support because of '{e}'",
+        )

implicit/gpu/als.py

@@ -55,6 +55,9 @@ def __init__(
         calculate_training_loss=False,
         random_state=None,
     ):
+        if not implicit.gpu.HAS_RMM:
+            raise ValueError("RMM isn't installed, can't train on GPU.")
+
         if not implicit.gpu.HAS_CUDA:
             raise ValueError("No CUDA extension has been built, can't train on GPU.")
 

implicit/gpu/bpr.py

@@ -56,6 +56,9 @@ def __init__(
         random_state=None,
     ):
         super().__init__()
+        if not implicit.gpu.HAS_RMM:
+            raise ValueError("RMM isn't installed, can't train on GPU.")
+
         if not implicit.gpu.HAS_CUDA:
             raise ValueError("No CUDA extension has been built, can't train on GPU.")
 

implicit/gpu/knn.cu

@@ -5,9 +5,9 @@
 #include <cub/device/device_segmented_radix_sort.cuh>
 #include <cuda_runtime.h>
 #include <rmm/cuda_stream_view.hpp>
-#include <rmm/mr/device/cuda_memory_resource.hpp>
-#include <rmm/mr/device/device_memory_resource.hpp>
-#include <rmm/mr/device/pool_memory_resource.hpp>
+#include <rmm/mr/cuda_memory_resource.hpp>
+#include <rmm/mr/device_memory_resource.hpp>
+#include <rmm/mr/pool_memory_resource.hpp>
 #include <thrust/device_ptr.h>
 #include <thrust/functional.h>
 #include <thrust/iterator/constant_iterator.h>

pyproject.toml

@@ -31,8 +31,11 @@ classifiers = [
 ]
 dependencies = ["numpy>=1.17.0", "scipy>=0.16", "tqdm>=4.27", "threadpoolctl"]
 
+[project.optional-dependencies]
+gpu = ["rmm-cu13"]
+
 [tool.scikit-build]
-wheel.exclude = ["**.pyx", "**.cmake", "**.cuh", "**.hpp", "**.h", "**.hpp.in"]
+wheel.exclude = ["**.pyx", "**.cmake", "**.cuh", "**.hpp", "**.h", "**.hpp.in", "**librapids_logger.so", "**librmm.so"]
 cmake.version = ">=3.21.0"
 ninja.version = ">=1.10"
 build-dir = "build"