microgemm 0.3.1

General matrix multiplication with custom configuration in Rust. Supports no_std and no_alloc environments.
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

use crate::std_prelude::*;
use crate::{MatMut, MatRef, Zero};
use core::ops::{Add, Mul};

pub fn naive_gemm<T>(alpha: T, a: MatRef<T>, b: MatRef<T>, beta: T, c: &mut MatMut<T>)
where
    T: Copy + Add<Output = T> + Mul<Output = T> + Zero,
{
    assert_eq!(a.nrows(), c.nrows());
    assert_eq!(b.ncols(), c.ncols());
    assert_eq!(a.ncols(), b.nrows());

    let k = a.ncols();

    for i in 0..a.nrows() {
        for j in 0..b.ncols() {
            let dot = (0..k)
                .map(|h| a.get(i, h) * b.get(h, j))
                .reduce(|accum, x| accum + x)
                .unwrap_or(T::zero());
            let z = c.get_mut(i, j);
            *z = alpha * dot + beta * *z;
        }
    }
}

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

    #[rustfmt::skip]
    #[test]
    fn fixed_1() {
        let a = [
            1, 2, 3,
            4, 5, 6,
        ];
        let b = [
            10, 11,
            20, 21,
            30, 31,
        ];
        let a = MatRef::row_major(2, 3, &a);
        let b = MatRef::row_major(3, 2, &b);

        let mut c = [-1; 4];
        let mut c = MatMut::row_major(2, 2, c.as_mut());
        let expect = [
            140, 146,
            320, 335,
        ];
        naive_gemm(1, a, b, 0, c.as_mut());
        assert_eq!(c.as_slice(), expect);

        let mut c = [-1; 4];
        let mut c = MatMut::col_major(2, 2, c.as_mut());
        let expect = [
            140, 320,
            146, 335,
        ];
        naive_gemm(1, a, b, 0, c.as_mut());
        assert_eq!(c.as_slice(), expect);
    }

    #[rustfmt::skip]
    #[test]
    fn fixed_2() {
        let alpha = 3;
        let beta = -4;

        let a = [
            1, 2, 3,
            4, 5, 6,
        ];
        let b = [
            2, 3, 4,
            5, 6, 7,
        ];
        let mut c = [
            -4, 1,
            -5, -6,
        ];
        let expect = [
            beta * c[0] + alpha * (2 + 2 * 3 + 3 * 4),
            beta * c[1] + alpha * (5 + 2 * 6 + 3 * 7),
            beta * c[2] + alpha * (4 * 2 + 5 * 3 + 6 * 4),
            beta * c[3] + alpha * (4 * 5 + 5 * 6 + 6 * 7),
        ];

        let a = MatRef::row_major(2, 3, a.as_ref());
        let b = MatRef::col_major(3, 2, b.as_ref());
        let mut c = MatMut::row_major(2, 2, c.as_mut());

        naive_gemm(alpha, a, b, beta, c.as_mut());
        assert_eq!(c.as_slice(), expect);
    }

    #[test]
    #[rustfmt::skip]
    fn fixed_3() {
        let a = [
            1, 0, 2,
            0, -1, 3,
        ];
        let a = MatRef::row_major(2, 3, a.as_ref());
        let b = [
            2, -1,
            0, 5,
            1, 1,
        ];
        let b = MatRef::row_major(3, 2, b.as_ref());

        let mut c = [-9; 2 * 2];
        let c = &mut MatMut::row_major(2, 2, c.as_mut());
        let expect = [
            4, 1,
            3, -2,
        ];
        naive_gemm(1, a, b, 0, c);
        assert_eq!(c.as_slice(), expect);
    }
}