matmul 0.1.0

Fast matrix multiplication in Rust with AVX2/AVX-512 SIMD, built from scratch
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

//! Fast matrix multiplication in Rust, built from scratch.
//!
//! I built this to understand what makes BLAS fast. Turns out it's:
//! cache blocking, SIMD intrinsics, and FMA instructions. This crate
//! implements all of that, achieving ~62% of NumPy/OpenBLAS performance.
//!
//! ## Usage
//!
//! ```
//! use matmul::multiply;
//!
//! let a = vec![1.0f64; 256 * 256];
//! let b = vec![1.0f64; 256 * 256];
//! let mut c = vec![0.0f64; 256 * 256];
//!
//! multiply(&a, &b, &mut c, 256, 256, 256);
//! ```
//!
//! For large matrices, use the multi-threaded version:
//!
//! ```
//! use matmul::multiply_parallel;
//!
//! let a = vec![1.0f64; 1024 * 1024];
//! let b = vec![1.0f64; 1024 * 1024];
//! let mut c = vec![0.0f64; 1024 * 1024];
//!
//! multiply_parallel(&a, &b, &mut c, 1024, 1024, 1024, 4);
//! ```
//!
//! ## What's inside
//!
//! - 4x4, 12x4 AVX2 kernels
//! - 8x8 AVX-512 kernel
//! - Cache blocking tuned for L1/L2
//! - Adaptive multi-threading (scales down for small matrices)

pub mod blocked;
pub mod kernels;
pub mod matrix;
pub mod threaded;

pub use matrix::naive_ijk::matmul_naive_ijk;
pub use matrix::naive_ikj::matmul_naive_ikj;

/// Matrix multiply: C += A * B
///
/// Picks the fastest available kernel for your CPU (AVX-512 > AVX2 > scalar).
/// Matrices are row-major: A is m×k, B is k×n, C is m×n.
///
/// # Panics
///
/// Panics if the slice sizes don't match m, n, k.
pub fn multiply(a: &[f64], b: &[f64], c: &mut [f64], m: usize, n: usize, k: usize) {
    assert_eq!(a.len(), m * k, "A: expected {}x{}={} elements", m, k, m * k);
    assert_eq!(b.len(), k * n, "B: expected {}x{}={} elements", k, n, k * n);
    assert_eq!(c.len(), m * n, "C: expected {}x{}={} elements", m, n, m * n);

    #[cfg(target_arch = "x86_64")]
    {
        if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("fma") {
            unsafe { blocked::gemm_8x8::matmul_blocked_8x8(a, b, c, m, n, k, None, None) };
            return;
        }
        if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
            unsafe { blocked::gemm_12x4::matmul_blocked_12x4(a, b, c, m, n, k, None, None) };
            return;
        }
    }

    matrix::naive_ikj::matmul_naive_ikj(a, b, c, m, n, k);
}

/// Same as [`multiply`] but uses multiple threads.
///
/// Thread count adapts to matrix size - small matrices use fewer threads
/// because the overhead isn't worth it.
pub fn multiply_parallel(
    a: &[f64],
    b: &[f64],
    c: &mut [f64],
    m: usize,
    n: usize,
    k: usize,
    num_threads: usize,
) {
    assert_eq!(a.len(), m * k, "A: expected {}x{}={} elements", m, k, m * k);
    assert_eq!(b.len(), k * n, "B: expected {}x{}={} elements", k, n, k * n);
    assert_eq!(c.len(), m * n, "C: expected {}x{}={} elements", m, n, m * n);

    #[cfg(target_arch = "x86_64")]
    {
        if is_x86_feature_detected!("avx512f") && is_x86_feature_detected!("fma") {
            threaded::gemm_8x8_mt::matmul_blocked_8x8_mt(a, b, c, m, n, k, num_threads);
            return;
        }
        if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
            threaded::gemm_12x4_mt::matmul_blocked_12x4_mt(a, b, c, m, n, k, num_threads);
            return;
        }
    }

    matrix::naive_ikj::matmul_naive_ikj(a, b, c, m, n, k);
}