//! Criterion bench for `matmul_q8_0_scalar`.
//!
//! Measures steady-state Q8_0 matmul throughput at sizes
//! representative of a small Llama layer's attention + FFN
//! projections. Numbers are absolute (ns / GFLOPS) — the
//! wave-12 prefetch wins show up as steady-state throughput
//! that's stable as `n` grows past the L2 cache.

// Benches and integration tests are their own crates, so the decision
// recorded in src/lib.rs does not reach them: these loops fill arrays by
// index, which is what the index is for.
#![allow(clippy::needless_range_loop)]

use std::hint::black_box;

use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion, Throughput};
use jouleclaw_loader_gguf::matmul_q8_0::{
    matmul_q8_0, matmul_q8_0_scalar, Q8_0_BLOCK_BYTES, Q8_0_ELEMS,
};

/// Build a synthetic Q8_0 weight buffer of shape `[n, k]`.
/// Bytes are bit-pattern-derived from indices so output is
/// deterministic.
fn synth_weights(n: usize, k: usize) -> Vec<u8> {
    let blocks_per_row = k / Q8_0_ELEMS;
    let row_bytes = blocks_per_row * Q8_0_BLOCK_BYTES;
    let total = n * row_bytes;
    let mut w = vec![0u8; total];
    for i in 0..total {
        w[i] = ((i.wrapping_mul(2654435761)) >> 24) as u8;
    }
    w
}

/// Build a synthetic activation buffer of shape `[m, k]`
/// (i8) + per-row scales.
fn synth_activations(m: usize, k: usize) -> (Vec<i8>, Vec<f32>) {
    let mut q = vec![0i8; m * k];
    for i in 0..(m * k) {
        q[i] = (((i as i32).wrapping_mul(48271)) >> 16) as i8;
    }
    let scales = vec![0.01_f32; m];
    (q, scales)
}

fn bench_matmul_q8_0_sizes(c: &mut Criterion) {
    let mut g = c.benchmark_group("matmul_q8_0_scalar");
    // Sizes inspired by a small Llama / Gemma layer: hidden
    // dim ≈ k, output dim ≈ n. m=1 mimics single-token
    // generation (the latency-sensitive case).
    let cases = [
        // (m, n, k) — m rows of activations, n output cols, k inner dim.
        (1, 256, 256),
        (1, 1024, 1024),
        (1, 2048, 2048),
        (1, 4096, 4096),
        // Batched: m > 1 stresses the row outer loop.
        (8, 1024, 1024),
        (32, 1024, 1024),
    ];
    for &(m, n, k) in &cases {
        let w = synth_weights(n, k);
        let (qa, a_s) = synth_activations(m, k);
        let mut out = vec![0f32; m * n];
        let bias: Vec<f32> = Vec::new();
        let label = format!("m={m},n={n},k={k}");
        // FLOPs ≈ 2 * m * n * k (one multiply + one add per
        // inner-product element).
        let flops = 2.0 * (m * n * k) as f64;
        g.throughput(Throughput::Elements(flops as u64));
        g.bench_with_input(
            BenchmarkId::from_parameter(label),
            &(m, n, k),
            |b, _| {
                b.iter(|| {
                    matmul_q8_0_scalar(
                        black_box(&qa),
                        black_box(&a_s),
                        black_box(&w),
                        black_box(&bias),
                        m,
                        n,
                        k,
                        black_box(&mut out),
                    );
                });
            },
        );
    }
    g.finish();
}

/// Bench the dispatching `matmul_q8_0` entry point — on
/// aarch64 this is the NEON fast path; on x86_64 it falls
/// back to the scalar kernel. Side-by-side with
/// `matmul_q8_0_scalar` reveals the NEON win directly.
///
/// Note this entry point also pays for `quantize_a_per_row`
/// (fp32 activation → i8 + per-row scale) per call, while
/// the scalar bench above takes pre-quantized inputs. The
/// quantize cost is real and should be included in
/// production comparisons.
fn bench_matmul_q8_0_dispatch_sizes(c: &mut Criterion) {
    let mut g = c.benchmark_group("matmul_q8_0_dispatch");
    let cases = [
        (1, 256, 256),
        (1, 1024, 1024),
        (1, 2048, 2048),
        (1, 4096, 4096),
        (8, 1024, 1024),
        (32, 1024, 1024),
    ];
    for &(m, n, k) in &cases {
        let w = synth_weights(n, k);
        // Build fp32 activations directly — the dispatch
        // entry point quantizes internally.
        let mut a = vec![0f32; m * k];
        for i in 0..(m * k) {
            a[i] = (((i as i32).wrapping_mul(48271)) >> 16) as f32 * 0.01;
        }
        let mut out = vec![0f32; m * n];
        let bias: Vec<f32> = Vec::new();
        let label = format!("m={m},n={n},k={k}");
        let flops = 2.0 * (m * n * k) as f64;
        g.throughput(Throughput::Elements(flops as u64));
        g.bench_with_input(
            BenchmarkId::from_parameter(label),
            &(m, n, k),
            |b, _| {
                b.iter(|| {
                    matmul_q8_0(
                        black_box(&a),
                        black_box(&w),
                        black_box(&bias),
                        m, n, k,
                        black_box(&mut out),
                    );
                });
            },
        );
    }
    g.finish();
}

