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

// wai-conformance: int-kernels
//! The declared cost of an integer decode (`int_payload::declared_cost`,
//! `extensions/integer-payloads.md` §7, *Memory and work*) against what the decode
//! holds, measured with a global allocator that tracks the bytes live and their peak.
//!
//! A sink reads the declared cost before it decodes and applies its own size policy,
//! so the cost must be one it can rely on: the decode, on the reference kernels, holds
//! at least the output and at most the declared working bytes beside the payload and
//! the parameter set as parsed: for a capability that decodes through the integer
//! operator set, parsed, lowered and loaded, which is the same for every payload the
//! set decodes. For a `WIV1` with keyframe 0, which runs no network,
//! the declared working bytes are what the decode holds, to within a few bytes.
//!
//! These checks run with the store of loaded sets (`int_set_cache`) off, so each decode
//! lowers and loads its set, as it did before there was a store. Then with the store on:
//! a decode that loads its set holds within the same bound, and leaves the set in the
//! store, which counts at least the bytes it holds for it (so its budget bounds what it
//! holds); a decode that finds its set held holds within the same bound, less than the
//! decode that loaded it, and leaves nothing more behind. A set kept in the store is the
//! set's, the same for every payload, and outside every decode's declared working bytes.
//!
//! The clip a few kilobytes can code: sixteen frames of 4729 × 4729 whose stream is
//! one symbol of frequency 65535/65536 over and over, about 3 KB, decoding to a clip at
//! the 2^30-byte cap. It is decoded in release builds only (CI's release step); a debug
//! build decodes a smaller clip of the same kind.
//!
//! `wai.neural.int_mlicpp` declares no cost, so its decode through the integer operator set
//! is held to the direct decode it replaced (feature `legacy_int`): with its bundle loaded,
//! it holds no more than the direct decode on the same payload.
//!
//! Its own test binary with a single test, so that no other test's allocations land in
//! the window it measures.
#![cfg(all(feature = "video_int", feature = "audio_int", feature = "splat_int", feature = "legacy_int"))]

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

use wai::int_kernels::{with_policy, Policy};
use wai::int_payload::{declared_cost, decode_int, resolve, wire, IntParams};
use wai::int_set_cache::LoadedSets;
use wai::pin::HeldBytes;

/// 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
/// (still held when the measure is taken).
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 read(rel: &str) -> Vec<u8> {
    let p = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/vectors").join(rel);
    std::fs::read(&p).unwrap_or_else(|e| panic!("{}: {e}", p.display()))
}

/// A rANS stream of `n` symbols, symbol `i` coded with `(start, freq)` from `sym(i)`,
/// from the state 2^31, so it ends canonically: what an encoder writes.
fn encode(n: u64, sym: impl Fn(u64) -> (u64, u64)) -> Vec<u8> {
    const L: u64 = 1 << 31;
    let mut state = L;
    let mut renorm: Vec<u32> = Vec::new();
    let mut i = n;
    while i > 0 {
        i -= 1;
        let (start, freq) = sym(i);
        let x_max = ((L >> 16) << 32) * freq;
        while state >= x_max {
            renorm.push(state as u32);
            state >>= 32;
        }
        state = ((state / freq) << 16) + (state % freq) + start;
    }
    renorm.reverse();
    let mut words = vec![state as u32, (state >> 32) as u32];
    words.extend(renorm);
    words.iter().flat_map(|w| w.to_le_bytes()).collect()
}

/// A keyframe-0 clip of `n_frames` frames of `h × w`, one block per frame, whose every
/// residual is 0 at frequency 65535/65536: a valid stream of a few bytes per 2^20
/// symbols.
fn constant_clip(n_frames: u32, h: u32, w: u32) -> Vec<u8> {
    let rans = encode(n_frames as u64 * h as u64 * w as u64 * 3, |_| (0, 65_535));
    wire::write_wiv1(&wire::Wiv1 {
        mv_precision: 1, keyframe: 0, n_frames, h, w, block: h.max(w), q: 1, offset: 0, cdf: vec![0, 65_535, 65_536],
        mvs: vec![(0, 0); n_frames as usize - 1], rans: &rans, wih1: None,
    })
}

