Remove RMM dependency and prepare for RAFT removal, Part 1

README.md

@@ -17,6 +17,30 @@ fm = nvforest.load_model("/path/to/xgboost_model.ubj", device="gpu",
 pred = fm.predict(X)
 ```
 
+By default, nvForest creates a new
+[CUDA stream](https://docs.nvidia.com/cuda/cuda-programming-guide/02-basics/asynchronous-execution.html#cuda-streams)
+for the model and retains it. To use a custom stream, pass in a :external+cuda-python:py:class:`cuda.core.Stream` or
+other stream-like objects (*):
+
+(*) A stream-like object is an object that exposes a method named `__cuda_stream__`. See [Stream Protocol](
+https://nvidia.github.io/cuda-python/cuda-core/latest/interoperability.html#cuda-stream-protocol) for more details.
+
+```python
+from cuda.core import Device
+
+device = Device(0)
+device.set_current()
+stream = device.create_stream()
+
+fm = nvforest.load_model(
+    "/path/to/xgboost_model.ubj",
+    device="gpu",
+    device_id=0,
+    stream=stream,
+)
+pred = fm.predict(X)  # Note: automatically syncs the stream
+```
+
 Load a scikit-learn random forest model and get class probabilities:
 
 ```python

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

@@ -11,6 +11,7 @@ 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,6 +11,7 @@ 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,6 +11,7 @@ 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,6 +11,7 @@ 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

@@ -208,7 +208,6 @@ 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)
@@ -261,7 +260,8 @@ if(NVFOREST_ENABLE_GPU)
             src/infer8.cu
             src/infer9.cu
             src/infer10.cu
-            src/infer11.cu)
+            src/infer11.cu
+            src/detail/device_buffer.cu)
 endif()
 target_sources(
   ${NVFOREST_CPP_TARGET}
@@ -321,7 +321,7 @@ elseif(NVFOREST_EXPORT_TREELITE_LINKAGE)
 endif()
 
 if(NVFOREST_ENABLE_GPU)
-  list(APPEND _nvforest_cpp_public_libs rmm::rmm raft::raft CUDA::cudart_static)
+  list(APPEND _nvforest_cpp_public_libs raft::raft CUDA::cudart_static CCCL::CCCL)
 endif()
 
 # These are always private:

cpp/include/nvforest/README.md

@@ -94,25 +94,24 @@ auto output = static_cast<float*>(nullptr);  // Loaded as single
 cudaMalloc((void**)&output, num_rows * num_outputs * sizeof(float));
 
 // Assuming that input is a float* pointing to data already located on-device
-
-auto handle = nvforest::handle_t{};
-
 nvforest_model.predict(
-  handle,
+  stream,
   output,
   input,
   num_rows,
   nvforest::device_type::gpu,  // out_mem_type
   nvforest::device_type::gpu,  // in_mem_type
+  nvforest::infer_kind::default_kind,
   4  // chunk_size
 );
 ```
 
-**handle**: To provide a unified interface on CPU and GPU, we introduce
-`nvforest::handle_t` as a wrapper for `raft::handle_t`. This is currently just a
-placeholder in CPU-only builds, and using it does not require any CUDA
-functionality. For GPU-enabled builds, you can construct a
-`nvforest::handle_t` directly from the `raft::handle_t` you wish to use.
+**stream**: The CUDA stream on which inference and any required data copies
+will run. It must be associated with the same GPU device as the model. Reusing
+the stream passed during model import preserves the ordering between model
+initialization and inference. Synchronize the stream before consuming the
+output on the host or releasing resources used by the operation. Keep the stream
+alive until the model is destroyed and all queued work has completed.
 
 **output**: Pointer to pre-allocated buffer where results should be
 written. If the model has been loaded at single precision, this should be a
@@ -145,4 +144,4 @@ will be selected based on heuristics.
 ## Learning More
 While the above usage summary should be enough to get started using nvForest in
 another C++ application, you can learn more about the details of this
-implementation by reading TODO(wphicks).
+implementation in [Implementation.md](Implementation.md).

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

@@ -7,23 +7,20 @@
 #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 = value_type{};
+        auto raw_id = int{};
         cuda_check(cudaGetDevice(&raw_id));
         return raw_id;
       }()} {};
-  device_id(value_type dev_id) noexcept : id_{dev_id} {};
+  device_id(int dev_id) noexcept : id_{dev_id} {};
 
-  auto value() const noexcept { return id_.value(); }
+  auto value() const noexcept { return id_; }
 
  private:
-  rmm::cuda_device_id id_;
+  int id_;
 };
 }  // namespace nvforest::detail

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

@@ -8,10 +8,22 @@
 #include <nvforest/detail/device_setter/base.hpp>
 #include <nvforest/device_type.hpp>
 
-#include <raft/util/cudart_utils.hpp>
-
 #include <cuda_runtime_api.h>
 
+#include <cstdio>
+
+#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)
+
 namespace nvforest::detail {
 
 /** Struct for setting current device within a code block */
@@ -27,7 +39,7 @@ struct device_setter<device_type::gpu> {
     cuda_check(cudaSetDevice(device.value()));
   }
 
-  ~device_setter() { RAFT_CUDA_TRY_NO_THROW(cudaSetDevice(prev_device_.value())); }
+  ~device_setter() { NVFOREST_CUDA_TRY_NO_THROW(cudaSetDevice(prev_device_.value())); }
 
  private:
   device_id<device_type::gpu> prev_device_;

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

@@ -21,6 +21,8 @@
 #include <nvforest/exceptions.hpp>
 #include <nvforest/infer_kind.hpp>
 
+#include <cuda/stream>
+
 #include <cstddef>
 #include <cstdint>
 #include <optional>
@@ -41,6 +43,19 @@ 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
  *
@@ -99,6 +114,8 @@ 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,38 +3,51 @@
  * 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 <rmm/device_buffer.hpp>
-
 #include <cuda/stream>
 #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() : data_{} {}
-
+  owning_buffer()  = default;
   owning_buffer(device_id<device_type::gpu> device_id,
                 std::size_t size,
-                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), cuda::stream_ref{stream}};
-      }()}
+                cuda_stream stream) noexcept(false)
+    : data_{device_id, size * sizeof(value_type), cuda::stream_ref{stream}}
   {
   }
 
-  auto* get() const { return reinterpret_cast<T*>(data_.data()); }
+  auto* get() const { return reinterpret_cast<T*>(data_.get()); }
 
  private:
-  mutable rmm::device_buffer data_;
+  mutable owning_device_buffer_type_erased data_;
 };
 }  // namespace nvforest::detail