/// Bench the f32 matmul against Apple Accelerate (AMX) when
/// the `accelerate` feature is on. Skips cleanly otherwise.
/// Direct apples-to-apples vs `matmul_f32_scalar` at the same
/// shapes — reveals the AMX win without the Q8_0 quantization
/// cost confounding the measurement.
fn bench_matmul_f32_paths(c: &mut Criterion) {
    use jouleclaw_loader_gguf::matmul_f::{
        matmul_f32_accelerate, matmul_f32_scalar,
    };
    let cases = [
        (1, 256, 256),
        (1, 1024, 1024),
        (1, 2048, 2048),
        (1, 4096, 4096),
        (8, 1024, 1024),
        (32, 1024, 1024),
    ];
    let mut g_scalar = c.benchmark_group("matmul_f32_scalar");
    for &(m, n, k) in &cases {
        let a: Vec<f32> = (0..(m * k)).map(|i| (i % 17) as f32 * 0.01).collect();
        let w: Vec<f32> =
            (0..(n * k)).map(|i| ((i + 99) % 13) as f32 * 0.01).collect();
        let mut out = vec![0f32; m * n];
        let bias: Vec<f32> = Vec::new();
        let label = format!("m={m},n={n},k={k}");
        let flops = 2.0 * (m * n * k) as f64;
        g_scalar.throughput(Throughput::Elements(flops as u64));
        g_scalar.bench_with_input(
            BenchmarkId::from_parameter(label),
            &(m, n, k),
            |b, _| {
                b.iter(|| {
                    matmul_f32_scalar(
                        black_box(&a), black_box(&w), black_box(&bias),
                        m, n, k, black_box(&mut out),
                    );
                });
            },
        );
    }
    g_scalar.finish();
    // Accelerate path — short-circuits on non-aarch64-macOS /
    // no-feature builds via the Err(Unsupported) return.
    let mut probe = vec![0f32; 1];
    let probe_ok = matmul_f32_accelerate(
        &[0f32; 1], &[0f32; 1], &[], 1, 1, 1, &mut probe,
    )
    .is_ok();
    if !probe_ok {
        return;
    }
    let mut g_acc = c.benchmark_group("matmul_f32_accelerate");
    for &(m, n, k) in &cases {
        let a: Vec<f32> = (0..(m * k)).map(|i| (i % 17) as f32 * 0.01).collect();
        let w: Vec<f32> =
            (0..(n * k)).map(|i| ((i + 99) % 13) as f32 * 0.01).collect();
        let mut out = vec![0f32; m * n];
        let bias: Vec<f32> = Vec::new();
        let label = format!("m={m},n={n},k={k}");
        let flops = 2.0 * (m * n * k) as f64;
        g_acc.throughput(Throughput::Elements(flops as u64));
        g_acc.bench_with_input(
            BenchmarkId::from_parameter(label),
            &(m, n, k),
            |b, _| {
                b.iter(|| {
                    let _ = matmul_f32_accelerate(
                        black_box(&a), black_box(&w), black_box(&bias),
                        m, n, k, black_box(&mut out),
                    );
                });
            },
        );
    }
    g_acc.finish();
}

/// Bench the wave-24 AMX-tiled Q8_0 matmul against the
/// dispatching NEON Q8_0 entry point. Reveals whether the
/// dequant-into-scratch cost is amortized by the AMX speedup.
fn bench_matmul_q8_0_amx_tiled(c: &mut Criterion) {
    use jouleclaw_loader_gguf::matmul_q8_0::matmul_q8_0_amx_tiled;
    // Skip cleanly when not on macOS aarch64 + accelerate.
    // Probe: m=1, n=1, k=32 → a len 32, w len 34 (1 block), out len 1.
    let mut probe_out = vec![0f32; 1];
    let probe_a = vec![0f32; 32];
    let probe_w = vec![0u8; Q8_0_BLOCK_BYTES];
    if matmul_q8_0_amx_tiled(
        &probe_a, &probe_w, &[], 1, 1, 32, &mut probe_out,
    )
    .is_err()
    {
        return;
    }
    let mut g = c.benchmark_group("matmul_q8_0_amx_tiled");
    let cases = [
        (1, 256, 256),
        (1, 1024, 1024),
        (1, 2048, 2048),
        (1, 4096, 4096),
        (8, 1024, 1024),
        (32, 1024, 1024),
    ];
    for &(m, n, k) in &cases {
        let w = synth_weights(n, k);
        let a: Vec<f32> = (0..(m * k)).map(|i| (i % 17) as f32 * 0.01).collect();
        let mut out = vec![0f32; m * n];
        let bias: Vec<f32> = Vec::new();
        let label = format!("m={m},n={n},k={k}");
        let flops = 2.0 * (m * n * k) as f64;
        g.throughput(Throughput::Elements(flops as u64));
        g.bench_with_input(
            BenchmarkId::from_parameter(label),
            &(m, n, k),
            |b, _| {
                b.iter(|| {
                    let _ = matmul_q8_0_amx_tiled(
                        black_box(&a),
                        black_box(&w),
                        black_box(&bias),
                        m, n, k,
                        black_box(&mut out),
                    );
                });
            },
        );
    }
    g.finish();
}

