gemmkit 0.1.1

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
81
82
83
84
85
86
87
88
89
90

//! Shared prepacked-i8 bit-parity check, used by the per-ISA pin binaries (`env_isa_avx512f`,
//! `env_isa_vnni`, `env_isa_scalar`)
//!
//! Each binary pins `GEMMKIT_REQUIRE_ISA` once, in its own process, before any dispatch runs,
//! then calls [`check`], which asserts `gemm_i8_packed_b` reproduces plain `gemm_i8`
//! **bit-identical** across a shape/layout/alpha-beta/thread-count sweep. The shape list mixes
//! in a `k` that is not a multiple of 4 and an `n` that is not a multiple of a kernel's panel
//! width, so the VNNI pin's depth-padding and the widen pin's panel-padding each get exercised
//! at their own edge
#![allow(dead_code)]

use gemmkit::{MatMut, MatRef, Parallelism};

fn rand_i8(n: usize, seed: u64) -> Vec<i8> {
    let mut s = seed.wrapping_add(0x9E3779B97F4A7C15);
    (0..n)
        .map(|_| {
            s ^= s << 13;
            s ^= s >> 7;
            s ^= s << 17;
            ((s >> 24) as i64 % 201 - 100) as i8
        })
        .collect()
}

/// Assert `gemm_i8_packed_b` reproduces plain `gemm_i8` bit-for-bit under whichever ISA the
/// calling binary pinned. Wrapping i32 accumulation is associative, so an exact match, not a
/// tolerance, is the right bar across every shape/layout/alpha-beta/thread count below
pub fn check() {
    for (m, k, n) in [
        (200usize, 130, 175),
        (65, 64, 64),
        (64, 65, 100), // k not a multiple of 4, n not a multiple of a panel width
        (33, 17, 19),
        (256, 257, 129),
        (8, 1024, 96), // small m, long k: the fixed-weights inference shape
        (2, 1023, 12),
    ] {
        let a = rand_i8(m * k, 0x51 + (m * 7 + n) as u64);
        let b_rm = rand_i8(k * n, 0x62 + (n * 3 + k) as u64); // logical k x n, row-major
        let mut b_cm = vec![0i8; k * n];
        for i in 0..k {
            for j in 0..n {
                b_cm[j * k + i] = b_rm[i * n + j];
            }
        }
        let c0: Vec<i32> = (0..m * n).map(|x| (x as i32 % 7) - 3).collect();

        for lb in 0..2 {
            let (bbuf, rsb, csb): (&[i8], isize, isize) = if lb == 0 {
                (&b_rm, n as isize, 1)
            } else {
                (&b_cm, 1, k as isize)
            };
            let bview = MatRef::new(bbuf, k, n, rsb, csb);
            let packed = gemmkit::prepack_rhs_i8(bview);

            for &(alpha, beta) in &[(1i32, 0i32), (2, 3), (0, 5)] {
                let mut c_ref = c0.clone();
                gemmkit::gemm_i8(
                    alpha,
                    MatRef::from_col_major(&a, m, k),
                    bview,
                    beta,
                    MatMut::from_col_major(&mut c_ref, m, n),
                    Parallelism::Serial,
                );
                for par in [
                    Parallelism::Serial,
                    Parallelism::Rayon(2),
                    Parallelism::Rayon(8),
                ] {
                    let mut c_pk = c0.clone();
                    gemmkit::gemm_i8_packed_b(
                        alpha,
                        MatRef::from_col_major(&a, m, k),
                        &packed,
                        beta,
                        MatMut::from_col_major(&mut c_pk, m, n),
                        par,
                    );
                    assert_eq!(
                        c_ref, c_pk,
                        "prepack_i8 != gemm_i8 for {m}x{k}x{n} lb={lb} a={alpha} b={beta} par={par:?}"
                    );
                }
            }
        }
    }
}