Revert "Remove RMM/RAFT dependencies, Part 1 (#195)"

conda/environments/all_cuda-129_arch-aarch64.yaml

@@ -11,7 +11,6 @@ dependencies:
 - click
 - cmake>=4.0
 - cuda-bindings>=12.9.2,<13.0
-- cuda-core>=1.0.0,<2.0.0
 - cuda-cudart-dev
 - cuda-nvcc
 - cuda-profiler-api

conda/environments/all_cuda-129_arch-x86_64.yaml

@@ -11,7 +11,6 @@ dependencies:
 - click
 - cmake>=4.0
 - cuda-bindings>=12.9.2,<13.0
-- cuda-core>=1.0.0,<2.0.0
 - cuda-cudart-dev
 - cuda-nvcc
 - cuda-profiler-api

conda/environments/all_cuda-133_arch-aarch64.yaml

@@ -11,7 +11,6 @@ dependencies:
 - click
 - cmake>=4.0
 - cuda-bindings>=13.0.1,<14.0
-- cuda-core>=1.0.0,<2.0.0
 - cuda-cudart-dev
 - cuda-nvcc
 - cuda-profiler-api

conda/environments/all_cuda-133_arch-x86_64.yaml

@@ -11,7 +11,6 @@ dependencies:
 - click
 - cmake>=4.0
 - cuda-bindings>=13.0.1,<14.0
-- cuda-core>=1.0.0,<2.0.0
 - cuda-cudart-dev
 - cuda-nvcc
 - cuda-profiler-api

cpp/CMakeLists.txt

@@ -199,6 +199,7 @@ endif()
 # CCCL before RMM, and RMM before RAFT
 include(cmake/thirdparty/get_cccl.cmake)
 if(NVFOREST_ENABLE_GPU)
+  include(cmake/thirdparty/get_rmm.cmake)
   include(cmake/thirdparty/get_raft.cmake)
 endif()
 include(cmake/thirdparty/get_treelite.cmake)
@@ -251,8 +252,7 @@ if(NVFOREST_ENABLE_GPU)
             src/infer8.cu
             src/infer9.cu
             src/infer10.cu
