spark_storage/
cuda_module.rs

1// SPDX-License-Identifier: AGPL-3.0-only
2//
3// CUDA module loading + event primitives. Split out of `cuda_min.rs` to
4// keep that file focused on the core context/buffer/copy primitives.
5
6use anyhow::{Context, Result, bail};
7use std::ffi::c_void;
8
9unsafe extern "C" {
10    fn cuModuleLoadData(module: *mut u64, image: *const c_void) -> i32;
11    fn cuModuleUnload(module: u64) -> i32;
12    fn cuModuleGetFunction(func: *mut u64, module: u64, name: *const std::ffi::c_char) -> i32;
13    fn cuLaunchKernel(
14        func: u64,
15        grid_x: u32,
16        grid_y: u32,
17        grid_z: u32,
18        block_x: u32,
19        block_y: u32,
20        block_z: u32,
21        shared_bytes: u32,
22        stream: u64,
23        kernel_params: *mut *mut c_void,
24        extra: *mut *mut c_void,
25    ) -> i32;
26    fn cuEventCreate(event: *mut u64, flags: u32) -> i32;
27    fn cuEventRecord(event: u64, stream: u64) -> i32;
28    fn cuEventSynchronize(event: u64) -> i32;
29    fn cuEventDestroy_v2(event: u64) -> i32;
30}
31
32/// Loaded CUmodule with cached function lookups. Drop unloads.
33pub struct CudaModule {
34    handle: u64,
35}
36
37impl CudaModule {
38    pub fn from_ptx(ptx: &str) -> Result<Self> {
39        let cstr = std::ffi::CString::new(ptx).context("PTX contains NUL")?;
40        let mut h = 0u64;
41        let s = unsafe { cuModuleLoadData(&mut h, cstr.as_ptr() as *const c_void) };
42        if s != 0 {
43            bail!("cuModuleLoadData failed: {s}");
44        }
45        Ok(Self { handle: h })
46    }
47
48    pub fn function(&self, name: &str) -> Result<u64> {
49        let cstr = std::ffi::CString::new(name).context("function name has NUL")?;
50        let mut f = 0u64;
51        let s = unsafe { cuModuleGetFunction(&mut f, self.handle, cstr.as_ptr()) };
52        if s != 0 {
53            bail!("cuModuleGetFunction({name}) failed: {s}");
54        }
55        Ok(f)
56    }
57}
58
59impl Drop for CudaModule {
60    fn drop(&mut self) {
61        unsafe {
62            let _ = cuModuleUnload(self.handle);
63        }
64    }
65}
66
67/// CUDA event for cross-stream / host-side completion tracking.
68pub struct CudaEvent {
69    pub handle: u64,
70}
71
72impl CudaEvent {
73    pub fn new() -> Result<Self> {
74        let mut h = 0u64;
75        // CU_EVENT_DISABLE_TIMING (0x2): we never query elapsed time.
76        const CU_EVENT_DISABLE_TIMING: u32 = 0x2;
77        let s = unsafe { cuEventCreate(&mut h, CU_EVENT_DISABLE_TIMING) };
78        if s != 0 {
79            bail!("cuEventCreate failed: {s}");
80        }
81        Ok(Self { handle: h })
82    }
83    pub fn record(&self, stream: u64) -> Result<()> {
84        let s = unsafe { cuEventRecord(self.handle, stream) };
85        if s != 0 {
86            bail!("cuEventRecord failed: {s}");
87        }
88        Ok(())
89    }
90    pub fn sync(&self) -> Result<()> {
91        let s = unsafe { cuEventSynchronize(self.handle) };
92        if s != 0 {
93            bail!("cuEventSynchronize failed: {s}");
94        }
95        Ok(())
96    }
97}
98
99impl Drop for CudaEvent {
100    fn drop(&mut self) {
101        unsafe {
102            let _ = cuEventDestroy_v2(self.handle);
103        }
104    }
105}
106
107/// Direct cuLaunchKernel wrapper. Caller is responsible for ensuring the
108/// `params` pointer array stays valid until the launch returns.
109pub fn launch_kernel(
110    func: u64,
111    grid: (u32, u32, u32),
112    block: (u32, u32, u32),
113    shared_bytes: u32,
114    stream: u64,
115    params: &mut [*mut c_void],
116) -> Result<()> {
117    let s = unsafe {
118        cuLaunchKernel(
119            func,
120            grid.0,
121            grid.1,
122            grid.2,
123            block.0,
124            block.1,
125            block.2,
126            shared_bytes,
127            stream,
128            params.as_mut_ptr(),
129            std::ptr::null_mut(),
130        )
131    };
132    if s != 0 {
133        bail!("cuLaunchKernel failed: status {s} grid={grid:?} block={block:?}");
134    }
135    Ok(())
136}
spark_storage/cuda_module.rs

spark_storage/
cuda_module.rs
