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

//! Wave 67: measure THIS machine's fabric, in joules.
//!
//! Runs the same BF16 matmul workloads on every serving
//! local member — NEON, AMX (Accelerate), Metal GPU
//! (resident-weight MPS) — under an energy meter, and
//! writes the measured µJ/GFLOP table the placer loads at
//! discovery (`jouleclaw_fabric::calibration_file_path`:
//! `JOULECLAW_FABRIC_CALIBRATION` where it is set, else
//! `~/.jouleclaw/fabric-calibration.json`).
//!
//! Two meters, chosen as the live autotuner chooses
//! (`fabric_dispatch::tune_shape_live`):
//!
//! - **The IOReport SoC counter, where its CPU energy rail
//!   is live.** It is machine-wide, so marginal energy per
//!   op = ΔE_busy − P_idle·Δt, with P_idle sampled
//!   immediately before the runs — an idle-subtraction MODEL
//!   on top of a hardware reading, the same caveat the
//!   joule-fix fleet dataset records. Run this on a quiet
//!   machine or the numbers are someone else's workload.
//!   Every figure must be a measurement or nothing is
//!   written: the idle baseline must advance and every
//!   member's marginal energy must be positive
//!   (`marginal_energy`); a zero, clamped or not, would be
//!   loaded as a measured placement cost.
//! - **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 member's loop is one
//!   `ProcessSpanMeter` span, its figure this process's CPU
//!   energy plus its GPU energy, with no idle to subtract
//!   (`marginal_energy::measure_member_process`). A member is
//!   written only from a span metered whole: a CPU member's
//!   GPU share an exact zero, a GPU member's the GPU rail
//!   while this process was the GPU's sole user. Any other
//!   member is refused, saying why, and the rest are still
//!   written. The domain is this process's CPU and GPU, not
//!   the SoC: DRAM is outside it, and AMX work through
//!   Accelerate is counted as the kernel attributes it to the
//!   calling threads' CPU energy, which is not checked
//!   against a rail here — the row's method says so.
//!
//! Each row it writes declares its grade — the meter's, so a
//! live reading (`HwShunt`) is `Measured` — and records its
//! method, its date, and the OS and chip it was taken on
//! (`marginal_energy::calibration_row`); the fabric loads it
//! at that grade and carries the rest along. The file's rows
//! for the three members this example measures
//! (`marginal_energy::FABRIC_CALIBRATE_MEMBERS`) are all
//! replaced (`marginal_energy::merge_calibration`): a member
//! written now by its new row, and a member refused or
//! skipped by none, so the fabric places it on its seeded
//! figure — never on a row an earlier run took, which may be
//! of another domain, grade or shape than the rows written
//! beside it. Rows of members it does not measure
//! (`local/metal-mpp`, which `mpp_energy` writes) stay as
//! they were.
//!
//! Exit status: 0 when every member was written; 2 when
//! some were and others were refused or skipped (absent at
//! warm-up); 1, having written nothing, when none could be.
//!
//! ```bash
//! cargo run --release --features "rayon fabric-meter" \
//!     --example fabric_calibrate
//! ```

use jouleclaw_energy::apple_ioreport::AppleIoReportCounter;
use jouleclaw_energy::span::ProcessSpanMeter;
use jouleclaw_energy::EnergyCounter;
use jouleclaw_loader_gguf::marginal_energy::{
    calibration_row, idle_draw_uw, marginal_uj, measure_member_process, merge_calibration, process_method, refusal,
    MemberRuns, NotMeasured, FABRIC_CALIBRATE_MEMBERS,
};
use std::collections::BTreeMap;
use std::time::Instant;

/// 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);
}

/// One cumulative reading, µJ.
fn read_uj(counter: &AppleIoReportCounter) -> Result<u64, NotMeasured> {
    counter.read().map(|r| r.uj).map_err(|e| NotMeasured::ReadFailed(e.to_string()))
}

/// The meter every member is priced on.
enum Meter {
    /// The machine-wide IOReport SoC counter, its CPU rail live, and the idle draw
    /// sampled before the runs, µW.
    IoReport { counter: AppleIoReportCounter, idle_uw: f64 },
    /// This process's own CPU and GPU counters.
    Process(ProcessSpanMeter),
}

struct Shape {
    label: &'static str,
    m: usize,
    n: usize,
    k: usize,
}

