mbtree_bench/
mbtree_bench.rs

1//! mb-tree lookahead cost/benefit instrument: encode the same clip with the
2//! lookahead OFF and ON, reporting each arm's wall (best-of-N) and an FNV hash
3//! of the bitstream.
4//!
5//! The hash is the correctness gate for lookahead *speed* work: a faster cost
6//! kernel that yields the same SATD values must leave the QP map — and therefore
7//! the bitstream — byte-identical. The wall ratio is the overhead mb-tree is
8//! asking us to pay for its BD win.
9//!
10//!   mbtree_bench <clip.y4m>     env: MB_FRAMES (48) MB_QP (27) MB_GOP (30) MB_REPS (3)
11//!                                    MB_LA (full|hybrid|half — RFF_MBTREE_LA also works)
12
13use rusty_h264_common::types::YuvFrame;
14use rusty_h264_encoder::{Encoder, EncoderConfig, Preset};
15
16fn fnv1a(d: &[u8]) -> u64 {
17    let mut h = 0xcbf29ce484222325u64;
18    for &b in d {
19        h ^= b as u64;
20        h = h.wrapping_mul(0x100000001b3);
21    }
22    h
23}
24
25fn read_y4m(path: &str, max: usize) -> (usize, usize, Vec<YuvFrame>) {
26    let raw = std::fs::read(path).expect("read clip");
27    let e = raw.iter().position(|&b| b == b'\n').expect("y4m header");
28    let hdr = std::str::from_utf8(&raw[..e]).expect("utf8 header");
29    let (mut w, mut h) = (0usize, 0usize);
30    for t in hdr.split_whitespace() {
31        match t.as_bytes().first() {
32            Some(b'W') => w = t[1..].parse().expect("W"),
33            Some(b'H') => h = t[1..].parse().expect("H"),
34            _ => {}
35        }
36    }
37    let (ys, cs) = (w * h, (w / 2) * (h / 2));
38    let (mut f, mut p) = (Vec::new(), e + 1);
39    while f.len() < max {
40        let Some(r) = raw[p..].iter().position(|&b| b == b'\n') else { break };
41        p += r + 1;
42        if p + ys + 2 * cs > raw.len() {
43            break;
44        }
45        f.push(YuvFrame {
46            width: w,
47            height: h,
48            y: raw[p..p + ys].to_vec(),
49            u: raw[p + ys..p + ys + cs].to_vec(),
50            v: raw[p + ys + cs..p + ys + 2 * cs].to_vec(),
51        });
52        p += ys + 2 * cs;
53    }
54    (w, h, f)
55}
56
57fn main() {
58    let path = std::env::args().nth(1).expect("usage: mbtree_bench <clip.y4m>");
59    let env = |k: &str, d: usize| -> usize {
60        std::env::var(k).ok().and_then(|v| v.parse().ok()).unwrap_or(d)
61    };
62    let (n, qp, gop, reps) = (env("MB_FRAMES", 48), env("MB_QP", 27), env("MB_GOP", 30), env("MB_REPS", 3));
63    if let Ok(la) = std::env::var("MB_LA") {
64        std::env::set_var("RFF_MBTREE_LA", la);
65    }
66    let (w, h, frames) = read_y4m(&path, n);
67    println!("mbtree_bench {}x{} x{} qp{qp} gop{gop} (best-of-{reps})", w, h, frames.len());
68
69    let mut off_ms = f64::MAX;
70    let mut on_ms = f64::MAX;
71    let (mut off_h, mut on_h, mut off_b, mut on_b) = (0u64, 0u64, 0usize, 0usize);
72    let mut calls = 0u64;
73    // Alternate the arms so thermal drift hits both equally.
74    for _ in 0..reps {
75        for on in [false, true] {
76            let mut cfg = EncoderConfig::new(w, h);
77            cfg.qp = qp as u8;
78            cfg.gop_size = gop as u32;
79            cfg.preset = Preset::Quality;
80            cfg.mbtree = on;
81            let enc = Encoder::new(cfg).expect("cfg");
82            rusty_h264_encoder::mbtree_satd_reset();
83            let t = std::time::Instant::now();
84            let out: Vec<u8> = enc.encode_all(&frames).expect("encode").concat();
85            let ms = t.elapsed().as_secs_f64() * 1e3;
86            if on {
87                on_ms = on_ms.min(ms);
88                on_h = fnv1a(&out);
89                on_b = out.len();
90                calls = rusty_h264_encoder::mbtree_satd_calls();
91            } else {
92                off_ms = off_ms.min(ms);
93                off_h = fnv1a(&out);
94                off_b = out.len();
95            }
96        }
97    }
98    println!("  mbtree OFF: {off_ms:8.1} ms  {off_b:>8} bytes  hash {off_h:016x}");
99    println!("  mbtree ON : {on_ms:8.1} ms  {on_b:>8} bytes  hash {on_h:016x}");
100    println!(
101        "  lookahead work: {calls} candidate evals ({:.0}/MB/frame, DETERMINISTIC)",
102        calls as f64 / ((w / 16 * (h / 16)) as f64 * frames.len() as f64)
103    );
104    println!(
105        "  lookahead overhead: {:+.1}%   size {:+.2}%",
106        100.0 * (on_ms / off_ms - 1.0),
107        100.0 * (on_b as f64 / off_b as f64 - 1.0)
108    );
109}
mbtree_bench/mbtree_bench.rs

mbtree_bench/
mbtree_bench.rs
