icicle-cuda-runtime 1.3.0

Ingonyama's Rust wrapper of CUDA runtime
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
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84

use crate::bindings::{
    cudaStreamCreate, cudaStreamCreateWithFlags, cudaStreamDefault, cudaStreamDestroy, cudaStreamNonBlocking,
    cudaStreamSynchronize, cudaStream_t,
};
use crate::error::{CudaResult, CudaResultWrap};
use bitflags::bitflags;
use std::mem::{forget, MaybeUninit};

#[repr(transparent)]
#[derive(Debug)]
pub struct CudaStream {
    pub(crate) handle: cudaStream_t,
}

unsafe impl Sync for CudaStream {}

bitflags! {
    pub struct CudaStreamCreateFlags: u32 {
        const DEFAULT = cudaStreamDefault;
        const NON_BLOCKING = cudaStreamNonBlocking;
    }
}

impl CudaStream {
    pub(crate) fn from_handle(handle: cudaStream_t) -> Self {
        Self { handle }
    }

    pub fn create() -> CudaResult<Self> {
        let mut handle = MaybeUninit::<cudaStream_t>::uninit();
        unsafe {
            cudaStreamCreate(handle.as_mut_ptr())
                .wrap_maybe_uninit(handle)
                .map(CudaStream::from_handle)
        }
    }

    pub fn create_with_flags(flags: CudaStreamCreateFlags) -> CudaResult<Self> {
        let mut handle = MaybeUninit::<cudaStream_t>::uninit();
        unsafe {
            cudaStreamCreateWithFlags(handle.as_mut_ptr(), flags.bits)
                .wrap_maybe_uninit(handle)
                .map(CudaStream::from_handle)
        }
    }

    pub fn destroy(self) -> CudaResult<()> {
        let handle = self.handle;
        forget(self);
        if handle.is_null() {
            Ok(())
        } else {
            unsafe { cudaStreamDestroy(handle).wrap() }
        }
    }

    pub fn synchronize(&self) -> CudaResult<()> {
        unsafe { cudaStreamSynchronize(self.handle).wrap() }
    }
}

impl Default for CudaStream {
    fn default() -> Self {
        Self {
            handle: std::ptr::null_mut(),
        }
    }
}

impl Drop for CudaStream {
    fn drop(&mut self) {
        let handle = self.handle;
        if handle.is_null() {
            return;
        }
        let _ = unsafe { cudaStreamDestroy(handle) };
    }
}

impl From<&CudaStream> for cudaStream_t {
    fn from(stream: &CudaStream) -> Self {
        stream.handle
    }
}