Use the new reference-based RMM API to set RMM pools in device_resources_snmg (#2972)

This PR updates the `device_resources_snmg` utility so that it uses the new reference-based RMM API to set RMM pools.

Authors:
  - Victor Lafargue (https://github.com/viclafargue)

Approvers:
  - Jinsol Park (https://github.com/jinsolp)
  - Artem M. Chirkin (https://github.com/achirkin)

URL: #2972

Author: viclafargue

cpp/include/raft/core/device_resources_snmg.hpp

@@ -80,12 +80,8 @@ class device_resources_snmg : public device_resources {
   device_resources_snmg& operator=(device_resources_snmg&&) = delete;
   ~device_resources_snmg()
   {
-    // Restore original device memory resources
-    if (!device_original_mrs_.empty()) {
-      for (const auto& [device_id, original_mr] : device_original_mrs_) {
-        rmm::cuda_device_id id(device_id);
-        rmm::mr::set_per_device_resource(id, original_mr);
-      }
+    for (int device_id : pool_device_ids_) {
+      rmm::mr::reset_per_device_resource_ref(rmm::cuda_device_id{device_id});
     }
   }
 
@@ -94,33 +90,27 @@ class device_resources_snmg : public device_resources {
    */
   void set_memory_pool(int percent_of_free_memory)
   {
-    // Protect against repeated calls - restore original resources and clear pools
     if (!per_device_pools_.empty()) {
-      for (const auto& [device_id, original_mr] : device_original_mrs_) {
-        rmm::cuda_device_id id(device_id);
-        rmm::mr::set_per_device_resource(id, original_mr);
+      for (int device_id : pool_device_ids_) {
+        rmm::mr::reset_per_device_resource_ref(rmm::cuda_device_id{device_id});
       }
       per_device_pools_.clear();
-      device_original_mrs_.clear();
+      pool_device_ids_.clear();
     }
 
     int world_size = raft::resource::get_num_ranks(*this);
     for (int rank = 0; rank < world_size; rank++) {
       const raft::resources& dev_res = raft::resource::set_current_device_to_rank(*this, rank);
 
-      // Get the actual device ID for this rank
       int device_id = raft::resource::get_device_id(dev_res);
+      pool_device_ids_.push_back(device_id);
 
-      // Store the original memory resource before replacing it
-      auto old_mr = rmm::mr::get_current_device_resource();
-      device_original_mrs_.push_back({device_id, old_mr});
-
-      // create a pool memory resource for each device
       per_device_pools_.push_back(
         std::make_unique<rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource>>(
-          old_mr, rmm::percent_of_free_device_memory(percent_of_free_memory)));
-      rmm::cuda_device_id id(device_id);
-      rmm::mr::set_per_device_resource(id, per_device_pools_.back().get());
+          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());
     }
     RAFT_CUDA_TRY(cudaSetDevice(main_gpu_id_));
   }
@@ -163,7 +153,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::pair<int, rmm::mr::device_memory_resource*>> device_original_mrs_;
+  std::vector<int> pool_device_ids_;
 };  // class device_resources_snmg
 
 }  // namespace raft