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

//! 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 `WIV1` with keyframe 0, which runs no network,
//! the declared working bytes are what the decode holds, to within a few 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.
//!
//! 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"))]

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::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;
    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 {
            c::CAP_NEURAL_INT_HYPER | c::CAP_VIDEO_INT_HYPER => {
                let e: serde_json::Value = serde_json::from_slice(&hyp["entropy.json"]).unwrap();
                let t = wai::int_hyper_synth::CondEntropy::try_from_value(&e).unwrap();
                drop(e);
                Box::new((wai::int_hyper_synth::IntModel::from_bytes(&hyp["g_s.whm"]).unwrap(),
                          wai::int_hyper_synth::IntModel::from_bytes(&hyp["h_s.whm"]).unwrap(), t)) as Box<dyn std::any::Any>
            }
            c::CAP_AUDIO_INT_CODEC => Box::new(wai::int_audio::AudioModel::from_bytes(&model).unwrap()) as Box<dyn std::any::Any>,
            c::CAP_SPLAT_INT_CODEC => Box::new(wai::int_hyper_synth::IntModel::from_bytes(&gs).unwrap()) 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));
    }
    for (what, cap, payload, params) in cases {
        let cost = declared_cost(cap, &payload, params).unwrap_or_else(|e| panic!("{what}: {e}"));
        // 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);
    }
}