/// A member's row, for the AMX member's method: what this process's CPU counter
/// covers of work it hands to Accelerate.
const AMX_NOTE: &str = "; AMX work through Accelerate counted as the kernel attributes it to the calling \
                        threads' CPU energy, not checked against a rail";

fn main() {
    // Decode GEMV at the lm_head class (wave 66's open
    // question) and a prefill batch shape.
    let shapes = [
        Shape { label: "decode-gemv (m=1, lm_head class)", m: 1, n: 98304, k: 2304 },
        Shape { label: "prefill-batch (m=16)", m: 16, n: 4096, k: 2304 },
    ];

    // The file the placer loads: the variable's where it is set.
    let file = jouleclaw_fabric::calibration_file_path();
    let path = file.display().to_string();

    let counter = AppleIoReportCounter::new();
    let meter = if counter.is_available() {
        // Idle baseline: 3 s of nothing.
        let idle_uw = {
            let e0 = read_uj(&counter).unwrap_or_else(|e| refuse(&path, e));
            let t0 = Instant::now();
            std::thread::sleep(std::time::Duration::from_secs(3));
            let e1 = read_uj(&counter).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);
        Meter::IoReport { counter, idle_uw }
    } else {
        let span = ProcessSpanMeter::new(jouleclaw_backend_apple::gpu_liveness_load);
        if !span.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 ({:?}): pricing each member on this process's own CPU and GPU counters",
            jouleclaw_energy::ioreport_status()
        );
        Meter::Process(span)
    };

    // Per member, the row from its LAST shape (the batch shape for
    // matrix members reflects their serving regime; the decode
    // shape stays in the printout for the record). A member whose
    // last shape was refused or skipped has no row: neither an
    // earlier shape's figure nor an earlier run's stands in for it.
    let mut latest: BTreeMap<String, jouleclaw_fabric::CalibrationEntry> = BTreeMap::new();
    let mut refused: BTreeMap<String, String> = BTreeMap::new();

    for shape in &shapes {
        let (m, n, k) = (shape.m, shape.n, shape.k);
        let gflop = 2.0 * m as f64 * n as f64 * k as f64 / 1e9;
        println!("\n── {} ({:.2} GFLOP/op) ──", shape.label, gflop);

        // Synthetic BF16 weight + activation.
        let mut w = vec![0u8; 2 * n * k];
        for (i, c) in w.as_chunks_mut::<2>().0.iter_mut().enumerate() {
            let v = (((i * 40503) % 2000) as f32 / 1000.0 - 1.0) * 0.05;
            c.copy_from_slice(
                &jouleclaw_loader_gguf::matmul_bf16::f32_to_bf16_rne(v).to_le_bytes(),
            );
        }
        let a: Vec<f32> = (0..m * k)
            .map(|i| ((i * 2654435761 % 1000) as f32 / 500.0) - 1.0)
            .collect();
        let mut out = vec![0f32; m * n];

        type Runner<'a> = Box<dyn Fn(&mut [f32]) -> bool + 'a>;
        let members: Vec<(&str, MemberRuns, Runner)> = vec![
            (
                "local/cpu-neon",
                MemberRuns::Cpu,
                Box::new(|out: &mut [f32]| {
                    jouleclaw_loader_gguf::matmul_bf16::matmul_bf16_neon(&a, &w, &[], m, n, k, out);
                    true
                }),
            ),
            (
                "local/amx",
                MemberRuns::Cpu,
                Box::new(|out: &mut [f32]| {
                    jouleclaw_loader_gguf::fabric_dispatch::matmul_bf16_amx(&a, &w, &[], m, n, k, out)
                }),
            ),
            (
                "local/metal-gpu",
                MemberRuns::Gpu,
                Box::new(|out: &mut [f32]| {
                    jouleclaw_loader_gguf::fabric_dispatch::matmul_bf16_gpu(&a, &w, &[], m, n, k, out)
                }),
            ),
        ];

        for (member, runs, run) in &members {
            // Warm up (and for the GPU: pay the residency
            // upload OUTSIDE the measured window — calibration
            // prices the steady state; the upload amortizes).
            if !run(&mut out) {
                println!("{member:>18}: honestly absent / refused — skipped");
                latest.remove(*member);
                refused.insert(member.to_string(), format!("{}: absent or refused at warm-up", shape.label));
                continue;
            }
            run(&mut out);
            match &meter {
                Meter::IoReport { counter, idle_uw } => {
                    // Measure ≥ 2 s of work.
                    let mut iters = 0u64;
                    let e0 = read_uj(counter).unwrap_or_else(|e| refuse(&path, e));
                    let t0 = Instant::now();
                    while t0.elapsed().as_secs_f64() < 2.0 {
                        run(&mut out);
                        iters += 1;
                    }
                    let dt = t0.elapsed().as_secs_f64();
                    let e1 = read_uj(counter).unwrap_or_else(|e| refuse(&path, e));
                    let marginal = marginal_uj(e0, e1, dt, *idle_uw)
                        .unwrap_or_else(|e| refuse(&path, format!("{member} on {}: {e}", shape.label)));
                    let uj_per_op = marginal / iters as f64;
                    let uj_per_gflop = uj_per_op / gflop;
                    println!(
                        "{member:>18}: {iters:>5} ops in {dt:.2}s → {:.0} µJ/op, {:.0} µJ/GFLOP ({:.1} ms/op)",
                        uj_per_op, uj_per_gflop, dt * 1e3 / iters as f64,
                    );
                    refused.remove(*member);
                    latest.insert(
                        member.to_string(),
                        calibration_row(
                            member,
                            "matmul",
                            uj_per_gflop,
                            counter.provenance(),
                            format!(
                                "IOReport SoC ΔE − idle({:.2} W)·Δt over {} ops of {} — \
                                 idle-subtraction model on a hardware reading; IOReport {:?}",
                                idle_uw / 1e6, iters, shape.label, jouleclaw_energy::ioreport_status(),
                            ),
                        ),
                    );
                }
                Meter::Process(span) => {
                    // ≥ 2 s of work, one span.
                    match measure_member_process(span, *runs, gflop, 2.0, || {
                        run(&mut out);
                    }) {
                        Ok(f) => {
                            println!(
                                "{member:>18}: {:>5} ops in {:.2}s → {:.0} µJ/op, {:.0} µJ/GFLOP ({:.1} ms/op; \
                                 CPU {} µJ + GPU {} µJ)",
                                f.ops,
                                f.secs,
                                f.uj_per_gflop * gflop,
                                f.uj_per_gflop,
                                f.secs * 1e3 / f.ops as f64,
                                f.cpu_uj,
                                f.gpu_uj,
                            );
                            let mut method = process_method(&f, *runs, shape.label);
                            if *member == "local/amx" {
                                method.push_str(AMX_NOTE);
                            }
                            refused.remove(*member);
                            latest.insert(
                                member.to_string(),
                                calibration_row(member, "matmul", f.uj_per_gflop, f.provenance, method),
                            );
                        }
                        Err(why) => {
                            println!("{member:>18}: refused — {why}");
                            latest.remove(*member);
                            refused.insert(member.to_string(), format!("{}: {why}", shape.label));
                        }
                    }
                }
            }
        }
    }

    if latest.is_empty() {
        let why: Vec<String> = refused.iter().map(|(m, w)| format!("{m} ({w})")).collect();
        refuse(&path, format!("no member was measured whole: {}", why.join("; ")));
    }
    // Every row of a member this example measures is replaced: by
    // its row from this run, or by none. Other members' rows stay.
    let existing: Vec<jouleclaw_fabric::CalibrationEntry> = std::fs::read(&file)
        .ok()
        .and_then(|b| serde_json::from_slice(&b).ok())
        .unwrap_or_default();
    let written: Vec<String> = latest.keys().cloned().collect();
    let rows = merge_calibration(existing, &FABRIC_CALIBRATE_MEMBERS, latest.into_values().collect());
    if let Some(dir) = file.parent().filter(|d| !d.as_os_str().is_empty()) {
        std::fs::create_dir_all(dir).expect("mkdir");
    }
    std::fs::write(&file, serde_json::to_vec_pretty(&rows).unwrap()).expect("write");
    println!("\nwrote {} to {path} — Fabric::discover() places them on these figures.", written.join(", "));
    if !refused.is_empty() {
        for (member, why) in &refused {
            eprintln!("not written: {member} — {why}; its row was dropped, so it is placed on its seeded figure");
        }
        std::process::exit(2);
    }
}