-            src/infer11.cu
-            src/detail/device_buffer.cu)
+            src/infer11.cu)
 endif()
 target_sources(
   ${NVFOREST_CPP_TARGET}
@@ -312,7 +312,7 @@ elseif(NVFOREST_EXPORT_TREELITE_LINKAGE)
 endif()
 
 if(NVFOREST_ENABLE_GPU)
-  list(APPEND _nvforest_cpp_public_libs raft::raft CUDA::cudart_static CCCL::CCCL)
+  list(APPEND _nvforest_cpp_public_libs rmm::rmm raft::raft CUDA::cudart_static)
 endif()
 
 # These are always private:

cpp/cmake/thirdparty/get_rmm.cmake

@@ -0,0 +1,13 @@
+# =============================================================================
+# cmake-format: off
+# SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+# cmake-format: on
+# =============================================================================
+
+function(find_and_configure_rmm)
+  include(${rapids-cmake-dir}/cpm/rmm.cmake)
+  rapids_cpm_rmm(BUILD_EXPORT_SET nvforest-exports INSTALL_EXPORT_SET nvforest-exports)
+endfunction()
+
+find_and_configure_rmm()

cpp/include/nvforest/detail/device_id/gpu.hpp

@@ -7,20 +7,23 @@
 #include <nvforest/detail/device_id/base.hpp>
 #include <nvforest/device_type.hpp>
 
+#include <rmm/cuda_device.hpp>
+
 namespace nvforest::detail {
 template <>
 struct device_id<device_type::gpu> {
+  using value_type = typename rmm::cuda_device_id::value_type;
   device_id() noexcept(false)
     : id_{[]() {
-        auto raw_id = int{};
+        auto raw_id = value_type{};
         cuda_check(cudaGetDevice(&raw_id));
         return raw_id;
       }()} {};
-  device_id(int dev_id) noexcept : id_{dev_id} {};
+  device_id(value_type dev_id) noexcept : id_{dev_id} {};
 
-  auto value() const noexcept { return id_; }
+  auto value() const noexcept { return id_.value(); }
 
  private:
-  int id_;
+  rmm::cuda_device_id id_;
 };
 }  // namespace nvforest::detail

cpp/include/nvforest/detail/device_setter/gpu.hpp

@@ -8,21 +8,9 @@
 #include <nvforest/detail/device_setter/base.hpp>
 #include <nvforest/device_type.hpp>
 
-#include <cuda_runtime_api.h>
-
-#include <cstdio>
+#include <raft/util/cudart_utils.hpp>
 
-#define NVFOREST_CUDA_TRY_NO_THROW(call)                           \
-  do {                                                             \
-    cudaError_t const status = call;                               \
-    if (cudaSuccess != status) {                                   \
-      printf("CUDA call='%s' at file=%s line=%d failed with %s\n", \
-             #call,                                                \
-             __FILE__,                                             \
-             __LINE__,                                             \
-             cudaGetErrorString(status));                          \
-    }                                                              \
-  } while (0)
+#include <cuda_runtime_api.h>
 
 namespace nvforest::detail {
 
@@ -39,7 +27,7 @@ struct device_setter<device_type::gpu> {
     cuda_check(cudaSetDevice(device.value()));
   }
 
-  ~device_setter() { NVFOREST_CUDA_TRY_NO_THROW(cudaSetDevice(prev_device_.value())); }
+  ~device_setter() { RAFT_CUDA_TRY_NO_THROW(cudaSetDevice(prev_device_.value())); }
 
  private:
   device_id<device_type::gpu> prev_device_;

cpp/include/nvforest/detail/infer/gpu.cuh

@@ -21,8 +21,6 @@
 #include <nvforest/exceptions.hpp>
 #include <nvforest/infer_kind.hpp>
 
-#include <cuda/stream>
-
 #include <cstddef>
 #include <cstdint>
 #include <optional>
@@ -43,19 +41,6 @@ inline auto compute_output_size(index_type row_output_size,
   return result;
 }
 
-// If a non-default stream is provided, it must reside on the correct device.
-inline void validate_stream(device_id<device_type::gpu> device, cuda::stream_ref stream)
-{
-  cuda::stream_ref legacy_default_stream{cudaStreamLegacy};
-  cuda::stream_ref per_thread_default_stream{cudaStreamPerThread};
-  if (stream != legacy_default_stream && stream != per_thread_default_stream &&
-      stream.device().get() != device.value()) {
-    throw std::runtime_error{std::string("Stream on the wrong device. ") +
-                             "Expected: " + std::to_string(device.value()) +
-                             ", Actual: " + std::to_string(stream.device().get())};
-  }
-}
-
 /* A wrapper around the underlying inference kernels to support dispatching to
  * the right kernel
  *
@@ -114,8 +99,6 @@ std::enable_if_t<D == device_type::gpu, void> infer(
 {
   using output_t = typename forest_t::template raw_output_type<vector_output_t>;
 
-  validate_stream(device, cuda::stream_ref{stream});
-
   auto sm_count                       = get_sm_count(device);
   auto const max_shared_mem_per_block = get_max_shared_mem_per_block(device);
   auto const max_shared_mem_per_sm    = get_max_shared_mem_per_sm(device);

cpp/include/nvforest/detail/owning_buffer/gpu.hpp

@@ -3,51 +3,37 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
-#include <nvforest/cuda_stream.hpp>
 #include <nvforest/detail/device_id.hpp>
+#include <nvforest/detail/device_setter.hpp>
 #include <nvforest/detail/owning_buffer/base.hpp>
 #include <nvforest/device_type.hpp>
 
-#include <cuda/stream>
+#include <rmm/device_buffer.hpp>
+
 #include <cuda_runtime_api.h>
 
-#include <cstddef>
-#include <memory>
 #include <type_traits>
 
 namespace nvforest::detail {
-
-struct owning_device_buffer_type_erased_impl;
-
-struct owning_device_buffer_type_erased {
-  owning_device_buffer_type_erased();
-  owning_device_buffer_type_erased(device_id<device_type::gpu> device_id,
-                                   std::size_t size,
-                                   cuda::stream_ref stream);
-  owning_device_buffer_type_erased(owning_device_buffer_type_erased&& other) noexcept;
-  owning_device_buffer_type_erased& operator=(owning_device_buffer_type_erased&& other) noexcept;
-  ~owning_device_buffer_type_erased();
-  std::byte* get();
-
- private:
-  std::unique_ptr<owning_device_buffer_type_erased_impl> impl_;
-};
-
 template <typename T>
 struct owning_buffer<device_type::gpu, T> {
   // TODO(wphicks): Assess need for buffers of const T
   using value_type = std::remove_const_t<T>;
-  owning_buffer()  = default;
+  owning_buffer() : data_{} {}
+
   owning_buffer(device_id<device_type::gpu> device_id,
                 std::size_t size,
-                cuda_stream stream) noexcept(false)
-    : data_{device_id, size * sizeof(value_type), cuda::stream_ref{stream}}
+                cudaStream_t stream) noexcept(false)
+    : data_{[&device_id, &size, &stream]() {
+        auto device_context = device_setter{device_id};
+        return rmm::device_buffer{size * sizeof(value_type), rmm::cuda_stream_view{stream}};
+      }()}
   {
   }
 
-  auto* get() const { return reinterpret_cast<T*>(data_.get()); }
+  auto* get() const { return reinterpret_cast<T*>(data_.data()); }
 
  private:
-  mutable owning_device_buffer_type_erased data_;
+  mutable rmm::device_buffer data_;
 };
 }  // namespace nvforest::detail