trueno 0.16.4

High-performance SIMD compute library with GPU support for matrix operations
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

use super::super::*;

// ========================================================================
// Phase 4: BLIS GEMM Tests
// ========================================================================

#[test]
fn test_gemm_blis_small() {
    let a = vec![1.0, 2.0, 3.0, 4.0];
    let b = vec![5.0, 6.0, 7.0, 8.0];
    let mut c = vec![0.0; 4];

    gemm_blis(2, 2, 2, &a, &b, &mut c, None).unwrap();

    assert_eq!(c, vec![19.0, 22.0, 43.0, 50.0]);
}

#[test]
fn test_gemm_blis_medium() {
    let n = 64;
    let a: Vec<f32> = (0..n * n).map(|i| (i % 10) as f32).collect();
    let b: Vec<f32> = (0..n * n).map(|i| ((i + 3) % 10) as f32).collect();
    let mut c_ref = vec![0.0; n * n];
    let mut c_blis = vec![0.0; n * n];

    gemm_reference(n, n, n, &a, &b, &mut c_ref).unwrap();
    gemm_blis(n, n, n, &a, &b, &mut c_blis, None).unwrap();

    for i in 0..n * n {
        let diff = (c_ref[i] - c_blis[i]).abs();
        assert!(diff < 1e-3, "Mismatch at {}: ref={}, blis={}", i, c_ref[i], c_blis[i]);
    }
}

#[test]
fn test_gemm_blis_large() {
    let n = 256;
    let a: Vec<f32> = (0..n * n).map(|i| ((i % 7) as f32) * 0.1).collect();
    let b: Vec<f32> = (0..n * n).map(|i| ((i % 11) as f32) * 0.1).collect();
    let mut c_ref = vec![0.0; n * n];
    let mut c_blis = vec![0.0; n * n];

    gemm_reference(n, n, n, &a, &b, &mut c_ref).unwrap();
    gemm_blis(n, n, n, &a, &b, &mut c_blis, None).unwrap();

    let mut max_diff = 0.0f32;
    for i in 0..n * n {
        let diff = (c_ref[i] - c_blis[i]).abs();
        max_diff = max_diff.max(diff);
    }

    assert!(max_diff < 1e-2, "Max diff: {}", max_diff);
}

#[test]
fn test_gemm_blis_rectangular() {
    // Common ML shape: 32 x 4096 @ 4096 x 11008
    let m = 32;
    let k = 128;
    let n = 256;

    let a: Vec<f32> = (0..m * k).map(|i| ((i % 5) as f32) * 0.1).collect();
    let b: Vec<f32> = (0..k * n).map(|i| ((i % 7) as f32) * 0.1).collect();
    let mut c_ref = vec![0.0; m * n];
    let mut c_blis = vec![0.0; m * n];

    gemm_reference(m, n, k, &a, &b, &mut c_ref).unwrap();
    gemm_blis(m, n, k, &a, &b, &mut c_blis, None).unwrap();

    let mut max_diff = 0.0f32;
    for i in 0..m * n {
        let diff = (c_ref[i] - c_blis[i]).abs();
        max_diff = max_diff.max(diff);
    }

    assert!(max_diff < 1e-3, "Max diff: {}", max_diff);
}

#[test]
fn test_gemm_blis_edge_m_not_divisible_by_mr() {
    let m = 13; // Not divisible by MR=8
    let n = 16;
    let k = 16;

    let a: Vec<f32> = (0..m * k).map(|i| (i as f32) * 0.01).collect();
    let b: Vec<f32> = (0..k * n).map(|i| (i as f32) * 0.01).collect();
    let mut c_ref = vec![0.0; m * n];
    let mut c_blis = vec![0.0; m * n];

    gemm_reference(m, n, k, &a, &b, &mut c_ref).unwrap();
    gemm_blis(m, n, k, &a, &b, &mut c_blis, None).unwrap();

    for i in 0..m * n {
        let diff = (c_ref[i] - c_blis[i]).abs();
        assert!(diff < 1e-3, "Mismatch at {}: {} vs {}", i, c_ref[i], c_blis[i]);
    }
}

#[test]
fn test_gemm_blis_edge_n_not_divisible_by_nr() {
    let m = 16;
    let n = 17; // Not divisible by NR=6
    let k = 16;

    let a: Vec<f32> = (0..m * k).map(|i| (i as f32) * 0.01).collect();
    let b: Vec<f32> = (0..k * n).map(|i| (i as f32) * 0.01).collect();
    let mut c_ref = vec![0.0; m * n];
    let mut c_blis = vec![0.0; m * n];

    gemm_reference(m, n, k, &a, &b, &mut c_ref).unwrap();
    gemm_blis(m, n, k, &a, &b, &mut c_blis, None).unwrap();

    for i in 0..m * n {
        let diff = (c_ref[i] - c_blis[i]).abs();
        assert!(diff < 1e-3, "Mismatch at {}: {} vs {}", i, c_ref[i], c_blis[i]);
    }
}