Migrate RMM usage to CCCL memory resource design

Remove device_memory_resource base class usage, de-template all resource
and adaptor types, replace pointer-based per-device resource APIs with
ref-based equivalents, and update all call sites for the new signatures.

Part of rapidsai/rmm#2011.

Author: bdice

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
  */
 
@@ -17,7 +17,7 @@
 #include <rmm/cuda_stream.hpp>
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_buffer.hpp>
-#include <rmm/mr/device_memory_resource.hpp>
+#include <rmm/mr/cuda_memory_resource.hpp>
 #include <rmm/mr/per_device_resource.hpp>
 #include <rmm/mr/pool_memory_resource.hpp>
 
@@ -33,26 +33,24 @@ namespace raft::bench {
  */
 struct using_pool_memory_res {
  private:
-  rmm::mr::device_memory_resource* orig_res_;
   rmm::mr::cuda_memory_resource cuda_res_{};
-  rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_res_;
+  rmm::mr::pool_memory_resource pool_res_;
+  cuda::mr::any_resource<cuda::mr::device_accessible> prev_res_;
 
  public:
   using_pool_memory_res(size_t initial_size, size_t max_size)
-    : orig_res_(rmm::mr::get_current_device_resource()),
-      pool_res_(&cuda_res_, initial_size, max_size)
+    : pool_res_(cuda_res_, initial_size, max_size),
+      prev_res_(rmm::mr::set_current_device_resource_ref(pool_res_))
   {
-    rmm::mr::set_current_device_resource(&pool_res_);
   }
 
   using_pool_memory_res()
-    : orig_res_(rmm::mr::get_current_device_resource()),
-      pool_res_(&cuda_res_, rmm::percent_of_free_device_memory(50))
+    : pool_res_(cuda_res_, rmm::percent_of_free_device_memory(50)),
+      prev_res_(rmm::mr::set_current_device_resource_ref(pool_res_))
   {
-    rmm::mr::set_current_device_resource(&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(prev_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
  */
 
@@ -13,7 +13,7 @@
 #include <raft/util/itertools.hpp>
 
 #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>
 
 namespace raft::bench::matrix {
@@ -35,18 +35,17 @@ 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)),
+      pool_mr(rmm::mr::get_current_device_resource_ref(), 2 * (1ULL << 30)),
+      prev_res_(rmm::mr::set_current_device_resource_ref(pool_mr)),
       matrix(this->handle),
       map(this->handle),
       out(this->handle),
       stencil(this->handle),
       matrix_h(this->handle)
   {
-    rmm::mr::set_current_device_resource(&pool_mr);
   }
 
-  ~Gather() { rmm::mr::set_current_device_resource(old_mr); }
+  ~Gather() { rmm::mr::set_current_device_resource_ref(prev_res_); }
 
   void allocate_data(const ::benchmark::State& state) override
   {
@@ -107,8 +106,8 @@ struct Gather : public fixture {
 
  private:
   GatherParams<IdxT> params;
-  rmm::mr::device_memory_resource* old_mr;
-  rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_mr;
+  rmm::mr::pool_memory_resource pool_mr;
+  cuda::mr::any_resource<cuda::mr::device_accessible> prev_res_;
   raft::device_matrix<T, IdxT> matrix, out;
   raft::host_matrix<T, IdxT> matrix_h;
   raft::device_vector<T, IdxT> stencil;

cpp/bench/prims/random/subsample.cu

@@ -50,16 +50,15 @@ 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),
+      pool_mr(rmm::mr::get_current_device_resource_ref(), 2 * GiB),
+      prev_mr(rmm::mr::set_current_device_resource_ref(pool_mr)),
       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::random::RngState r(123456ULL);
   }
 
-  ~sample() { rmm::mr::set_current_device_resource(old_mr); }
+  ~sample() { rmm::mr::set_current_device_resource_ref(prev_mr); }
   void run_benchmark(::benchmark::State& state) override
   {
     std::ostringstream label_stream;
@@ -81,8 +80,8 @@ struct sample : public fixture {
  private:
   float GiB = 1073741824.0f;
   raft::device_resources res;
-  rmm::mr::device_memory_resource* old_mr;
-  rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_mr;
+  rmm::mr::pool_memory_resource pool_mr;
+  cuda::mr::any_resource<cuda::mr::device_accessible> prev_mr;
   sample_inputs params;
   raft::device_vector<T, int64_t> out, in;
 };  // struct sample

cpp/include/raft/core/device_resources.hpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -26,7 +26,7 @@
 
 #include <rmm/cuda_stream_pool.hpp>
 #include <rmm/exec_policy.hpp>
-#include <rmm/mr/device_memory_resource.hpp>
+#include <rmm/resource_ref.hpp>
 
 #include <cuda_runtime.h>
 
@@ -51,15 +51,6 @@ namespace raft {
  */
 class device_resources : public resources {
  public:
-  device_resources(const device_resources& handle,
-                   std::shared_ptr<rmm::mr::device_memory_resource> workspace_resource,
-                   std::optional<std::size_t> allocation_limit = std::nullopt)
-    : resources{handle}
-  {
-    // replace the resource factory for the workspace_resources
-    resource::set_workspace_resource(*this, workspace_resource, allocation_limit);
-  }
-
   device_resources(const device_resources& handle) : resources{handle} {}
   device_resources(device_resources&&)            = delete;
   device_resources& operator=(device_resources&&) = delete;
@@ -70,25 +61,16 @@ class device_resources : public resources {
    * @param[in] stream_view the default stream (which has the default per-thread stream if
    * unspecified)
    * @param[in] stream_pool the stream pool used (which has default of nullptr if unspecified)
-   * @param[in] workspace_resource an optional resource used by some functions for allocating
-   *            temporary workspaces.
-   * @param[in] allocation_limit the total amount of memory in bytes available to the temporary
-   *            workspace resources.
    */
   device_resources(rmm::cuda_stream_view stream_view                  = rmm::cuda_stream_per_thread,
-                   std::shared_ptr<rmm::cuda_stream_pool> stream_pool = {nullptr},
-                   std::shared_ptr<rmm::mr::device_memory_resource> workspace_resource = {nullptr},
-                   std::optional<std::size_t> allocation_limit = std::nullopt)
+                   std::shared_ptr<rmm::cuda_stream_pool> stream_pool = {nullptr})
     : resources{}
   {
     resources::add_resource_factory(std::make_shared<resource::device_id_resource_factory>());
     resources::add_resource_factory(
       std::make_shared<resource::cuda_stream_resource_factory>(stream_view));
     resources::add_resource_factory(
       std::make_shared<resource::cuda_stream_pool_resource_factory>(stream_pool));
-    if (workspace_resource) {
-      resource::set_workspace_resource(*this, workspace_resource, allocation_limit);
-    }
   }
 
   /** Destroys all held-up resources */
@@ -214,7 +196,7 @@ class device_resources : public resources {
     return resource::get_subcomm(*this, key);
   }
 
-  rmm::mr::device_memory_resource* get_workspace_resource() const
+  rmm::mr::limiting_resource_adaptor* get_workspace_resource() const
   {
     return resource::get_workspace_resource(*this);
   }

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
  */
 
@@ -12,6 +12,7 @@
 #include <rmm/cuda_stream_pool.hpp>
 #include <rmm/mr/cuda_memory_resource.hpp>
 #include <rmm/mr/per_device_resource.hpp>
+#include <rmm/mr/pool_memory_resource.hpp>
 
 #include <algorithm>
 #include <memory>
@@ -114,12 +115,6 @@ struct device_resources_manager {
     std::optional<std::size_t> max_mem_pool_size{std::size_t{}};
     // Limit on workspace memory for the returned device_resources object
     std::optional<std::size_t> workspace_allocation_limit{std::nullopt};
-    // Optional specification of separate workspace memory resources for each
-    // device. The integer in each pair indicates the device for this memory
-    // resource.
-    std::vector<std::pair<std::shared_ptr<rmm::mr::device_memory_resource>, int>> workspace_mrs{};
-
-    auto get_workspace_memory_resource(int device_id) {}
   } params_;
 
   // This struct stores the underlying resources to be shared among
@@ -152,35 +147,18 @@ struct device_resources_manager {
         }()},
         pool_mr_{[&params, this]() {
           auto scoped_device = device_setter{device_id_};
-          auto result =
-            std::shared_ptr<rmm::mr::pool_memory_resource<rmm::mr::cuda_memory_resource>>{nullptr};
+          auto result        = std::optional<rmm::mr::pool_memory_resource>{};
           // If max_mem_pool_size is nullopt or non-zero, create a pool memory
           // 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());
-            if (upstream != nullptr) {
-              result =
-                std::make_shared<rmm::mr::pool_memory_resource<rmm::mr::cuda_memory_resource>>(
-                  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());
-            } else {
-              RAFT_LOG_WARN(
-                "Pool allocation requested, but other memory resource has already been set and "
-                "will not be overwritten");
-            }
+            auto upstream = rmm::mr::get_current_device_resource_ref();
+            result.emplace(
+              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_ref(*result);
           }
           return result;
-        }()},
-        workspace_mr_{[&params, this]() {
-          auto result = std::shared_ptr<rmm::mr::device_memory_resource>{nullptr};
-          auto iter   = std::find_if(std::begin(params.workspace_mrs),
-                                   std::end(params.workspace_mrs),
-                                   [this](auto&& pair) { return pair.second == device_id_; });
-          if (iter != std::end(params.workspace_mrs)) { result = iter->first; }
-          return result;
         }()}
     {
     }
@@ -216,27 +194,14 @@ struct device_resources_manager {
       if (pool_count() != 0) { result = pools_[get_thread_id() % pool_count()]; }
       return result;
     }
-    // Return a (possibly null) shared_ptr to the pool memory resource
-    // created for this device by the manager
-    [[nodiscard]] auto get_pool_memory_resource() const { return pool_mr_; }
-    // Return the RAFT workspace allocation limit that will be used by
-    // `device_resources` returned from this manager
-    [[nodiscard]] auto get_workspace_allocation_limit() const
-    {
-      return workspace_allocation_limit_;
-    }
-    // Return a (possibly null) shared_ptr to the memory resource that will
-    // be used for workspace allocations by `device_resources` returned from
-    // this manager
-    [[nodiscard]] auto get_workspace_memory_resource() { return workspace_mr_; }
+    // Return the pool memory resource created for this device by the manager (if any)
+    [[nodiscard]] auto& get_pool_memory_resource() { return pool_mr_; }
 
    private:
     int device_id_;
     std::unique_ptr<rmm::cuda_stream_pool> streams_;
     std::vector<std::shared_ptr<rmm::cuda_stream_pool>> pools_;
-    std::shared_ptr<rmm::mr::pool_memory_resource<rmm::mr::cuda_memory_resource>> pool_mr_;
-    std::shared_ptr<rmm::mr::device_memory_resource> workspace_mr_;
-    std::optional<std::size_t> workspace_allocation_limit_{std::nullopt};
+    std::optional<rmm::mr::pool_memory_resource> pool_mr_;
   };
 
   // Mutex used to lock access to shared data until after the first
@@ -290,10 +255,7 @@ struct device_resources_manager {
       auto scoped_device = device_setter(device_id);
       // Build the device_resources object for this thread out of shared
       // components
-      thread_resources[device_id].emplace(component_iter->get_stream(),
-                                          component_iter->get_pool(),
-                                          component_iter->get_workspace_memory_resource(),
-                                          component_iter->get_workspace_allocation_limit());
+      thread_resources[device_id].emplace(component_iter->get_stream(), component_iter->get_pool());
     }
 
     return thread_resources[device_id].value();
@@ -373,27 +335,6 @@ struct device_resources_manager {
     }
   }
 
-  // Thread-safe setter for workspace memory resources
-  void set_workspace_memory_resource_(std::shared_ptr<rmm::mr::device_memory_resource> mr,
-                                      int device_id)
-  {
-    auto lock = get_lock();
-    if (params_finalized_) {
-      RAFT_LOG_WARN(
-        "Attempted to set device_resources_manager properties after resources have already been "
-        "retrieved");
-    } else {
-      auto iter = std::find_if(std::begin(params_.workspace_mrs),
-                               std::end(params_.workspace_mrs),
-                               [device_id](auto&& pair) { return pair.second == device_id; });
-      if (iter != std::end(params_.workspace_mrs)) {
-        iter->first = mr;
-      } else {
-        params_.workspace_mrs.emplace_back(mr, device_id);
-      }
-    }
-  }
-
   // Retrieve the instance of this singleton
   static auto& get_manager()
   {
@@ -543,24 +484,5 @@ struct device_resources_manager {
     set_init_mem_pool_size(init_mem);
     set_max_mem_pool_size(max_mem);
   }
-
-  /**
-   * @brief Set the workspace memory resource to be used on a specific device
-   *
-   * RAFT device_resources objects can be built with a separate memory
-   * resource for allocating temporary workspaces. If a (non-nullptr) memory
-   * resource is provided by this setter, it will be used as the
-   * workspace memory resource for all `device_resources` returned for the
-   * indicated device.
-   *
-   * If called after the first call to
-   * `raft::device_resources_manager::get_device_resources`, no change will be made,
-   * and a warning will be emitted.
-   */
-  static void set_workspace_memory_resource(std::shared_ptr<rmm::mr::device_memory_resource> mr,
-                                            int device_id = device_setter::get_current_device())
-  {
-    get_manager().set_workspace_memory_resource_(mr, device_id);
-  }
 };
 }  // namespace raft

cpp/include/raft/core/device_resources_snmg.hpp

@@ -10,7 +10,6 @@
 #include <raft/core/resource/resource_types.hpp>
 
 #include <rmm/cuda_device.hpp>
-#include <rmm/mr/device_memory_resource.hpp>
 #include <rmm/mr/per_device_resource.hpp>
 #include <rmm/mr/pool_memory_resource.hpp>
 
@@ -105,10 +104,9 @@ class device_resources_snmg : public device_resources {
       int device_id = raft::resource::get_device_id(dev_res);
       pool_device_ids_.push_back(device_id);
 
-      per_device_pools_.push_back(
-        std::make_unique<rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource>>(
-          rmm::mr::get_current_device_resource_ref(),
-          rmm::percent_of_free_device_memory(percent_of_free_memory)));
+      per_device_pools_.push_back(std::make_unique<rmm::mr::pool_memory_resource>(
+        rmm::mr::get_current_device_resource_ref(),
+        rmm::percent_of_free_device_memory(percent_of_free_memory)));
       rmm::mr::set_per_device_resource_ref(rmm::cuda_device_id{device_id},
                                            *per_device_pools_.back());
     }
@@ -151,8 +149,7 @@ class device_resources_snmg : public device_resources {
     }
   }
   int main_gpu_id_;
-  std::vector<std::unique_ptr<rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource>>>
-    per_device_pools_;
+  std::vector<std::unique_ptr<rmm::mr::pool_memory_resource>> per_device_pools_;
   std::vector<int> pool_device_ids_;
 };  // class device_resources_snmg