Update RMM memory resource APIs to ref-based equivalents

cpp/bench/prims/common/benchmark.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2022-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2022-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -20,6 +20,7 @@
 #include <rmm/mr/device_memory_resource.hpp>
 #include <rmm/mr/per_device_resource.hpp>
 #include <rmm/mr/pool_memory_resource.hpp>
+#include <rmm/resource_ref.hpp>
 
 #include <benchmark/benchmark.h>
 
@@ -33,26 +34,26 @@ namespace raft::bench {
  */
 struct using_pool_memory_res {
  private:
-  rmm::mr::device_memory_resource* orig_res_;
+  rmm::device_async_resource_ref orig_res_;
   rmm::mr::cuda_memory_resource cuda_res_{};
   rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_res_;
 
  public:
   using_pool_memory_res(size_t initial_size, size_t max_size)
-    : orig_res_(rmm::mr::get_current_device_resource()),
+    : orig_res_(rmm::mr::get_current_device_resource_ref()),
       pool_res_(&cuda_res_, initial_size, max_size)
   {
-    rmm::mr::set_current_device_resource(&pool_res_);
+    rmm::mr::set_current_device_resource_ref(&pool_res_);
   }
 
   using_pool_memory_res()
-    : orig_res_(rmm::mr::get_current_device_resource()),
+    : orig_res_(rmm::mr::get_current_device_resource_ref()),
       pool_res_(&cuda_res_, rmm::percent_of_free_device_memory(50))
   {
-    rmm::mr::set_current_device_resource(&pool_res_);
+    rmm::mr::set_current_device_resource_ref(&pool_res_);
   }
 
-  ~using_pool_memory_res() { rmm::mr::set_current_device_resource(orig_res_); }
+  ~using_pool_memory_res() { rmm::mr::set_current_device_resource_ref(orig_res_); }
 };
 
 /**

cpp/bench/prims/matrix/gather.cu

@@ -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
  */
 
@@ -14,7 +14,9 @@
 
 #include <rmm/device_uvector.hpp>
 #include <rmm/mr/device_memory_resource.hpp>
+#include <rmm/mr/per_device_resource.hpp>
 #include <rmm/mr/pool_memory_resource.hpp>
+#include <rmm/resource_ref.hpp>
 
 namespace raft::bench::matrix {
 
@@ -35,18 +37,18 @@ 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()),
+      old_mr(rmm::mr::get_current_device_resource_ref()),
       pool_mr(rmm::mr::get_current_device_resource(), 2 * (1ULL << 30)),
       matrix(this->handle),
       map(this->handle),
       out(this->handle),
       stencil(this->handle),
       matrix_h(this->handle)
   {
-    rmm::mr::set_current_device_resource(&pool_mr);
+    rmm::mr::set_current_device_resource_ref(&pool_mr);
   }
 
-  ~Gather() { rmm::mr::set_current_device_resource(old_mr); }
+  ~Gather() { rmm::mr::set_current_device_resource_ref(old_mr); }
 
   void allocate_data(const ::benchmark::State& state) override
   {
@@ -107,7 +109,7 @@ struct Gather : public fixture {
 
  private:
   GatherParams<IdxT> params;
-  rmm::mr::device_memory_resource* old_mr;
+  rmm::device_async_resource_ref old_mr;
   rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_mr;
   raft::device_matrix<T, IdxT> matrix, out;
   raft::host_matrix<T, IdxT> matrix_h;

cpp/bench/prims/random/subsample.cu

@@ -17,6 +17,7 @@
 #include <rmm/device_scalar.hpp>
 #include <rmm/mr/per_device_resource.hpp>
 #include <rmm/mr/pool_memory_resource.hpp>
+#include <rmm/resource_ref.hpp>
 
 namespace raft::bench::random {
 
@@ -50,16 +51,16 @@ template <typename T>
 struct sample : public fixture {
   sample(const sample_inputs& p)
     : params(p),
-      old_mr(rmm::mr::get_current_device_resource()),
+      old_mr(rmm::mr::get_current_device_resource_ref()),
       pool_mr(rmm::mr::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))
   {
-    rmm::mr::set_current_device_resource(&pool_mr);
+    rmm::mr::set_current_device_resource_ref(&pool_mr);
     raft::random::RngState r(123456ULL);
   }
 
-  ~sample() { rmm::mr::set_current_device_resource(old_mr); }
+  ~sample() { rmm::mr::set_current_device_resource_ref(old_mr); }
   void run_benchmark(::benchmark::State& state) override
   {
     std::ostringstream label_stream;
@@ -81,7 +82,7 @@ struct sample : public fixture {
  private:
   float GiB = 1073741824.0f;
   raft::device_resources res;
-  rmm::mr::device_memory_resource* old_mr;
+  rmm::device_async_resource_ref old_mr;
   rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_mr;
   sample_inputs params;
   raft::device_vector<T, int64_t> out, in;

cpp/include/raft/core/device_resources_manager.hpp

@@ -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
  */
 
@@ -165,7 +165,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());
+              rmm::mr::set_current_device_resource_ref(result.get());
             } else {
               RAFT_LOG_WARN(
                 "Pool allocation requested, but other memory resource has already been set and "

cpp/include/raft/random/detail/multi_variable_gaussian.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -22,6 +22,8 @@
 #include <rmm/device_uvector.hpp>
 #include <rmm/resource_ref.hpp>
 
+#include <cuda/memory_resource>
+
 #include <cmath>
 #include <cstdio>
 #include <memory>
@@ -365,7 +367,7 @@ class multi_variable_gaussian_setup_token {
  private:
   std::unique_ptr<multi_variable_gaussian_impl<ValueType>> impl_;
   raft::resources const& handle_;
-  rmm::device_async_resource_ref mem_resource_;
+  mutable cuda::mr::any_resource<cuda::mr::device_accessible> mem_resource_;
   int dim_ = 0;
 
   auto allocate_workspace() const

cpp/tests/mr/device/buffer.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2020-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -53,7 +53,7 @@ TEST(Raft, DeviceBufferZeroResize)
     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());
+  rmm::mr::set_current_device_resource_ref(limit_mr.get());
 
   cudaStream_t stream;
   RAFT_CUDA_TRY(cudaStreamCreate(&stream));
@@ -73,7 +73,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);
+  rmm::mr::set_current_device_resource_ref(curr_mr);
 
   RAFT_CUDA_TRY(cudaStreamSynchronize(stream));
   RAFT_CUDA_TRY(cudaStreamDestroy(stream));