Migrate RMM usage to CCCL MR design

c/src/core/c_api.cpp

@@ -14,12 +14,12 @@
 #include <raft/util/cudart_utils.hpp>
 #include <rapids_logger/logger.hpp>
 #include <rmm/cuda_stream_view.hpp>
-#include <rmm/mr/device_memory_resource.hpp>
+#include <rmm/mr/cuda_memory_resource.hpp>
 #include <rmm/mr/managed_memory_resource.hpp>
-#include <rmm/mr/owning_wrapper.hpp>
 #include <rmm/mr/per_device_resource.hpp>
-#include <rmm/mr/pool_memory_resource.hpp>
 #include <rmm/mr/pinned_host_memory_resource.hpp>
+#include <rmm/mr/pool_memory_resource.hpp>
+#include <rmm/resource_ref.hpp>
 
 #include "../core/exceptions.hpp"
 
@@ -132,62 +132,41 @@ extern "C" cuvsError_t cuvsRMMAlloc(cuvsResources_t res, void** ptr, size_t byte
 {
   return cuvs::core::translate_exceptions([=] {
     auto res_ptr = reinterpret_cast<raft::resources*>(res);
-    auto mr      = rmm::mr::get_current_device_resource();
-    *ptr         = mr->allocate(raft::resource::get_cuda_stream(*res_ptr), bytes);
+    auto mr      = rmm::mr::get_current_device_resource_ref();
+    *ptr         = mr.allocate(raft::resource::get_cuda_stream(*res_ptr), bytes);
   });
 }
 
 extern "C" cuvsError_t cuvsRMMFree(cuvsResources_t res, void* ptr, size_t bytes)
 {
   return cuvs::core::translate_exceptions([=] {
     auto res_ptr = reinterpret_cast<raft::resources*>(res);
-    auto mr      = rmm::mr::get_current_device_resource();
-    mr->deallocate(raft::resource::get_cuda_stream(*res_ptr), ptr, bytes);
+    auto mr      = rmm::mr::get_current_device_resource_ref();
+    mr.deallocate(raft::resource::get_cuda_stream(*res_ptr), ptr, bytes);
   });
 }
 
-thread_local std::shared_ptr<
-  rmm::mr::owning_wrapper<rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource>,
-                          rmm::mr::device_memory_resource>>
-  pool_mr;
-
 extern "C" cuvsError_t cuvsRMMPoolMemoryResourceEnable(int initial_pool_size_percent,
                                                        int max_pool_size_percent,
                                                        bool managed)
 {
   return cuvs::core::translate_exceptions([=] {
-    // Upstream memory resource needs to be a cuda_memory_resource
-    auto cuda_mr         = rmm::mr::get_current_device_resource();
-    auto* cuda_mr_casted = dynamic_cast<rmm::mr::cuda_memory_resource*>(cuda_mr);
-    if (cuda_mr_casted == nullptr) {
-      throw std::runtime_error("Current memory resource is not a cuda_memory_resource");
-    }
-
     auto initial_size = rmm::percent_of_free_device_memory(initial_pool_size_percent);
     auto max_size     = rmm::percent_of_free_device_memory(max_pool_size_percent);
 
-    auto mr = std::shared_ptr<rmm::mr::device_memory_resource>();
     if (managed) {
-      mr = std::static_pointer_cast<rmm::mr::device_memory_resource>(
-        std::make_shared<rmm::mr::managed_memory_resource>());
+      rmm::mr::set_current_device_resource(
+        rmm::mr::pool_memory_resource{rmm::mr::managed_memory_resource{}, initial_size, max_size});
     } else {
-      mr = std::static_pointer_cast<rmm::mr::device_memory_resource>(
-        std::make_shared<rmm::mr::cuda_memory_resource>());
+      rmm::mr::set_current_device_resource(
+        rmm::mr::pool_memory_resource{rmm::mr::cuda_memory_resource{}, initial_size, max_size});
     }
-
-    pool_mr =
-      rmm::mr::make_owning_wrapper<rmm::mr::pool_memory_resource>(mr, initial_size, max_size);
-
-    rmm::mr::set_current_device_resource(pool_mr.get());
   });
 }
 
 extern "C" cuvsError_t cuvsRMMMemoryResourceReset()
 {
-  return cuvs::core::translate_exceptions([=] {
-    rmm::mr::set_current_device_resource(rmm::mr::detail::initial_resource());
-    pool_mr.reset();
-  });
+  return cuvs::core::translate_exceptions([=] { rmm::mr::reset_current_device_resource(); });
 }
 
 thread_local std::unique_ptr<rmm::mr::pinned_host_memory_resource> pinned_mr;

cpp/bench/ann/src/common/cuda_huge_page_resource.hpp

@@ -1,14 +1,16 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
 
 #include <raft/core/error.hpp>
 #include <raft/core/logger_macros.hpp>
 
-#include <rmm/cuda_stream_view.hpp>
-#include <rmm/mr/device_memory_resource.hpp>
+#include <rmm/detail/aligned.hpp>
+#include <rmm/resource_ref.hpp>
+
+#include <cuda/memory_resource>
 
 #include <sys/mman.h>
 
@@ -17,37 +19,25 @@
 
 namespace raft::mr {
 /**
- * @brief `device_memory_resource` derived class that uses mmap to allocate memory.
- * This class enables memory allocation using huge pages.
+ * @brief Memory resource that uses mmap to allocate memory with huge pages.
  * It is assumed that the allocated memory is directly accessible on device. This currently only
  * works on GH systems.
  *
  * TODO(tfeher): consider improving or removing this helper once we made progress with
  * https://github.com/rapidsai/raft/issues/1819
  */
-class cuda_huge_page_resource final : public rmm::mr::device_memory_resource {
+class cuda_huge_page_resource {
  public:
   cuda_huge_page_resource()                                                  = default;
-  ~cuda_huge_page_resource() override                                        = default;
+  ~cuda_huge_page_resource()                                                 = default;
   cuda_huge_page_resource(cuda_huge_page_resource const&)                    = default;
   cuda_huge_page_resource(cuda_huge_page_resource&&)                         = default;
   auto operator=(cuda_huge_page_resource const&) -> cuda_huge_page_resource& = default;
   auto operator=(cuda_huge_page_resource&&) -> cuda_huge_page_resource&      = default;
 
- private:
-  /**
-   * @brief Allocates memory of size at least `bytes` using cudaMalloc.
-   *
-   * The returned pointer has at least 256B alignment.
-   *
-   * @note Stream argument is ignored
-   *
-   * @throws `rmm::bad_alloc` if the requested allocation could not be fulfilled
-   *
-   * @param bytes The size, in bytes, of the allocation
-   * @return void* Pointer to the newly allocated memory
-   */
-  auto do_allocate(std::size_t bytes, rmm::cuda_stream_view) -> void* override
+  void* allocate(cuda::stream_ref,
+                 std::size_t bytes,
+                 std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT)
   {
     void* addr{nullptr};
     addr = mmap(nullptr, bytes, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
@@ -60,36 +50,40 @@ class cuda_huge_page_resource final : public rmm::mr::device_memory_resource {
     return addr;
   }
 
-  /**
-   * @brief Deallocate memory pointed to by \p p.
-   *
-   * @note Stream argument is ignored.
-   *
-   * @throws Nothing.
-   *
-   * @param p Pointer to be deallocated
-   */
-  void do_deallocate(void* ptr, std::size_t size, rmm::cuda_stream_view) noexcept override
+  void deallocate(cuda::stream_ref,
+                  void* ptr,
+                  std::size_t size,
+                  std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept
   {
     if (munmap(ptr, size) == -1) { RAFT_LOG_ERROR("huge_page_resource::munmap failed"); }
   }
 
-  /**
-   * @brief Compare this resource to another.
-   *
-   * Two cuda_huge_page_resources always compare equal, because they can each
-   * deallocate memory allocated by the other.
-   *
-   * @throws Nothing.
-   *
-   * @param other The other resource to compare to
-   * @return true If the two resources are equivalent
-   * @return false If the two resources are not equal
-   */
-  [[nodiscard]] auto do_is_equal(device_memory_resource const& other) const noexcept
-    -> bool override
+  void* allocate_sync(std::size_t bytes, std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT)
+  {
+    return allocate(cuda::stream_ref{cudaStream_t{nullptr}}, bytes, alignment);
+  }
+
+  void deallocate_sync(void* ptr,
+                       std::size_t bytes,
+                       std::size_t alignment = rmm::CUDA_ALLOCATION_ALIGNMENT) noexcept
+  {
+    deallocate(cuda::stream_ref{cudaStream_t{nullptr}}, ptr, bytes, alignment);
+  }
+
+  [[nodiscard]] bool operator==(cuda_huge_page_resource const&) const noexcept { return true; }
+  [[nodiscard]] bool operator!=(cuda_huge_page_resource const&) const noexcept { return false; }
+
+  constexpr friend void get_property(cuda_huge_page_resource const&,
+                                     cuda::mr::device_accessible) noexcept
+  {
+  }
+  constexpr friend void get_property(cuda_huge_page_resource const&,
+                                     cuda::mr::host_accessible) noexcept
   {
-    return dynamic_cast<cuda_huge_page_resource const*>(&other) != nullptr;
   }
 };
+static_assert(cuda::mr::synchronous_resource<cuda_huge_page_resource>);
+static_assert(cuda::mr::resource<cuda_huge_page_resource>);
+static_assert(cuda::mr::resource_with<cuda_huge_page_resource, cuda::mr::device_accessible>);
+static_assert(cuda::mr::resource_with<cuda_huge_page_resource, cuda::mr::host_accessible>);
 }  // namespace raft::mr

cpp/bench/ann/src/cuvs/cuvs_ann_bench_utils.h

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -15,14 +15,16 @@
 #include <raft/core/host_mdspan.hpp>
 #include <raft/core/logger.hpp>
 #include <raft/core/operators.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/util/cudart_utils.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_uvector.hpp>
-#include <rmm/mr/device_memory_resource.hpp>
 #include <rmm/mr/failure_callback_resource_adaptor.hpp>
 #include <rmm/mr/managed_memory_resource.hpp>
+#include <rmm/mr/per_device_resource.hpp>
 #include <rmm/mr/pool_memory_resource.hpp>
+#include <rmm/resource_ref.hpp>
 
 #include <memory>
 #include <type_traits>
@@ -65,17 +67,16 @@ inline auto rmm_oom_callback(std::size_t bytes, void*) -> bool
  */
 class shared_raft_resources {
  public:
-  using pool_mr_type  = rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource>;
-  using mr_type       = rmm::mr::failure_callback_resource_adaptor<pool_mr_type>;
   using large_mr_type = rmm::mr::managed_memory_resource;
 
   shared_raft_resources()
-  try
-    : orig_resource_{rmm::mr::get_current_device_resource()},
-      pool_resource_(orig_resource_, 1024 * 1024 * 1024ull),
-      resource_(&pool_resource_, rmm_oom_callback, nullptr),
-      large_mr_() {
-    rmm::mr::set_current_device_resource(&resource_);
+  try : large_mr_() {
+    orig_resource_ =
+      rmm::mr::set_current_device_resource(rmm::mr::failure_callback_resource_adaptor<>{
+        rmm::mr::pool_memory_resource{rmm::mr::get_current_device_resource_ref(),
+                                      1024 * 1024 * 1024ull},
+        rmm_oom_callback,
+        nullptr});
   } catch (const std::exception& e) {
     auto cuda_status = cudaGetLastError();
     size_t free      = 0;
@@ -97,15 +98,10 @@ class shared_raft_resources {
 
   ~shared_raft_resources() noexcept { rmm::mr::set_current_device_resource(orig_resource_); }
 
-  auto get_large_memory_resource() noexcept
-  {
-    return static_cast<rmm::mr::device_memory_resource*>(&large_mr_);
-  }
+  auto get_large_memory_resource() noexcept -> rmm::device_async_resource_ref { return large_mr_; }
 
  private:
-  rmm::mr::device_memory_resource* orig_resource_;
-  pool_mr_type pool_resource_;
-  mr_type resource_;
+  cuda::mr::any_resource<cuda::mr::device_accessible> orig_resource_;
   large_mr_type large_mr_;
 };
 
@@ -129,12 +125,8 @@ class configured_raft_resources {
       res_{std::make_unique<raft::device_resources>(
         rmm::cuda_stream_view(get_stream_from_global_pool()))}
   {
-    // set the large workspace resource to the raft handle, but without the deleter
-    // (this resource is managed by the shared_res).
     raft::resource::set_large_workspace_resource(
-      *res_,
-      std::shared_ptr<rmm::mr::device_memory_resource>(shared_res_->get_large_memory_resource(),
-                                                       raft::void_op{}));
+      *res_, raft::mr::device_resource{shared_res_->get_large_memory_resource()});
   }
 
   /** Default constructor creates all resources anew. */

cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h

@@ -166,9 +166,9 @@ class cuvs_cagra : public algo<T>, public algo_gpu {
   inline rmm::device_async_resource_ref get_mr(AllocatorType mem_type)
   {
     switch (mem_type) {
-      case (AllocatorType::kHostPinned): return &mr_pinned_;
-      case (AllocatorType::kHostHugePage): return &mr_huge_page_;
-      default: return rmm::mr::get_current_device_resource();
+      case (AllocatorType::kHostPinned): return mr_pinned_;
+      case (AllocatorType::kHostHugePage): return mr_huge_page_;
+      default: return rmm::mr::get_current_device_resource_ref();
     }
   }
 };