metrale_storage/
cuda_min.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: Minimal CUDA driver-API FFI for metrale-storage: context and
4//! stream, device and pinned host buffers, async copies, stream sync.
5//!
6//! Owner: metrale-storage.
7//! Invariants:
8//! - A non-zero driver status becomes an `Err` naming the call, except in the
9//!   `Drop` impls, which ignore the status of the free or destroy call.
10
11use anyhow::{Result, bail};
12use std::ffi::c_void;
13
14// 2026-09-25: Re-exported so callers reach the module and event helpers of
15// `cuda_module.rs` through `cuda_min::` as well.
16pub use crate::cuda_module::{CudaEvent, CudaModule, launch_kernel};
17
18unsafe extern "C" {
19    fn cuInit(flags: u32) -> i32;
20    fn cuDeviceGet(device: *mut i32, ordinal: i32) -> i32;
21    fn cuCtxCreate_v2(pctx: *mut u64, flags: u32, dev: i32) -> i32;
22    fn cuCtxDestroy_v2(ctx: u64) -> i32;
23    fn cuMemAlloc_v2(dptr: *mut u64, bytesize: usize) -> i32;
24    fn cuMemFree_v2(dptr: u64) -> i32;
25    fn cuMemAllocHost_v2(pp: *mut *mut c_void, bytesize: usize) -> i32;
26    fn cuMemFreeHost(p: *mut c_void) -> i32;
27    fn cuMemHostGetDevicePointer_v2(pdptr: *mut u64, p: *mut c_void, flags: u32) -> i32;
28    fn cuMemcpyHtoDAsync_v2(dst: u64, src: *const c_void, bytes: usize, stream: u64) -> i32;
29    fn cuMemcpyDtoHAsync_v2(dst: *mut c_void, src: u64, bytes: usize, stream: u64) -> i32;
30    fn cuMemGetInfo_v2(free: *mut usize, total: *mut usize) -> i32;
31    fn cuStreamCreate(phStream: *mut u64, flags: u32) -> i32;
32    fn cuStreamDestroy_v2(stream: u64) -> i32;
33    fn cuStreamSynchronize(stream: u64) -> i32;
34}
35
36/// 2026-09-25: The current context's free and total device memory in bytes
37/// (`cuMemGetInfo`).
38pub fn mem_info() -> Result<(usize, usize)> {
39    let mut free = 0usize;
40    let mut total = 0usize;
41    let s = unsafe { cuMemGetInfo_v2(&mut free, &mut total) };
42    if s != 0 {
43        bail!("cuMemGetInfo_v2 failed: {s}");
44    }
45    Ok((free, total))
46}
47
48pub struct CudaCtx {
49    pub ctx: u64,
50    pub stream: u64,
51}
52
53impl CudaCtx {
54    pub fn new(ordinal: i32) -> Result<Self> {
55        unsafe {
56            let s = cuInit(0);
57            if s != 0 {
58                bail!("cuInit failed: {s}");
59            }
60            let mut dev = 0i32;
61            let s = cuDeviceGet(&mut dev, ordinal);
62            if s != 0 {
63                bail!("cuDeviceGet({ordinal}) failed: {s}");
64            }
65            let mut ctx = 0u64;
66            let s = cuCtxCreate_v2(&mut ctx, 0, dev);
67            if s != 0 {
68                bail!("cuCtxCreate failed: {s}");
69            }
70            let mut stream = 0u64;
71            let s = cuStreamCreate(&mut stream, 0);
72            if s != 0 {
73                cuCtxDestroy_v2(ctx);
74                bail!("cuStreamCreate failed: {s}");
75            }
76            Ok(Self { ctx, stream })
77        }
78    }
79}
80
81impl Drop for CudaCtx {
82    fn drop(&mut self) {
83        unsafe {
84            let _ = cuStreamDestroy_v2(self.stream);
85            let _ = cuCtxDestroy_v2(self.ctx);
86        }
87    }
88}
89
90pub struct DeviceBuffer {
91    pub ptr: u64,
92    pub bytes: usize,
93}
94
95impl DeviceBuffer {
96    pub fn new(bytes: usize) -> Result<Self> {
97        let mut p = 0u64;
98        let s = unsafe { cuMemAlloc_v2(&mut p, bytes) };
99        if s != 0 {
100            bail!("cuMemAlloc_v2({bytes}) failed: {s}");
101        }
102        Ok(Self { ptr: p, bytes })
103    }
104}
105
106impl Drop for DeviceBuffer {
107    fn drop(&mut self) {
108        unsafe {
109            let _ = cuMemFree_v2(self.ptr);
110        }
111    }
112}
113
114pub struct PinnedBuffer {
115    pub ptr: *mut c_void,
116    pub bytes: usize,
117}
118
119// 2026-09-25: SAFETY: `cuMemAllocHost` memory keeps its address for the
120// buffer's lifetime, so moving a `PinnedBuffer` moves only a pointer and a
121// length. `ptr` is a public field: code that shares the memory across threads
122// must order its own accesses.
123unsafe impl Send for PinnedBuffer {}
124unsafe impl Sync for PinnedBuffer {}
125
126impl PinnedBuffer {
127    pub fn new(bytes: usize) -> Result<Self> {
128        let mut p: *mut c_void = std::ptr::null_mut();
129        let s = unsafe { cuMemAllocHost_v2(&mut p, bytes) };
130        if s != 0 {
131            bail!("cuMemAllocHost_v2({bytes}) failed: {s}");
132        }
133        Ok(Self { ptr: p, bytes })
134    }
135
136    /// 2026-09-25: The device pointer for this pinned host allocation
137    /// (`cuMemHostGetDevicePointer`, flags 0). `expert_arena.rs` uses it as the
138    /// arena's device base.
139    pub fn device_ptr(&self) -> Result<u64> {
140        let mut dptr = 0u64;
141        let s = unsafe { cuMemHostGetDevicePointer_v2(&mut dptr, self.ptr, 0) };
142        if s != 0 {
143            bail!("cuMemHostGetDevicePointer_v2 failed: {s}");
144        }
145        Ok(dptr)
146    }
147}
148
149impl Drop for PinnedBuffer {
150    fn drop(&mut self) {
151        unsafe {
152            let _ = cuMemFreeHost(self.ptr);
153        }
154    }
155}
156
157#[allow(clippy::not_unsafe_ptr_arg_deref)]
158pub fn copy_h_to_d_async(dst: u64, src: *const c_void, bytes: usize, stream: u64) -> Result<()> {
159    let s = unsafe { cuMemcpyHtoDAsync_v2(dst, src, bytes, stream) };
160    if s != 0 {
161        bail!("cuMemcpyHtoDAsync_v2 failed: {s}");
162    }
163    Ok(())
164}
165
166#[allow(clippy::not_unsafe_ptr_arg_deref)]
167pub fn copy_d_to_h_async(dst: *mut c_void, src: u64, bytes: usize, stream: u64) -> Result<()> {
168    let s = unsafe { cuMemcpyDtoHAsync_v2(dst, src, bytes, stream) };
169    if s != 0 {
170        bail!("cuMemcpyDtoHAsync_v2 failed: {s}");
171    }
172    Ok(())
173}
174
175pub fn stream_sync(stream: u64) -> Result<()> {
176    let s = unsafe { cuStreamSynchronize(stream) };
177    if s != 0 {
178        bail!("cuStreamSynchronize failed: {s}");
179    }
180    Ok(())
181}
metrale_storage/cuda_min.rs

metrale_storage/
cuda_min.rs
