From 7d660431b9b41f089fecaa80613f4ecb6d51ad3f Mon Sep 17 00:00:00 2001
From: Bradley Dice <bdice@bradleydice.com>
Date: Sun, 17 May 2026 11:22:51 -0500
Subject: [PATCH 1/2] Restore resource_cast memory limit check

---
 cpp/src/utilities/cuda_helpers.cuh | 19 ++++++++++++++++++-
 1 file changed, 18 insertions(+), 1 deletion(-)

diff --git a/cpp/src/utilities/cuda_helpers.cuh b/cpp/src/utilities/cuda_helpers.cuh
index eccf8e1538..379e6c162b 100644
--- a/cpp/src/utilities/cuda_helpers.cuh
+++ b/cpp/src/utilities/cuda_helpers.cuh
@@ -9,13 +9,18 @@
 
 #include <utilities/macros.cuh>
 
+#include <cuda/memory_resource>
+
 #include <thrust/host_vector.h>
 #include <thrust/tuple.h>
+#include <algorithm>
 #include <mutex>
 #include <raft/core/device_span.hpp>
 #include <raft/util/cuda_utils.cuh>
 #include <raft/util/cudart_utils.hpp>
 #include <rmm/device_uvector.hpp>
+#include <rmm/mr/limiting_resource_adaptor.hpp>
+#include <rmm/mr/per_device_resource.hpp>
 #include <shared_mutex>
 #include <unordered_map>
 
@@ -242,7 +247,19 @@ inline size_t get_device_memory_size()
 {
   size_t free_mem, total_mem;
   RAFT_CUDA_TRY(cudaMemGetInfo(&free_mem, &total_mem));
-  // TODO (bdice): Restore limiting adaptor check after updating CCCL to support resource_cast
+
+  auto res              = rmm::mr::get_current_device_resource_ref();
+  auto limiting_adaptor = cuda::mr::resource_cast<rmm::mr::limiting_resource_adaptor>(&res);
+  if (limiting_adaptor) {
+    printf("limiting_adaptor->get_allocation_limit(): %fMiB\n",
+           limiting_adaptor->get_allocation_limit() / (double)1e6);
+    printf("used_mem: %fMiB\n", limiting_adaptor->get_allocated_bytes() / (double)1e6);
+    printf("free_mem: %fMiB\n",
+           (limiting_adaptor->get_allocation_limit() - limiting_adaptor->get_allocated_bytes()) /
+             (double)1e6);
+    return std::min(total_mem, limiting_adaptor->get_allocation_limit());
+  }
+
   return total_mem;
 }
 

From aa2d6656dce4f5635294dacf1b0fe1c04a973aee Mon Sep 17 00:00:00 2001
From: Bradley Dice <bdice@bradleydice.com>
Date: Thu, 28 May 2026 10:05:14 -0500
Subject: [PATCH 2/2] Remove printf

---
 cpp/src/utilities/cuda_helpers.cuh | 10 +---------
 1 file changed, 1 insertion(+), 9 deletions(-)

diff --git a/cpp/src/utilities/cuda_helpers.cuh b/cpp/src/utilities/cuda_helpers.cuh
index 379e6c162b..778dcd5b60 100644
--- a/cpp/src/utilities/cuda_helpers.cuh
+++ b/cpp/src/utilities/cuda_helpers.cuh
@@ -250,15 +250,7 @@ inline size_t get_device_memory_size()
 
   auto res              = rmm::mr::get_current_device_resource_ref();
   auto limiting_adaptor = cuda::mr::resource_cast<rmm::mr::limiting_resource_adaptor>(&res);
-  if (limiting_adaptor) {
-    printf("limiting_adaptor->get_allocation_limit(): %fMiB\n",
-           limiting_adaptor->get_allocation_limit() / (double)1e6);
-    printf("used_mem: %fMiB\n", limiting_adaptor->get_allocated_bytes() / (double)1e6);
-    printf("free_mem: %fMiB\n",
-           (limiting_adaptor->get_allocation_limit() - limiting_adaptor->get_allocated_bytes()) /
-             (double)1e6);
-    return std::min(total_mem, limiting_adaptor->get_allocation_limit());
-  }
+  if (limiting_adaptor) { return std::min(total_mem, limiting_adaptor->get_allocation_limit()); }
 
   return total_mem;
 }