sim_lib_compute_cuda/
runtime.rs

1//! Narrow runtime-loaded CUDA execution boundary.
2
3use std::{ffi::c_void, sync::Arc};
4
5use libloading::Library;
6
7use crate::{CudaLibrarySet, CudaLoadError};
8
9const CUDA_MEMCPY_HOST_TO_DEVICE: i32 = 1;
10const CUDA_MEMCPY_DEVICE_TO_HOST: i32 = 2;
11const CUBLAS_OP_N: i32 = 0;
12
13type CudaMalloc = unsafe extern "C" fn(*mut *mut c_void, usize) -> i32;
14type CudaFree = unsafe extern "C" fn(*mut c_void) -> i32;
15type CudaMemcpy = unsafe extern "C" fn(*mut c_void, *const c_void, usize, i32) -> i32;
16type CudaDeviceSynchronize = unsafe extern "C" fn() -> i32;
17type CublasCreate = unsafe extern "C" fn(*mut *mut c_void) -> i32;
18type CublasDestroy = unsafe extern "C" fn(*mut c_void) -> i32;
19type CublasSgemm = unsafe extern "C" fn(
20    *mut c_void,
21    i32,
22    i32,
23    i32,
24    i32,
25    i32,
26    *const f32,
27    *const f32,
28    i32,
29    *const f32,
30    i32,
31    *const f32,
32    *mut f32,
33    i32,
34) -> i32;
35
36pub(crate) struct CudaDeviceBuffer {
37    runtime: Arc<CudaLibrarySet>,
38    address: usize,
39    len: usize,
40}
41
42// SAFETY: CUDA device allocations are opaque integer handles. CUDA runtime and
43// cuBLAS calls are thread-safe, and the allocation is freed exactly once by
44// `Drop` after the last Arc owner is released.
45unsafe impl Send for CudaDeviceBuffer {}
46// SAFETY: See the `Send` justification; no host dereference of the device
47// address is ever performed.
48unsafe impl Sync for CudaDeviceBuffer {}
49
50impl CudaDeviceBuffer {
51    pub(crate) fn len(&self) -> usize {
52        self.len
53    }
54
55    pub(crate) fn runtime(&self) -> &Arc<CudaLibrarySet> {
56        &self.runtime
57    }
58
59    fn pointer(&self) -> *mut c_void {
60        self.address as *mut c_void
61    }
62
63    pub(crate) fn read(&self) -> Result<Vec<f32>, CudaLoadError> {
64        self.runtime.download(self)
65    }
66}
67
68impl Drop for CudaDeviceBuffer {
69    fn drop(&mut self) {
70        let _ = self.runtime.free(self.pointer());
71    }
72}
73
74impl CudaLibrarySet {
75    pub(crate) fn upload(
76        self: &Arc<Self>,
77        values: &[f32],
78    ) -> Result<Arc<CudaDeviceBuffer>, CudaLoadError> {
79        let bytes = byte_count(values.len())?;
80        let (runtime, _) = self.execution_handles();
81        let malloc = symbol::<CudaMalloc>(runtime, b"cudaMalloc\0")?;
82        let copy = symbol::<CudaMemcpy>(runtime, b"cudaMemcpy\0")?;
83        let mut pointer = std::ptr::null_mut();
84        // SAFETY: Function pointers use the documented CUDA runtime ABI.
85        // `pointer` is an initialized out parameter and `values` covers `bytes`.
86        let status = unsafe { malloc(&mut pointer, bytes) };
87        check_cuda(status, "cudaMalloc")?;
88        // SAFETY: CUDA owns `pointer` for `bytes`, and the host slice is valid
89        // for the same byte count.
90        let status = unsafe {
91            copy(
92                pointer,
93                values.as_ptr().cast::<c_void>(),
94                bytes,
95                CUDA_MEMCPY_HOST_TO_DEVICE,
96            )
97        };
98        if let Err(error) = check_cuda(status, "cudaMemcpy host-to-device") {
99            let _ = self.free(pointer);
100            return Err(error);
101        }
102        Ok(Arc::new(CudaDeviceBuffer {
103            runtime: Arc::clone(self),
104            address: pointer as usize,
105            len: values.len(),
106        }))
107    }
108
109    pub(crate) fn matmul(
110        self: &Arc<Self>,
111        left: &CudaDeviceBuffer,
112        right: &CudaDeviceBuffer,
113        rows: usize,
114        inner: usize,
115        cols: usize,
116    ) -> Result<Arc<CudaDeviceBuffer>, CudaLoadError> {
117        validate_matrix_lengths(left.len(), right.len(), rows, inner, cols)?;
118        if !Arc::ptr_eq(self, left.runtime()) || !Arc::ptr_eq(self, right.runtime()) {
119            return Err(error("CUDA inputs belong to another runtime"));
120        }
121        let output_len = rows
122            .checked_mul(cols)
123            .ok_or_else(|| error("CUDA output length overflowed"))?;
124        let zeros = vec![0.0_f32; output_len];
125        let output = self.upload(&zeros)?;
126        let (_, cublas) = self.execution_handles();
127        let create = symbol::<CublasCreate>(cublas, b"cublasCreate_v2\0")?;
128        let destroy = symbol::<CublasDestroy>(cublas, b"cublasDestroy_v2\0")?;
129        let sgemm = symbol::<CublasSgemm>(cublas, b"cublasSgemm_v2\0")?;
130        let mut handle = std::ptr::null_mut();
131        // SAFETY: Function pointers use the documented cuBLAS ABI and the
132        // handle out parameter is initialized.
133        check_cublas(unsafe { create(&mut handle) }, "cublasCreate_v2")?;
134        let dimensions = matrix_dimensions(rows, inner, cols)?;
135        let alpha = 1.0_f32;
136        let beta = 0.0_f32;
137        // Row-major C=A*B is column-major C^T=B^T*A^T.
138        // SAFETY: All device buffers cover the validated dimensions; cuBLAS
139        // receives the documented column-major SGEMM arguments.
140        let status = unsafe {
141            sgemm(
142                handle,
143                CUBLAS_OP_N,
144                CUBLAS_OP_N,
145                dimensions.cols,
146                dimensions.rows,
147                dimensions.inner,
148                &alpha,
149                right.pointer().cast::<f32>(),
150                dimensions.cols,
151                left.pointer().cast::<f32>(),
152                dimensions.inner,
153                &beta,
154                output.pointer().cast::<f32>(),
155                dimensions.cols,
156            )
157        };
158        let gemm_result = check_cublas(status, "cublasSgemm_v2");
159        // SAFETY: `handle` came from the successful create call.
160        let destroy_result = check_cublas(unsafe { destroy(handle) }, "cublasDestroy_v2");
161        gemm_result?;
162        destroy_result?;
163        Ok(output)
164    }
165
166    /// Runs one row-major dense f32 matrix multiplication through cuBLAS and
167    /// returns a synchronized host result.
168    pub fn matmul_f32(
169        self: &Arc<Self>,
170        left: &[f32],
171        right: &[f32],
172        rows: usize,
173        inner: usize,
174        cols: usize,
175    ) -> Result<Vec<f32>, CudaLoadError> {
176        let left = self.upload(left)?;
177        let right = self.upload(right)?;
178        self.matmul(&left, &right, rows, inner, cols)?.read()
179    }
180
181    fn download(&self, buffer: &CudaDeviceBuffer) -> Result<Vec<f32>, CudaLoadError> {
182        let bytes = byte_count(buffer.len())?;
183        let (runtime, _) = self.execution_handles();
184        let copy = symbol::<CudaMemcpy>(runtime, b"cudaMemcpy\0")?;
185        let synchronize = symbol::<CudaDeviceSynchronize>(runtime, b"cudaDeviceSynchronize\0")?;
186        let mut values = vec![0.0_f32; buffer.len()];
187        // SAFETY: The output vector covers `bytes`, and the CUDA allocation was
188        // created for the same logical f32 length.
189        check_cuda(
190            unsafe {
191                copy(
192                    values.as_mut_ptr().cast::<c_void>(),
193                    buffer.pointer(),
194                    bytes,
195                    CUDA_MEMCPY_DEVICE_TO_HOST,
196                )
197            },
198            "cudaMemcpy device-to-host",
199        )?;
200        // SAFETY: The function has no arguments and uses the documented ABI.
201        check_cuda(unsafe { synchronize() }, "cudaDeviceSynchronize")?;
202        Ok(values)
203    }
204
205    fn free(&self, pointer: *mut c_void) -> Result<(), CudaLoadError> {
206        if pointer.is_null() {
207            return Ok(());
208        }
209        let (runtime, _) = self.execution_handles();
210        let free = symbol::<CudaFree>(runtime, b"cudaFree\0")?;
211        // SAFETY: The pointer was returned by `cudaMalloc` and is freed once.
212        check_cuda(unsafe { free(pointer) }, "cudaFree")
213    }
214}
215
216struct MatrixDimensions {
217    rows: i32,
218    inner: i32,
219    cols: i32,
220}
221
222fn matrix_dimensions(
223    rows: usize,
224    inner: usize,
225    cols: usize,
226) -> Result<MatrixDimensions, CudaLoadError> {
227    Ok(MatrixDimensions {
228        rows: i32::try_from(rows).map_err(|_| error("CUDA row count exceeds i32"))?,
229        inner: i32::try_from(inner).map_err(|_| error("CUDA inner count exceeds i32"))?,
230        cols: i32::try_from(cols).map_err(|_| error("CUDA column count exceeds i32"))?,
231    })
232}
233
234fn validate_matrix_lengths(
235    left: usize,
236    right: usize,
237    rows: usize,
238    inner: usize,
239    cols: usize,
240) -> Result<(), CudaLoadError> {
241    if rows.checked_mul(inner) != Some(left) || inner.checked_mul(cols) != Some(right) {
242        return Err(error("CUDA matmul shape does not match input lengths"));
243    }
244    Ok(())
245}
246
247fn byte_count(len: usize) -> Result<usize, CudaLoadError> {
248    len.checked_mul(std::mem::size_of::<f32>())
249        .ok_or_else(|| error("CUDA byte count overflowed"))
250}
251
252fn symbol<'library, T>(
253    library: &'library Library,
254    name: &[u8],
255) -> Result<libloading::Symbol<'library, T>, CudaLoadError> {
256    // SAFETY: Every caller supplies the official function signature for the
257    // named CUDA/cuBLAS symbol. Library handles outlive returned symbols.
258    unsafe { library.get(name) }.map_err(|load| error(load.to_string()))
259}
260
261fn check_cuda(status: i32, operation: &str) -> Result<(), CudaLoadError> {
262    (status == 0)
263        .then_some(())
264        .ok_or_else(|| error(format!("{operation} failed with CUDA status {status}")))
265}
266
267fn check_cublas(status: i32, operation: &str) -> Result<(), CudaLoadError> {
268    (status == 0)
269        .then_some(())
270        .ok_or_else(|| error(format!("{operation} failed with cuBLAS status {status}")))
271}
272
273fn error(message: impl Into<String>) -> CudaLoadError {
274    CudaLoadError {
275        message: message.into(),
276    }
277}
sim_lib_compute_cuda/runtime.rs

sim_lib_compute_cuda/
runtime.rs
