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

//! Honest µJ/GFLOP of the Metal-4 matmul2d prefill path vs the
//! threaded-NEON CPU path.
//! `cargo run --release --features "rayon fabric fabric-meter" --example mpp_energy`
//!
//! Two meters, chosen as the live autotuner chooses (`fabric_dispatch::tune_shape_live`):
//!
//! - **The IOReport SoC counter, where its CPU energy rail is live.** Run on a QUIET
//!   box: marginal energy = ΔE_busy − P_idle·Δt. It writes nothing, and exits non-zero
//!   saying why, unless both figures are measurements: an idle baseline that advances,
//!   and a positive marginal for the GPU **and** the CPU run. A CPU run that draws
//!   nothing above idle means the meter cannot see CPU work, so its SoC figure for the
//!   GPU run is not a whole-SoC figure either.
//! - **This process's own CPU and GPU counters, where it is not** (macOS 27.0.1: the
//!   rail is frozen for an unprivileged reader): each run is one `ProcessSpanMeter`
//!   span, this process's CPU energy plus its GPU energy, no idle to subtract
//!   (`marginal_energy::measure_member_process`). The GPU run's figure is written only
//!   from a span metered whole — the GPU rail while this process was the GPU's sole
//!   user — else nothing is written. The CPU run (its GPU share must be an exact zero)
//!   stands on its own span and gives the ratio; its refusal is printed and does not
//!   void the GPU figure, which needs no idle baseline to be this process's alone.
#[cfg(all(feature = "fabric", feature = "fabric-meter"))]
fn main() {
    use jouleclaw_energy::span::ProcessSpanMeter;
    use jouleclaw_energy::EnergyCounter;
    use jouleclaw_loader_gguf::marginal_energy::{
        calibration_row, energy_ratio, idle_draw_uw, marginal_uj, measure_member_process, process_method, refusal,
        MemberRuns, NotMeasured,
    };
    use std::time::{Duration, Instant};
    fn f32_to_bf16(v: f32) -> u16 { let b = v.to_bits(); ((b + (((b >> 16) & 1) + 0x7fff)) >> 16) as u16 }
    /// Write nothing: say why, with the IOReport status, on stderr and exit non-zero.
    fn refuse(path: &str, why: impl std::fmt::Display) -> ! {
        eprintln!("{}", refusal(path, why, jouleclaw_energy::ioreport_status()));
        std::process::exit(1);
    }

    let path = jouleclaw_fabric::calibration_file_path().display().to_string();
    let counter = jouleclaw_energy::apple_ioreport::AppleIoReportCounter::new();
    if !counter.is_available() {
        return process_counters(&path);
    }
    let read_uj = || counter.read().map(|r| r.uj).map_err(|e| NotMeasured::ReadFailed(e.to_string()));
    let idle_uw = {
        let e0 = read_uj().unwrap_or_else(|e| refuse(&path, e));
        let t0 = Instant::now();
        std::thread::sleep(Duration::from_secs(3));
        let e1 = read_uj().unwrap_or_else(|e| refuse(&path, e));
        idle_draw_uw(e0, e1, t0.elapsed().as_secs_f64()).unwrap_or_else(|e| refuse(&path, e))
    };
    println!("idle draw: {:.2} W (machine-wide SoC)", idle_uw / 1e6);

    // Prefill projection shape.
    let (m, n, k) = (256usize, 4096usize, 4096usize);
    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;

    // run a workload for ~`secs` and return (µJ/GFLOP, GFLOP/s), or why the meter
    // produced no figure for it.
    let measure = |secs: f64, mut run: Box<dyn FnMut()>| -> Result<(f64, f64), NotMeasured> {
        run(); // warm
        let e0 = read_uj()?;
        let t0 = Instant::now();
        let mut ops = 0u64;
        while t0.elapsed().as_secs_f64() < secs {
            run();
            ops += 1;
        }
        let dt = t0.elapsed().as_secs_f64();
        let e1 = read_uj()?;
        let total_gf = gflop * ops as f64;
        let marginal = marginal_uj(e0, e1, dt, idle_uw)?;
        Ok((marginal / total_gf, total_gf / dt))
    };

    let mut out = vec![0f32; m * n];
    let optr = out.as_mut_ptr();
    let (a_p, w_p) = (&a, &w);
    let gpu = measure(3.0, Box::new(move || {
        let o = unsafe { std::slice::from_raw_parts_mut(optr, m * n) };
        let ok = jouleclaw_loader_gguf::fabric_dispatch::matmul_bf16_mpp_bt(a_p, w_p, m, n, k, o);
        assert!(ok, "mpp_bt fell back — set JOULECLAW_FABRIC=1 on a Metal-4 box");
    }));

    let mut out2 = vec![0f32; m * n];
    let o2ptr = out2.as_mut_ptr();
    let cpu = measure(3.0, Box::new(move || {
        let o = unsafe { std::slice::from_raw_parts_mut(o2ptr, m * n) };
        // Assert it ran, as the GPU run does: a refused shape would meter an empty loop.
        jouleclaw_loader_gguf::matmul_bf16::matmul_bf16_threaded(a_p, w_p, &[], m, n, k, o)
            .expect("threaded-NEON matmul refused this shape");
    }));

    println!("shape {m}x{n}x{k}");
    for (label, r) in [("matmul2d (metal-gpu)", &gpu), ("threaded-NEON (cpu) ", &cpu)] {
        match r {
            Ok((ujpgf, gfps)) => println!("  {label}: {ujpgf:>9.0} µJ/GFLOP   {gfps:>7.1} GFLOP/s"),
            Err(e) => println!("  {label}: not measured — {e}"),
        }
    }
    let ratio = match (&gpu, &cpu) {
        (Ok((gpu_ujpgf, _)), Ok((cpu_ujpgf, _))) => energy_ratio(*cpu_ujpgf, *gpu_ujpgf),
        _ => None,
    };
    match ratio {
        Some(ratio) => println!("  energy ratio (CPU/GPU per GFLOP): {ratio:.2}×"),
        None => println!("  energy ratio (CPU/GPU per GFLOP): n/a — both runs must be measured"),
    }
    let gpu_ujpgf = match (gpu, cpu) {
        (Ok((gpu_ujpgf, _)), Ok(_)) => gpu_ujpgf,
        (Err(e), _) => refuse(&path, format!("the matmul2d (metal-gpu) run: {e}")),
        (_, Err(e)) => refuse(&path, format!("the threaded-NEON (cpu) run: {e}")),
    };

    // Overlay the measured matmul2d µJ/GFLOP into this box's
    // calibration file under local/metal-mpp, so the placer +
    // ledger price the prefill path at its TRUE cost (not the
    // borrowed MPSMatrix metal-gpu figure). Merge, don't clobber.
    let row = calibration_row(
        "local/metal-mpp",
        "matmul",
        gpu_ujpgf,
        counter.provenance(),
        format!(
            "IOReport SoC ΔE − idle({:.2} W)·Δt over matmul2d {m}×{n}×{k} on a quiet box; IOReport {:?}",
            idle_uw / 1e6,
            jouleclaw_energy::ioreport_status()
        ),
    );
    write_row(&path, row);

    /// Merge `row` into the calibration file under its member, replacing that member's
    /// row and leaving every other row as it was (merge, don't clobber).
    fn write_row(path: &str, row: jouleclaw_fabric::CalibrationEntry) {
        let mut entries: Vec<jouleclaw_fabric::CalibrationEntry> = std::fs::read(path)
            .ok()
            .and_then(|b| serde_json::from_slice(&b).ok())
            .unwrap_or_default();
        entries.retain(|e| e.member_id != row.member_id);
        let (member, ujpgf, grade) = (row.member_id.clone(), row.uj_per_gflop, row.loaded_provenance());
        entries.push(row);
        if let Some(dir) = std::path::Path::new(path).parent() {
            let _ = std::fs::create_dir_all(dir);
        }
        if std::fs::write(path, serde_json::to_vec_pretty(&entries).unwrap()).is_ok() {
            println!("wrote {member} = {ujpgf:.0} µJ/GFLOP ({grade:?}) → {path}");
        }
    }

    /// The same two runs on this process's own CPU and GPU counters, each one span
    /// ([`measure_member_process`]): the GPU run's figure is written only from a span
    /// metered whole, and the CPU run, on its own span, gives the ratio.
    fn process_counters(path: &str) {
        let meter = ProcessSpanMeter::new(jouleclaw_backend_apple::gpu_liveness_load);
        if !meter.is_metered() {
            refuse(
                path,
                NotMeasured::CounterUnavailable(
                    "apple-ioreport declined, and this process's own CPU and GPU counters are not both live".into(),
                ),
            );
        }
        println!(
            "IOReport declined ({:?}): each run metered on this process's own CPU and GPU counters",
            jouleclaw_energy::ioreport_status()
        );
        let (m, n, k) = (256usize, 4096usize, 4096usize);
        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 out = vec![0f32; m * n];
        let gpu_run = |o: &mut [f32]| {
            let ok = jouleclaw_loader_gguf::fabric_dispatch::matmul_bf16_mpp_bt(&a, &w, m, n, k, o);
            assert!(ok, "mpp_bt fell back — set JOULECLAW_FABRIC=1 on a Metal-4 box");
        };
        gpu_run(&mut out); // warm, outside the span
        let gpu = measure_member_process(&meter, MemberRuns::Gpu, gflop, 3.0, || gpu_run(&mut out));
        let cpu_run = |o: &mut [f32]| {
            // Assert it ran, as the GPU run does: a refused shape would meter an empty loop.
            jouleclaw_loader_gguf::matmul_bf16::matmul_bf16_threaded(&a, &w, &[], m, n, k, o)
                .expect("threaded-NEON matmul refused this shape");
        };
        cpu_run(&mut out);
        let cpu = measure_member_process(&meter, MemberRuns::Cpu, gflop, 3.0, || cpu_run(&mut out));

        println!("shape {m}x{n}x{k}");
        for (label, r) in [("matmul2d (metal-gpu)", &gpu), ("threaded-NEON (cpu) ", &cpu)] {
            match r {
                Ok(f) => println!(
                    "  {label}: {:>9.0} µJ/GFLOP   {:>7.1} GFLOP/s   (CPU {} µJ + GPU {} µJ)",
                    f.uj_per_gflop,
                    gflop * f.ops as f64 / f.secs,
                    f.cpu_uj,
                    f.gpu_uj
                ),
                Err(e) => println!("  {label}: not measured — {e}"),
            }
        }
        match (&gpu, &cpu) {
            (Ok(g), Ok(c)) => match energy_ratio(c.uj_per_gflop, g.uj_per_gflop) {
                Some(ratio) => println!("  energy ratio (CPU/GPU per GFLOP): {ratio:.2}×"),
                None => println!("  energy ratio (CPU/GPU per GFLOP): n/a"),
            },
            _ => println!("  energy ratio (CPU/GPU per GFLOP): n/a — both runs must be measured"),
        }
        let g = gpu.unwrap_or_else(|e| refuse(path, format!("the matmul2d (metal-gpu) run: {e}")));
        let method = process_method(&g, MemberRuns::Gpu, &format!("matmul2d {m}×{n}×{k}"));
        write_row(path, calibration_row("local/metal-mpp", "matmul", g.uj_per_gflop, g.provenance, method));
    }
}
#[cfg(not(all(feature = "fabric", feature = "fabric-meter")))]
fn main() { println!("build with --features \"rayon fabric fabric-meter\""); }