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] 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;
 }