rocmrc 0.4.0

Minimal safe ROCm bindings (HIP, hipRTC), modeled after cudarc
Documentation 
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
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99

use super::{result, result::RocblasError, sys};
use crate::hip::HipStream;
use std::sync::Arc;

mod asum;
mod axpy;
mod copy;
mod dot;
mod gemm;
mod gemv;
mod nrm2;
mod scal;

pub use asum::*;
pub use axpy::*;
pub use copy::*;
pub use dot::*;
pub use gemm::*;
pub use gemv::*;
pub use nrm2::*;
pub use scal::*;

pub struct RocBlas {
    pub(crate) handle: sys::rocblas_handle,
    pub(crate) stream: Arc<HipStream>,
}

unsafe impl Send for RocBlas {}
unsafe impl Sync for RocBlas {}

impl RocBlas {
    pub fn new(stream: Arc<HipStream>) -> Result<Self, RocblasError> {
        let ctx = stream.context();
        ctx.record_err(ctx.bind_to_thread());
        let handle = result::create_handle()?;
        unsafe { result::set_stream(handle, stream.hip_stream() as _) }?;
        let blas = Self { handle, stream };
        Ok(blas)
    }

    pub fn handle(&self) -> &sys::rocblas_handle {
        &self.handle
    }

    pub fn set_stream(&mut self, stream: Arc<HipStream>) -> Result<(), RocblasError> {
        self.stream = stream;
        unsafe { result::set_stream(self.handle, self.stream.hip_stream() as _) }
    }

    pub fn set_pointer_mode(
        &self,
        pointer_mode: sys::rocblas_pointer_mode,
    ) -> Result<(), RocblasError> {
        unsafe { sys::rocblas_set_pointer_mode(self.handle, pointer_mode).result() }
    }

    pub fn get_pointer_mode(&self) -> Result<sys::rocblas_pointer_mode, RocblasError> {
        unsafe {
            let mut mode = ::core::mem::MaybeUninit::uninit();
            sys::rocblas_get_pointer_mode(self.handle, mode.as_mut_ptr()).result()?;
            Ok(mode.assume_init())
        }
    }
}

impl Drop for RocBlas {
    fn drop(&mut self) {
        let handle = std::mem::replace(&mut self.handle, std::ptr::null_mut());
        if !handle.is_null() {
            unsafe { result::destroy_handle(handle) }.unwrap();
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::hip::HipContext;

    #[test]
    fn rocblas_pointer_mode() {
        let ctx = HipContext::new(0).unwrap();
        let stream = ctx.default_stream();
        let blas = RocBlas::new(stream.clone()).unwrap();
        assert_eq!(
            blas.get_pointer_mode().unwrap(),
            sys::rocblas_pointer_mode::rocblas_pointer_mode_host,
            "The default pointer mode uses host pointers"
        );

        blas.set_pointer_mode(sys::rocblas_pointer_mode::rocblas_pointer_mode_device)
            .unwrap();
        assert_eq!(
            blas.get_pointer_mode().unwrap(),
            sys::rocblas_pointer_mode::rocblas_pointer_mode_device,
            "We have set the mode to use device pointers"
        );
    }
}