#[test]
fn the_declared_cost_bounds_what_the_decode_holds() {
    use wai::codecs as c;
    // The store off: every decode lowers and loads its set.
    LoadedSets::global().set_budget(0);
    let hyp: BTreeMap<String, Vec<u8>> = [("entropy.json", "int_hyper/tiny_entropy.json"), ("g_s.whm", "int_hyper/tiny_g_s.bin"), ("h_s.whm", "int_hyper/tiny_h_s.bin")]
        .into_iter().map(|(k, v)| (k.to_string(), read(v))).collect();
    let hyp_view: BTreeMap<String, &[u8]> = hyp.iter().map(|(k, v)| (k.clone(), v.as_slice())).collect();
    let (model, gs) = (read("int_audio/audio_model.bin"), read("int_splat/splat_gs.bin"));
    let unpinned = |cap, held| resolve(cap, None, &[held], None).expect("an unpinned set resolves");
    let none = IntParams::none();
    let hyp_p = unpinned(c::CAP_NEURAL_INT_HYPER, HeldBytes::Set { listing: b"", files: &hyp_view });
    let (aud, spl) = (unpinned(c::CAP_AUDIO_INT_CODEC, HeldBytes::File(&model)), unpinned(c::CAP_SPLAT_INT_CODEC, HeldBytes::File(&gs)));
    // The parameter set as a decode parses it, which the declared cost leaves out.
    let set_bytes = |cap: &str| -> u64 {
        let (r, peak) = peak_during(|| match cap {
            // The conditional codec's set: its tables, and its two networks parsed, then
            // lowered and loaded (`int_hyper_synth::ops::load`).
            c::CAP_NEURAL_INT_HYPER | c::CAP_VIDEO_INT_HYPER => {
                let t = wai::int_hyper_synth::CondEntropy::read(&hyp["entropy.json"]).unwrap();
                let (g, h) = (wai::int_hyper_synth::IntModel::from_bytes(&hyp["g_s.whm"]).unwrap(),
                              wai::int_hyper_synth::IntModel::from_bytes(&hyp["h_s.whm"]).unwrap());
                let n = wai::int_hyper_synth::ops::load(&g, &h).unwrap();
                Box::new((g, h, t, n)) as Box<dyn std::any::Any>
            }
            c::CAP_AUDIO_INT_CODEC => {
                let m = wai::int_audio::AudioModel::from_bytes(&model).unwrap();
                let g = wai::int_audio::ops::load(&m).unwrap();
                Box::new((m, g)) as Box<dyn std::any::Any>
            }
            // A decode through the integer operator set holds its set parsed and its
            // lowered graph loaded (`int_splat::ops::load`), and the lowering forms each
            // in turn.
            c::CAP_SPLAT_INT_CODEC => {
                let m = wai::int_hyper_synth::IntModel::from_bytes(&gs).unwrap();
                let g = wai::int_splat::ops::load(&m).unwrap();
                Box::new((m, g)) as Box<dyn std::any::Any>
            }
            _ => Box::new(()) as Box<dyn std::any::Any>,
        });
        drop(r);
        peak
    };
    let small_clip = constant_clip(4, 300, 200);
    let mut cases: Vec<(&str, &str, Vec<u8>, &IntParams<'_>)> = vec![
        ("int_video", c::CAP_VIDEO_INT_MOTION, read("int_video/int_video_bitstream.wiv1"), &none),
        ("video_real", c::CAP_VIDEO_INT_MOTION, read("int_video/video_real_bitstream.wiv1"), &none),
        ("4 frames of 300 × 200, a constant stream", c::CAP_VIDEO_INT_MOTION, small_clip, &none),
        ("video_nk", c::CAP_VIDEO_INT_HYPER, read("int_video/video_nk_bitstream.wiv1"), &hyp_p),
        ("tiny", c::CAP_NEURAL_INT_HYPER, read("int_hyper/tiny_bitstream.wih1"), &hyp_p),
        ("audio", c::CAP_AUDIO_INT_CODEC, read("int_audio/audio_bitstream.wia1"), &aud),
        ("splat", c::CAP_SPLAT_INT_CODEC, read("int_splat/splat.wis1"), &spl),
    ];
    if !cfg!(debug_assertions) {
        let big = constant_clip(16, 4729, 4729);
        assert!(big.len() < 4096, "the clip at the cap is coded in {} bytes", big.len());
        cases.push(("16 frames of 4729 × 4729 in about 3 KB", c::CAP_VIDEO_INT_MOTION, big, &none));
    }
    let mut costs = BTreeMap::new();
    for (what, cap, payload, params) in cases {
        let cost = declared_cost(cap, &payload, params).unwrap_or_else(|e| panic!("{what}: {e}"));
        costs.insert(what, cost);
        // The header-only reading agrees where the header fixes the output.
        let size = wai::container::declared_size(cap, &payload).unwrap().unwrap();
        assert_eq!(size.symbols, cost.symbols, "{what}: symbols");
        if let Some(o) = size.output_bytes {
            assert_eq!(o, cost.output_bytes, "{what}: output bytes");
        }
        if let (Some(w), Some(m)) = (size.working_bytes, size.macs) {
            assert_eq!((w, m), (cost.working_bytes, cost.macs), "{what}: working bytes and multiply-accumulates");
        }
        let set = set_bytes(cap);
        let (r, peak) = with_policy(Policy::reference(), || peak_during(|| decode_int(cap, &payload, params)));
        r.unwrap_or_else(|e| panic!("{what}: {e}"));
        assert!(peak >= cost.output_bytes, "{what}: the decode held {peak} bytes, under its {} output bytes", cost.output_bytes);
        assert!(peak <= cost.working_bytes + set, "{what}: the decode held {peak} bytes, over the declared {} and the set's {set}",
                cost.working_bytes);
        if cap == c::CAP_VIDEO_INT_MOTION {
            assert!(cost.working_bytes - peak <= 4096, "{what}: declared {}, held {peak}", cost.working_bytes);
            assert_eq!(cost.macs, 0, "{what}: a keyframe-0 clip runs no network");
        } else {
            assert!(cost.macs > 0, "{what}: its networks run");
        }
        println!("{what}: declared {} working bytes ({} output, {} symbols, {} multiply-accumulates); held {peak}",
                 cost.working_bytes, cost.output_bytes, cost.symbols, cost.macs);
    }

    // The store on, and empty. For each payload that decodes against a set: the decode
    // that loads it, then the decode that finds it held.
    let store = LoadedSets::global();
    store.set_budget(LoadedSets::DEFAULT_BUDGET);
    for (what, cap, rel, params) in [
        ("video_nk", c::CAP_VIDEO_INT_HYPER, "int_video/video_nk_bitstream.wiv1", &hyp_p),
        ("tiny", c::CAP_NEURAL_INT_HYPER, "int_hyper/tiny_bitstream.wih1", &hyp_p),
        ("audio", c::CAP_AUDIO_INT_CODEC, "int_audio/audio_bitstream.wia1", &aud),
        ("splat", c::CAP_SPLAT_INT_CODEC, "int_splat/splat.wis1", &spl),
    ] {
        let (payload, cost, set) = (read(rel), costs[what], set_bytes(cap));
        let counted0 = store.stats().resident_bytes;
        let base = CUR.load(Relaxed);
        let (r, miss) = with_policy(Policy::reference(), || peak_during(|| decode_int(cap, &payload, params)));
        drop(r.unwrap_or_else(|e| panic!("{what}: {e}")));
        let held = (CUR.load(Relaxed) - base) as u64;
        let counted = store.stats().resident_bytes - counted0;
        assert!(miss <= cost.working_bytes + set, "{what}, its set loaded and kept: held {miss}, over the declared {} and the set's {set}", cost.working_bytes);
        assert!(held <= counted && counted - held <= LoadedSets::INDEX_BYTES_PER_ENTRY, "{what}: the store holds {held} bytes for the set and counts {counted}");
        let (r, hit) = with_policy(Policy::reference(), || peak_during(|| decode_int(cap, &payload, params)));
        drop(r.unwrap_or_else(|e| panic!("{what}: {e}")));
        assert_eq!((CUR.load(Relaxed) - base) as u64, held, "{what}: a decode that finds its set held keeps nothing more");
        assert!(hit <= cost.working_bytes + set && hit < miss, "{what}, its set held: held {hit}; loading it held {miss}");
        println!("{what}, store on: the decode that loads the set held {miss}, the one that finds it held {hit}; the store holds {held} and counts {counted}");
    }
    store.set_budget(0);

    // The decode through the integer operator set, with its set already parsed and
    // loaded: what it then holds for the payload (the latent, each operator's input,
    // output and uncropped plane, the attribute bytes) is within the declared working
    // bytes alone. The set's own lowering, above, holds more for a moment; that is the
    // set's, the same for every payload.
    let p = read("int_splat/splat.wis1");
    let cost = declared_cost(c::CAP_SPLAT_INT_CODEC, &p, &spl).unwrap();
    let v = wire::parse_wis1(&p).unwrap();
    let st = wai::int_splat::SplatStream {
        afrac: v.afrac as u32, c: v.c as usize, sh: v.sh as usize, sw: v.sw as usize, s: v.s as usize,
        n_attr: v.attr as usize, cdfs: &v.cdfs, offsets: &v.offsets, median_q: &v.median_q, rans: v.rans,
    };
    let m = wai::int_hyper_synth::IntModel::from_bytes(&gs).unwrap();
    let g = wai::int_splat::ops::load(&m).unwrap();
    let (r, peak) = with_policy(Policy::reference(), || peak_during(|| wai::int_splat::ops::decode_loaded(&g, m.afrac, &st)));
    r.unwrap();
    assert!(peak >= cost.output_bytes && peak <= cost.working_bytes, "splat, set loaded: held {peak}, declared {}", cost.working_bytes);
    println!("splat, its set loaded beforehand: declared {} working bytes; held {peak}", cost.working_bytes);

    // The same for the audio synthesis.
    let p = read("int_audio/audio_bitstream.wia1");
    let cost = declared_cost(c::CAP_AUDIO_INT_CODEC, &p, &aud).unwrap();
    let a = wire::parse_wia1(&p).unwrap();
    let st = wai::int_audio::AudioStream {
        afrac: a.afrac as u32, n_ch: a.n_ch as usize, lat_len: a.lat_len as usize, t_out: a.t_out as usize,
        offset: a.offset, median_q: &a.median_q, cdfs: &a.cdfs, rans: a.rans,
    };
    let m = wai::int_audio::AudioModel::from_bytes(&model).unwrap();
    let g = wai::int_audio::ops::load(&m).unwrap();
    let (r, peak) = with_policy(Policy::reference(), || peak_during(|| wai::int_audio::ops::decode_loaded(&g, m.afrac, &st)));
    r.unwrap();
    assert!(peak >= cost.output_bytes && peak <= cost.working_bytes, "audio, set loaded: held {peak}, declared {}", cost.working_bytes);
    println!("audio, its set loaded beforehand: declared {} working bytes; held {peak}", cost.working_bytes);

    // And for the conditional codec, its two networks loaded beforehand.
    let p = read("int_hyper/tiny_bitstream.wih1");
    let cost = declared_cost(c::CAP_NEURAL_INT_HYPER, &p, &hyp_p).unwrap();
    let k = wire::parse_wih1(&p).unwrap();
    let st = wai::int_hyper_synth::CondStream {
        m: k.m as usize, sh: k.sh as usize, sw: k.sw as usize, n_z: k.n_z as usize, zh: k.zh as usize, zw: k.zw as usize,
        z_rans: k.z_rans, rans: k.rans,
    };
    let (g, h) = (wai::int_hyper_synth::IntModel::from_bytes(&hyp["g_s.whm"]).unwrap(), wai::int_hyper_synth::IntModel::from_bytes(&hyp["h_s.whm"]).unwrap());
    let ent = wai::int_hyper_synth::CondEntropy::read(&hyp["entropy.json"]).unwrap();
    let nets = wai::int_hyper_synth::ops::load(&g, &h).unwrap();
    let (r, peak) = with_policy(Policy::reference(), || peak_during(|| wai::int_hyper_synth::ops::decode_loaded(&nets, g.afrac, &ent, &st)));
    r.unwrap();
    assert!(peak >= cost.output_bytes && peak <= cost.working_bytes, "int_hyper, set loaded: held {peak}, declared {}", cost.working_bytes);
    println!("int_hyper, its set loaded beforehand: declared {} working bytes; held {peak}", cost.working_bytes);

    // wai.neural.int_mlicpp declares no cost: integer-payloads §7 declares one for the
    // formats SPEC §5 pins, and `check_cost` passes this one. So what its decode holds is
    // held to the direct decode it replaced (`int_mlicpp::legacy`, test-only): with the
    // bundle parsed and its three networks loaded beforehand, the decode through the operator
    // set holds at least the image and no more than the direct decode holds on the same
    // payload, both on the reference kernels. Through the sink, with the store on, the
    // decode that loads the bundle holds within what the direct decode and the set as parsed,
    // lowered and loaded hold, and the store counts at least what it keeps; the decode that
    // finds the bundle held holds less, and keeps nothing more behind.
    {
        use wai::int_mlicpp::{legacy, ops, IntMlicModel};
        let bundle = read("int_mlicpp/model.wmb");
        let p = wai::container::Wai::from_bytes(&read("int_mlicpp/stream.wai")).unwrap().payload;
        let (m, parse) = peak_during(|| IntMlicModel::from_bundle(&bundle).unwrap());
        let (want, direct) = with_policy(Policy::reference(), || peak_during(|| legacy::decode(&m, &p)));
        let want = want.unwrap();
        let hd = ops::check(&m, &p).unwrap();
        let (nets, set) = peak_during(|| ops::load(&m).unwrap());
        ops::fit(&nets, &hd).unwrap();
        let (got, held) = with_policy(Policy::reference(), || peak_during(|| ops::decode_loaded(&nets, &m.tables, m.afrac(), &hd, &p, None)));
        assert!(got.unwrap() == want, "int_mlicpp: the image through the lowering is the direct decode's");
        let out = 3 * (hd.height * hd.width) as u64;
        assert!(held >= out && held <= direct, "int_mlicpp, set loaded: held {held}; the direct decode held {direct}");
        println!("int_mlicpp, its set loaded beforehand: held {held}; the direct decode held {direct}; loading the networks held {set}");
        let params = unpinned(c::CAP_NEURAL_INT_MLICPP, HeldBytes::File(&bundle));
        let (r, off) = with_policy(Policy::reference(), || peak_during(|| decode_int(c::CAP_NEURAL_INT_MLICPP, &p, &params)));
        drop(r.unwrap());
        store.set_budget(LoadedSets::DEFAULT_BUDGET);
        let counted0 = store.stats().resident_bytes;
        let base = CUR.load(Relaxed);
        let (r, miss) = with_policy(Policy::reference(), || peak_during(|| decode_int(c::CAP_NEURAL_INT_MLICPP, &p, &params)));
        drop(r.unwrap());
        let kept = (CUR.load(Relaxed) - base) as u64;
        let counted = store.stats().resident_bytes - counted0;
        assert!(kept <= counted && counted - kept <= LoadedSets::INDEX_BYTES_PER_ENTRY, "int_mlicpp: the store holds {kept} bytes for the set and counts {counted}");
        let (r, hit) = with_policy(Policy::reference(), || peak_during(|| decode_int(c::CAP_NEURAL_INT_MLICPP, &p, &params)));
        drop(r.unwrap());
        assert_eq!((CUR.load(Relaxed) - base) as u64, kept, "int_mlicpp: a decode that finds its set held keeps nothing more");
        assert!(miss <= parse + set + direct, "int_mlicpp, store on: the decode that loads the set held {miss}, over the bundle parsed ({parse}), loaded ({set}) and the direct decode's {direct}");
        assert!(hit < miss && hit <= parse + direct, "int_mlicpp, store on: held {hit} finding the set, over the bundle parsed ({parse}) and the direct decode's {direct}; loading it held {miss}");
        println!("int_mlicpp, through the sink: store off held {off}; store on, the decode that loads the set held {miss}, the one that finds it held {hit}; the store holds {kept} and counts {counted}");
        store.set_budget(0);
    }
}