openvm-cuda-common 2.0.0

CUDA common utils
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

use std::{
    ffi::c_void,
    sync::atomic::{AtomicU64, Ordering},
};

use crate::error::{check, CudaError};

#[link(name = "cudart")]
extern "C" {
    fn cudaFree(dev_ptr: *mut c_void) -> i32;
    fn cudaGetDevice(device: *mut i32) -> i32;
    fn cudaSetDevice(device: i32) -> i32;
    fn cudaDeviceReset() -> i32;
}

static DEVICE_RESET_EPOCH: AtomicU64 = AtomicU64::new(0);

pub fn get_device() -> Result<i32, CudaError> {
    let mut device = 0;
    unsafe {
        check(cudaGetDevice(&mut device))?;
    }
    assert!(device >= 0);
    Ok(device)
}

pub fn device_reset_epoch() -> u64 {
    DEVICE_RESET_EPOCH.load(Ordering::Acquire)
}

pub fn set_device_by_id(device: i32) -> Result<(), CudaError> {
    assert!(device >= 0);
    unsafe {
        check(cudaSetDevice(device))?;
        // Force primary-context initialization on the selected device.
        check(cudaFree(std::ptr::null_mut()))?;
    }
    Ok(())
}

pub fn set_device() -> Result<i32, CudaError> {
    let device = get_device()?;
    set_device_by_id(device)?;
    Ok(device)
}

pub fn reset_device() -> Result<(), CudaError> {
    check(unsafe { cudaDeviceReset() })?;
    DEVICE_RESET_EPOCH.fetch_add(1, Ordering::AcqRel);
    Ok(())
}