open-standards/jouleclaw/jouleclaw-rs/crates/jouleclaw-loader-gguf/examples/mpp_bt_bench.rs

//! Prefill GEMM throughput: the square-tiling Metal-4 matmul2d
//! wrapper (`matmul_bf16_mpp_bt`) vs the threaded-NEON CPU path,
//! at realistic projection shapes. Run with JOULECLAW_FABRIC=1.
#[cfg(feature = "fabric")]
fn main() {
    use std::time::Instant;
    fn f32_to_bf16(v: f32) -> u16 {
        let b = v.to_bits();
        ((b + (((b >> 16) & 1) + 0x7fff)) >> 16) as u16
    }
    // (M tokens, N out, K in)
    let shapes = [
        (128usize, 4096usize, 4096usize),
        (256, 4096, 4096),
        (256, 11008, 4096),
        (512, 4096, 4096),
    ];
    println!("shape (MxNxK)        mpp_bt ms   GFLOP/s    neon ms   GFLOP/s   speedup  max_rel");
    for &(m, n, k) in &shapes {
        let a: Vec<f32> = (0..m * k).map(|i| ((i % 101) as f32 / 50.0) - 1.0).collect();
        let w: Vec<u8> = (0..n * k)
            .flat_map(|i| f32_to_bf16(((i % 97) as f32 / 48.0 - 1.0) * 0.1).to_le_bytes())
            .collect();
        let gflop = 2.0 * m as f64 * n as f64 * k as f64 / 1e9;

        let mut got = vec![0f32; m * n];
        if !jouleclaw_loader_gguf::fabric_dispatch::matmul_bf16_mpp_bt(&a, &w, m, n, k, &mut got) {
            println!("  {m}x{n}x{k}: mpp_bt fell back (fabric off / pre-Metal-4)");
            continue;
        }
        let iters = 10;
        let t0 = Instant::now();
        for _ in 0..iters {
            jouleclaw_loader_gguf::fabric_dispatch::matmul_bf16_mpp_bt(&a, &w, m, n, k, &mut got);
        }
        let mpp_ms = t0.elapsed().as_secs_f64() * 1e3 / iters as f64;

        let mut neon = vec![0f32; m * n];
        let t1 = Instant::now();
        for _ in 0..iters {
            jouleclaw_loader_gguf::matmul_bf16::matmul_bf16_threaded(&a, &w, &[], m, n, k, &mut neon)
                .expect("threaded bf16 matmul");
        }
        let neon_ms = t1.elapsed().as_secs_f64() * 1e3 / iters as f64;

        let mut max_rel = 0f32;
        for (g, r) in got.iter().zip(neon.iter()) {
            max_rel = max_rel.max((g - r).abs() / (r.abs() + 1e-2));
        }
        println!(
            "  {m:>4}x{n:>5}x{k:>5}   {mpp_ms:8.3}  {:8.1}   {neon_ms:8.3}  {:8.1}   {:5.2}x   {max_rel:.2e}",
            gflop / (mpp_ms / 1e3),
            gflop / (neon_ms / 1e3),
            neon_ms / mpp_ms,
        );
    }
}
#[cfg(not(feature = "fabric"))]
fn main() { println!("build with --features fabric"); }