From e29efc7da34d9b418ed6c7bfe5d15fc76e39d0ae Mon Sep 17 00:00:00 2001
From: Mark Harris <mharris@nvidia.com>
Date: Wed, 28 Aug 2024 08:49:35 +0000
Subject: [PATCH 01/14] Add current_resource wrappers.

---
 .../core/resource/device_memory_resource.hpp  | 105 ++++++++++++++++--
 1 file changed, 98 insertions(+), 7 deletions(-)

diff --git a/cpp/include/raft/core/resource/device_memory_resource.hpp b/cpp/include/raft/core/resource/device_memory_resource.hpp
index b785010a0a..99fb770bb2 100644
--- a/cpp/include/raft/core/resource/device_memory_resource.hpp
+++ b/cpp/include/raft/core/resource/device_memory_resource.hpp
@@ -24,6 +24,7 @@
 #include <rmm/mr/device/limiting_resource_adaptor.hpp>
 #include <rmm/mr/device/per_device_resource.hpp>
 #include <rmm/mr/device/pool_memory_resource.hpp>
+#include <rmm/resource_ref.hpp>
 
 #include <cstddef>
 #include <optional>
@@ -35,6 +36,96 @@ namespace raft::resource {
  * @{
  */
 
+/**
+ * @brief Alias for a `cuda::mr::resource_ref` with the property
+ * `cuda::mr::device_accessible`.
+ */
+using device_resource_ref = rmm::device_resource_ref;
+
+/**
+ * @brief Alias for a `cuda::mr::async_resource_ref` with the property
+ * `cuda::mr::device_accessible`.
+ */
+using device_async_resource_ref = rmm::device_async_resource_ref;
+
+/**
+ * @brief Alias for a `cuda::mr::resource_ref` with the property
+ * `cuda::mr::host_accessible`.
+ */
+using host_resource_ref = rmm::host_resource_ref;
+
+/**
+ * @brief Alias for a `cuda::mr::async_resource_ref` with the property
+ * `cuda::mr::host_accessible`.
+ */
+using host_async_resource_ref = rmm::host_async_resource_ref;
+
+/**
+ * @brief Alias for a `cuda::mr::resource_ref` with the properties
+ * `cuda::mr::host_accessible` and `cuda::mr::device_accessible`.
+ */
+using host_device_resource_ref = rmm::host_device_resource_ref;
+
+/**
+ * @brief Alias for a `cuda::mr::async_resource_ref` with the properties
+ * `cuda::mr::host_accessible` and `cuda::mr::device_accessible`.
+ */
+using host_device_async_resource_ref = rmm::host_device_async_resource_ref;
+
+/**
+ * @brief Get the current device memory resource.
+ *
+ * @return The current device memory resource.
+ */
+inline rmm::mr::device_memory_resource* get_current_device_resource()
+{
+  return rmm::mr::get_current_device_resource();
+}
+
+/**
+ * @brief Get the current device memory resource reference.
+ *
+ * @return The current device memory resource reference.
+ */
+inline device_async_resource_ref get_current_device_resource_ref()
+{
+  // For now, match current behavior which is to return current resource pointer
+  return rmm::mr::get_current_device_resource();
+}
+
+/**
+ * @brief Set the current device memory resource.
+ *
+ * @param mr The new device memory resource.
+ * @return The previous device memory resource.
+ */
+inline rmm::mr::device_memory_resource* set_current_device_resource(
+  rmm::mr::device_memory_resource* mr)
+{
+  return rmm::mr::set_current_device_resource(mr);
+}
+
+/**
+ * @brief Set the current device memory resource reference.
+ *
+ * @param mr The new device memory resource reference.
+ * @return The previous device memory resource reference.
+ */
+inline device_async_resource_ref set_current_device_resource_ref(device_async_resource_ref mr)
+{
+  return rmm::mr::set_current_device_resource_ref(mr);
+}
+
+/**
+ * @brief Reset the current device memory resource reference to the initial resource.
+ *
+ * @return The previous device memory resource reference.
+ */
+inline device_async_resource_ref reset_current_device_resource_ref()
+{
+  return rmm::mr::reset_current_device_resource_ref();
+}
+
 class device_memory_resource : public resource {
  public:
   explicit device_memory_resource(std::shared_ptr<rmm::mr::device_memory_resource> mr) : mr_(mr) {}
@@ -85,8 +176,8 @@ class large_workspace_resource_factory : public resource_factory {
   explicit large_workspace_resource_factory(
     std::shared_ptr<rmm::mr::device_memory_resource> mr = {nullptr})
     : mr_{mr ? mr
-             : std::shared_ptr<rmm::mr::device_memory_resource>{
-                 rmm::mr::get_current_device_resource(), void_op{}}}
+             : std::shared_ptr<rmm::mr::device_memory_resource>{get_current_device_resource(),
+                                                                void_op{}}}
   {
   }
   auto get_resource_type() -> resource_type override
@@ -140,7 +231,7 @@ class workspace_resource_factory : public resource_factory {
     //      resource adaptor bad_alloc error than into the pool bad_alloc error.
     //   2) The pool doesn't grab too much memory on top of the 'limit'.
     auto max_size = std::min<std::size_t>(limit + kOneGb / 2lu, limit * 3lu / 2lu);
-    auto upstream = rmm::mr::get_current_device_resource();
+    auto upstream = get_current_device_resource();
     RAFT_LOG_DEBUG(
       "Setting the workspace pool resource; memory limit = %zu, initial pool size = %zu, max pool "
       "size = %zu.",
@@ -154,13 +245,13 @@ class workspace_resource_factory : public resource_factory {
   /**
    * Get the global memory resource wrapped into an unmanaged shared_ptr (with no deleter).
    *
-   * Note: the lifetime of the underlying `rmm::mr::get_current_device_resource()` is managed
+   * Note: the lifetime of the underlying `get_current_device_resource()` is managed
    * somewhere else, since it's passed by a raw pointer. Hence, this shared_ptr wrapper is not
    * allowed to delete the pointer on destruction.
    */
   static inline auto default_plain_resource() -> std::shared_ptr<rmm::mr::device_memory_resource>
   {
-    return std::shared_ptr<rmm::mr::device_memory_resource>{rmm::mr::get_current_device_resource(),
+    return std::shared_ptr<rmm::mr::device_memory_resource>{get_current_device_resource(),
                                                             void_op{}};
   }
 
@@ -238,7 +329,7 @@ inline void set_workspace_resource(resources const& res,
 
 /**
  * Set the temporary workspace resource to a pool on top of the global memory resource
- * (`rmm::mr::get_current_device_resource()`.
+ * (`raft::resource::get_current_device_resource()`.
  *
  * @param res raft resources object for managing resources
  * @param allocation_limit
@@ -258,7 +349,7 @@ inline void set_workspace_to_pool_resource(
 
 /**
  * Set the temporary workspace resource the same as the global memory resource
- * (`rmm::mr::get_current_device_resource()`.
+ * (`raft::resource::get_current_device_resource()`.
  *
  * Note, the workspace resource is always limited; the limit here defines how much of the global
  * memory resource can be consumed by the workspace allocations.

From 3f6dd48fb0cbce68b201c4f1b6c73cd6277fb3b2 Mon Sep 17 00:00:00 2001
From: Mark Harris <mharris@nvidia.com>
Date: Wed, 28 Aug 2024 09:00:50 +0000
Subject: [PATCH 02/14] Use wrappers

---
 cpp/bench/ann/src/raft/raft_ann_bench_utils.h |  2 +-
 cpp/bench/ann/src/raft/raft_cagra_wrapper.h   |  2 +-
 cpp/bench/prims/common/benchmark.hpp          |  4 ++--
 cpp/bench/prims/matrix/gather.cu              |  4 ++--
 cpp/bench/prims/neighbors/knn.cuh             | 22 +++++++++----------
 cpp/bench/prims/neighbors/refine.cuh          |  4 ++--
 cpp/bench/prims/random/subsample.cu           |  6 ++---
 .../raft/core/device_container_policy.hpp     |  4 +---
 .../raft/core/device_resources_manager.hpp    |  4 ++--
 .../raft/neighbors/detail/cagra/utils.hpp     |  9 ++++----
 .../neighbors/detail/ivf_flat_search-inl.cuh  | 21 +++++++++---------
 .../raft/neighbors/detail/vpq_dataset.cuh     |  2 +-
 cpp/include/raft/neighbors/ivf_flat-inl.cuh   |  4 ++--
 .../raft/random/multi_variable_gaussian.cuh   |  2 +-
 .../raft/spatial/knn/detail/ann_utils.cuh     | 13 ++++++-----
 cpp/template/src/cagra_example.cu             |  4 ++--
 cpp/template/src/ivf_flat_example.cu          |  4 ++--
 cpp/template/src/ivf_pq_example.cu            |  4 ++--
 cpp/test/core/device_resources_manager.cpp    |  2 +-
 cpp/test/core/handle.cpp                      |  4 ++--
 cpp/test/mr/device/buffer.cpp                 |  7 +++---
 cpp/test/random/multi_variable_gaussian.cu    |  2 +-
 22 files changed, 64 insertions(+), 66 deletions(-)

diff --git a/cpp/bench/ann/src/raft/raft_ann_bench_utils.h b/cpp/bench/ann/src/raft/raft_ann_bench_utils.h
index 9b086fdb23..1a70aa9657 100644
--- a/cpp/bench/ann/src/raft/raft_ann_bench_utils.h
+++ b/cpp/bench/ann/src/raft/raft_ann_bench_utils.h
@@ -80,7 +80,7 @@ class shared_raft_resources {
   using large_mr_type = rmm::mr::managed_memory_resource;
 
   shared_raft_resources()
-  try : orig_resource_{rmm::mr::get_current_device_resource()},
+  try : orig_resource_{raft::resource::get_current_device_resource_ref()},
     pool_resource_(orig_resource_, 1024 * 1024 * 1024ull),
     resource_(&pool_resource_, rmm_oom_callback, nullptr), large_mr_() {
     rmm::mr::set_current_device_resource(&resource_);
diff --git a/cpp/bench/ann/src/raft/raft_cagra_wrapper.h b/cpp/bench/ann/src/raft/raft_cagra_wrapper.h
index b03f875a8e..081fae1a94 100644
--- a/cpp/bench/ann/src/raft/raft_cagra_wrapper.h
+++ b/cpp/bench/ann/src/raft/raft_cagra_wrapper.h
@@ -153,7 +153,7 @@ class RaftCagra : public ANN<T>, public AnnGPU {
     switch (mem_type) {
       case (AllocatorType::HostPinned): return &mr_pinned_;
       case (AllocatorType::HostHugePage): return &mr_huge_page_;
-      default: return rmm::mr::get_current_device_resource();
+      default: return raft::resource::get_current_device_resource_ref();
     }
   }
 };
diff --git a/cpp/bench/prims/common/benchmark.hpp b/cpp/bench/prims/common/benchmark.hpp
index 3ce43cc1e7..f2ad980e9e 100644
--- a/cpp/bench/prims/common/benchmark.hpp
+++ b/cpp/bench/prims/common/benchmark.hpp
@@ -50,14 +50,14 @@ struct using_pool_memory_res {
 
  public:
   using_pool_memory_res(size_t initial_size, size_t max_size)
-    : orig_res_(rmm::mr::get_current_device_resource()),
+    : orig_res_(raft::resource::get_current_device_resource_ref()),
       pool_res_(&cuda_res_, initial_size, max_size)
   {
     rmm::mr::set_current_device_resource(&pool_res_);
   }
 
   using_pool_memory_res()
-    : orig_res_(rmm::mr::get_current_device_resource()),
+    : orig_res_(raft::resource::get_current_device_resource_ref()),
       pool_res_(&cuda_res_, rmm::percent_of_free_device_memory(50))
   {
     rmm::mr::set_current_device_resource(&pool_res_);
diff --git a/cpp/bench/prims/matrix/gather.cu b/cpp/bench/prims/matrix/gather.cu
index 876e47525c..a12792d774 100644
--- a/cpp/bench/prims/matrix/gather.cu
+++ b/cpp/bench/prims/matrix/gather.cu
@@ -46,8 +46,8 @@ template <typename T, typename MapT, typename IdxT, bool Conditional = false>
 struct Gather : public fixture {
   Gather(const GatherParams<IdxT>& p)
     : params(p),
-      old_mr(rmm::mr::get_current_device_resource()),
-      pool_mr(rmm::mr::get_current_device_resource(), 2 * (1ULL << 30)),
+      old_mr(raft::resource::get_current_device_resource_ref()),
+      pool_mr(raft::resource::get_current_device_resource_ref(), 2 * (1ULL << 30)),
       matrix(this->handle),
       map(this->handle),
       out(this->handle),
diff --git a/cpp/bench/prims/neighbors/knn.cuh b/cpp/bench/prims/neighbors/knn.cuh
index 6499078623..f8b837d77f 100644
--- a/cpp/bench/prims/neighbors/knn.cuh
+++ b/cpp/bench/prims/neighbors/knn.cuh
@@ -20,6 +20,7 @@
 
 #include <raft/core/bitset.cuh>
 #include <raft/core/resource/device_id.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/neighbors/ivf_flat.cuh>
 #include <raft/neighbors/ivf_pq.cuh>
 #include <raft/neighbors/sample_filter.cuh>
@@ -29,13 +30,12 @@
 
 #include <rmm/mr/device/device_memory_resource.hpp>
 #include <rmm/mr/device/managed_memory_resource.hpp>
-#include <rmm/mr/device/per_device_resource.hpp>
 #include <rmm/mr/host/new_delete_resource.hpp>
 #include <rmm/mr/host/pinned_memory_resource.hpp>
-#include <rmm/resource_ref.hpp>
 
 #include <thrust/sequence.h>
 
+#include <memory>
 #include <optional>
 
 namespace raft::bench::spatial {
@@ -89,25 +89,23 @@ inline auto operator<<(std::ostream& os, const Scope& s) -> std::ostream&
 
 struct device_resource {
  public:
-  explicit device_resource(bool managed) : managed_(managed)
+  explicit device_resource(bool managed)
+    : managed_(managed ? std::make_shared<rmm::mr::managed_memory_resource>() : nullptr)
   {
-    if (managed_) {
-      res_ = new rmm::mr::managed_memory_resource();
+    if (managed) {
+      res_ = managed.get();
     } else {
-      res_ = rmm::mr::get_current_device_resource();
+      res_ = raft::resource::get_current_device_resource_ref();
     }
   }
 
-  ~device_resource()
-  {
-    if (managed_) { delete res_; }
-  }
+  ~device_resource() {}
 
   [[nodiscard]] auto get() const -> rmm::device_async_resource_ref { return res_; }
 
  private:
-  const bool managed_;
-  rmm::mr::device_memory_resource* res_;
+  std::shared_ptr<rmm::mr::device_memory_resource> managed_;
+  raft::device_async_resource_ref res_;
 };
 
 template <typename T>
diff --git a/cpp/bench/prims/neighbors/refine.cuh b/cpp/bench/prims/neighbors/refine.cuh
index 0360babd82..3e8721f353 100644
--- a/cpp/bench/prims/neighbors/refine.cuh
+++ b/cpp/bench/prims/neighbors/refine.cuh
@@ -58,7 +58,7 @@ class RefineAnn : public fixture {
     label_stream << data.p;
     state.SetLabel(label_stream.str());
 
-    auto old_mr = rmm::mr::get_current_device_resource();
+    auto old_mr = raft::resource::get_current_device_resource_ref();
     rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_mr(
       old_mr, rmm::percent_of_free_device_memory(50));
     rmm::mr::set_current_device_resource(&pool_mr);
@@ -84,7 +84,7 @@ class RefineAnn : public fixture {
                                                               data.p.metric);
       });
     }
-    rmm::mr::set_current_device_resource(old_mr);
+    raft::set_current_device_resource(old_mr);
   }
 
  private:
diff --git a/cpp/bench/prims/random/subsample.cu b/cpp/bench/prims/random/subsample.cu
index 70a9c65e0d..58cd35276c 100644
--- a/cpp/bench/prims/random/subsample.cu
+++ b/cpp/bench/prims/random/subsample.cu
@@ -64,8 +64,8 @@ template <typename T>
 struct sample : public fixture {
   sample(const sample_inputs& p)
     : params(p),
-      old_mr(rmm::mr::get_current_device_resource()),
-      pool_mr(rmm::mr::get_current_device_resource(), 2 * GiB),
+      old_mr(raft::resource::get_current_device_resource()),
+      pool_mr(raft::resource::get_current_device_resource(), 2 * GiB),
       in(make_device_vector<T, int64_t>(res, p.n_samples)),
       out(make_device_vector<T, int64_t>(res, p.n_train))
   {
@@ -73,7 +73,7 @@ struct sample : public fixture {
     raft::random::RngState r(123456ULL);
   }
 
-  ~sample() { rmm::mr::set_current_device_resource(old_mr); }
+  ~sample() { raft::set_current_device_resource(old_mr); }
   void run_benchmark(::benchmark::State& state) override
   {
     std::ostringstream label_stream;
diff --git a/cpp/include/raft/core/device_container_policy.hpp b/cpp/include/raft/core/device_container_policy.hpp
index 18d8b77364..904baef5d2 100644
--- a/cpp/include/raft/core/device_container_policy.hpp
+++ b/cpp/include/raft/core/device_container_policy.hpp
@@ -31,8 +31,6 @@
 
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_uvector.hpp>
-#include <rmm/mr/device/per_device_resource.hpp>
-#include <rmm/resource_ref.hpp>
 
 #include <thrust/device_ptr.h>
 
@@ -185,7 +183,7 @@ class device_uvector_policy {
   [[nodiscard]] auto make_accessor_policy() const noexcept { return const_accessor_policy{}; }
 
  private:
-  rmm::device_async_resource_ref mr_{rmm::mr::get_current_device_resource()};
+  rmm::device_async_resource_ref mr_{raft::resource::get_current_device_resource_ref()};
 };
 
 }  // namespace raft
diff --git a/cpp/include/raft/core/device_resources_manager.hpp b/cpp/include/raft/core/device_resources_manager.hpp
index 1c4bee15a6..bef70a8740 100644
--- a/cpp/include/raft/core/device_resources_manager.hpp
+++ b/cpp/include/raft/core/device_resources_manager.hpp
@@ -17,12 +17,12 @@
 #pragma once
 #include <raft/core/device_resources.hpp>
 #include <raft/core/device_setter.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
 
 #include <rmm/cuda_device.hpp>
 #include <rmm/cuda_stream.hpp>
 #include <rmm/cuda_stream_pool.hpp>
 #include <rmm/mr/device/cuda_memory_resource.hpp>
-#include <rmm/mr/device/per_device_resource.hpp>
 
 #include <algorithm>
 #include <memory>
@@ -169,7 +169,7 @@ struct device_resources_manager {
           // resource
           if (params.max_mem_pool_size.value_or(1) != 0) {
             auto* upstream =
-              dynamic_cast<rmm::mr::cuda_memory_resource*>(rmm::mr::get_current_device_resource());
+              dynamic_cast<rmm::mr::cuda_memory_resource*>(resource::get_current_device_resource());
             if (upstream != nullptr) {
               result =
                 std::make_shared<rmm::mr::pool_memory_resource<rmm::mr::cuda_memory_resource>>(
diff --git a/cpp/include/raft/neighbors/detail/cagra/utils.hpp b/cpp/include/raft/neighbors/detail/cagra/utils.hpp
index ece95a7cb7..4592726ff5 100644
--- a/cpp/include/raft/neighbors/detail/cagra/utils.hpp
+++ b/cpp/include/raft/neighbors/detail/cagra/utils.hpp
@@ -258,10 +258,11 @@ class host_matrix_view_from_device {
 
 // Copy matrix src to dst. pad rows with 0 if necessary to make them 16 byte aligned.
 template <typename T, typename data_accessor>
-void copy_with_padding(raft::resources const& res,
-                       raft::device_matrix<T, int64_t, row_major>& dst,
-                       mdspan<const T, matrix_extent<int64_t>, row_major, data_accessor> src,
-                       rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource())
+void copy_with_padding(
+  raft::resources const& res,
+  raft::device_matrix<T, int64_t, row_major>& dst,
+  mdspan<const T, matrix_extent<int64_t>, row_major, data_accessor> src,
+  rmm::device_async_resource_ref mr = raft::resource::get_current_device_resource_ref())
 {
   size_t padded_dim = round_up_safe<size_t>(src.extent(1) * sizeof(T), 16) / sizeof(T);
 
diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_search-inl.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_search-inl.cuh
index 388dd60f14..9fa03ea5a3 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_search-inl.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_search-inl.cuh
@@ -268,16 +268,17 @@ void search_impl(raft::resources const& handle,
 template <typename T,
           typename IdxT,
           typename IvfSampleFilterT = raft::neighbors::filtering::none_ivf_sample_filter>
-inline void search(raft::resources const& handle,
-                   const search_params& params,
-                   const index<T, IdxT>& index,
-                   const T* queries,
-                   uint32_t n_queries,
-                   uint32_t k,
-                   IdxT* neighbors,
-                   float* distances,
-                   rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource(),
-                   IvfSampleFilterT sample_filter    = IvfSampleFilterT())
+inline void search(
+  raft::resources const& handle,
+  const search_params& params,
+  const index<T, IdxT>& index,
+  const T* queries,
+  uint32_t n_queries,
+  uint32_t k,
+  IdxT* neighbors,
+  float* distances,
+  rmm::device_async_resource_ref mr = raft::resource::get_current_device_resource_ref(),
+  IvfSampleFilterT sample_filter    = IvfSampleFilterT())
 {
   common::nvtx::range<common::nvtx::domain::raft> fun_scope(
     "ivf_flat::search(k = %u, n_queries = %u, dim = %zu)", k, n_queries, index.dim());
diff --git a/cpp/include/raft/neighbors/detail/vpq_dataset.cuh b/cpp/include/raft/neighbors/detail/vpq_dataset.cuh
index f6cd2a1ceb..14e58315f6 100644
--- a/cpp/include/raft/neighbors/detail/vpq_dataset.cuh
+++ b/cpp/include/raft/neighbors/detail/vpq_dataset.cuh
@@ -365,7 +365,7 @@ auto process_and_fill_codes(const raft::resources& res,
                                                         dim,
                                                         max_batch_size,
                                                         stream,
-                                                        rmm::mr::get_current_device_resource())) {
+                                                        resource::get_current_device_resource())) {
     auto batch_view = raft::make_device_matrix_view(batch.data(), ix_t(batch.size()), dim);
     auto labels     = predict_vq<label_t>(res, batch_view, vq_centers);
     dim3 blocks(div_rounding_up_safe<ix_t>(n_rows, kBlockSize / threads_per_vec), 1, 1);
diff --git a/cpp/include/raft/neighbors/ivf_flat-inl.cuh b/cpp/include/raft/neighbors/ivf_flat-inl.cuh
index ea7cff7060..12e22fbe7a 100644
--- a/cpp/include/raft/neighbors/ivf_flat-inl.cuh
+++ b/cpp/include/raft/neighbors/ivf_flat-inl.cuh
@@ -416,7 +416,7 @@ void extend(raft::resources const& handle,
  *   ...
  *   // Create a pooling memory resource with a pre-defined initial size.
  *   rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> mr(
- *     rmm::mr::get_current_device_resource(), 1024 * 1024);
+ *     raft::resource::get_current_device_resource(), 1024 * 1024);
  *   // use default search parameters
  *   ivf_flat::search_params search_params;
  *   filtering::none_ivf_sample_filter filter;
@@ -482,7 +482,7 @@ void search_with_filtering(raft::resources const& handle,
  *   ...
  *   // Create a pooling memory resource with a pre-defined initial size.
  *   rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> mr(
- *     rmm::mr::get_current_device_resource(), 1024 * 1024);
+ *     raft::resource::get_current_device_resource(), 1024 * 1024);
  *   // use default search parameters
  *   ivf_flat::search_params search_params;
  *   // Use the same allocator across multiple searches to reduce the number of
diff --git a/cpp/include/raft/random/multi_variable_gaussian.cuh b/cpp/include/raft/random/multi_variable_gaussian.cuh
index 4b37e1ff65..014776d915 100644
--- a/cpp/include/raft/random/multi_variable_gaussian.cuh
+++ b/cpp/include/raft/random/multi_variable_gaussian.cuh
@@ -52,7 +52,7 @@ void multi_variable_gaussian(raft::resources const& handle,
                              const multi_variable_gaussian_decomposition_method method)
 {
   detail::compute_multi_variable_gaussian_impl(
-    handle, rmm::mr::get_current_device_resource(), x, P, X, method);
+    handle, raft::resource::get_current_device_resource_ref(), x, P, X, method);
 }
 
 /** @} */
diff --git a/cpp/include/raft/spatial/knn/detail/ann_utils.cuh b/cpp/include/raft/spatial/knn/detail/ann_utils.cuh
index 920249172f..da88a1ea91 100644
--- a/cpp/include/raft/spatial/knn/detail/ann_utils.cuh
+++ b/cpp/include/raft/spatial/knn/detail/ann_utils.cuh
@@ -498,12 +498,13 @@ struct batch_load_iterator {
    * @param stream the ordering for the host->device copies, if applicable.
    * @param mr a custom memory resource for the intermediate buffer, if applicable.
    */
-  batch_load_iterator(const T* source,
-                      size_type n_rows,
-                      size_type row_width,
-                      size_type batch_size,
-                      rmm::cuda_stream_view stream,
-                      rmm::device_async_resource_ref mr = rmm::mr::get_current_device_resource())
+  batch_load_iterator(
+    const T* source,
+    size_type n_rows,
+    size_type row_width,
+    size_type batch_size,
+    rmm::cuda_stream_view stream,
+    rmm::device_async_resource_ref mr = raft::resource::get_current_device_resource_ref())
     : cur_batch_(new batch(source, n_rows, row_width, batch_size, stream, mr)), cur_pos_(0)
   {
   }
diff --git a/cpp/template/src/cagra_example.cu b/cpp/template/src/cagra_example.cu
index 3c1be8b4f8..fdbfd5b238 100644
--- a/cpp/template/src/cagra_example.cu
+++ b/cpp/template/src/cagra_example.cu
@@ -67,8 +67,8 @@ int main()
 
   // Set pool memory resource with 1 GiB initial pool size. All allocations use the same pool.
   rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_mr(
-    rmm::mr::get_current_device_resource(), 1024 * 1024 * 1024ull);
-  rmm::mr::set_current_device_resource(&pool_mr);
+    raft::resource::get_current_device_resource(), 1024 * 1024 * 1024ull);
+  raft::resource::set_current_device_resource(&pool_mr);
 
   // Alternatively, one could define a pool allocator for temporary arrays (used within RAFT
   // algorithms). In that case only the internal arrays would use the pool, any other allocation
diff --git a/cpp/template/src/ivf_flat_example.cu b/cpp/template/src/ivf_flat_example.cu
index 60694aea0f..9fb0bfc708 100644
--- a/cpp/template/src/ivf_flat_example.cu
+++ b/cpp/template/src/ivf_flat_example.cu
@@ -132,8 +132,8 @@ int main()
 
   // Set pool memory resource with 1 GiB initial pool size. All allocations use the same pool.
   rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_mr(
-    rmm::mr::get_current_device_resource(), 1024 * 1024 * 1024ull);
-  rmm::mr::set_current_device_resource(&pool_mr);
+    raft::resource::get_current_device_resource(), 1024 * 1024 * 1024ull);
+  raft::set_current_device_resource(&pool_mr);
 
   // Alternatively, one could define a pool allocator for temporary arrays (used within RAFT
   // algorithms). In that case only the internal arrays would use the pool, any other allocation
diff --git a/cpp/template/src/ivf_pq_example.cu b/cpp/template/src/ivf_pq_example.cu
index 4bc0ba4348..8ded79ba2e 100644
--- a/cpp/template/src/ivf_pq_example.cu
+++ b/cpp/template/src/ivf_pq_example.cu
@@ -92,8 +92,8 @@ int main()
 
   // Set pool memory resource with 1 GiB initial pool size. All allocations use the same pool.
   rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_mr(
-    rmm::mr::get_current_device_resource(), 1024 * 1024 * 1024ull);
-  rmm::mr::set_current_device_resource(&pool_mr);
+    raft::resource::get_current_device_resource(), 1024 * 1024 * 1024ull);
+  raft::set_current_device_resource(&pool_mr);
 
   // Alternatively, one could define a pool allocator for temporary arrays (used within RAFT
   // algorithms). In that case only the internal arrays would use the pool, any other allocation
diff --git a/cpp/test/core/device_resources_manager.cpp b/cpp/test/core/device_resources_manager.cpp
index c63d5896e5..cb7465c887 100644
--- a/cpp/test/core/device_resources_manager.cpp
+++ b/cpp/test/core/device_resources_manager.cpp
@@ -114,7 +114,7 @@ TEST(DeviceResourcesManager, ObeysSetters)
     EXPECT_EQ(streams_per_pool, pool.get_pool_size());
 
     auto* mr = dynamic_cast<rmm::mr::pool_memory_resource<rmm::mr::cuda_memory_resource>*>(
-      rmm::mr::get_current_device_resource());
+      resource::get_current_device_resource());
 
     if (upstream_mrs[i % devices.size()] != nullptr) {
       // Expect that the current memory resource is a pool memory resource as requested
diff --git a/cpp/test/core/handle.cpp b/cpp/test/core/handle.cpp
index be18b0d5b4..6cee514a3e 100644
--- a/cpp/test/core/handle.cpp
+++ b/cpp/test/core/handle.cpp
@@ -287,7 +287,7 @@ TEST(Raft, WorkspaceResource)
 
   // Let's create a pooled resource
   auto pool_mr = std::shared_ptr<rmm::mr::device_memory_resource>{new rmm::mr::pool_memory_resource(
-    rmm::mr::get_current_device_resource(), rmm::percent_of_free_device_memory(50))};
+    raft::resource::get_current_device_resource(), rmm::percent_of_free_device_memory(50))};
 
   // A tiny workspace of 1MB
   size_t max_size = 1024 * 1024;
@@ -326,7 +326,7 @@ TEST(Raft, WorkspaceResourceCopy)
     resource::set_workspace_resource(
       tmp_res,
       std::shared_ptr<rmm::mr::device_memory_resource>{new rmm::mr::pool_memory_resource(
-        rmm::mr::get_current_device_resource(), rmm::percent_of_free_device_memory(50))},
+        raft::resource::get_current_device_resource(), rmm::percent_of_free_device_memory(50))},
       orig_size * 2);
 
     ASSERT_EQ(orig_mr, resource::get_workspace_resource(res));
diff --git a/cpp/test/mr/device/buffer.cpp b/cpp/test/mr/device/buffer.cpp
index d14aa09b7a..ec2fe698f4 100644
--- a/cpp/test/mr/device/buffer.cpp
+++ b/cpp/test/mr/device/buffer.cpp
@@ -58,13 +58,12 @@ TEST(Raft, DeviceBufferAlloc)
 TEST(Raft, DeviceBufferZeroResize)
 {
   // Create a limiting_resource_adaptor to track allocations
-  auto curr_mr =
-    dynamic_cast<rmm::mr::cuda_memory_resource*>(rmm::mr::get_current_device_resource());
+  auto curr_mr = dynamic_cast<rmm::mr::cuda_memory_resource*>(raft::resource::get_current_device_resource());
   auto limit_mr =
     std::make_shared<rmm::mr::limiting_resource_adaptor<rmm::mr::cuda_memory_resource>>(curr_mr,
                                                                                         1000);
 
-  rmm::mr::set_current_device_resource(limit_mr.get());
+  raft::set_current_device_resource(limit_mr.get());
 
   cudaStream_t stream;
   RAFT_CUDA_TRY(cudaStreamCreate(&stream));
@@ -84,7 +83,7 @@ TEST(Raft, DeviceBufferZeroResize)
   // Now check that there is no memory left. (Used to not be true)
   ASSERT_EQ(0, limit_mr->get_allocated_bytes());
 
-  rmm::mr::set_current_device_resource(curr_mr);
+  raft::set_current_device_resource(curr_mr);
 
   RAFT_CUDA_TRY(cudaStreamSynchronize(stream));
   RAFT_CUDA_TRY(cudaStreamDestroy(stream));
diff --git a/cpp/test/random/multi_variable_gaussian.cu b/cpp/test/random/multi_variable_gaussian.cu
index bed9515a53..d7ac5da8ae 100644
--- a/cpp/test/random/multi_variable_gaussian.cu
+++ b/cpp/test/random/multi_variable_gaussian.cu
@@ -289,7 +289,7 @@ class MVGMdspanTest : public ::testing::TestWithParam<MVGInputs<T>> {
     raft::device_matrix_view<T, int, raft::col_major> X_view(X_d.data(), dim, nPoints);
 
     raft::random::multi_variable_gaussian(
-      handle, rmm::mr::get_current_device_resource(), x_view, P_view, X_view, method);
+      handle, raft::resource::get_current_device_resource_ref(), x_view, P_view, X_view, method);
 
     // saving the mean of the randoms in Rand_mean
     //@todo can be swapped with a API that calculates mean

From 20cfc8513f055a890187b8202e956d5b29acb62a Mon Sep 17 00:00:00 2001
From: Mark Harris <mharris@nvidia.com>
Date: Thu, 29 Aug 2024 21:08:03 +0000
Subject: [PATCH 03/14] More wrapper use.

---
 cpp/bench/ann/src/faiss/faiss_gpu_benchmark.cu     | 4 ++--
 cpp/bench/ann/src/raft/raft_ann_bench_utils.h      | 6 +++---
 cpp/bench/ann/src/raft/raft_cagra_hnswlib.cu       | 4 ++--
 cpp/bench/prims/common/benchmark.hpp               | 7 ++++---
 cpp/bench/prims/matrix/gather.cu                   | 4 ++--
 cpp/bench/prims/neighbors/refine.cuh               | 2 +-
 cpp/bench/prims/random/subsample.cu                | 2 +-
 cpp/include/raft/core/device_resources_manager.hpp | 2 +-
 8 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/cpp/bench/ann/src/faiss/faiss_gpu_benchmark.cu b/cpp/bench/ann/src/faiss/faiss_gpu_benchmark.cu
index b47c497e3d..46e2eb2457 100644
--- a/cpp/bench/ann/src/faiss/faiss_gpu_benchmark.cu
+++ b/cpp/bench/ann/src/faiss/faiss_gpu_benchmark.cu
@@ -183,10 +183,10 @@ int main(int argc, char** argv)
   rmm::mr::pool_memory_resource<rmm::mr::cuda_memory_resource> pool_mr{
     &cuda_mr, rmm::percent_of_free_device_memory(50)};
   // Updates the current device resource pointer to `pool_mr`
-  auto old_mr = rmm::mr::set_current_device_resource(&pool_mr);
+  auto old_mr = raft::resource::set_current_device_resource(&pool_mr);
   auto ret    = raft::bench::ann::run_main(argc, argv);
   // Restores the current device resource pointer to its previous value
-  rmm::mr::set_current_device_resource(old_mr);
+  raft::resource::set_current_device_resource(old_mr);
   return ret;
 }
 #endif
diff --git a/cpp/bench/ann/src/raft/raft_ann_bench_utils.h b/cpp/bench/ann/src/raft/raft_ann_bench_utils.h
index 1a70aa9657..2797524eab 100644
--- a/cpp/bench/ann/src/raft/raft_ann_bench_utils.h
+++ b/cpp/bench/ann/src/raft/raft_ann_bench_utils.h
@@ -80,10 +80,10 @@ class shared_raft_resources {
   using large_mr_type = rmm::mr::managed_memory_resource;
 
   shared_raft_resources()
-  try : orig_resource_{raft::resource::get_current_device_resource_ref()},
+  try : orig_resource_{raft::resource::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_);
+    raft::resource::set_current_device_resource(&resource_);
   } catch (const std::exception& e) {
     auto cuda_status = cudaGetLastError();
     size_t free      = 0;
@@ -103,7 +103,7 @@ class shared_raft_resources {
   shared_raft_resources(const shared_raft_resources& res)              = delete;
   shared_raft_resources& operator=(const shared_raft_resources& other) = delete;
 
-  ~shared_raft_resources() noexcept { rmm::mr::set_current_device_resource(orig_resource_); }
+  ~shared_raft_resources() noexcept { raft::resource::set_current_device_resource(orig_resource_); }
 
   auto get_large_memory_resource() noexcept
   {
diff --git a/cpp/bench/ann/src/raft/raft_cagra_hnswlib.cu b/cpp/bench/ann/src/raft/raft_cagra_hnswlib.cu
index d9ef1d74a3..6c4a81c03a 100644
--- a/cpp/bench/ann/src/raft/raft_cagra_hnswlib.cu
+++ b/cpp/bench/ann/src/raft/raft_cagra_hnswlib.cu
@@ -91,10 +91,10 @@ int main(int argc, char** argv)
   rmm::mr::pool_memory_resource<rmm::mr::cuda_memory_resource> pool_mr{
     &cuda_mr, rmm::percent_of_free_device_memory(50)};
   // Updates the current device resource pointer to `pool_mr`
-  auto old_mr = rmm::mr::set_current_device_resource(&pool_mr);
+  auto old_mr = raft::resource::set_current_device_resource(&pool_mr);
   auto ret    = raft::bench::ann::run_main(argc, argv);
   // Restores the current device resource pointer to its previous value
-  rmm::mr::set_current_device_resource(old_mr);
+  raft::resource::set_current_device_resource(old_mr);
   return ret;
 }
 #endif
diff --git a/cpp/bench/prims/common/benchmark.hpp b/cpp/bench/prims/common/benchmark.hpp
index f2ad980e9e..db36d4961e 100644
--- a/cpp/bench/prims/common/benchmark.hpp
+++ b/cpp/bench/prims/common/benchmark.hpp
@@ -21,6 +21,7 @@
 #include <raft/core/device_resources.hpp>
 #include <raft/core/interruptible.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/random/make_blobs.cuh>
 #include <raft/util/cudart_utils.hpp>
 
@@ -53,17 +54,17 @@ struct using_pool_memory_res {
     : orig_res_(raft::resource::get_current_device_resource_ref()),
       pool_res_(&cuda_res_, initial_size, max_size)
   {
-    rmm::mr::set_current_device_resource(&pool_res_);
+    raft::resource::set_current_device_resource(&pool_res_);
   }
 
   using_pool_memory_res()
     : orig_res_(raft::resource::get_current_device_resource_ref()),
       pool_res_(&cuda_res_, rmm::percent_of_free_device_memory(50))
   {
-    rmm::mr::set_current_device_resource(&pool_res_);
+    raft::resource::set_current_device_resource(&pool_res_);
   }
 
-  ~using_pool_memory_res() { rmm::mr::set_current_device_resource(orig_res_); }
+  ~using_pool_memory_res() { raft::resource::set_current_device_resource(orig_res_); }
 };
 
 /**
diff --git a/cpp/bench/prims/matrix/gather.cu b/cpp/bench/prims/matrix/gather.cu
index a12792d774..0a2029b05d 100644
--- a/cpp/bench/prims/matrix/gather.cu
+++ b/cpp/bench/prims/matrix/gather.cu
@@ -54,10 +54,10 @@ struct Gather : public fixture {
       stencil(this->handle),
       matrix_h(this->handle)
   {
-    rmm::mr::set_current_device_resource(&pool_mr);
+    raft::resource::set_current_device_resource(&pool_mr);
   }
 
-  ~Gather() { rmm::mr::set_current_device_resource(old_mr); }
+  ~Gather() { raft::resource::set_current_device_resource(old_mr); }
 
   void allocate_data(const ::benchmark::State& state) override
   {
diff --git a/cpp/bench/prims/neighbors/refine.cuh b/cpp/bench/prims/neighbors/refine.cuh
index 3e8721f353..14a01321ac 100644
--- a/cpp/bench/prims/neighbors/refine.cuh
+++ b/cpp/bench/prims/neighbors/refine.cuh
@@ -61,7 +61,7 @@ class RefineAnn : public fixture {
     auto old_mr = raft::resource::get_current_device_resource_ref();
     rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_mr(
       old_mr, rmm::percent_of_free_device_memory(50));
-    rmm::mr::set_current_device_resource(&pool_mr);
+    raft::resource::set_current_device_resource(&pool_mr);
 
     if (data.p.host_data) {
       loop_on_state(state, [this]() {
diff --git a/cpp/bench/prims/random/subsample.cu b/cpp/bench/prims/random/subsample.cu
index 58cd35276c..7eef2b747c 100644
--- a/cpp/bench/prims/random/subsample.cu
+++ b/cpp/bench/prims/random/subsample.cu
@@ -69,7 +69,7 @@ struct sample : public fixture {
       in(make_device_vector<T, int64_t>(res, p.n_samples)),
       out(make_device_vector<T, int64_t>(res, p.n_train))
   {
-    rmm::mr::set_current_device_resource(&pool_mr);
+    raft::resource::set_current_device_resource(&pool_mr);
     raft::random::RngState r(123456ULL);
   }
 
diff --git a/cpp/include/raft/core/device_resources_manager.hpp b/cpp/include/raft/core/device_resources_manager.hpp
index bef70a8740..9c22cab6eb 100644
--- a/cpp/include/raft/core/device_resources_manager.hpp
+++ b/cpp/include/raft/core/device_resources_manager.hpp
@@ -176,7 +176,7 @@ struct device_resources_manager {
                   upstream,
                   params.init_mem_pool_size.value_or(rmm::percent_of_free_device_memory(50)),
                   params.max_mem_pool_size);
-              rmm::mr::set_current_device_resource(result.get());
+              raft::resource::set_current_device_resource(result.get());
             } else {
               RAFT_LOG_WARN(
                 "Pool allocation requested, but other memory resource has already been set and "

From 2131acb4e848a9505fc36347bc0d8873179657cb Mon Sep 17 00:00:00 2001
From: Mark Harris <mharris@nvidia.com>
Date: Tue, 3 Sep 2024 01:15:25 +0000
Subject: [PATCH 04/14] style

---
 cpp/test/mr/device/buffer.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/cpp/test/mr/device/buffer.cpp b/cpp/test/mr/device/buffer.cpp
index ec2fe698f4..2ba6ca112a 100644
--- a/cpp/test/mr/device/buffer.cpp
+++ b/cpp/test/mr/device/buffer.cpp
@@ -58,7 +58,8 @@ TEST(Raft, DeviceBufferAlloc)
 TEST(Raft, DeviceBufferZeroResize)
 {
   // Create a limiting_resource_adaptor to track allocations
-  auto curr_mr = dynamic_cast<rmm::mr::cuda_memory_resource*>(raft::resource::get_current_device_resource());
+  auto curr_mr =
+    dynamic_cast<rmm::mr::cuda_memory_resource*>(raft::resource::get_current_device_resource());
   auto limit_mr =
     std::make_shared<rmm::mr::limiting_resource_adaptor<rmm::mr::cuda_memory_resource>>(curr_mr,
                                                                                         1000);

From c7e91a000c5e367891dd2952b6aa9a2297ee7093 Mon Sep 17 00:00:00 2001
From: Mark Harris <mharris@nvidia.com>
Date: Tue, 3 Sep 2024 01:15:38 +0000
Subject: [PATCH 05/14] Not ready for _ref here yet.

---
 cpp/bench/prims/common/benchmark.hpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cpp/bench/prims/common/benchmark.hpp b/cpp/bench/prims/common/benchmark.hpp
index db36d4961e..ccbff22b13 100644
--- a/cpp/bench/prims/common/benchmark.hpp
+++ b/cpp/bench/prims/common/benchmark.hpp
@@ -51,14 +51,14 @@ struct using_pool_memory_res {
 
  public:
   using_pool_memory_res(size_t initial_size, size_t max_size)
-    : orig_res_(raft::resource::get_current_device_resource_ref()),
+    : orig_res_(raft::resource::get_current_device_resource()),
       pool_res_(&cuda_res_, initial_size, max_size)
   {
     raft::resource::set_current_device_resource(&pool_res_);
   }
 
   using_pool_memory_res()
-    : orig_res_(raft::resource::get_current_device_resource_ref()),
+    : orig_res_(raft::resource::get_current_device_resource()),
       pool_res_(&cuda_res_, rmm::percent_of_free_device_memory(50))
   {
     raft::resource::set_current_device_resource(&pool_res_);

From de3a1df4767da1435afc521bd0f1ed7ccad1ce74 Mon Sep 17 00:00:00 2001
From: Mark Harris <mharris@nvidia.com>
Date: Wed, 4 Sep 2024 02:04:51 +0000
Subject: [PATCH 06/14] Add missing namespace and include

---
 cpp/bench/prims/neighbors/refine.cuh | 3 ++-
 cpp/bench/prims/random/subsample.cu  | 3 ++-
 cpp/template/src/ivf_flat_example.cu | 3 ++-
 cpp/template/src/ivf_pq_example.cu   | 3 ++-
 cpp/test/mr/device/buffer.cpp        | 5 +++--
 5 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/cpp/bench/prims/neighbors/refine.cuh b/cpp/bench/prims/neighbors/refine.cuh
index 14a01321ac..c29349927a 100644
--- a/cpp/bench/prims/neighbors/refine.cuh
+++ b/cpp/bench/prims/neighbors/refine.cuh
@@ -20,6 +20,7 @@
 
 #include <raft/core/device_mdspan.hpp>
 #include <raft/core/device_resources.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/distance/distance_types.hpp>
 #include <raft/neighbors/detail/refine.cuh>
 #include <raft/neighbors/refine.cuh>
@@ -84,7 +85,7 @@ class RefineAnn : public fixture {
                                                               data.p.metric);
       });
     }
-    raft::set_current_device_resource(old_mr);
+    raft::resource::set_current_device_resource(old_mr);
   }
 
  private:
diff --git a/cpp/bench/prims/random/subsample.cu b/cpp/bench/prims/random/subsample.cu
index 7eef2b747c..ca0aa79101 100644
--- a/cpp/bench/prims/random/subsample.cu
+++ b/cpp/bench/prims/random/subsample.cu
@@ -20,6 +20,7 @@
 #include <raft/core/device_resources.hpp>
 #include <raft/core/host_mdarray.hpp>
 #include <raft/core/operators.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/random/permute.cuh>
 #include <raft/random/rng.cuh>
 #include <raft/random/sample_without_replacement.cuh>
@@ -73,7 +74,7 @@ struct sample : public fixture {
     raft::random::RngState r(123456ULL);
   }
 
-  ~sample() { raft::set_current_device_resource(old_mr); }
+  ~sample() { raft::resource::set_current_device_resource(old_mr); }
   void run_benchmark(::benchmark::State& state) override
   {
     std::ostringstream label_stream;
diff --git a/cpp/template/src/ivf_flat_example.cu b/cpp/template/src/ivf_flat_example.cu
index 9fb0bfc708..1f6fe7dff8 100644
--- a/cpp/template/src/ivf_flat_example.cu
+++ b/cpp/template/src/ivf_flat_example.cu
@@ -18,6 +18,7 @@
 
 #include <raft/core/device_mdarray.hpp>
 #include <raft/core/device_resources.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/core/resource/thrust_policy.hpp>
 #include <raft/neighbors/ivf_flat.cuh>
 #include <raft/util/cudart_utils.hpp>
@@ -133,7 +134,7 @@ int main()
   // Set pool memory resource with 1 GiB initial pool size. All allocations use the same pool.
   rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_mr(
     raft::resource::get_current_device_resource(), 1024 * 1024 * 1024ull);
-  raft::set_current_device_resource(&pool_mr);
+  raft::resource::set_current_device_resource(&pool_mr);
 
   // Alternatively, one could define a pool allocator for temporary arrays (used within RAFT
   // algorithms). In that case only the internal arrays would use the pool, any other allocation
diff --git a/cpp/template/src/ivf_pq_example.cu b/cpp/template/src/ivf_pq_example.cu
index 8ded79ba2e..98fdc76c12 100644
--- a/cpp/template/src/ivf_pq_example.cu
+++ b/cpp/template/src/ivf_pq_example.cu
@@ -18,6 +18,7 @@
 
 #include <raft/core/device_mdarray.hpp>
 #include <raft/core/device_resources.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/neighbors/ivf_pq.cuh>
 #include <raft/neighbors/refine.cuh>
 
@@ -93,7 +94,7 @@ int main()
   // Set pool memory resource with 1 GiB initial pool size. All allocations use the same pool.
   rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_mr(
     raft::resource::get_current_device_resource(), 1024 * 1024 * 1024ull);
-  raft::set_current_device_resource(&pool_mr);
+  raft::resource::set_current_device_resource(&pool_mr);
 
   // Alternatively, one could define a pool allocator for temporary arrays (used within RAFT
   // algorithms). In that case only the internal arrays would use the pool, any other allocation
diff --git a/cpp/test/mr/device/buffer.cpp b/cpp/test/mr/device/buffer.cpp
index 2ba6ca112a..6c9a968c86 100644
--- a/cpp/test/mr/device/buffer.cpp
+++ b/cpp/test/mr/device/buffer.cpp
@@ -14,6 +14,7 @@
  * limitations under the License.
  */
 
+#include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/util/cudart_utils.hpp>
 
 #include <rmm/device_uvector.hpp>
@@ -64,7 +65,7 @@ TEST(Raft, DeviceBufferZeroResize)
     std::make_shared<rmm::mr::limiting_resource_adaptor<rmm::mr::cuda_memory_resource>>(curr_mr,
                                                                                         1000);
 
-  raft::set_current_device_resource(limit_mr.get());
+  raft::resource::set_current_device_resource(limit_mr.get());
 
   cudaStream_t stream;
   RAFT_CUDA_TRY(cudaStreamCreate(&stream));
@@ -84,7 +85,7 @@ TEST(Raft, DeviceBufferZeroResize)
   // Now check that there is no memory left. (Used to not be true)
   ASSERT_EQ(0, limit_mr->get_allocated_bytes());
 
-  raft::set_current_device_resource(curr_mr);
+  raft::resource::set_current_device_resource(curr_mr);
 
   RAFT_CUDA_TRY(cudaStreamSynchronize(stream));
   RAFT_CUDA_TRY(cudaStreamDestroy(stream));

From 04ca772baf787ee56fe00938b86ba9a38e067662 Mon Sep 17 00:00:00 2001
From: Mark Harris <mharris@nvidia.com>
Date: Wed, 4 Sep 2024 02:09:04 +0000
Subject: [PATCH 07/14] No need to alias RMM's aliases.

---
 .../core/resource/device_memory_resource.hpp  | 43 ++-----------------
 1 file changed, 4 insertions(+), 39 deletions(-)

diff --git a/cpp/include/raft/core/resource/device_memory_resource.hpp b/cpp/include/raft/core/resource/device_memory_resource.hpp
index 99fb770bb2..5f0effb03b 100644
--- a/cpp/include/raft/core/resource/device_memory_resource.hpp
+++ b/cpp/include/raft/core/resource/device_memory_resource.hpp
@@ -36,42 +36,6 @@ namespace raft::resource {
  * @{
  */
 
-/**
- * @brief Alias for a `cuda::mr::resource_ref` with the property
- * `cuda::mr::device_accessible`.
- */
-using device_resource_ref = rmm::device_resource_ref;
-
-/**
- * @brief Alias for a `cuda::mr::async_resource_ref` with the property
- * `cuda::mr::device_accessible`.
- */
-using device_async_resource_ref = rmm::device_async_resource_ref;
-
-/**
- * @brief Alias for a `cuda::mr::resource_ref` with the property
- * `cuda::mr::host_accessible`.
- */
-using host_resource_ref = rmm::host_resource_ref;
-
-/**
- * @brief Alias for a `cuda::mr::async_resource_ref` with the property
- * `cuda::mr::host_accessible`.
- */
-using host_async_resource_ref = rmm::host_async_resource_ref;
-
-/**
- * @brief Alias for a `cuda::mr::resource_ref` with the properties
- * `cuda::mr::host_accessible` and `cuda::mr::device_accessible`.
- */
-using host_device_resource_ref = rmm::host_device_resource_ref;
-
-/**
- * @brief Alias for a `cuda::mr::async_resource_ref` with the properties
- * `cuda::mr::host_accessible` and `cuda::mr::device_accessible`.
- */
-using host_device_async_resource_ref = rmm::host_device_async_resource_ref;
-
 /**
  * @brief Get the current device memory resource.
  *
@@ -87,7 +51,7 @@ inline rmm::mr::device_memory_resource* get_current_device_resource()
  *
  * @return The current device memory resource reference.
  */
-inline device_async_resource_ref get_current_device_resource_ref()
+inline rmm::device_async_resource_ref get_current_device_resource_ref()
 {
   // For now, match current behavior which is to return current resource pointer
   return rmm::mr::get_current_device_resource();
@@ -111,7 +75,8 @@ inline rmm::mr::device_memory_resource* set_current_device_resource(
  * @param mr The new device memory resource reference.
  * @return The previous device memory resource reference.
  */
-inline device_async_resource_ref set_current_device_resource_ref(device_async_resource_ref mr)
+inline rmm::device_async_resource_ref set_current_device_resource_ref(
+  rmm::device_async_resource_ref mr)
 {
   return rmm::mr::set_current_device_resource_ref(mr);
 }
@@ -121,7 +86,7 @@ inline device_async_resource_ref set_current_device_resource_ref(device_async_re
  *
  * @return The previous device memory resource reference.
  */
-inline device_async_resource_ref reset_current_device_resource_ref()
+inline rmm::device_async_resource_ref reset_current_device_resource_ref()
 {
   return rmm::mr::reset_current_device_resource_ref();
 }

From 9f2faf0be4670d82d7c62c49dc4cf3df56f474fc Mon Sep 17 00:00:00 2001
From: Mark Harris <mharris@nvidia.com>
Date: Wed, 4 Sep 2024 07:44:05 +0000
Subject: [PATCH 08/14] doc

---
 docs/source/vector_search_tutorial.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/source/vector_search_tutorial.md b/docs/source/vector_search_tutorial.md
index d1d5c57700..a528a37d0c 100644
--- a/docs/source/vector_search_tutorial.md
+++ b/docs/source/vector_search_tutorial.md
@@ -375,7 +375,7 @@ rmm::mr::cuda_memory_resource cuda_mr;
 // set the initial size to half of the free device memory
 auto init_size = rmm::percent_of_free_device_memory(50);
 rmm::mr::pool_memory_resource<rmm::mr::cuda_memory_resource> pool_mr{&cuda_mr, init_size};
-rmm::mr::set_current_device_resource(&pool_mr); // Updates the current device resource pointer to `pool_mr`
+raft::resource::set_current_device_resource(&pool_mr); // Updates the current device resource pointer to `pool_mr`
 ```
 
 The `raft::device_resources` object will now also use the `rmm::current_device_resource`.  This isn't limited to C++, however. Often a user will be interacting with PyTorch, RAPIDS, or Tensorflow through Python and so they can set and use RMM's `current_device_resource` [right in Python](https://github.com/rapidsai/rmm#using-rmm-in-python-code).

From 0cd99c79040415ef39201679d1428805a905c67c Mon Sep 17 00:00:00 2001
From: Mark Harris <mharris@nvidia.com>
Date: Mon, 9 Sep 2024 23:46:31 +0000
Subject: [PATCH 09/14] Fix a couple of compilation errors.

---
 cpp/bench/prims/matrix/gather.cu  | 2 +-
 cpp/bench/prims/neighbors/knn.cuh | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/cpp/bench/prims/matrix/gather.cu b/cpp/bench/prims/matrix/gather.cu
index 0a2029b05d..c0b6711b77 100644
--- a/cpp/bench/prims/matrix/gather.cu
+++ b/cpp/bench/prims/matrix/gather.cu
@@ -46,7 +46,7 @@ template <typename T, typename MapT, typename IdxT, bool Conditional = false>
 struct Gather : public fixture {
   Gather(const GatherParams<IdxT>& p)
     : params(p),
-      old_mr(raft::resource::get_current_device_resource_ref()),
+      old_mr(raft::resource::get_current_device_resource()),
       pool_mr(raft::resource::get_current_device_resource_ref(), 2 * (1ULL << 30)),
       matrix(this->handle),
       map(this->handle),
diff --git a/cpp/bench/prims/neighbors/knn.cuh b/cpp/bench/prims/neighbors/knn.cuh
index f8b837d77f..176b10cee8 100644
--- a/cpp/bench/prims/neighbors/knn.cuh
+++ b/cpp/bench/prims/neighbors/knn.cuh
@@ -32,6 +32,7 @@
 #include <rmm/mr/device/managed_memory_resource.hpp>
 #include <rmm/mr/host/new_delete_resource.hpp>
 #include <rmm/mr/host/pinned_memory_resource.hpp>
+#include <rmm/resource_ref.hpp>
 
 #include <thrust/sequence.h>
 
@@ -105,7 +106,7 @@ struct device_resource {
 
  private:
   std::shared_ptr<rmm::mr::device_memory_resource> managed_;
-  raft::device_async_resource_ref res_;
+  rmm::device_async_resource_ref res_;
 };
 
 template <typename T>

From 44d9b11003d764db6d6bcc41280cdfe7e4a90244 Mon Sep 17 00:00:00 2001
From: Mark Harris <mharris@nvidia.com>
Date: Tue, 10 Sep 2024 05:15:06 +0000
Subject: [PATCH 10/14] Fix resource_ref init

---
 cpp/bench/prims/neighbors/knn.cuh | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/cpp/bench/prims/neighbors/knn.cuh b/cpp/bench/prims/neighbors/knn.cuh
index 176b10cee8..8dabe5dc46 100644
--- a/cpp/bench/prims/neighbors/knn.cuh
+++ b/cpp/bench/prims/neighbors/knn.cuh
@@ -91,13 +91,9 @@ inline auto operator<<(std::ostream& os, const Scope& s) -> std::ostream&
 struct device_resource {
  public:
   explicit device_resource(bool managed)
-    : managed_(managed ? std::make_shared<rmm::mr::managed_memory_resource>() : nullptr)
+    : managed_(managed ? std::make_shared<rmm::mr::managed_memory_resource>() : nullptr),
+      res_(managed ? managed.get() : raft::resource::get_current_device_resource_ref())
   {
-    if (managed) {
-      res_ = managed.get();
-    } else {
-      res_ = raft::resource::get_current_device_resource_ref();
-    }
   }
 
   ~device_resource() {}

From dfdfe9d147a80b7dd3d6a7eb59ff5d7e7ac5dced Mon Sep 17 00:00:00 2001
From: Mark Harris <mharris@nvidia.com>
Date: Tue, 10 Sep 2024 06:27:57 +0000
Subject: [PATCH 11/14] Typo

---
 cpp/bench/prims/neighbors/knn.cuh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/bench/prims/neighbors/knn.cuh b/cpp/bench/prims/neighbors/knn.cuh
index 8dabe5dc46..4a48855a82 100644
--- a/cpp/bench/prims/neighbors/knn.cuh
+++ b/cpp/bench/prims/neighbors/knn.cuh
@@ -92,7 +92,7 @@ struct device_resource {
  public:
   explicit device_resource(bool managed)
     : managed_(managed ? std::make_shared<rmm::mr::managed_memory_resource>() : nullptr),
-      res_(managed ? managed.get() : raft::resource::get_current_device_resource_ref())
+      res_(managed ? managed_.get() : raft::resource::get_current_device_resource_ref())
   {
   }
 

From 2fa415cfe3e73f07b6fb86de748d0b27ddb0d4dc Mon Sep 17 00:00:00 2001
From: Mark Harris <mharris@nvidia.com>
Date: Tue, 10 Sep 2024 20:42:15 +0000
Subject: [PATCH 12/14] Use =default;

---
 cpp/bench/prims/neighbors/knn.cuh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/bench/prims/neighbors/knn.cuh b/cpp/bench/prims/neighbors/knn.cuh
index 4a48855a82..02442dd8f6 100644
--- a/cpp/bench/prims/neighbors/knn.cuh
+++ b/cpp/bench/prims/neighbors/knn.cuh
@@ -96,7 +96,7 @@ struct device_resource {
   {
   }
 
-  ~device_resource() {}
+  ~device_resource() = default;
 
   [[nodiscard]] auto get() const -> rmm::device_async_resource_ref { return res_; }
 

From 9ced2434aabf3c1f96f6051be6c13bece6e59b68 Mon Sep 17 00:00:00 2001
From: Mark Harris <mharris@nvidia.com>
Date: Tue, 10 Sep 2024 20:42:50 +0000
Subject: [PATCH 13/14] set_current_device_resource() (not ref)

---
 cpp/bench/prims/neighbors/refine.cuh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/bench/prims/neighbors/refine.cuh b/cpp/bench/prims/neighbors/refine.cuh
index c29349927a..b2d85a8085 100644
--- a/cpp/bench/prims/neighbors/refine.cuh
+++ b/cpp/bench/prims/neighbors/refine.cuh
@@ -59,7 +59,7 @@ class RefineAnn : public fixture {
     label_stream << data.p;
     state.SetLabel(label_stream.str());
 
-    auto old_mr = raft::resource::get_current_device_resource_ref();
+    auto old_mr = raft::resource::get_current_device_resource();
     rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_mr(
       old_mr, rmm::percent_of_free_device_memory(50));
     raft::resource::set_current_device_resource(&pool_mr);

From 3c7259674d1f43106a4ab8e0a23f5b13bd4625ab Mon Sep 17 00:00:00 2001
From: Mark Harris <mharris@nvidia.com>
Date: Wed, 11 Sep 2024 09:49:13 +0000
Subject: [PATCH 14/14] Add include to doc example.

---
 docs/source/vector_search_tutorial.md | 1 +
 1 file changed, 1 insertion(+)

diff --git a/docs/source/vector_search_tutorial.md b/docs/source/vector_search_tutorial.md
index a528a37d0c..9ef1d8cd61 100644
--- a/docs/source/vector_search_tutorial.md
+++ b/docs/source/vector_search_tutorial.md
@@ -368,6 +368,7 @@ The RAPIDS software ecosystem makes heavy use of the [RAPIDS Memory Manager](htt
 
 As an example, the following code snippet creates a `pool_memory_resource` and sets it as the default memory resource, which means all other libraries that use RMM will now allocate their device memory from this same pool:
 ```c++
+#include <raft/core/resource/device_memory_resource.hpp>
 #include <rmm/mr/device/pool_memory_resource.hpp>
 
 rmm::mr::cuda_memory_resource cuda_mr;