a2f-mlx 0.2.0

//! Micro-benchmark for a2f-mlx: streaming per-window latency, batched
//! throughput, MLX memory, and full-utterance output dump for accuracy diffs.
//!
//!   cargo run --release -p a2f-mlx --example bench -- BUNDLE WAV [--batch N]
//!       [--precision f32|f16|bf16|q8|q4] [--compile] [--out coefs.npy] [--iters N]

use std::path::PathBuf;
use std::time::Instant;

use anyhow::{Context, Result};
use a2f_core::{read_wav_16k_mono, write_npy_f32, BUFFER_LEN, FRAME_HOP, HOP_LEN, OUT_COUNT, SR};
use a2f_mlx::{A2FMlx, A2FMlxCompiled, Options};

enum Model {
    Plain(A2FMlx),
    Compiled(A2FMlxCompiled),
}

impl Model {
    fn base(&self) -> &A2FMlx {
        match self { Model::Plain(m) => m, Model::Compiled(c) => c.model() }
    }
    fn coefficients_batch(&self, w: &[&[f32]], e: &a2f_core::Emotion) -> Result<Vec<[f32; OUT_COUNT]>> {
        match self { Model::Plain(m) => m.coefficients_batch(w, e), Model::Compiled(c) => c.coefficients_batch(w, e) }
    }
    fn coefficients(&self, w: &[f32], e: &a2f_core::Emotion) -> Result<Vec<f32>> {
        Ok(self.coefficients_batch(&[w], e)?.remove(0).to_vec())
    }
}

fn windows(audio: &[f32]) -> Vec<Vec<f32>> {
    let mut out = Vec::new();
    let mut target = 0usize;
    loop {
        let start = target as isize - HOP_LEN as isize;
        let mut w = vec![0f32; BUFFER_LEN];
        let lo = start.max(0) as usize;
        let hi = (start + BUFFER_LEN as isize).min(audio.len() as isize) as usize;
        if hi > lo {
            w[(lo as isize - start) as usize..(hi as isize - start) as usize]
                .copy_from_slice(&audio[lo..hi]);
        }
        out.push(w);
        if target >= audio.len().max(1) {
            break;
        }
        target += FRAME_HOP;
    }
    out
}

fn pct(v: &mut [f64], p: f64) -> f64 {
    v.sort_by(|a, b| a.partial_cmp(b).unwrap());
    v[((v.len() - 1) as f64 * p).round() as usize]
}

fn main() -> Result<()> {
    let args: Vec<String> = std::env::args().collect();
    let bundle = PathBuf::from(args.get(1).context("BUNDLE")?);
    let wav = PathBuf::from(args.get(2).context("WAV")?);
    let mut opts = Options::default();
    let mut compile = false;
    let mut batch = 1usize;
    let mut out = None;
    let mut iters = 3usize;
    let mut stages = false;
    let mut i = 3;
    while i < args.len() {
        match args[i].as_str() {
            "--batch" => { batch = args[i + 1].parse()?; i += 1; }
            "--precision" => { opts.precision = args[i + 1].parse()?; i += 1; }
            "--compile" => compile = true,
            "--stages" => stages = true,
            "--out" => { out = Some(PathBuf::from(&args[i + 1])); i += 1; }
            "--iters" => { iters = args[i + 1].parse()?; i += 1; }
            a => anyhow::bail!("unknown arg {a}"),
        }
        i += 1;
    }

    let audio = read_wav_16k_mono(&wav)?;
    let dur = audio.len() as f64 / SR as f64;
    mlx_rs::memory::reset_peak_memory()?;
    let t = Instant::now();
    let plain = A2FMlx::load_with(&bundle, None, opts)?;
    let model = if compile { Model::Compiled(plain.compiled()) } else { Model::Plain(plain) };
    let load = t.elapsed().as_secs_f64();
    let emo = model.base().default_emotion().clone();
    let wins = windows(&audio);
    let refs: Vec<&[f32]> = wins.iter().map(|w| w.as_slice()).collect();

    // warmup (first-call kernel JIT / compile tracing), one per batch shape used
    model.coefficients(refs[0], &emo)?;
    model.coefficients_batch(&refs[..batch.min(refs.len())], &emo)?;
    if refs.len() % batch != 0 {
        model.coefficients_batch(&refs[..refs.len() % batch], &emo)?;
    }

    if stages {
        let mut flat = Vec::new();
        for w in &refs[100..100 + batch] {
            flat.extend(w.iter().map(|s| s * model.base().input_strength()));
        }
        let a = mlx_rs::Array::from_slice(&flat, &[batch as i32, 1, BUFFER_LEN as i32]);
        let e: Vec<f32> = (0..batch).flat_map(|_| emo.0.clone()).collect();
        let e = mlx_rs::Array::from_slice(&e, &[batch as i32, 1, a2f_core::N_EMO as i32]);
        for (name, ms) in model.base().stage_times(&a, &e, 21)? {
            println!("  stage {name:<14} {ms:7.2} ms  (batch {batch})");
        }
        let mut v: Vec<f64> = (0..21).map(|_| {
            let t = Instant::now();
            model.coefficients_batch(&refs[100..100 + batch], &emo).unwrap();
            t.elapsed().as_secs_f64() * 1e3
        }).collect();
        println!("  stage {:<14} {:7.2} ms  (batch {batch})", "coefs_batch", pct(&mut v, 0.5));
        let mut v: Vec<f64> = (0..21).map(|_| {
            let t = Instant::now();
            model.base().forward(&a, &e).unwrap().eval().unwrap();
            t.elapsed().as_secs_f64() * 1e3
        }).collect();
        println!("  stage {:<14} {:7.2} ms  (batch {batch})", "forward_f32in", pct(&mut v, 0.5));
    }

    mlx_rs::memory::clear_cache()?;
    let weights_mb = mlx_rs::memory::active_memory()? as f64 / 1e6; // resident after warmup

    // full utterance, batched
    let mut walls = Vec::new();
    let mut coefs = Vec::new();
    for _ in 0..iters {
        coefs.clear();
        let t = Instant::now();
        for chunk in refs.chunks(batch) {
            for f in model.coefficients_batch(chunk, &emo)? {
                coefs.extend_from_slice(&f);
            }
        }
        walls.push(t.elapsed().as_secs_f64());
    }
    let wall = pct(&mut walls, 0.5);

    // streaming latency: one call of `batch` windows (what a consumer pays per tick)
    let mut lat = Vec::new();
    for k in 0..60 {
        let s = (k * 3) % (refs.len() - batch);
        let t = Instant::now();
        model.coefficients_batch(&refs[s..s + batch], &emo)?;
        lat.push(t.elapsed().as_secs_f64() * 1e3);
    }
    let (p50, p95) = (pct(&mut lat, 0.5), pct(&mut lat, 0.95));
    let peak_mb = mlx_rs::memory::peak_memory()? as f64 / 1e6;

    println!(
        "precision={:?} compile={} batch={batch} frames={} load={load:.2}s weights={weights_mb:.1}MB peak={peak_mb:.1}MB \
         wall={wall:.3}s RTF={:.4} perFrame={:.2}ms callLatency p50={p50:.2}ms p95={p95:.2}ms",
        opts.precision, compile, refs.len(), wall / dur, wall * 1e3 / refs.len() as f64,
    );
    if let Some(out) = out {
        write_npy_f32(&out, &[refs.len(), OUT_COUNT], &coefs)?;
    }
    Ok(())
}