open-standards/wai/wai-rs/tests/jpegai_memory.rs

//! What `wai.image.jpegai`'s exact decoder holds now that every integer stage runs
//! through the integer operator set, against the direct decode it replaced
//! (`jpegai_decode::legacy`, feature `legacy_int`), measured with a global allocator that
//! tracks the bytes live and their peak.
//!
//! The capability declares no cost. So what the decode holds is held to the direct
//! decode: with the bundle's scale path and networks loaded beforehand
//! (`jpegai_decode::load` and `jpegai_decode::load_networks`, once per bundle, as the
//! sink's store of loaded sets keeps them), a decode through the operator set holds no
//! more than the direct decode of the same codestream, on the reference kernels; and so
//! does each plane's front half (its integers through §13.6.1), measured alone. The
//! scale path (two graphs and the latent tables) and the networks (four graphs, §13.6.2)
//! are the set's, the same for every codestream, and are reported. The inputs are the `jpegai/v1` bundle and the test
//! codestream (`WAI_JPEGAI_CODESTREAM`, `WAI_JPEGAI_WEIGHTS_DIR`, `WAI_JPEGAI_NORM`), so
//! the test runs where they are fetched (`wai-jpegai-e2e.yml`) and skips loudly elsewhere.
//!
//! Its own test binary with a single test, so that no other test's allocations land in
//! the window it measures.
#![cfg(all(feature = "neural_int", feature = "legacy_int"))]

#[path = "support/require_vectors.rs"]
mod require_vectors;

use std::alloc::{GlobalAlloc, Layout, System};
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicUsize, Ordering::Relaxed};

use wai::int_kernels::{with_policy, Policy};
use wai::jpegai_decode::{self, legacy, DecoderWeights, NormTables, PlaneWeights};
use wai::jpegai_synthesis::{IntNet, Weights};

/// The system allocator, tracking the bytes live and the most live at once.
struct Live;
static CUR: AtomicUsize = AtomicUsize::new(0);
static PEAK: AtomicUsize = AtomicUsize::new(0);

unsafe impl GlobalAlloc for Live {
    unsafe fn alloc(&self, l: Layout) -> *mut u8 {
        PEAK.fetch_max(CUR.fetch_add(l.size(), Relaxed) + l.size(), Relaxed);
        unsafe { System.alloc(l) }
    }
    unsafe fn alloc_zeroed(&self, l: Layout) -> *mut u8 {
        PEAK.fetch_max(CUR.fetch_add(l.size(), Relaxed) + l.size(), Relaxed);
        unsafe { System.alloc_zeroed(l) }
    }
    unsafe fn realloc(&self, p: *mut u8, l: Layout, new: usize) -> *mut u8 {
        if new >= l.size() {
            PEAK.fetch_max(CUR.fetch_add(new - l.size(), Relaxed) + (new - l.size()), Relaxed);
        } else {
            CUR.fetch_sub(l.size() - new, Relaxed);
        }
        unsafe { System.realloc(p, l, new) }
    }
    unsafe fn dealloc(&self, p: *mut u8, l: Layout) {
        CUR.fetch_sub(l.size(), Relaxed);
        unsafe { System.dealloc(p, l) }
    }
}

#[global_allocator]
static ALLOC: Live = Live;

/// The most bytes `f` holds at once beyond what was live when it started, and its result.
fn peak_during<R>(f: impl FnOnce() -> R) -> (R, u64) {
    let base = CUR.load(Relaxed);
    PEAK.store(base, Relaxed);
    let r = f();
    (r, (PEAK.load(Relaxed) - base) as u64)
}

fn json(p: &Path) -> serde_json::Value {
    serde_json::from_slice(&std::fs::read(p).unwrap_or_else(|e| panic!("{}: {e}", p.display()))).unwrap()
}

