Remove RMM/RAFT dependencies, Part 1

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

@@ -198,10 +198,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)
 
 if(BUILD_NVFOREST_TESTS)
@@ -252,7 +248,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}
@@ -312,7 +309,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 CUDA::cudart_static CCCL::CCCL)
 endif()
 
 # These are always private:

cpp/include/nvforest/buffer.hpp

@@ -13,6 +13,8 @@
 #include <nvforest/detail/owning_buffer.hpp>
 #include <nvforest/device_type.hpp>
 
+#include <cuda/stream>
+
 #include <stdint.h>
 
 #include <cstddef>
@@ -57,7 +59,8 @@ struct buffer {
         switch (mem_type) {
           case device_type::cpu: result = detail::owning_buffer<device_type::cpu, T>{size}; break;
           case device_type::gpu:
-            result = detail::owning_buffer<device_type::gpu, T>{std::get<1>(device_), size, stream};
+            result = detail::owning_buffer<device_type::gpu, T>{
+              std::get<1>(device_), size, cuda::stream_ref{stream}};
             break;
         }
         return result;
@@ -138,8 +141,8 @@ struct buffer {
           result_data = buf.get();
           result      = std::move(buf);
         } else if (mem_type == device_type::gpu) {
-          auto buf =
-            detail::owning_buffer<device_type::gpu, T>(std::get<1>(device_), other.size(), stream);
+          auto buf = detail::owning_buffer<device_type::gpu, T>(
+            std::get<1>(device_), other.size(), cuda::stream_ref{stream});
           result_data = buf.get();
           result      = std::move(buf);
         }
@@ -212,7 +215,8 @@ struct buffer {
             result_data = buf.get();
             result      = std::move(buf);
           } else if (mem_type == device_type::gpu) {
-            auto buf    = detail::owning_buffer<device_type::gpu, T>{device, other.size(), stream};
+            auto buf = detail::owning_buffer<device_type::gpu, T>{
+              device, other.size(), cuda::stream_ref{stream}};
             result_data = buf.get();
             result      = std::move(buf);
           }

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,20 @@ 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 stream)
+{
+  cuda::stream_ref legacy_default_stream{cudaStreamLegacy};
+  cuda::stream_ref per_thread_default_stream{cudaStreamPerThread};
+  cuda::stream_ref stream_wrapped{stream};
+  if (stream_wrapped != legacy_default_stream && stream_wrapped != per_thread_default_stream &&
+      stream_wrapped.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_wrapped.device().get())};
+  }
+}
+
 /* A wrapper around the underlying inference kernels to support dispatching to
  * the right kernel
  *
@@ -99,6 +115,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, 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

@@ -4,36 +4,49 @@
  */
 #pragma once
 #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), rmm::cuda_stream_view{stream}};
-      }()}
+                cuda::stream_ref stream) noexcept(false)
+    : data_{device_id, size * sizeof(value_type), 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