yarnn 0.1.0

Yet Another rust Neural Network framework
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


fn gemm_nn(m: usize, n: usize, k: usize, alpha: f32, 
           a: &[f32], lda: usize, 
           b: &[f32], ldb: usize,
           c: &mut [f32], ldc: usize)
{
    let a = &a[0..m * k];
    let b = &b[0..n * k];
    let c = &mut c[0..m * n];

    for i_m in 0..m {
        for i_k in 0..k {
            let a_part = alpha * a[i_m * lda + i_k];
            for i_n in 0..n {
                c[i_m * ldc + i_n] += a_part * b[i_k * ldb + i_n];
            }
        }
    }
}

fn gemm_nt(m: usize, n: usize, k: usize, alpha: f32, 
           a: &[f32], lda: usize, 
           b: &[f32], ldb: usize,
           c: &mut [f32], ldc: usize)
{
    let a = &a[0..m * k];
    let b = &b[0..n * k];
    let c = &mut c[0..m * n];

    for i_m in 0..m {
        for i_n in 0..n {
            let mut sum = 0.0;

            for i_k in 0..k {
                sum += alpha * a[i_m * lda + i_k] * b[i_n * ldb + i_k];
            }

            c[i_m * ldc + i_n] += sum;
        }
    }
}

fn gemm_tn(m: usize, n: usize, k: usize, alpha: f32, 
           a: &[f32], lda: usize, 
           b: &[f32], ldb: usize,
           c: &mut [f32], ldc: usize)
{    
    let a = &a[0..m * k];
    let b = &b[0..n * k];
    let c = &mut c[0..m * n];

    for i_m in 0..m {
        for i_k in 0..k {
            let a_part = alpha * a[i_k * lda + i_m];

            for i_n in 0..n {
                c[i_m * ldc + i_n] += a_part * b[i_k * ldb + i_n];
            }
        }
    }
}

fn gemm_tt(m: usize, n: usize, k: usize, alpha: f32, 
           a: &[f32], lda: usize, 
           b: &[f32], ldb: usize,
           c: &mut [f32], ldc: usize)
{
    let a = &a[0..m * k];
    let b = &b[0..n * k];
    let c = &mut c[0..m * n];
    
    for i_m in 0..m {
        for i_n in 0..n {
            let mut sum = 0.0;
            
            for i_k in 0..k {
                sum += alpha * a[i_k * lda + i_m] * b[i_n * ldb + i_k];
            }

            c[i_m * ldc + i_n] += sum;
        }
    }
}

pub fn gemm(ta: bool, tb: bool, m: usize, n: usize, k: usize, alpha: f32, 
        a: &[f32], lda: usize, 
        b: &[f32], ldb: usize, beta: f32,
        c: &mut [f32], ldc: usize)
{
    for i in 0..m {
        for j in 0..n {
            c[i * ldc + j] *= beta;
        }
    }

    if !ta && !tb {
        gemm_nn(m, n, k, alpha, a, lda, b, ldb, c, ldc);
    } else if ta && !tb {
        gemm_tn(m, n, k, alpha, a, lda, b, ldb, c, ldc);
    } else if !ta && tb {
        gemm_nt(m, n, k, alpha, a, lda, b, ldb, c, ldc);
    } else {
        gemm_tt(m, n, k, alpha, a, lda, b, ldb, c, ldc);
    }
}