spark_storage/
cuda_min.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//
3// Minimal CUDA driver FFI for the Phase-0 probe. Intentionally separate from
4// spark-runtime's `cuda_backend.rs` so the probe binary doesn't pull in the
5// full Atlas runtime / kernel registry. Only the symbols the probe needs.
6
7use anyhow::{Result, bail};
8use std::ffi::c_void;
9
10// Re-export the helpers that used to live here, keeping
11// `cuda_min::{CudaModule, CudaEvent, launch_kernel}` paths working after the
12// split into `cuda_module.rs`.
13pub use crate::cuda_module::{CudaEvent, CudaModule, launch_kernel};
14
15unsafe extern "C" {
16    fn cuInit(flags: u32) -> i32;
17    fn cuDeviceGet(device: *mut i32, ordinal: i32) -> i32;
18    fn cuCtxCreate_v2(pctx: *mut u64, flags: u32, dev: i32) -> i32;
19    fn cuCtxDestroy_v2(ctx: u64) -> i32;
20    fn cuMemAlloc_v2(dptr: *mut u64, bytesize: usize) -> i32;
21    fn cuMemFree_v2(dptr: u64) -> i32;
22    fn cuMemAllocHost_v2(pp: *mut *mut c_void, bytesize: usize) -> i32;
23    fn cuMemFreeHost(p: *mut c_void) -> i32;
24    fn cuMemHostGetDevicePointer_v2(pdptr: *mut u64, p: *mut c_void, flags: u32) -> i32;
25    fn cuMemcpyHtoDAsync_v2(dst: u64, src: *const c_void, bytes: usize, stream: u64) -> i32;
26    fn cuMemcpyDtoHAsync_v2(dst: *mut c_void, src: u64, bytes: usize, stream: u64) -> i32;
27    fn cuMemGetInfo_v2(free: *mut usize, total: *mut usize) -> i32;
28    fn cuStreamCreate(phStream: *mut u64, flags: u32) -> i32;
29    fn cuStreamDestroy_v2(stream: u64) -> i32;
30    fn cuStreamSynchronize(stream: u64) -> i32;
31}
32
33/// Query the current context's free/total HBM in bytes. Used by HSS install
34/// preflight to fail fast with an actionable error before a multi-GB
35/// `cuMemAlloc` blows up cryptically. Phase-7 follow-up to PR #47.
36pub fn mem_info() -> Result<(usize, usize)> {
37    let mut free = 0usize;
38    let mut total = 0usize;
39    let s = unsafe { cuMemGetInfo_v2(&mut free, &mut total) };
40    if s != 0 {
41        bail!("cuMemGetInfo_v2 failed: {s}");
42    }
43    Ok((free, total))
44}
45
46pub struct CudaCtx {
47    pub ctx: u64,
48    pub stream: u64,
49}
50
51impl CudaCtx {
52    pub fn new(ordinal: i32) -> Result<Self> {
53        unsafe {
54            let s = cuInit(0);
55            if s != 0 {
56                bail!("cuInit failed: {s}");
57            }
58            let mut dev = 0i32;
59            let s = cuDeviceGet(&mut dev, ordinal);
60            if s != 0 {
61                bail!("cuDeviceGet({ordinal}) failed: {s}");
62            }
63            let mut ctx = 0u64;
64            let s = cuCtxCreate_v2(&mut ctx, 0, dev);
65            if s != 0 {
66                bail!("cuCtxCreate failed: {s}");
67            }
68            let mut stream = 0u64;
69            let s = cuStreamCreate(&mut stream, 0);
70            if s != 0 {
71                cuCtxDestroy_v2(ctx);
72                bail!("cuStreamCreate failed: {s}");
73            }
74            Ok(Self { ctx, stream })
75        }
76    }
77}
78
79impl Drop for CudaCtx {
80    fn drop(&mut self) {
81        unsafe {
82            let _ = cuStreamDestroy_v2(self.stream);
83            let _ = cuCtxDestroy_v2(self.ctx);
84        }
85    }
86}
87
88pub struct DeviceBuffer {
89    pub ptr: u64,
90    pub bytes: usize,
91}
92
93impl DeviceBuffer {
94    pub fn new(bytes: usize) -> Result<Self> {
95        let mut p = 0u64;
96        let s = unsafe { cuMemAlloc_v2(&mut p, bytes) };
97        if s != 0 {
98            bail!("cuMemAlloc_v2({bytes}) failed: {s}");
99        }
100        Ok(Self { ptr: p, bytes })
101    }
102}
103
104impl Drop for DeviceBuffer {
105    fn drop(&mut self) {
106        unsafe {
107            let _ = cuMemFree_v2(self.ptr);
108        }
109    }
110}
111
112pub struct PinnedBuffer {
113    pub ptr: *mut c_void,
114    pub bytes: usize,
115}
116
117// SAFETY: `cuMemAllocHost` returns a process-pinned allocation whose
118// virtual address is stable for the buffer's entire lifetime — moving the
119// `PinnedBuffer` between threads only transfers a pointer + length + the
120// CUcontext handle used by Drop, none of which alias mutable state. The
121// inner pointer never escapes through `&self` accessors; concurrent users
122// of the underlying memory must coordinate externally (Atlas does this
123// via the io_uring submission queue, which is single-threaded per rank).
124unsafe impl Send for PinnedBuffer {}
125unsafe impl Sync for PinnedBuffer {}
126
127impl PinnedBuffer {
128    pub fn new(bytes: usize) -> Result<Self> {
129        let mut p: *mut c_void = std::ptr::null_mut();
130        let s = unsafe { cuMemAllocHost_v2(&mut p, bytes) };
131        if s != 0 {
132            bail!("cuMemAllocHost_v2({bytes}) failed: {s}");
133        }
134        Ok(Self { ptr: p, bytes })
135    }
136
137    /// Device pointer the GPU uses to address this pinned host allocation.
138    ///
139    /// On GB10's unified LPDDR (unified addressing) this equals the host
140    /// pointer numerically — the property the UMA zero-copy expert arena relies
141    /// on. `cuMemAllocHost` memory is portable + page-locked + device-accessible,
142    /// so no `Mapped` flag is required.
143    pub fn device_ptr(&self) -> Result<u64> {
144        let mut dptr = 0u64;
145        let s = unsafe { cuMemHostGetDevicePointer_v2(&mut dptr, self.ptr, 0) };
146        if s != 0 {
147            bail!("cuMemHostGetDevicePointer_v2 failed: {s}");
148        }
149        Ok(dptr)
150    }
151}
152
153impl Drop for PinnedBuffer {
154    fn drop(&mut self) {
155        unsafe {
156            let _ = cuMemFreeHost(self.ptr);
157        }
158    }
159}
160
161#[allow(clippy::not_unsafe_ptr_arg_deref)]
162pub fn copy_h_to_d_async(dst: u64, src: *const c_void, bytes: usize, stream: u64) -> Result<()> {
163    let s = unsafe { cuMemcpyHtoDAsync_v2(dst, src, bytes, stream) };
164    if s != 0 {
165        bail!("cuMemcpyHtoDAsync_v2 failed: {s}");
166    }
167    Ok(())
168}
169
170#[allow(clippy::not_unsafe_ptr_arg_deref)]
171pub fn copy_d_to_h_async(dst: *mut c_void, src: u64, bytes: usize, stream: u64) -> Result<()> {
172    let s = unsafe { cuMemcpyDtoHAsync_v2(dst, src, bytes, stream) };
173    if s != 0 {
174        bail!("cuMemcpyDtoHAsync_v2 failed: {s}");
175    }
176    Ok(())
177}
178
179pub fn stream_sync(stream: u64) -> Result<()> {
180    let s = unsafe { cuStreamSynchronize(stream) };
181    if s != 0 {
182        bail!("cuStreamSynchronize failed: {s}");
183    }
184    Ok(())
185}
spark_storage/cuda_min.rs

spark_storage/
cuda_min.rs
