aria_kernel/
cuda_rt.rs

1//! Runtime CUDA/cuBLAS via libloading (no compile-time toolkit required).
2
3use crate::EngineError;
4use libloading::{Library, Symbol};
5use std::collections::HashMap;
6use std::ffi::{c_char, c_int, c_void, CStr};
7#[cfg(unix)]
8use std::path::Path;
9use std::sync::Mutex;
10
11const CUDA_SUCCESS: c_int = 0;
12const CUBLAS_SUCCESS: c_int = 0;
13const CUDA_MEMCPY_H2D: c_int = 1;
14const CUDA_MEMCPY_D2H: c_int = 2;
15const CUBLAS_OP_N: c_int = 0;
16const CUBLAS_OP_T: c_int = 1;
17
18type CudaGetDeviceCount = unsafe extern "C" fn(*mut c_int) -> c_int;
19type CudaMalloc = unsafe extern "C" fn(*mut *mut c_void, usize) -> c_int;
20type CudaFree = unsafe extern "C" fn(*mut c_void) -> c_int;
21type CudaMemcpy = unsafe extern "C" fn(*mut c_void, *const c_void, usize, c_int) -> c_int;
22type CudaGetErrorString = unsafe extern "C" fn(c_int) -> *const c_char;
23type CublasCreate = unsafe extern "C" fn(*mut *mut c_void) -> c_int;
24type CublasDestroy = unsafe extern "C" fn(*mut c_void) -> c_int;
25type CublasSgemm = unsafe extern "C" fn(
26    *mut c_void,
27    c_int,
28    c_int,
29    c_int,
30    c_int,
31    c_int,
32    *const f32,
33    *const f32,
34    c_int,
35    *const f32,
36    c_int,
37    *const f32,
38    *mut f32,
39    c_int,
40) -> c_int;
41
42fn open_first(candidates: &[&str]) -> Result<Library, EngineError> {
43    let mut last = String::from("no candidate loaded");
44    for c in candidates {
45        match unsafe { Library::new(c) } {
46            Ok(lib) => return Ok(lib),
47            Err(e) => last = format!("{c}: {e}"),
48        }
49        #[cfg(unix)]
50        if let Ok(lib) = unsafe { Library::new(Path::new(c)) } {
51            return Ok(lib);
52        }
53    }
54    Err(EngineError::Unsupported(format!("libloading: {last}")))
55}
56
57fn nvidia_smi_name() -> Option<String> {
58    let out = std::process::Command::new("nvidia-smi")
59        .args(["--query-gpu=name", "--format=csv,noheader"])
60        .output()
61        .ok()?;
62    if !out.status.success() {
63        return None;
64    }
65    let s = String::from_utf8_lossy(&out.stdout);
66    let line = s.lines().next()?.trim();
67    if line.is_empty() {
68        None
69    } else {
70        Some(line.to_string())
71    }
72}
73
74/// Probe CUDA runtime + cuBLAS. Does not allocate a persistent handle.
75pub fn device_info() -> Result<String, EngineError> {
76    let cudart = open_first(&[
77        "libcudart.so.12",
78        "libcudart.so",
79        "nvcudart.dll",
80        "libcudart.dylib",
81    ])?;
82    let cublas = open_first(&[
83        "libcublas.so.12",
84        "libcublas.so",
85        "cublas64_12.dll",
86        "libcublas.dylib",
87    ])?;
88    let get_count: Symbol<CudaGetDeviceCount> = unsafe {
89        cudart
90            .get(b"cudaGetDeviceCount")
91            .map_err(|e| EngineError::Unsupported(e.to_string()))?
92    };
93    let mut n: c_int = 0;
94    let st = unsafe { get_count(&mut n) };
95    if st != CUDA_SUCCESS || n <= 0 {
96        return Err(EngineError::Unsupported(format!(
97            "cudaGetDeviceCount status={st} count={n}"
98        )));
99    }
100    // Touch a cublas symbol so missing SONAME fails here, not at first GEMM.
101    let _create: Symbol<CublasCreate> = unsafe {
102        cublas
103            .get(b"cublasCreate_v2")
104            .map_err(|e| EngineError::Unsupported(e.to_string()))?
105    };
106    let name = nvidia_smi_name().unwrap_or_else(|| format!("{n} device(s)"));
107    Ok(name)
108}
109
110struct Fns {
111    malloc: CudaMalloc,
112    free: CudaFree,
113    memcpy: CudaMemcpy,
114    errstr: CudaGetErrorString,
115    sgemm: CublasSgemm,
116}
117
118fn cuda_err(errstr: CudaGetErrorString, st: c_int, what: &str) -> EngineError {
119    let msg = unsafe {
120        let p = errstr(st);
121        if p.is_null() {
122            format!("{what} cuda status={st}")
123        } else {
124            format!("{what}: {}", CStr::from_ptr(p).to_string_lossy())
125        }
126    };
127    EngineError::Unsupported(msg)
128}
129
130/// Persistent cuBLAS handle + device copies of host weight buffers.
131pub struct CudaContext {
132    _cudart: Library,
133    _cublas: Library,
134    handle: *mut c_void,
135    fns: Fns,
136    weights: Mutex<HashMap<usize, (*mut f32, usize)>>,
137}
138
139unsafe impl Send for CudaContext {}
140unsafe impl Sync for CudaContext {}
141
142impl CudaContext {
143    pub fn new() -> Result<Self, EngineError> {
144        let cudart = open_first(&[
145            "libcudart.so.12",
146            "libcudart.so",
147            "nvcudart.dll",
148            "libcudart.dylib",
149        ])?;
150        let cublas = open_first(&[
151            "libcublas.so.12",
152            "libcublas.so",
153            "cublas64_12.dll",
154            "libcublas.dylib",
155        ])?;
156        let malloc: Symbol<CudaMalloc> = unsafe {
157            cudart
158                .get(b"cudaMalloc")
159                .map_err(|e| EngineError::Unsupported(e.to_string()))?
160        };
161        let free: Symbol<CudaFree> = unsafe {
162            cudart
163                .get(b"cudaFree")
164                .map_err(|e| EngineError::Unsupported(e.to_string()))?
165        };
166        let memcpy: Symbol<CudaMemcpy> = unsafe {
167            cudart
168                .get(b"cudaMemcpy")
169                .map_err(|e| EngineError::Unsupported(e.to_string()))?
170        };
171        let errstr: Symbol<CudaGetErrorString> = unsafe {
172            cudart
173                .get(b"cudaGetErrorString")
174                .map_err(|e| EngineError::Unsupported(e.to_string()))?
175        };
176        let create: Symbol<CublasCreate> = unsafe {
177            cublas
178                .get(b"cublasCreate_v2")
179                .map_err(|e| EngineError::Unsupported(e.to_string()))?
180        };
181        let sgemm: Symbol<CublasSgemm> = unsafe {
182            cublas
183                .get(b"cublasSgemm_v2")
184                .map_err(|e| EngineError::Unsupported(e.to_string()))?
185        };
186        let fns = Fns {
187            malloc: *malloc,
188            free: *free,
189            memcpy: *memcpy,
190            errstr: *errstr,
191            sgemm: *sgemm,
192        };
193        let mut handle: *mut c_void = std::ptr::null_mut();
194        let st = unsafe { create(&mut handle) };
195        if st != CUBLAS_SUCCESS || handle.is_null() {
196            return Err(EngineError::Unsupported(format!(
197                "cublasCreate_v2 status={st}"
198            )));
199        }
200        Ok(Self {
201            _cudart: cudart,
202            _cublas: cublas,
203            handle,
204            fns,
205            weights: Mutex::new(HashMap::new()),
206        })
207    }
208
209    pub fn upload(&self, host: &[f32]) -> Result<(), EngineError> {
210        if host.is_empty() {
211            return Ok(());
212        }
213        let key = host.as_ptr() as usize;
214        {
215            let map = self
216                .weights
217                .lock()
218                .map_err(|e| EngineError::Unsupported(e.to_string()))?;
219            if let Some((_, n)) = map.get(&key) {
220                if *n == host.len() {
221                    return Ok(());
222                }
223            }
224        }
225        let bytes = host.len() * 4;
226        let mut dev: *mut c_void = std::ptr::null_mut();
227        let st = unsafe { (self.fns.malloc)(&mut dev, bytes) };
228        if st != CUDA_SUCCESS {
229            return Err(cuda_err(self.fns.errstr, st, "cudaMalloc"));
230        }
231        let st = unsafe { (self.fns.memcpy)(dev, host.as_ptr().cast(), bytes, CUDA_MEMCPY_H2D) };
232        if st != CUDA_SUCCESS {
233            unsafe {
234                (self.fns.free)(dev);
235            }
236            return Err(cuda_err(self.fns.errstr, st, "cudaMemcpy H2D"));
237        }
238        let mut map = self
239            .weights
240            .lock()
241            .map_err(|e| EngineError::Unsupported(e.to_string()))?;
242        if let Some((old, _)) = map.insert(key, (dev.cast(), host.len())) {
243            unsafe {
244                (self.fns.free)(old.cast());
245            }
246        }
247        Ok(())
248    }
249
250    /// y = W @ x for row-major W `[out_f, in_f]`, x `[batch, in_f]`.
251    pub fn linear(
252        &self,
253        x: &[f32],
254        w: &[f32],
255        out_f: usize,
256        in_f: usize,
257    ) -> Result<Vec<f32>, EngineError> {
258        if in_f == 0 || !x.len().is_multiple_of(in_f) {
259            return Err(EngineError::ShapeMismatch("cuda linear x".into()));
260        }
261        if w.len() != out_f * in_f {
262            return Err(EngineError::ShapeMismatch("cuda linear w".into()));
263        }
264        self.upload(w)?;
265        let batch = x.len() / in_f;
266        let key = w.as_ptr() as usize;
267        let map = self
268            .weights
269            .lock()
270            .map_err(|e| EngineError::Unsupported(e.to_string()))?;
271        let (w_dev, n) = map
272            .get(&key)
273            .copied()
274            .ok_or_else(|| EngineError::Unsupported("cuda weight not uploaded".into()))?;
275        if n != w.len() {
276            return Err(EngineError::Unsupported("cuda weight size drift".into()));
277        }
278        drop(map);
279
280        let x_bytes = x.len() * 4;
281        let y_bytes = batch * out_f * 4;
282        let mut x_dev: *mut c_void = std::ptr::null_mut();
283        let mut y_dev: *mut c_void = std::ptr::null_mut();
284        let st = unsafe { (self.fns.malloc)(&mut x_dev, x_bytes) };
285        if st != CUDA_SUCCESS {
286            return Err(cuda_err(self.fns.errstr, st, "cudaMalloc x"));
287        }
288        let st = unsafe { (self.fns.malloc)(&mut y_dev, y_bytes) };
289        if st != CUDA_SUCCESS {
290            unsafe {
291                (self.fns.free)(x_dev);
292            }
293            return Err(cuda_err(self.fns.errstr, st, "cudaMalloc y"));
294        }
295        let st = unsafe { (self.fns.memcpy)(x_dev, x.as_ptr().cast(), x_bytes, CUDA_MEMCPY_H2D) };
296        if st != CUDA_SUCCESS {
297            unsafe {
298                (self.fns.free)(x_dev);
299                (self.fns.free)(y_dev);
300            }
301            return Err(cuda_err(self.fns.errstr, st, "cudaMemcpy x"));
302        }
303
304        let alpha = 1.0f32;
305        let beta = 0.0f32;
306        // See plan: W row-major [out,in] == col-major [in,out]; y = W_rm @ x => SGEMM(T, N).
307        let m = out_f as c_int;
308        let n_b = batch as c_int;
309        let k = in_f as c_int;
310        let st = unsafe {
311            (self.fns.sgemm)(
312                self.handle,
313                CUBLAS_OP_T,
314                CUBLAS_OP_N,
315                m,
316                n_b,
317                k,
318                &alpha,
319                w_dev,
320                k,
321                x_dev.cast(),
322                k,
323                &beta,
324                y_dev.cast(),
325                m,
326            )
327        };
328        if st != CUBLAS_SUCCESS {
329            unsafe {
330                (self.fns.free)(x_dev);
331                (self.fns.free)(y_dev);
332            }
333            return Err(EngineError::Unsupported(format!(
334                "cublasSgemm_v2 status={st}"
335            )));
336        }
337        let mut y = vec![0.0f32; batch * out_f];
338        let st = unsafe { (self.fns.memcpy)(y.as_mut_ptr().cast(), y_dev, y_bytes, CUDA_MEMCPY_D2H) };
339        unsafe {
340            (self.fns.free)(x_dev);
341            (self.fns.free)(y_dev);
342        }
343        if st != CUDA_SUCCESS {
344            return Err(cuda_err(self.fns.errstr, st, "cudaMemcpy y"));
345        }
346        let _ = CUBLAS_OP_N; // keep const used in docs
347        Ok(y)
348    }
349}
350
351impl Drop for CudaContext {
352    fn drop(&mut self) {
353        if let Ok(mut map) = self.weights.lock() {
354            for (_, (ptr, _)) in map.drain() {
355                unsafe {
356                    (self.fns.free)(ptr.cast());
357                }
358            }
359        }
360        if let Ok(destroy) = unsafe { self._cublas.get::<CublasDestroy>(b"cublasDestroy_v2") } {
361            if !self.handle.is_null() {
362                unsafe {
363                    destroy(self.handle);
364                }
365            }
366        }
367    }
368}
369
370#[cfg(test)]
371mod tests {
372    use super::*;
373    use crate::linear;
374
375    #[test]
376    fn cuda_linear_matches_cpu_if_available() {
377        let Ok(ctx) = CudaContext::new() else {
378            return;
379        };
380        let out_f = 4usize;
381        let in_f = 3usize;
382        let w: Vec<f32> = (0..out_f * in_f).map(|i| i as f32 * 0.1 - 0.2).collect();
383        let x: Vec<f32> = (0..in_f).map(|i| i as f32 * 0.25).collect();
384        let cpu = linear(&x, &w, out_f, in_f).unwrap();
385        let gpu = ctx.linear(&x, &w, out_f, in_f).unwrap();
386        for (a, b) in cpu.iter().zip(gpu.iter()) {
387            assert!((a - b).abs() < 1e-3, "{a} vs {b}");
388        }
389        let xb: Vec<f32> = (0..in_f * 3).map(|i| i as f32 * 0.1).collect();
390        let cpu_b = linear(&xb, &w, out_f, in_f).unwrap();
391        let gpu_b = ctx.linear(&xb, &w, out_f, in_f).unwrap();
392        for (a, b) in cpu_b.iter().zip(gpu_b.iter()) {
393            assert!((a - b).abs() < 1e-3, "batch {a} vs {b}");
394        }
395    }
396}
aria_kernel/cuda_rt.rs

aria_kernel/
cuda_rt.rs
