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
100
101
102
103
104
105
106
107
108
109
110
111
112
113

use super::{result, result::RocblasError, sys};
use crate::hip::DevicePtrMut;
use crate::rocblas::RocBlas;

#[derive(Debug, Copy, Clone)]
pub struct ScalConfig<T> {
    pub n: sys::rocblas_int,
    pub alpha: T,
    pub incx: sys::rocblas_int,
}

pub trait Scal<T> {
    /// `x := alpha*x`.
    fn scal<X: DevicePtrMut<T>>(&self, cfg: ScalConfig<T>, x: &mut X) -> Result<(), RocblasError>;
}

impl Scal<f32> for RocBlas {
    fn scal<X: DevicePtrMut<f32>>(
        &self,
        cfg: ScalConfig<f32>,
        x: &mut X,
    ) -> Result<(), RocblasError> {
        let (x, _record_x) = x.device_ptr_mut(&self.stream);
        unsafe {
            result::sscal(
                self.handle,
                cfg.n,
                (&cfg.alpha) as *const _,
                x as *mut _,
                cfg.incx,
            )
        }
    }
}

impl Scal<f64> for RocBlas {
    fn scal<X: DevicePtrMut<f64>>(
        &self,
        cfg: ScalConfig<f64>,
        x: &mut X,
    ) -> Result<(), RocblasError> {
        let (x, _record_x) = x.device_ptr_mut(&self.stream);
        unsafe {
            result::dscal(
                self.handle,
                cfg.n,
                (&cfg.alpha) as *const _,
                x as *mut _,
                cfg.incx,
            )
        }
    }
}

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

    #[test]
    fn test_sscal() {
        let ctx = HipContext::new(0).unwrap();
        let stream = ctx.default_stream();
        let blas = RocBlas::new(stream.clone()).unwrap();

        let alpha = -1.75f32;
        let x: [f32; 5] = [-0.5944882, 1.8055636, 0.52204555, -0.00397902, 0.39125723];
        let truth: Vec<f32> = x.iter().map(|v| alpha * v).collect();

        let mut x_dev = stream.clone_htod(&x).unwrap();
        blas.scal(
            ScalConfig {
                n: 5,
                alpha,
                incx: 1,
            },
            &mut x_dev,
        )
        .unwrap();

        let found = stream.clone_dtoh(&x_dev).unwrap();
        for i in 0..5 {
            assert!((found[i] - truth[i]).abs() <= 1e-5);
        }
    }

    #[test]
    fn test_dscal() {
        let ctx = HipContext::new(0).unwrap();
        let stream = ctx.default_stream();
        let blas = RocBlas::new(stream.clone()).unwrap();

        let alpha = -1.75f64;
        let x: [f64; 5] = [-0.5944882, 1.8055636, 0.52204555, -0.00397902, 0.39125723];
        let truth: Vec<f64> = x.iter().map(|v| alpha * v).collect();

        let mut x_dev = stream.clone_htod(&x).unwrap();
        blas.scal(
            ScalConfig {
                n: 5,
                alpha,
                incx: 1,
            },
            &mut x_dev,
        )
        .unwrap();

        let found = stream.clone_dtoh(&x_dev).unwrap();
        for i in 0..5 {
            assert!((found[i] - truth[i]).abs() <= 1e-12);
        }
    }
}