#[test]
fn the_decode_holds_no_more_than_the_direct_decode() {
    let Some([cs, dir, nf]) = require_vectors::env_paths(
        "jpegai_memory::the_decode_holds_no_more_than_the_direct_decode",
        ["WAI_JPEGAI_CODESTREAM", "WAI_JPEGAI_WEIGHTS_DIR", "WAI_JPEGAI_NORM"],
        "tools/wai_jpegai_fetch_e2e_inputs.sh <dir> (pinned by tools/jpegai-e2e-v1.sha256)",
    ) else {
        return;
    };
    let d = PathBuf::from(dir);
    let w = |n: &str| Weights::from_json(&json(&d.join(n))).unwrap();
    let net = |n: &str| IntNet::from_json(&json(&d.join(n))).unwrap();
    let weights = DecoderWeights {
        luma: PlaneWeights { hyper_scale: net("hs_y.json"), hyper_decoder: w("hyp_y.json"), synthesis: w("synth_y.json") },
        chroma: PlaneWeights { hyper_scale: net("hs_uv.json"), hyper_decoder: w("hyp_uv.json"), synthesis: w("synth_uv.json") },
        mcm_y: [w("mcm0_y.json"), w("mcm1_y.json"), w("mcm2_y.json"), w("mcm3_y.json")],
        icci: wai::jpegai_icci::IcciNet::from_json(&json(&d.join("icci.json"))).unwrap(),
    };
    let nv = json(Path::new(&nf));
    let ints = |k: &str| -> Vec<i64> { nv[k].as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect() };
    let (thr, mag) = (nv["lef_thr0"].as_array().unwrap().clone(), nv["lef_mag0"].as_array().unwrap().clone());
    let norm = NormTables {
        pmf: nv["pmf"].as_array().unwrap().iter().map(|r| r.as_array().unwrap().iter().map(|x| x.as_i64().unwrap()).collect()).collect(),
        bounds: ints("bounds"),
        per_channel: nv["per_channel"].as_u64().unwrap() as usize,
        cdfs_y: ints("cdfs_y"),
        cdfs_uv: ints("cdfs_uv"),
        gain_y: ints("gain_y"),
        gain_uv: ints("gain_uv"),
        lsbs_t0: ints("lsbs_t0").iter().map(|&x| x as i32).collect(),
        lsbs_t1: ints("lsbs_t1").iter().map(|&x| x as i32).collect(),
        lef_thr: [0, 1, 2].map(|i| thr[i].as_i64().unwrap()),
        lef_mag: [0, 1, 2].map(|i| mag[i].as_f64().unwrap() as f32),
    };
    let data = std::fs::read(&cs).unwrap();

    with_policy(Policy::reference(), || {
        let (direct, direct_peak) = peak_during(|| legacy::decode(&data, &weights, &norm).expect("the direct decode"));
        drop(direct);
        let (sp, set) = peak_during(|| jpegai_decode::load(&weights, &norm).expect("the scale path loads"));
        let (nets, nets_load) = peak_during(|| jpegai_decode::load_networks(&weights).expect("the networks load"));
        let (img, held) = peak_during(|| jpegai_decode::decode_loaded(&data, &weights.icci, &norm, &sp, &nets, &mut |_, _| {}).expect("the decode"));
        drop(img);
        let (img, whole) = peak_during(|| jpegai_decode::decode(&data, &weights, &norm).expect("the decode"));
        drop(img);
        eprintln!(
            "jpegai/v1, reference kernels: the direct decode held {direct_peak} bytes; with the scale path ({set} bytes to load, {} held) and the networks ({nets_load} bytes to load, {} held) loaded beforehand, the decode through the operator set held {held}; loading and decoding in one call held {whole}",
            sp.heap_bytes(),
            nets.heap_bytes()
        );
        assert!(held <= direct_peak, "the decode through the operator set held {held} bytes; the direct decode held {direct_peak}");

        // The front half alone, each plane: its integers through the operator set (the scale
        // path loaded beforehand) against the direct computation (which builds its tables
        // on each call).
        use wai::jpegai_bitstream::{find_segment, parse_container, parse_picture_header, Marker};
        use wai::jpegai_ops::{self, Plane};
        let segs = parse_container(&data);
        let pay = |m| { let s = find_segment(&segs, m).unwrap(); &data[s.start..s.start + s.len] };
        let ph = parse_picture_header(&data, find_segment(&segs, Marker::Pih).unwrap().start_bit());
        for (k, (p, m)) in [(Plane::Luma, Marker::Sorp), (Plane::Chroma, Marker::Sors)].into_iter().enumerate() {
            let grfs = &ph.components[k].grfs_channel_flags;
            let (pw, gain) = if k == 0 { (&weights.luma, &norm.gain_y) } else { (&weights.chroma, &norm.gain_uv) };
            let (a, ops) = peak_during(|| jpegai_ops::plane(&sp, p, pay(Marker::Soz), pay(m), grfs).unwrap());
            drop(a);
            let (b, direct) = peak_during(|| jpegai_ops::legacy::plane(&norm, &pw.hyper_scale, gain, p, pay(Marker::Soz), pay(m), grfs));
            drop(b);
            eprintln!("{p:?}: the front half held {ops} bytes through the operator set, {direct} directly");
            assert!(ops <= direct, "{p:?}: the front half through the operator set held {ops} bytes; directly, {direct}");
        }
    });
}