rlst 0.10.0

A Rust native linear algebra library.
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

//! Implementation of the Gemm trait for matrix multiplication.

use crate::base_types::TransMode;
use crate::base_types::{c32, c64};
use crate::traits::linalg::base::Gemm;
use blas::{cgemm, dgemm, sgemm, zgemm};

/// Compute expected size of a data slice from stride and shape.
fn get_expected_data_size(stride: [usize; 2], shape: [usize; 2]) -> usize {
    if shape[0] == 0 || shape[1] == 0 {
        return 0;
    }

    1 + (shape[0] - 1) * stride[0] + (shape[1] - 1) * stride[1]
}

/// Panic if expected data size is not identical to actual data size.
fn assert_data_size(nelems: usize, stride: [usize; 2], shape: [usize; 2]) {
    let expected = get_expected_data_size(stride, shape);

    assert_eq!(
        expected, nelems,
        "Wrong size for data slice. Actual size {nelems}. Expected size {expected}."
    );
}

macro_rules! impl_gemm {
    ($scalar:ty, $gemm:ident) => {
        impl Gemm for $scalar {
            fn gemm(
                transa: TransMode,
                transb: TransMode,
                m: usize,
                n: usize,
                k: usize,
                alpha: Self,
                a: &[Self],
                rsa: usize,
                csa: usize,
                b: &[Self],
                rsb: usize,
                csb: usize,
                beta: Self,
                c: &mut [Self],
                rsc: usize,
                csc: usize,
            ) {
                match transa {
                    TransMode::NoTrans => assert_data_size(a.len(), [rsa, csa], [m, k]),
                    TransMode::ConjNoTrans => assert_data_size(a.len(), [rsa, csa], [m, k]),
                    TransMode::Trans => assert_data_size(a.len(), [rsa, csa], [k, m]),
                    TransMode::ConjTrans => assert_data_size(a.len(), [rsa, csa], [k, m]),
                }

                match transb {
                    TransMode::NoTrans => assert_data_size(b.len(), [rsb, csb], [k, n]),
                    TransMode::ConjNoTrans => assert_data_size(b.len(), [rsb, csb], [k, n]),
                    TransMode::Trans => assert_data_size(b.len(), [rsb, csb], [n, k]),
                    TransMode::ConjTrans => assert_data_size(b.len(), [rsb, csb], [n, k]),
                }

                assert_data_size(c.len(), [rsc, csc], [m, n]);

                assert_eq!(
                    rsa, 1,
                    "Input matrix for dgemm must have rsa=1. Actual {}",
                    rsa
                );
                assert_eq!(
                    rsa, 1,
                    "Input matrix for dgemm must have rsb=1. Actual {}",
                    rsb
                );
                assert_eq!(
                    rsa, 1,
                    "Input matrix for dgemm must have rsc=1. Actual {}",
                    rsc
                );

                let lda = csa as i32;
                let ldb = csb as i32;
                let ldc = csc as i32;

                let transa = match transa {
                    TransMode::NoTrans => b'N',
                    TransMode::ConjNoTrans => {
                        panic!("TransMode::ConjNoTrans not supported for gemm implementation.")
                    }
                    TransMode::Trans => b'T',
                    TransMode::ConjTrans => b'C',
                };

                let transb = match transb {
                    TransMode::NoTrans => b'N',
                    TransMode::ConjNoTrans => {
                        panic!("TransMode::ConjNoTrans not supported for gemm implementation.")
                    }
                    TransMode::Trans => b'T',
                    TransMode::ConjTrans => b'C',
                };

                unsafe {
                    $gemm(
                        transa as u8,
                        transb as u8,
                        m as i32,
                        n as i32,
                        k as i32,
                        alpha,
                        a,
                        lda,
                        b,
                        ldb,
                        beta,
                        c,
                        ldc,
                    )
                }
            }
        }
    };
}

impl_gemm!(f32, sgemm);
impl_gemm!(f64, dgemm);
impl_gemm!(c32, cgemm);
impl_gemm!(c64, zgemm);