metrale_storage/
cuda_module.rs

1// SPDX-License-Identifier: MIT OR Apache-2.0
2
3//! 2026-09-25: CUDA driver-API module loading, kernel launch and events.
4//!
5//! Owner: metrale-storage.
6//! Invariants:
7//! - A non-zero driver status becomes an `Err` naming the call, except in the
8//!   `Drop` impls, which ignore the status of the unload or destroy call.
9
10use anyhow::{Context, Result, bail};
11use std::ffi::c_void;
12
13unsafe extern "C" {
14    fn cuModuleLoadData(module: *mut u64, image: *const c_void) -> i32;
15    fn cuModuleUnload(module: u64) -> i32;
16    fn cuModuleGetFunction(func: *mut u64, module: u64, name: *const std::ffi::c_char) -> i32;
17    fn cuLaunchKernel(
18        func: u64,
19        grid_x: u32,
20        grid_y: u32,
21        grid_z: u32,
22        block_x: u32,
23        block_y: u32,
24        block_z: u32,
25        shared_bytes: u32,
26        stream: u64,
27        kernel_params: *mut *mut c_void,
28        extra: *mut *mut c_void,
29    ) -> i32;
30    fn cuEventCreate(event: *mut u64, flags: u32) -> i32;
31    fn cuEventRecord(event: u64, stream: u64) -> i32;
32    fn cuEventSynchronize(event: u64) -> i32;
33    fn cuEventDestroy_v2(event: u64) -> i32;
34}
35
36/// 2026-09-25: A loaded `CUmodule`; `Drop` unloads it. `function` looks the
37/// name up on every call.
38pub struct CudaModule {
39    handle: u64,
40}
41
42impl CudaModule {
43    pub fn from_ptx(ptx: &str) -> Result<Self> {
44        let cstr = std::ffi::CString::new(ptx).context("PTX contains NUL")?;
45        let mut h = 0u64;
46        let s = unsafe { cuModuleLoadData(&mut h, cstr.as_ptr() as *const c_void) };
47        if s != 0 {
48            bail!("cuModuleLoadData failed: {s}");
49        }
50        Ok(Self { handle: h })
51    }
52
53    pub fn function(&self, name: &str) -> Result<u64> {
54        let cstr = std::ffi::CString::new(name).context("function name has NUL")?;
55        let mut f = 0u64;
56        let s = unsafe { cuModuleGetFunction(&mut f, self.handle, cstr.as_ptr()) };
57        if s != 0 {
58            bail!("cuModuleGetFunction({name}) failed: {s}");
59        }
60        Ok(f)
61    }
62}
63
64impl Drop for CudaModule {
65    fn drop(&mut self) {
66        unsafe {
67            let _ = cuModuleUnload(self.handle);
68        }
69    }
70}
71
72pub struct CudaEvent {
73    pub handle: u64,
74}
75
76impl CudaEvent {
77    pub fn new() -> Result<Self> {
78        let mut h = 0u64;
79        // 2026-09-25: No method reads elapsed time, so events skip timing.
80        const CU_EVENT_DISABLE_TIMING: u32 = 0x2;
81        let s = unsafe { cuEventCreate(&mut h, CU_EVENT_DISABLE_TIMING) };
82        if s != 0 {
83            bail!("cuEventCreate failed: {s}");
84        }
85        Ok(Self { handle: h })
86    }
87    pub fn record(&self, stream: u64) -> Result<()> {
88        let s = unsafe { cuEventRecord(self.handle, stream) };
89        if s != 0 {
90            bail!("cuEventRecord failed: {s}");
91        }
92        Ok(())
93    }
94    pub fn sync(&self) -> Result<()> {
95        let s = unsafe { cuEventSynchronize(self.handle) };
96        if s != 0 {
97            bail!("cuEventSynchronize failed: {s}");
98        }
99        Ok(())
100    }
101}
102
103impl Drop for CudaEvent {
104    fn drop(&mut self) {
105        unsafe {
106            let _ = cuEventDestroy_v2(self.handle);
107        }
108    }
109}
110
111/// 2026-09-25: `cuLaunchKernel` with `extra = NULL`. Each `params` entry must
112/// point at a live argument of the type the kernel expects, until this
113/// returns.
114pub fn launch_kernel(
115    func: u64,
116    grid: (u32, u32, u32),
117    block: (u32, u32, u32),
118    shared_bytes: u32,
119    stream: u64,
120    params: &mut [*mut c_void],
121) -> Result<()> {
122    let s = unsafe {
123        cuLaunchKernel(
124            func,
125            grid.0,
126            grid.1,
127            grid.2,
128            block.0,
129            block.1,
130            block.2,
131            shared_bytes,
132            stream,
133            params.as_mut_ptr(),
134            std::ptr::null_mut(),
135        )
136    };
137    if s != 0 {
138        bail!("cuLaunchKernel failed: status {s} grid={grid:?} block={block:?}");
139    }
140    Ok(())
141}
metrale_storage/cuda_module.rs

metrale_storage/
cuda_module.rs
