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

//! Wave 46 real-checkpoint bring-up: run the BF16 kernel
//! ladder (scalar / NEON / bulk-dequant→f32) against REAL
//! weight bytes from a BF16 safetensors checkpoint.
//!
//! ```bash
//! cargo run --release --example bf16_bringup -- path/to/model-00001-of-N.safetensors
//! ```
//!
//! For each of the largest BF16 matrices in the file:
//!   1. slice the raw bytes out of the mmap
//!   2. run a [1, k] × [n, k] matmul through all three routes
//!   3. cross-check the routes against each other (1e-4 rel)
//!   4. report wall-clock per route
//!
//! Exit code 1 if any route disagrees — this is a
//! verification harness, not a benchmark.

use jouleclaw_loader_gguf::dispatch::{read_any_model, AnyModel};
use jouleclaw_loader_gguf::matmul_bf16::{
    dequant_bf16_to_f32, matmul_bf16_neon, matmul_bf16_scalar,
};
use jouleclaw_loader_gguf::matmul_f::matmul_f32_scalar;
use jouleclaw_loader_gguf::safetensors_reader::SafetensorsDtype;
use std::path::PathBuf;
use std::time::Instant;

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = std::env::args().collect();
    if args.len() < 2 {
        eprintln!("usage: bf16_bringup <path.safetensors>");
        std::process::exit(1);
    }
    let path = PathBuf::from(&args[1]);
    println!("checkpoint: {}", path.display());

    let t0 = Instant::now();
    let model = read_any_model(&path).map_err(|e| format!("dispatch: {e}"))?;
    let st = match model {
        AnyModel::Safetensors(s) => s,
        other => {
            return Err(format!(
                "expected a safetensors checkpoint, dispatcher said {other:?}"
            ).into());
        }
    };
    println!("loaded via read_any_model in {:.3}s — {} tensors, {}",
        t0.elapsed().as_secs_f64(),
        st.tensors.len(),
        humanize(st.mmap.len() as u64),
    );

    // Pick the largest 2-D BF16 tensors.
    let mut bf16_mats: Vec<_> = st.tensors.iter()
        .filter(|t| t.dtype == SafetensorsDtype::BF16 && t.shape.len() == 2)
        .collect();
    bf16_mats.sort_by_key(|t| std::cmp::Reverse(t.data_len));
    if bf16_mats.is_empty() {
        return Err("no 2-D BF16 tensors in this file".into());
    }
    println!("\n{} 2-D BF16 matrices; verifying the 3 largest:\n", bf16_mats.len());

    let mut failures = 0u32;
    for t in bf16_mats.iter().take(3) {
        let n = t.shape[0] as usize;
        let k = t.shape[1] as usize;
        let w = st.tensor_bytes(t);
        println!("== {}  [{} x {}]  {}", t.name, n, k, humanize(t.data_len));

        // Synthetic but deterministic activation row.
        let m = 1usize;
        let a: Vec<f32> = (0..k).map(|i| ((i % 31) as f32 - 15.0) * 0.05).collect();

        // Route 1: fused scalar.
        let mut out_scalar = vec![0f32; m * n];
        let t1 = Instant::now();
        matmul_bf16_scalar(&a, w, &[], m, n, k, &mut out_scalar);
        let d_scalar = t1.elapsed();

        // Route 2: fused NEON.
        let mut out_neon = vec![0f32; m * n];
        let t2 = Instant::now();
        matmul_bf16_neon(&a, w, &[], m, n, k, &mut out_neon);
        let d_neon = t2.elapsed();

        // Route 3: bulk dequant → f32 matmul.
        let mut w_f32 = vec![0f32; n * k];
        let t3 = Instant::now();
        dequant_bf16_to_f32(w, &mut w_f32);
        let mut out_dequant = vec![0f32; m * n];
        matmul_f32_scalar(&a, &w_f32, &[], m, n, k, &mut out_dequant);
        let d_dequant = t3.elapsed();

        // Cross-check all three routes.
        let mut max_rel = 0f32;
        for j in 0..n {
            let s = out_scalar[j];
            for (label, v) in [("neon", out_neon[j]), ("dequant", out_dequant[j])] {
                let denom = s.abs().max(1.0);
                let rel = ((s - v) / denom).abs();
                if rel > max_rel {
                    max_rel = rel;
                }
                if rel > 1e-4 {
                    eprintln!("  MISMATCH col {j} {label}: scalar={s}, {label}={v}");
                    failures += 1;
                }
            }
        }
        println!("  scalar  {:>9.2} ms", d_scalar.as_secs_f64() * 1e3);
        println!("  neon    {:>9.2} ms   ({:.1}x vs scalar)",
            d_neon.as_secs_f64() * 1e3,
            d_scalar.as_secs_f64() / d_neon.as_secs_f64().max(1e-9));
        println!("  dequant {:>9.2} ms   (incl. widen of {})",
            d_dequant.as_secs_f64() * 1e3, humanize((n * k * 4) as u64));
        println!("  max relative divergence across routes: {max_rel:.2e}");
        // Sanity: output is finite and non-degenerate.
        let finite = out_scalar.iter().all(|v| v.is_finite());
        let nonzero = out_scalar.iter().any(|v| *v != 0.0);
        println!("  output finite: {finite}, non-degenerate: {nonzero}\n");
        if !finite || !nonzero {
            failures += 1;
        }
    }

    if failures > 0 {
        eprintln!("FAILED: {failures} mismatches");
        std::process::exit(1);
    }
    println!("ALL ROUTES AGREE on real checkpoint bytes.");
    Ok(())
}

fn humanize(n: u64) -> String {
    const KB: u64 = 1024;
    const MB: u64 = KB * 1024;
    const GB: u64 = MB * 1024;
    if n >= GB {
        format!("{:.2} GiB", n as f64 / GB as f64)
    } else if n >= MB {
        format!("{:.2} MiB", n as f64 / MB as f64)
    } else if n >= KB {
        format!("{:.2} KiB", n as f64 / KB as f64)
    } else {
        format!("{} B", n)
    }
}