From 3320304dfe0639ae1ff142c6ff853f53eb289e56 Mon Sep 17 00:00:00 2001
From: midagedev <midagedev@gmail.com>
Date: Wed, 23 Sep 2026 22:57:43 +0900
Subject: [PATCH] feat(cuda-core): CudaContext::mem_info, the (free, total)
 device memory query

cuMemGetInfo had no public wrapper: the shim's mem_get_info sits in a
pub(crate) module, so callers reached for the raw binding. Add the
method on CudaContext next to the other device queries; it binds the
context first. The leak test used the raw call and now uses the method.

Signed-off-by: midagedev <midagedev@gmail.com>
---
 cuda-core/src/simt/context.rs               | 14 ++++++++++
 cuda-core/tests/simt_device_buffer_leaks.rs | 29 +++++++++++++++------
 cutile-rs/CHANGELOG.md                      |  5 ++++
 3 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/cuda-core/src/simt/context.rs b/cuda-core/src/simt/context.rs
index 2a76c5a389..87529aae24 100644
--- a/cuda-core/src/simt/context.rs
+++ b/cuda-core/src/simt/context.rs
@@ -548,6 +548,20 @@ impl CudaContext {
         }
     }
 
+    /// Queries the device memory of this context as `(free, total)` bytes.
+    ///
+    /// Binds the context first, then calls `cuMemGetInfo`. `total` is the
+    /// device's memory; `free` is what the driver can hand out at this
+    /// moment, so allocations on other streams and driver-side caches move
+    /// it between two calls.
+    pub fn mem_info(&self) -> Result<(usize, usize), DriverError> {
+        self.bind_to_thread()?;
+        let mut free = 0usize;
+        let mut total = 0usize;
+        unsafe { cuda_bindings::cuMemGetInfo_v2(&mut free, &mut total) }.result()?;
+        Ok((free, total))
+    }
+
     /// Queries dimension, thread-count, and portable shared-memory launch
     /// limits for this device.
     ///
diff --git a/cuda-core/tests/simt_device_buffer_leaks.rs b/cuda-core/tests/simt_device_buffer_leaks.rs
index 2e51c654ea..72e0706b98 100644
--- a/cuda-core/tests/simt_device_buffer_leaks.rs
+++ b/cuda-core/tests/simt_device_buffer_leaks.rs
@@ -67,19 +67,13 @@ fn ctx_strong_count_returns_to_baseline_after_buffer_lifecycle() {
 /// This cannot deterministically force the async-enqueue failure that
 /// triggered the original leak (no public API constructs an invalid
 /// `CudaStream`), but it pins the happy-path accounting with
-/// `cuMemGetInfo` before and after the cycles.
+/// [`CudaContext::mem_info`] before and after the cycles.
 #[test]
 fn vram_returns_to_baseline_after_buffer_cycles() {
     let ctx = CudaContext::new(0).expect("failed to create CUDA context");
     let stream = ctx.new_stream().expect("failed to create CUDA stream");
 
-    fn free_mem() -> usize {
-        let mut free = 0usize;
-        let mut total = 0usize;
-        let rc = unsafe { cuda_bindings::cuMemGetInfo_v2(&mut free, &mut total) };
-        assert_eq!(rc, 0, "cuMemGetInfo failed: {rc}");
-        free
-    }
+    let free_mem = || ctx.mem_info().expect("mem_info failed").0;
 
     // Warm up driver allocator caches so the measured window is stable.
     for _ in 0..4 {
@@ -135,3 +129,22 @@ fn zero_length_construction_succeeds_for_both_constructors() {
         "empty buffers must not leak a ctx strong count"
     );
 }
+
+/// `mem_info` reports the device's memory: `total` matches the device's
+/// total memory attribute and `free` never exceeds it.
+#[test]
+fn mem_info_reports_free_within_the_device_total() {
+    let ctx = CudaContext::new(0).expect("failed to create CUDA context");
+    let (free, total) = ctx.mem_info().expect("mem_info failed");
+
+    let mut device_total = 0usize;
+    let rc = unsafe { cuda_bindings::cuDeviceTotalMem_v2(&mut device_total, ctx.cu_device()) };
+    assert_eq!(rc, 0, "cuDeviceTotalMem failed: {rc}");
+
+    assert_eq!(total, device_total, "total must be the device's memory");
+    assert!(
+        free <= total,
+        "free ({free}) must not exceed total ({total})"
+    );
+    assert!(total > 0, "a device has memory");
+}
diff --git a/cutile-rs/CHANGELOG.md b/cutile-rs/CHANGELOG.md
index 4b7ec6eff9..c463099d58 100644
--- a/cutile-rs/CHANGELOG.md
+++ b/cutile-rs/CHANGELOG.md
@@ -7,6 +7,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/).
 
 ## [Unreleased]
 
+### Added
+
+- `CudaContext::mem_info`, the `(free, total)` device memory query
+  (`cuMemGetInfo`) for the context.
+
 ### Changed
 
 - cuTile Rust now lives in the [NVIDIA/cuda-rust](https://github.com/NVIDIA/cuda-rust)