gemmkit 0.1.2

A clean, extensible, high-performance GEMM (general matrix multiply) engine
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

//! `GEMMKIT_REQUIRE_ISA=avx512bf16` pin: forces the `vdpbf16ps` dot kernel `Bf16DotGemm`, so its
//! multi-slice deep-k twin `Bf16DotGemmF32` gets exercised (`kc` rounded up to `DEPTH_MULTIPLE =
//! 2` there, so a k-pair never straddles a slice boundary). A separate test binary from the
//! other `env_isa_*` pins, since a memoized ISA choice needs its own process to force. Pins
//! through [`env_isa_common::pin`] (a single `set_var` under a `Once`, before any dispatch; the
//! shared write also overrides an inherited pin). The multi-slice route must be byte-for-byte
//! the single depth panel for `beta in {0, 1}`. Skips when the host lacks `avx512bf16`
#![cfg(all(
    feature = "half",
    feature = "std",
    any(target_arch = "x86", target_arch = "x86_64"),
    not(miri)
))]

// Shared GEMMKIT_REQUIRE_ISA pin helper; this binary's only test pins `avx512bf16` with it
mod env_isa_common;

use gemmkit::{MatMut, MatRef, Parallelism, bf16, gemm, tuning};

fn fill(n: usize, seed: u64) -> Vec<bf16> {
    let mut s = seed | 1;
    (0..n)
        .map(|_| {
            s ^= s << 13;
            s ^= s >> 7;
            s ^= s << 17;
            bf16::from_f32((s >> 40) as f32 / (1u64 << 24) as f32 - 0.5)
        })
        .collect()
}

// bf16 dot deep-k route matches the single depth panel for beta in {0, 1}, checked at k on
// and 1 past a slice boundary, serial and parallel

#[test]
fn deep_k_bf16_dot_under_avx512bf16_pin() {
    env_isa_common::pin("avx512bf16");
    if !(is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("avx512bf16")) {
        eprintln!("skipping: host does not report avx512f+avx512bf16");
        return;
    }
    let restore = tuning::deep_kc_bytes();

    // 4097 is odd, so both a slice boundary and the final tail land on odd depth: the
    // DEPTH_MULTIPLE rounding must keep every interior slice even and pad only the tail
    for &k in &[4096usize, 4097] {
        let (m, n) = (40usize, 50);
        let a = fill(m * k, 0x1);
        let b = fill(k * n, 0x2);
        let c0 = fill(m * n, 0x3);
        for &beta in &[0.0f32, 1.0] {
            let run = |engage: bool, par| -> Vec<u16> {
                tuning::set_deep_kc_bytes(if engage { 1 } else { usize::MAX });
                let mut c = c0.clone();
                gemm(
                    bf16::from_f32(1.25),
                    MatRef::from_col_major(&a, m, k),
                    MatRef::from_col_major(&b, k, n),
                    bf16::from_f32(beta),
                    MatMut::from_col_major(&mut c, m, n),
                    par,
                );
                c.iter().map(|v| v.to_bits()).collect()
            };
            let single = run(false, Parallelism::Serial);
            let deep_ser = run(true, Parallelism::Serial);
            let deep_par = run(true, Parallelism::Rayon(0));
            assert_eq!(
                deep_ser, single,
                "bf16 dot: deep-k != single panel (k={k} beta={beta})"
            );
            assert_eq!(
                deep_par, deep_ser,
                "bf16 dot: deep-k serial != parallel (k={k} beta={beta})"
            );
        }
    }

    tuning::set_deep_kc_bytes(restore);
}