sp1_gpu_cudart/
device.rs

1use std::{
2    ffi::{c_char, CStr},
3    sync::OnceLock,
4};
5
6use crate::{CudaError, TaskScope};
7use slop_alloc::{mem::CopyError, CopyIntoBackend, CopyToBackend, CpuBackend};
8use sp1_gpu_sys::runtime::{cuda_get_device_name, cuda_mem_get_info};
9
10pub trait DeviceCopy: Copy + 'static + Sized {}
11
12impl<T: Copy + 'static + Sized> DeviceCopy for T {}
13
14/// Returns a pair `(free, total)` of the amount of free and total memory on the device.
15pub fn cuda_memory_info() -> Result<(usize, usize), CudaError> {
16    let mut free: usize = 0;
17    let mut total: usize = 0;
18    CudaError::result_from_ffi(unsafe { cuda_mem_get_info(&mut free, &mut total) })?;
19    Ok((free, total))
20}
21
22/// Returns the name of the CUDA device the calling thread is bound to, e.g. `"NVIDIA L4"`.
23///
24/// The name is queried once and cached for the lifetime of the process, because the underlying
25/// `cudaGetDeviceProperties` call materializes the whole device property struct and costs tens of
26/// milliseconds on the first call. Errors are never cached, so a failed query may be retried.
27pub fn cuda_device_name() -> Result<&'static str, CudaError> {
28    static DEVICE_NAME: OnceLock<String> = OnceLock::new();
29    if let Some(name) = DEVICE_NAME.get() {
30        return Ok(name);
31    }
32
33    // `cudaDeviceProp::name` is a 256 byte NUL terminated buffer, and the shim NUL terminates
34    // whatever it writes, so a NUL is always present within the buffer.
35    let mut name = [0u8; 256];
36    CudaError::result_from_ffi(unsafe {
37        cuda_get_device_name(name.as_mut_ptr().cast::<c_char>(), name.len())
38    })?;
39    let name = CStr::from_bytes_until_nul(&name)
40        .map(|name| name.to_string_lossy().into_owned())
41        .unwrap_or_default();
42
43    Ok(DEVICE_NAME.get_or_init(|| name))
44}
45
46pub trait IntoDevice: CopyIntoBackend<TaskScope, CpuBackend> + Sized {
47    fn into_device_in(self, backend: &TaskScope) -> Result<Self::Output, CopyError> {
48        self.copy_into_backend(backend)
49    }
50}
51
52impl<T> IntoDevice for T where T: CopyIntoBackend<TaskScope, CpuBackend> + Sized {}
53
54pub trait ToDevice: CopyToBackend<TaskScope, CpuBackend> + Sized {
55    fn to_device_in(&self, backend: &TaskScope) -> Result<Self::Output, CopyError> {
56        self.copy_to_backend(backend)
57    }
58}
59
60impl<T> ToDevice for T where T: CopyToBackend<TaskScope, CpuBackend> + Sized {}
61
62#[macro_export]
63macro_rules! args {
64    ($($arg:expr),*) => {
65        [
66            $(
67                &$arg as *const _ as *mut std::ffi::c_void
68            ),*
69        ]
70    };
71}
sp1_gpu_cudart/device.rs

sp1_gpu_cudart/
device.rs
