mdarray-linalg-blas 0.2.0

BLAS backend for mdarray-linalg
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
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138

//! Simple function-based interface to linear algebra libraries

use std::mem::MaybeUninit;

use cblas_sys::{CBLAS_LAYOUT, CBLAS_TRANSPOSE};
use mdarray::{Array, Dense, Dim, Layout, Slice};
use mdarray_linalg::utils::{dims3, into_i32};
use num_complex::ComplexFloat;

use super::scalar::BlasScalar;
use crate::trans_stride;

pub(super) fn gemm<T, La, Lb, Lc, D0, D1, D2>(
    alpha: T,
    a: &Slice<T, (D0, D1), La>,
    b: &Slice<T, (D1, D2), Lb>,
    beta: T,
    c: &mut Slice<T, (D0, D2), Lc>,
) where
    T: BlasScalar + ComplexFloat,
    La: Layout,
    Lb: Layout,
    Lc: Layout,
    D0: Dim,
    D1: Dim,
    D2: Dim,
{
    let (m, n, k) = dims3(*a.shape(), *b.shape(), *c.shape());

    if a.stride(0) != 1 && a.stride(1) != 1 {
        let a_owned: Array<T, (D0, D1)> = a.to_array();
        return gemm(alpha, &a_owned, b, beta, c);
    }

    if b.stride(0) != 1 && b.stride(1) != 1 {
        let b_owned: Array<T, (D1, D2)> = b.to_array();
        return gemm(alpha, a, &b_owned, beta, c);
    }

    let row_major = c.stride(1) == 1;
    assert!(
        row_major || c.stride(0) == 1,
        "c must be contiguous in one dimension"
    );

    let (same_order, other_order) = if row_major {
        (CBLAS_TRANSPOSE::CblasNoTrans, CBLAS_TRANSPOSE::CblasTrans)
    } else {
        (CBLAS_TRANSPOSE::CblasTrans, CBLAS_TRANSPOSE::CblasNoTrans)
    };
    let (a_trans, a_stride) = trans_stride(a, same_order, other_order);
    let (b_trans, b_stride) = trans_stride(b, same_order, other_order);

    let c_stride = into_i32(c.stride(if row_major { 0 } else { 1 }));

    unsafe {
        T::cblas_gemm(
            if row_major {
                CBLAS_LAYOUT::CblasRowMajor
            } else {
                CBLAS_LAYOUT::CblasColMajor
            },
            a_trans,
            b_trans,
            m,
            n,
            k,
            alpha,
            a.as_ptr(),
            a_stride,
            b.as_ptr(),
            b_stride,
            beta,
            c.as_mut_ptr(),
            c_stride,
        )
    }
}

pub(super) fn gemm_uninit<T, La, Lb, Lc, D0, D1, D2>(
    alpha: T,
    a: &Slice<T, (D0, D1), La>,
    b: &Slice<T, (D1, D2), Lb>,
    beta: T,
    mut c: Array<MaybeUninit<T>, (D0, D2)>,
) -> Array<T, (D0, D2)>
where
    T: BlasScalar + ComplexFloat,
    La: Layout,
    Lb: Layout,
    Lc: Layout,
    D0: Dim,
    D1: Dim,
    D2: Dim,
{
    let (m, n, k) = dims3(*a.shape(), *b.shape(), *c.shape());

    if a.stride(0) != 1 && a.stride(1) != 1 {
        let a_owned: Array<T, (D0, D1)> = a.to_array();
        return gemm_uninit::<T, Dense, Lb, Lc, D0, D1, D2>(alpha, &a_owned, b, beta, c);
    }

    if b.stride(0) != 1 && b.stride(1) != 1 {
        let b_owned: Array<T, (D1, D2)> = b.to_array();
        return gemm_uninit::<T, La, Dense, Lc, D0, D1, D2>(alpha, a, &b_owned, beta, c);
    }

    debug_assert!(c.stride(1) == 1);

    let same_order = CBLAS_TRANSPOSE::CblasNoTrans;
    let other_order = CBLAS_TRANSPOSE::CblasTrans;

    let (a_trans, a_stride) = trans_stride(a, same_order, other_order);
    let (b_trans, b_stride) = trans_stride(b, same_order, other_order);

    let c_stride = into_i32(c.stride(0));

    unsafe {
        T::cblas_gemm(
            CBLAS_LAYOUT::CblasRowMajor,
            a_trans,
            b_trans,
            m,
            n,
            k,
            alpha,
            a.as_ptr(),
            a_stride,
            b.as_ptr(),
            b_stride,
            beta,
            c.as_mut_ptr() as *mut T,
            c_stride,
        );

        c.assume_init()
    }
}