/// Bench the wave-25 multi-thread Q8_0 matmul. Compares
/// matmul_q8_0_threaded against the single-thread NEON
/// dispatch.
fn bench_matmul_q8_0_threaded(c: &mut Criterion) {
    use jouleclaw_loader_gguf::matmul_q8_0::{matmul_q8_0, matmul_q8_0_threaded};
    // Skip cleanly when rayon feature is off.
    let mut probe_out = vec![0f32; 1];
    if matmul_q8_0_threaded(
        &[0f32; 32], &[0u8; Q8_0_BLOCK_BYTES], &[],
        1, 1, 32, &mut probe_out,
    )
    .is_err()
    {
        return;
    }
    let mut g = c.benchmark_group("matmul_q8_0_threaded");
    let cases = [
        (1, 256, 256),
        (1, 1024, 1024),
        (1, 2048, 2048),
        (1, 4096, 4096),
        (8, 1024, 1024),
        (32, 1024, 1024),
    ];
    for &(m, n, k) in &cases {
        let w = synth_weights(n, k);
        let a: Vec<f32> = (0..(m * k)).map(|i| (i % 17) as f32 * 0.01).collect();
        let mut out = vec![0f32; m * n];
        let bias: Vec<f32> = Vec::new();
        let label = format!("m={m},n={n},k={k}");
        let flops = 2.0 * (m * n * k) as f64;
        g.throughput(Throughput::Elements(flops as u64));
        g.bench_with_input(
            BenchmarkId::from_parameter(label),
            &(m, n, k),
            |b, _| {
                b.iter(|| {
                    let _ = matmul_q8_0_threaded(
                        black_box(&a),
                        black_box(&w),
                        black_box(&bias),
                        m, n, k,
                        black_box(&mut out),
                    );
                });
            },
        );
    }
    // Silence unused-import warning when matmul_q8_0 isn't called from any cfg.
    let _ = matmul_q8_0;
    g.finish();
}

/// Bench the threaded AMX-tiled Q8_0 matmul (rayon over chunks).
fn bench_matmul_q8_0_amx_tiled_threaded(c: &mut Criterion) {
    use jouleclaw_loader_gguf::matmul_q8_0::matmul_q8_0_amx_tiled_threaded;
    let mut probe_out = vec![0f32; 1];
    if matmul_q8_0_amx_tiled_threaded(
        &[0f32; 32], &[0u8; Q8_0_BLOCK_BYTES], &[],
        1, 1, 32, &mut probe_out,
    )
    .is_err()
    {
        return;
    }
    let mut g = c.benchmark_group("matmul_q8_0_amx_tiled_threaded");
    let cases = [
        (1, 256, 256),
        (1, 1024, 1024),
        (1, 2048, 2048),
        (1, 4096, 4096),
        (8, 1024, 1024),
        (32, 1024, 1024),
    ];
    for &(m, n, k) in &cases {
        let w = synth_weights(n, k);
        let a: Vec<f32> = (0..(m * k)).map(|i| (i % 17) as f32 * 0.01).collect();
        let mut out = vec![0f32; m * n];
        let bias: Vec<f32> = Vec::new();
        let label = format!("m={m},n={n},k={k}");
        let flops = 2.0 * (m * n * k) as f64;
        g.throughput(Throughput::Elements(flops as u64));
        g.bench_with_input(
            BenchmarkId::from_parameter(label),
            &(m, n, k),
            |b, _| {
                b.iter(|| {
                    let _ = matmul_q8_0_amx_tiled_threaded(
                        black_box(&a),
                        black_box(&w),
                        black_box(&bias),
                        m, n, k,
                        black_box(&mut out),
                    );
                });
            },
        );
    }
    g.finish();
}

criterion_group!(
    benches,
    bench_matmul_q8_0_sizes,
    bench_matmul_q8_0_dispatch_sizes,
    bench_matmul_f32_paths,
    bench_matmul_q8_0_amx_tiled,
    bench_matmul_q8_0_threaded,
    bench_matmul_q8_0_amx_tiled_threaded,
);
criterion_main!(benches);