rusty_h264_encoder/mbtree.rs
1//! Macroblock-tree lookahead adaptive quantization (temporal AQ).
2//!
3//! A cheap forward pass over a GOP's SOURCE frames estimates, per macroblock, how
4//! much of the *future's* coding cost depends on it, then lowers the QP of
5//! heavily-referenced macroblocks — investing bits where they pay off across many
6//! later frames (a sharp reference makes every frame that predicts from it cheaper).
7//! This is the temporal complement to the spatial AQ (`aq_qp_map`): AQ moves bits
8//! by texture *within* a frame; mb-tree moves them by reference *importance across*
9//! frames.
10//!
11//! Method (x264's mb-tree, adapted to our CQP GOP):
12//! 1. Per frame, per MB: `intra` = spatial AC SATD; `inter` = best small-search
13//! motion-compensated residual SATD to the previous SOURCE frame (capped at
14//! `intra`), plus the winning MV. Source-domain (like x264's lowres lookahead)
15//! so no reconstruction is needed — it's a pure pre-pass.
16//! 2. Backward propagation: walk frames last→first. Each MB's total importance is
17//! `intra + propagate_in`; the fraction its predictor earned — `(intra-inter)/
18//! intra` — is credited to the reference MBs it points to (bilinear by MV,
19//! area-weighted over the up-to-4 overlapped MBs) in the previous frame.
20//! 3. QP offset `= -strength · log2((intra + propagate_in) / intra)` (≤ 0:
21//! heavily-referenced MBs get finer QP; leaves get 0). CENTERED per GOP
22//! (subtract the GOP-mean offset) so the average QP — hence the rate — is
23//! preserved and the effect is a pure redistribution of bits toward the MBs
24//! the future depends on.
25
26use crate::config::{EncoderConfig, LookaheadMode};
27use rusty_h264_common::inter::mc_luma;
28use rusty_h264_common::transform::hadamard_4x4;
29use rusty_h264_common::YuvFrame;
30
31/// Per-MB lookahead cost + motion for one frame.
32#[derive(Clone, Copy)]
33struct MbCost {
34 intra: i32, // spatial AC SATD, >= 1
35 inter: i32, // best MC-residual SATD to the previous frame, capped at `intra`
36 mv: (i32, i32), // winning MV (quarter-pel) — propagation direction
37}
38
39/// SATD of a 4×4 residual (sum of |Hadamard coeffs|).
40fn satd4(res: &[i32; 16]) -> i64 {
41 hadamard_4x4(res).iter().map(|&v| v.unsigned_abs() as i64).sum()
42}
43
44/// Edge-clamped coded-size luma (matches the encoder's source preparation).
45pub(crate) fn coded_luma(cfg: &EncoderConfig, frame: &YuvFrame) -> Vec<u8> {
46 let (cw, ch) = (cfg.mb_width() * 16, cfg.mb_height() * 16);
47 let (w, h) = (frame.width, frame.height);
48 let mut y = vec![0u8; cw * ch];
49 // Row-slice copy + right-edge fill instead of a per-pixel double-`min`
50 // gather — the `load_mb` shape (write into a PRE-SIZED destination), which
51 // is the panic-round shape that works, as distinct from the append shape
52 // that regressed. Interior rows become a `memcpy` + `memset`; the bottom
53 // padding re-copies row `h-1`. Bit-identical to the per-pixel form by the
54 // `coded_luma_matches_per_pixel` oracle.
55 let wc = w.min(cw);
56 for j in 0..ch {
57 let src = &frame.y[j.min(h - 1) * w..][..w];
58 let dst = &mut y[j * cw..][..cw];
59 dst[..wc].copy_from_slice(&src[..wc]);
60 dst[wc..].fill(src[w - 1]);
61 }
62 y
63}
64
65/// 2×2-average downsample of a luma plane to half resolution (both dims are MB
66/// multiples → stay even). The Hybrid/HalfRes lookahead runs the MV search on this:
67/// 4× fewer pixels, ~4× cheaper. The MV direction survives; only the COST accuracy
68/// suffers on blurred detail — which the Hybrid mode fixes by re-scoring at full-res.
69fn downsample2x(y: &[u8], cw: usize, ch: usize) -> (Vec<u8>, usize, usize) {
70 let (hw, hh) = (cw / 2, ch / 2);
71 let mut out = vec![0u8; hw * hh];
72 // Row-pair slices + `chunks_exact(2)` instead of five multiplied indexings
73 // per output pixel: the iterator shape carries the extents, so the four
74 // source loads and the store lose their per-pixel bounds checks and the
75 // row-base multiplies hoist to once per row. Same sums, same `(s+2)/4`
76 // rounding: BIT-IDENTICAL.
77 for j in 0..hh {
78 let r0 = &y[2 * j * cw..][..cw];
79 let r1 = &y[(2 * j + 1) * cw..][..cw];
80 let dst = &mut out[j * hw..][..hw];
81 for ((o, p0), p1) in dst.iter_mut().zip(r0.chunks_exact(2)).zip(r1.chunks_exact(2)) {
82 let s = p0[0] as u32 + p0[1] as u32 + p1[0] as u32 + p1[1] as u32;
83 *o = ((s + 2) / 4) as u8;
84 }
85 }
86 (out, hw, hh)
87}
88
89/// Spatial AC SATD of a `bs`×`bs` block at pixel `(bx0, by0)` (DC excluded, summed
90/// over 4×4 sub-blocks). The intra "cost" floor — how expensive with no prediction.
91fn intra_cost(sy: &[u8], cw: usize, bx0: usize, by0: usize, bs: usize) -> i32 {
92 let mut s = 0i64;
93 for by in 0..bs / 4 {
94 for bx in 0..bs / 4 {
95 let mut blk = [0i32; 16];
96 for dy in 0..4 {
97 for dx in 0..4 {
98 blk[dy * 4 + dx] = sy[(by0 + by * 4 + dy) * cw + bx0 + bx * 4 + dx] as i32;
99 }
100 }
101 let h = hadamard_4x4(&blk);
102 s += h[1..].iter().map(|&v| v.unsigned_abs() as i64).sum::<i64>();
103 }
104 }
105 (s.min(i32::MAX as i64) as i32).max(1)
106}
107
108/// Full-pel MC-residual SATD of a `bs`×`bs` block at a given (plane) quarter-pel MV.
109/// DETERMINISTIC cost instrument for the lookahead (H-36). Wall-clock on this box
110/// swings ±40 points run-to-run on an IDENTICAL config, which is far larger than
111/// the content effect being measured — so the lookahead's cost is judged by its
112/// WORK COUNT (candidate evaluations), which is exactly reproducible.
113pub(crate) static SATD_CALLS: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0);
114
115fn mc_satd(sy: &[u8], cw: usize, ch: usize, ref_y: &[u8], bx0: usize, by0: usize, bs: usize, mv: (i32, i32)) -> i64 {
116 SATD_CALLS.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
117 // H-35: the diamond below only ever probes FULL-PEL vectors (`dx * 4` in
118 // quarter-pel units), so nearly every call can read the reference IN PLACE and
119 // hand both planes to the vendored asm SATD — instead of copying a bs×bs block
120 // out of `mc_luma` and then running a SCALAR per-4×4 Hadamard. ("Exported ≠
121 // wired": the main ME has used the asm kernel for months; the lookahead had its
122 // own scalar twin, which is why mb-tree cost far more than its half-res search
123 // should.) Identical value: `satd_px`'s scalar arm IS this function's old sum,
124 // and its asm arm is pinned byte-exact to that by the accel oracles.
125 let (ix, iy) = (bx0 as isize + (mv.0 >> 2) as isize, by0 as isize + (mv.1 >> 2) as isize);
126 if mv.0 & 3 == 0
127 && mv.1 & 3 == 0
128 && ix >= 0
129 && iy >= 0
130 && ix as usize + bs <= cw
131 && iy as usize + bs <= ch
132 && by0 + bs <= ch
133 && bx0 + bs <= cw
134 {
135 return crate::mb16::satd_px(
136 &sy[by0 * cw + bx0..],
137 cw,
138 &ref_y[iy as usize * cw + ix as usize..],
139 cw,
140 bs,
141 bs,
142 );
143 }
144 // Sub-pel seed probe or an edge-overhanging vector: the general path.
145 let mut pred = [0u8; 256]; // bs ≤ 16 → fits; stride = bs
146 mc_luma(ref_y, cw, ch, bx0, by0, bs, bs, mv.0, mv.1, &mut pred);
147 let mut s = 0i64;
148 for by in 0..bs / 4 {
149 for bx in 0..bs / 4 {
150 let mut res = [0i32; 16];
151 for dy in 0..4 {
152 for dx in 0..4 {
153 res[dy * 4 + dx] = sy[(by0 + by * 4 + dy) * cw + bx0 + bx * 4 + dx] as i32
154 - pred[(by * 4 + dy) * bs + (bx * 4 + dx)] as i32;
155 }
156 }
157 s += satd4(&res);
158 }
159 }
160 s
161}
162
163/// Best MC-residual SATD of a `bs`×`bs` block and its winning (plane) MV, via a
164/// full-pel diamond search SEEDED from a predictor (the neighbour's MV, for pan
165/// coherence). The diamond (step 8→1 full-pel) tracks large motion a fixed ±2px set
166/// missed — a wrong MV gives mb-tree a wrong propagation DIRECTION (misdirects bits).
167fn inter_cost(sy: &[u8], cw: usize, ch: usize, ref_y: &[u8], bx0: usize, by0: usize, bs: usize, seed: (i32, i32), max_step: i32) -> (i32, (i32, i32)) {
168 let mut best_mv = (0, 0);
169 let mut best = mc_satd(sy, cw, ch, ref_y, bx0, by0, bs, (0, 0));
170 // H-45: a PROVABLY byte-identical early-out. SATD is a sum of absolute values,
171 // so every candidate is ≥ 0; once `best == 0` the guard `s < best` can never
172 // fire again, and both the winning MV and the cost are already final. Without
173 // this the diamond still pays its fixed floor — 4 probes at each of the 4 step
174 // levels (8→4→2→1), because the round that TERMINATES a level costs 4 probes
175 // that do not move — which is why the lookahead's eval count came out nearly
176 // content-invariant (16–18 /MB/frame on flat AND busy clips) and why mb-tree's
177 // relative cost was WORST on the static content it helps most (+18.8% akiyo).
178 if best == 0 {
179 return (0, best_mv);
180 }
181 if seed != (0, 0) {
182 let s = mc_satd(sy, cw, ch, ref_y, bx0, by0, bs, seed);
183 if s < best {
184 best = s;
185 best_mv = seed;
186 }
187 }
188 // `max_step` bounds the initial diamond hop: 8 for a from-scratch search, small
189 // (2) for the hybrid's full-res refine around an already-good coarse MV.
190 let mut step = max_step;
191 while step >= 1 {
192 loop {
193 let mut moved = false;
194 for &(dx, dy) in &[(step, 0), (-step, 0), (0, step), (0, -step)] {
195 let mv = (best_mv.0 + dx * 4, best_mv.1 + dy * 4); // quarter-pel units
196 let s = mc_satd(sy, cw, ch, ref_y, bx0, by0, bs, mv);
197 if s < best {
198 best = s;
199 best_mv = mv;
200 moved = true;
201 }
202 }
203 if !moved {
204 break;
205 }
206 }
207 step >>= 1;
208 }
209 (best.min(i32::MAX as i64) as i32, best_mv)
210}
211
212/// Per-MB costs for one frame, at the lookahead `mode`'s resolution(s). MVs are
213/// always returned in FULL-res quarter-pel (propagation is resolution-independent).
214/// - `FullRes`: search + score at full-res (16×16).
215/// - `HalfRes`: search + score at half-res (8×8) — MV scaled ×2.
216/// - `Hybrid`: search the MV on half-res, then REFINE + score intra/inter at
217/// full-res (the cost accuracy the pure-half-res path lost, at ~its speed).
218///
219/// `ref_full`/`ref_half` are `None` for the IDR (intra-only, nothing to propagate).
220#[allow(clippy::too_many_arguments)]
221fn frame_costs(
222 full: &[u8],
223 cwf: usize,
224 chf: usize,
225 half: &[u8],
226 cwh: usize,
227 chh: usize,
228 mb_w: usize,
229 mb_h: usize,
230 ref_full: Option<&[u8]>,
231 ref_half: Option<&[u8]>,
232 mode: LookaheadMode,
233) -> Vec<MbCost> {
234 let mut out: Vec<MbCost> = Vec::with_capacity(mb_w * mb_h);
235 for mb_y in 0..mb_h {
236 for mb_x in 0..mb_w {
237 // Neighbour MV seed (full-res quarter-pel), for pan coherence.
238 let seed_full = if mb_x > 0 {
239 out[mb_y * mb_w + mb_x - 1].mv
240 } else if mb_y > 0 {
241 out[(mb_y - 1) * mb_w + mb_x].mv
242 } else {
243 (0, 0)
244 };
245 // Intra cost at the scoring resolution (full for FullRes/Hybrid, half for HalfRes).
246 let intra = if mode == LookaheadMode::HalfRes {
247 intra_cost(half, cwh, mb_x * 8, mb_y * 8, 8)
248 } else {
249 intra_cost(full, cwf, mb_x * 16, mb_y * 16, 16)
250 };
251 let (inter, mv) = match (mode, ref_full, ref_half) {
252 (LookaheadMode::FullRes, Some(rf), _) => {
253 let (ic, mv) = inter_cost(full, cwf, chf, rf, mb_x * 16, mb_y * 16, 16, seed_full, 8);
254 (ic.min(intra), mv)
255 }
256 (LookaheadMode::HalfRes, _, Some(rh)) => {
257 let seed = (seed_full.0 / 2, seed_full.1 / 2);
258 let (ic, mvp) = inter_cost(half, cwh, chh, rh, mb_x * 8, mb_y * 8, 8, seed, 8);
259 (ic.min(intra), (mvp.0 * 2, mvp.1 * 2))
260 }
261 (LookaheadMode::Hybrid, Some(rf), Some(rh)) => {
262 // Coarse MV from the cheap half-res search…
263 let seed = (seed_full.0 / 2, seed_full.1 / 2);
264 let (_, mvp) = inter_cost(half, cwh, chh, rh, mb_x * 8, mb_y * 8, 8, seed, 8);
265 let coarse = (mvp.0 * 2, mvp.1 * 2); // → full-res quarter-pel
266 // …then a SMALL full-res refine that also gives the accurate cost.
267 let (ic, mv) = inter_cost(full, cwf, chf, rf, mb_x * 16, mb_y * 16, 16, coarse, 2);
268 (ic.min(intra), mv)
269 }
270 _ => (intra, (0, 0)), // IDR (no reference)
271 };
272 out.push(MbCost { intra, inter, mv });
273 }
274 }
275 out
276}
277
278/// Distribute `amount` from frame `f`'s MB (referencing the previous frame at MV
279/// `mv`) into `prev`'s per-MB propagation accumulator, area-weighted over the up-to-4
280/// macroblocks the referenced 16×16 block overlaps (edge-clamped).
281fn propagate_to(prev: &mut [f64], mb_w: usize, mb_h: usize, mb_x: usize, mb_y: usize, mv: (i32, i32), amount: f64) {
282 if amount <= 0.0 {
283 return;
284 }
285 // Referenced block top-left in pixels (integer part of the quarter-pel MV),
286 // clamped so it stays inside the frame.
287 let rx = (mb_x as i32 * 16 + (mv.0 >> 2)).clamp(0, (mb_w as i32 - 1) * 16);
288 let ry = (mb_y as i32 * 16 + (mv.1 >> 2)).clamp(0, (mb_h as i32 - 1) * 16);
289 let cx0 = (rx / 16) as usize;
290 let cy0 = (ry / 16) as usize;
291 // Overlap widths with the left/top MB column/row (the remaining area spills into
292 // the right/bottom neighbour when the block isn't MB-aligned).
293 let fx = (rx % 16) as f64;
294 let fy = (ry % 16) as f64;
295 let wl = 16.0 - fx; // area in column cx0
296 let wt = 16.0 - fy; // area in row cy0
297 for (dy, wy) in [(0usize, wt), (1, fy)] {
298 if wy <= 0.0 {
299 continue;
300 }
301 let cy = (cy0 + dy).min(mb_h - 1);
302 for (dx, wx) in [(0usize, wl), (1, fx)] {
303 if wx <= 0.0 {
304 continue;
305 }
306 let cx = (cx0 + dx).min(mb_w - 1);
307 prev[cy * mb_w + cx] += amount * (wx * wy) / 256.0;
308 }
309 }
310}
311
312/// One frame's detector preparation: coded-size luma + its half-res plane —
313/// the per-FRAME half of the scene-cut pair ratio, split out so pair scorers
314/// can prepare each frame ONCE. Scoring a batch through `windows(2)` prepped
315/// every interior frame TWICE (as one pair's `cur`, then the next pair's
316/// `prev`); the rolling caches in `lookahead::segment_gops` and the streaming
317/// encoder hand the previous pair's `cur` prep across, halving the dominant
318/// detector cost. Same planes from the same functions: BIT-IDENTICAL ratios.
319pub(crate) struct PairPrep {
320 full: Vec<u8>,
321 half: Vec<u8>,
322}
323
324// Manual Debug: the derived form would dump both planes byte-by-byte through
325// the `Encoder` derive that requires this.
326impl std::fmt::Debug for PairPrep {
327 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
328 write!(f, "PairPrep({} + {} bytes)", self.full.len(), self.half.len())
329 }
330}
331
332pub(crate) fn pair_prep(cfg: &EncoderConfig, f: &YuvFrame) -> PairPrep {
333 let (mb_w, mb_h) = (cfg.mb_width(), cfg.mb_height());
334 let (cwf, chf) = (mb_w * 16, mb_h * 16);
335 let full = coded_luma(cfg, f);
336 let (half, _, _) = downsample2x(&full, cwf, chf);
337 PairPrep { full, half }
338}
339
340/// Scene-cut pair ratio `Σ min(inter, intra) / Σ intra` on the LOOKAHEAD
341/// estimator: half-res intra SATD vs half-res diamond-searched MC residual.
342/// The diamond (step 8→1) is the load-bearing part — the cheap ±2px activity
343/// set was already refuted IN THIS FILE for exactly the failure a scene-cut
344/// detector cannot afford: a fast pan reads as "unpredictable" when the
345/// probe cannot reach the true vector, and every pair looks like a cut.
346pub(crate) fn pair_ratio_prepped(cfg: &EncoderConfig, cur: &PairPrep, prev: &PairPrep) -> f64 {
347 let (mb_w, mb_h) = (cfg.mb_width(), cfg.mb_height());
348 let (cwf, chf) = (mb_w * 16, mb_h * 16);
349 let (cwh, chh) = (cwf / 2, chf / 2);
350 let costs = frame_costs(
351 &cur.full, cwf, chf, &cur.half, cwh, chh, mb_w, mb_h,
352 Some(&prev.full), Some(&prev.half), LookaheadMode::HalfRes,
353 );
354 let (mut num, mut den) = (0i64, 0i64);
355 for c in &costs {
356 num += c.inter.min(c.intra) as i64;
357 den += c.intra as i64;
358 }
359 num as f64 / den.max(1) as f64
360}
361
362/// mb-tree per-frame per-MB QP offsets for a GOP of SOURCE frames (display order,
363/// the IDR first). `strength <= 0` returns all-zero (no-op / byte-identical). The
364/// offsets are centered per GOP so the mean QP — hence the rate — is preserved.
365/// Per-GOP gate telemetry — the Front-B harvest seam.
366///
367/// The mb-tree latch decides PER GOP, so fitting a law on per-clip signals is a
368/// unit mismatch: within one clip the GOPs straddle the threshold (football
369/// measured 0.747 / 0.480 / 0.402). This records one row per GOP so the
370/// refinery can pair signals with a per-GOP objective.
371///
372/// Observe-only and off unless `RFF_MBTREE_GOPSTATS` is set.
373pub mod gopstats {
374 use std::sync::Mutex;
375 // GLOBAL, not thread_local: `encode_all` encodes GOPs IN PARALLEL, each on
376 // its own worker thread with a fresh encoder. Thread-local rows land on the
377 // worker and are invisible to the caller — the first version of this seam
378 // silently harvested ZERO rows for exactly that reason.
379 //
380 // Rows therefore arrive in worker-completion order, NOT GOP order. Harvest
381 // with `RUSTY_THREADS=1` so the order is the GOP order the objective is
382 // keyed by; `take()` refuses to guess otherwise.
383 static ROWS: Mutex<Vec<GopRow>> = Mutex::new(Vec::new());
384 /// One GOP's gate inputs and its latch decision.
385 #[derive(Debug, Clone, Copy)]
386 pub struct GopRow {
387 /// Strength-invariant offset dispersion (the latch signal).
388 pub sd: f64,
389 /// Raw RMS before normalisation (scales with mbtree_strength).
390 pub sd_raw: f64,
391 /// 1 - mean predictability; the back-off axis.
392 pub residual_frac: f64,
393 /// Effective strength after the back-off latch (0 = latched off).
394 pub eff_strength: f64,
395 /// Did the differentiation latch ZERO this GOP's offsets?
396 pub latched_off: bool,
397 }
398 pub fn on() -> bool {
399 static V: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
400 *V.get_or_init(|| std::env::var_os("RFF_MBTREE_GOPSTATS").is_some())
401 }
402 pub(crate) fn push(r: GopRow) {
403 if on() {
404 if let Ok(mut g) = ROWS.lock() {
405 g.push(r);
406 }
407 }
408 }
409 /// Drain the rows recorded so far.
410 ///
411 /// ⚠ In GOP order ONLY when the encode ran single-threaded
412 /// (`RUSTY_THREADS=1`). With parallel GOP encoding the rows interleave by
413 /// completion, and pairing them positionally with a per-GOP objective would
414 /// silently mismatch signals to outcomes.
415 pub fn take() -> Vec<GopRow> {
416 ROWS.lock().map(|mut g| std::mem::take(&mut *g)).unwrap_or_default()
417 }
418}
419
420/// Minimum propagation-offset dispersion for mb-tree to apply at all.
421/// `RFF_MBTREE_SDMIN=0` restores the ungated behaviour exactly.
422fn mbtree_spread_min(cfg: &EncoderConfig) -> f64 {
423 static V: std::sync::OnceLock<Option<f64>> = std::sync::OnceLock::new();
424 V.get_or_init(|| {
425 std::env::var("RFF_MBTREE_SDMIN")
426 .ok()
427 .and_then(|v| v.parse().ok())
428 })
429 .unwrap_or(cfg.mbtree_spread_min)
430}
431
432pub fn gop_qp_offsets(cfg: &EncoderConfig, frames: &[YuvFrame], strength: f64) -> Vec<Vec<i32>> {
433 gop_qp_offsets_refs(cfg, &frames.iter().collect::<Vec<_>>(), strength)
434}
435
436/// [`gop_qp_offsets`] over BORROWED frames. The bframes path's mb-tree runs on
437/// each GOP's anchor SUB-SEQUENCE — non-contiguous display indices — and used
438/// to materialize every window by DEEP-CLONING the anchor frames (a full
439/// Y+U+V copy per anchor per window, ~150 KB each at CIF, ~3 MB at 1080p).
440/// A slice of references carries the same frames with zero copies; the owned
441/// wrapper above keeps the contiguous callers unchanged.
442pub fn gop_qp_offsets_refs(cfg: &EncoderConfig, frames: &[&YuvFrame], strength: f64) -> Vec<Vec<i32>> {
443 let (mb_w, mb_h) = (cfg.mb_width(), cfg.mb_height());
444 let n = frames.len();
445 if strength <= 0.0 || n == 0 || mb_w * mb_h == 0 {
446 gopstats::push(gopstats::GopRow {
447 sd: 0.0, sd_raw: 0.0, residual_frac: 0.0, eff_strength: 0.0, latched_off: true,
448 });
449 return vec![vec![0i32; mb_w * mb_h]; n];
450 }
451 // Lookahead resolution mode (Hybrid default: half-res MV search + full-res cost
452 // scoring). `RFF_MBTREE_LA=full|hybrid|half` overrides for A/B.
453 let mode = match std::env::var("RFF_MBTREE_LA").as_deref() {
454 Ok("full") => LookaheadMode::FullRes,
455 Ok("hybrid") => LookaheadMode::Hybrid,
456 Ok("half") => LookaheadMode::HalfRes,
457 _ => cfg.mbtree_lookahead,
458 };
459 let (cwf, chf) = (mb_w * 16, mb_h * 16);
460 let (cwh, chh) = (mb_w * 8, mb_h * 8);
461 let full: Vec<Vec<u8>> = frames.iter().map(|f| coded_luma(cfg, f)).collect();
462 // GRAIN LATCH (Great Gate P3 item 1 — docs/gate-ledger.md mbtree-grain-veto):
463 // propagation credit is FICTION on noise (nothing persists), so mb-tree
464 // redistributes on false gradients — measured +4.41% BD-SSIM on grain once
465 // the AQ grain veto stopped masking it. Reuse the aq-grain-veto conjunction
466 // ("unexplained temporal residual: not texture, not motion → noise"),
467 // SOURCE-vs-SOURCE on the GOP's first frame pair — grain is stationary
468 // (per-frame floor spread was 7.7–8.3 over 120 frames), and both probe
469 // frames sit INSIDE the GOP, so a boundary scene cut cannot sit between
470 // them. Fires → zero offsets, byte-identical to mb-tree off for the GOP.
471 // Thresholds transfer from the AQ fit: its IDR arm already validated this
472 // exact conjunction on source-vs-source signals. `RFF_MBTREE_GRAIN=0`
473 // disables (bisection anchor). Single-frame GOPs fail open (no pair).
474 let grain_veto = {
475 static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
476 *ON.get_or_init(|| std::env::var("RFF_MBTREE_GRAIN").map(|s| s != "0").unwrap_or(true))
477 };
478 if grain_veto && n >= 2 {
479 let sig = crate::signals::FrameSignals::new(&full[1], cwf, mb_w, mb_h, Some(&full[0]));
480 let grain = sig.grain_signature();
481 crate::signals::census::bump(crate::signals::census::MBTREE_GRAIN, grain);
482 if grain {
483 if std::env::var("RFF_MBTREE_DBG").is_ok() {
484 eprintln!("MBTREE_DBG grain latch: eff=0.000 (zero offsets)");
485 }
486 // The grain veto IS a gate decision — record it, or the harvest
487 // drops the GOP and every later row pairs with the wrong objective.
488 gopstats::push(gopstats::GopRow {
489 sd: 0.0, sd_raw: 0.0, residual_frac: 0.0, eff_strength: 0.0, latched_off: true,
490 });
491 return vec![vec![0i32; mb_w * mb_h]; n];
492 }
493 }
494 // Half-res planes needed for Hybrid + HalfRes (the MV search); FullRes skips them.
495 let need_half = mode != LookaheadMode::FullRes;
496 let half: Vec<Vec<u8>> = if need_half {
497 full.iter().map(|f| downsample2x(f, cwf, chf).0).collect()
498 } else {
499 Vec::new()
500 };
501 let empty: Vec<u8> = Vec::new();
502 // 1. per-frame per-MB costs (frame 0 = IDR, intra-only).
503 let costs: Vec<Vec<MbCost>> = (0..n)
504 .map(|f| {
505 let ref_full = if f == 0 { None } else { Some(full[f - 1].as_slice()) };
506 let ref_half = if f == 0 || !need_half { None } else { Some(half[f - 1].as_slice()) };
507 let hf = if need_half { half[f].as_slice() } else { &empty[..] };
508 frame_costs(&full[f], cwf, chf, hf, cwh, chh, mb_w, mb_h, ref_full, ref_half, mode)
509 })
510 .collect();
511 // 2. backward propagation: each MB credits the fraction its predictor earned to
512 // the previous frame's referenced MBs.
513 let mbs = mb_w * mb_h;
514 // ONE flat accumulator instead of `Vec<Vec<f64>>` (n+1 allocations → 1);
515 // the same `split_at_mut` cur/prev discipline, the same cells, the same
516 // values. `frac_buf` records each MB's `(intra−inter)/intra` as it is
517 // computed here, because the residual-fraction pass below was recomputing
518 // the IDENTICAL divide for every MB of every frame — reading the stored
519 // value in the same forward order keeps `fsum` bit-identical while
520 // removing (n−1)·mbs float divides. (Rows are frames 1..n.)
521 let mut propagate: Vec<f64> = vec![0.0; n * mbs];
522 let mut frac_buf: Vec<f64> = vec![0.0; n.saturating_sub(1) * mbs];
523 for f in (1..n).rev() {
524 let (head, tail) = propagate.split_at_mut(f * mbs);
525 let cur = &tail[..mbs];
526 let prev = &mut head[(f - 1) * mbs..];
527 for mb_y in 0..mb_h {
528 for mb_x in 0..mb_w {
529 let m = mb_y * mb_w + mb_x;
530 let c = costs[f][m];
531 let total = c.intra as f64 + cur[m];
532 // Fraction of this MB's cost the previous frame's reference "carries".
533 let frac = (c.intra - c.inter) as f64 / c.intra as f64; // in [0,1]
534 frac_buf[(f - 1) * mbs + m] = frac;
535 propagate_to(prev, mb_w, mb_h, mb_x, mb_y, c.mv, total * frac);
536 }
537 }
538 }
539 // CONTENT-ADAPTIVE STRENGTH (codec-content-adaptive-dispatch): mb-tree's benefit
540 // scales with how many future bits it can redistribute, ∝ the mean residual
541 // fraction `1 − pred` (pred = mean predictability `(intra−inter)/intra` over inter
542 // frames). When prediction is near-perfect (pred → 1, frames near-free/mostly
543 // skip — a slow, smooth pan) there is nothing to gain and QP perturbation only adds
544 // noise, so mb-tree REGRESSES; ramp strength to 0 as the residual fraction falls
545 // below `MBTREE_RES_MIN`. Natural detailed/mixed content sits well above it.
546 // PREDICTABILITY BACK-OFF, re-fitted 2026-08-06 (Great Gate P2 —
547 // docs/gate-ledger.md mbtree-backoff-refit). The original linear ramp
548 // (eff = strength·min(rf/0.10, 1)) had the right AXIS and the wrong SHAPE
549 // AND THRESHOLD: it throttled mb-tree's biggest real-content WINNERS
550 // (akiyo_qcif rf 0.046 → eff 0.44, forgoing −5.09→−9.24% BD-SSIM;
551 // screen_text rf 0.04 → eff 0.35, forgoing −4.53→−7.06) while the one
552 // class that genuinely regresses at full strength — tsrc-class synthetic
553 // (rf 0.023–0.025, +3.53% BD-SSIM unthrottled) — still leaked +0.50
554 // through the ramp's partial strength. The measured populations are
555 // disjoint with a 1.56× natural gap (tsrc ≤ 0.025 | winners ≥ 0.039), so
556 // the honest form is the single-sided LATCH: OFF below `res_min` (a zero
557 // qpo — byte-identical to mb-tree off for the GOP), FULL strength above.
558 // Per-GOP steps are safe: each GOP's offsets are independent and centered.
559 // `RFF_MBTREE_RESMIN` overrides (0 = no back-off, always full strength).
560 let res_min: f64 = {
561 static E: std::sync::OnceLock<f64> = std::sync::OnceLock::new();
562 *E.get_or_init(|| {
563 std::env::var("RFF_MBTREE_RESMIN").ok().and_then(|v| v.parse().ok()).unwrap_or(0.03)
564 })
565 };
566 // `frac_buf` holds frames 1..n in exactly this pass's iteration order, so
567 // the running sum sees the SAME summands in the SAME order — bit-identical
568 // `fsum` with zero divides. `fc` was `+= 1.0` per iteration: a count, and
569 // every increment is exact (integers < 2^53), so the closed form is the
570 // same number.
571 let fsum: f64 = frac_buf.iter().sum();
572 let fc = (n.saturating_sub(1) * mbs) as f64;
573 let residual_frac = 1.0 - if fc > 0.0 { fsum / fc } else { 0.0 };
574 let eff_strength = if res_min > 0.0 && residual_frac < res_min { 0.0 } else { strength };
575 crate::signals::census::bump(
576 crate::signals::census::MBTREE_BACKOFF,
577 eff_strength == 0.0,
578 );
579 if eff_strength == 0.0 {
580 gopstats::push(gopstats::GopRow {
581 sd: 0.0, sd_raw: 0.0, residual_frac, eff_strength: 0.0, latched_off: true,
582 });
583 // Latched off: zero offsets are byte-identical to mb-tree off.
584 if std::env::var("RFF_MBTREE_DBG").is_ok() {
585 eprintln!("MBTREE_DBG spread=0.000 residual_frac={residual_frac:.3} eff=0.000 (latched off)");
586 }
587 return vec![vec![0i32; mb_w * mb_h]; n];
588 }
589 // 3. QP offset per MB (≤ 0), then center per GOP to preserve the mean QP.
590 //
591 // A2 (fast-transcendentals addendum): `propagate == 0` ⇒ `total == intra`
592 // ⇒ the ratio is EXACTLY 1.0 (`intra >= 1` by construction, and `x/x` is
593 // exact) ⇒ `log2(1.0)` is exactly +0.0 — so the libm `log2` call AND the
594 // divide are skipped and the surviving multiply is the ORIGINAL expression
595 // evaluated at its exact value (`-eff_strength * 0.0` keeps the -0.0 the
596 // original produced). The whole LAST frame of every GOP takes this path —
597 // backward propagation never credits it — plus every MB the future never
598 // references; that is ≥ 1/n of all calls by construction, more on content
599 // with dead regions. Flat `offs` (n allocations → 1) with the same
600 // frame-major order everywhere downstream.
601 // Site 6's ★★ arm (Round 10): poly `log2` behind the same switch as the
602 // AQ pipeline; the zero-propagate shortcut is exact either way.
603 let poly = crate::fastmath::polytier_on();
604 let mut offs: Vec<f64> = Vec::with_capacity(n * mbs);
605 for f in 0..n {
606 for m in 0..mbs {
607 let p = propagate[f * mbs + m];
608 offs.push(if p == 0.0 {
609 -eff_strength * 0.0
610 } else {
611 let intra = costs[f][m].intra as f64;
612 let total = intra + p;
613 let l = if poly {
614 crate::fastmath::log2_poly(total / intra)
615 } else {
616 (total / intra).log2()
617 };
618 -eff_strength * l
619 });
620 }
621 }
622 // Per-GOP CENTERING: subtract the GOP-mean offset so the average QP — hence the
623 // rate — is preserved. This is the right rate-neutralization in BOTH modes: in CQP
624 // it holds the fixed QP; in RC mode (mb-tree runs per-GOP over the anchor chain, the
625 // controller picks the frame base) it keeps the offsets rate-neutral per GOP so the
626 // controller's model is undisturbed. MEASURED: routing the cross-frame allocation
627 // through the RC's `complexity` instead (uncentered per-MB + a per-frame multiplier)
628 // was WORSE — it destroyed the cross-frame redistribution the centered offsets carry
629 // (tsrc −1.5% → +1.9%). Centering stays; the RC just supplies the base QP.
630 let cnt = (n * mbs) as f64;
631 // Same frame-major summation order as the nested `flatten` had.
632 let mean: f64 = offs.iter().sum::<f64>() / cnt;
633 for o in offs.iter_mut() {
634 *o -= mean;
635 }
636 // DIFFERENTIATION LATCH (Great Gate P3 item 4 — the pan loser).
637 //
638 // mb-tree lowers QP on blocks whose quality PROPAGATES to the frames that
639 // reference them. That only buys anything when propagation is DIFFERENTIAL
640 // — some blocks matter much more than others (a static scene with a moving
641 // subject; screen content with dead regions). On a smooth pan EVERY block
642 // propagates about equally, so these centered offsets carry no information
643 // and mb-tree just redistributes rate for no perceptual reason. It is the
644 // clip class mb-tree has always lost on: stockholm +3.10, ducks +1.20,
645 // shields +1.06, bus +0.68, in_to_tree +0.69, crowd_run +0.35, city +0.27
646 // BD-SSIM, all REGRESSIONS IN THE SHIPPED DEFAULT before this latch.
647 //
648 // `sd` is the dispersion of mb-tree's OWN output. That is the whole point:
649 // it is not a content proxy standing in for the phenomenon (the fit that
650 // was refused here twice used `lv_spread`/`flat_run`, spatial statistics
651 // that merely correlate with panning on this corpus), it is the tool
652 // measuring whether it has anything to say. Below the line its offsets are
653 // noise around a centered mean, and zeroing them is EXACTLY mb-tree off.
654 //
655 // Measured over 21 clips: fires on 5, net +26.82 with ZERO regressions,
656 // versus +23.25 with SEVEN for always-on. Costs four modest forgone wins
657 // (tempete -1.61, football -1.39, soccer -0.59, crew -0.19).
658 //
659 // The refutation arm that killed the alternative: a second clause
660 // (`headroom>10 && tdecay>1.3`) recovered football/soccer at perfect fit on
661 // the 16-clip table, then fired on bus_cif -- a FAST PAN -- and LOST +0.68.
662 // Dropped. See gate-ledger `mbtree-dispatch`.
663 // STRENGTH-INVARIANT. Each offset is `-eff_strength * log2(total/intra)`, so
664 // a raw RMS scales LINEARLY with `mbtree_strength` — and the threshold was
665 // fitted at the default 0.9. At `--mbtree-strength 0.5` every sd shrinks 44%
666 // and clips that should fire latch OFF; at 2.0 the losers start firing.
667 // Dividing by eff_strength measures the DIFFERENTIATION itself (RMS of the
668 // log2 importance ratio) — what the gate is actually about, and invariant to
669 // how hard the offsets are then applied.
670 //
671 // Same defect class as the CAVLC bits/MB bug: a threshold on a signal whose
672 // SCALE depends on an axis the fitting corpus never varied.
673 let sd_raw = (offs.iter().map(|o| o * o).sum::<f64>() / cnt).sqrt();
674 let sd = sd_raw / eff_strength.max(1e-9);
675 if std::env::var("RFF_MBTREE_DBG").is_ok() {
676 eprintln!("MBTREE_DBG spread={sd:.3} raw={sd_raw:.3} residual_frac={residual_frac:.3} eff={eff_strength:.3}");
677 }
678 let sd_min = mbtree_spread_min(cfg);
679 gopstats::push(gopstats::GopRow {
680 sd, sd_raw, residual_frac, eff_strength, latched_off: sd < sd_min,
681 });
682 crate::signals::census::bump(crate::signals::census::MBTREE_SPREAD_LATCH, sd < sd_min);
683 if sd < sd_min {
684 return vec![vec![0i32; mb_w * mb_h]; n];
685 }
686 // Round + clamp to a sane per-MB QP swing. (Indexed slicing, not
687 // `chunks_exact` — the latter computes `len / mbs`, a runtime-divisor
688 // `div` this campaign exists to remove.) Site 8's ★★ arm: magic-number
689 // round behind the poly-tier switch.
690 const MBTREE_DQP_MAX: i32 = 6;
691 (0..n)
692 .map(|f| {
693 offs[f * mbs..][..mbs]
694 .iter()
695 .map(|&o| {
696 let r = if poly { crate::fastmath::round_ties_even_fast(o) } else { o.round() };
697 (r as i32).clamp(-MBTREE_DQP_MAX, MBTREE_DQP_MAX)
698 })
699 .collect()
700 })
701 .collect()
702}
703
704#[cfg(test)]
705mod tests {
706 use super::*;
707
708 /// Deterministic synthetic GOP: textured static background, a textured
709 /// 16x16 block translating 4px/frame, and a strip whose texture scrolls
710 /// 1px/frame — motion the diamond can track, with enough residual that the
711 /// predictability back-off would not zero it even if it were live.
712 fn synth_frames(w: usize, h: usize, n: usize) -> Vec<YuvFrame> {
713 (0..n)
714 .map(|f| {
715 let mut y = vec![0u8; w * h];
716 for j in 0..h {
717 for i in 0..w {
718 // Static checker+gradient background.
719 let mut v = (((i / 4 + j / 4) % 2) as i32 * 60 + (i as i32 * 3 + j as i32 * 2) % 90 + 40) as i32;
720 // Scrolling strip (1px/frame): partial predictability.
721 if (16..24).contains(&j) {
722 v = 30 + (((i + f) * 13) % 200) as i32;
723 }
724 // Moving textured square (4px/frame).
725 let sx = 4 + f * 4;
726 if (sx..sx + 16).contains(&i) && (28..44).contains(&j) {
727 v = 200 - ((i * 7 + j * 11) % 120) as i32;
728 }
729 y[j * w + i] = v.clamp(0, 255) as u8;
730 }
731 }
732 YuvFrame { width: w, height: h, y, u: vec![128; (w / 2) * (h / 2)], v: vec![128; (w / 2) * (h / 2)] }
733 })
734 .collect()
735 }
736
737 fn fnv1a_i32s(offs: &[Vec<i32>]) -> u64 {
738 let mut hsh: u64 = 0xcbf2_9ce4_8422_2325;
739 for fr in offs {
740 for &v in fr {
741 for b in v.to_le_bytes() {
742 hsh ^= b as u64;
743 hsh = hsh.wrapping_mul(0x100_0000_01b3);
744 }
745 }
746 }
747 hsh
748 }
749
750 /// The row-slice `coded_luma` must equal the original per-pixel
751 /// double-`min` gather (kept here as the oracle), including on frames
752 /// whose display size is not an MB multiple (live right/bottom padding).
753 #[test]
754 fn coded_luma_matches_per_pixel() {
755 for (w, h) in [(64usize, 48usize), (20, 13), (33, 17), (16, 16)] {
756 let cfg = EncoderConfig::new(w, h);
757 let frame = &synth_frames(w, h, 1)[0];
758 let (cw, ch) = (cfg.mb_width() * 16, cfg.mb_height() * 16);
759 let mut want = vec![0u8; cw * ch];
760 for j in 0..ch {
761 for i in 0..cw {
762 want[j * cw + i] = frame.y[j.min(h - 1) * w + i.min(w - 1)];
763 }
764 }
765 assert_eq!(coded_luma(&cfg, frame), want, "{w}x{h}");
766 }
767 }
768
769 /// End-to-end golden gate for `gop_qp_offsets`: the exact `Vec<Vec<i32>>`
770 /// output is pinned by hash across all three lookahead modes, so any edit
771 /// claiming bit-identity has a whole-pipeline oracle (costs → propagation
772 /// → offsets → centering → latch → rounding). The grain veto and the
773 /// residual back-off are pinned OPEN via their documented env anchors
774 /// (`RFF_MBTREE_GRAIN=0`, `RFF_MBTREE_RESMIN=0` — both OnceLock-cached, so
775 /// they are set before the first call) so the golden pins the ARITHMETIC,
776 /// not a veto's early-out zeros.
777 #[test]
778 fn gop_qp_offsets_golden() {
779 std::env::set_var("RFF_MBTREE_GRAIN", "0");
780 std::env::set_var("RFF_MBTREE_RESMIN", "0");
781 let (w, h, n) = (64usize, 48usize, 6usize);
782 let frames = synth_frames(w, h, n);
783 let mut cfg = EncoderConfig::new(w, h);
784
785 crate::fastmath::TEST_POLYTIER.with(|c| c.set(Some(false)));
786 cfg.mbtree_lookahead = LookaheadMode::HalfRes;
787 let a = gop_qp_offsets(&cfg, &frames, 0.9);
788 // The gate must prove the tool ran: an all-zero output would pin only
789 // a latch, not the arithmetic under edit.
790 assert!(a.iter().flatten().any(|&v| v != 0), "HalfRes offsets all zero");
791 assert_eq!(fnv1a_i32s(&a), 1359955132549194384, "HalfRes golden");
792
793 cfg.mbtree_lookahead = LookaheadMode::Hybrid;
794 let b = gop_qp_offsets(&cfg, &frames, 0.9);
795 assert!(b.iter().flatten().any(|&v| v != 0), "Hybrid offsets all zero");
796 assert_eq!(fnv1a_i32s(&b), 7391805242828194773, "Hybrid golden");
797
798 cfg.mbtree_lookahead = LookaheadMode::FullRes;
799 let c = gop_qp_offsets(&cfg, &frames, 2.0);
800 assert!(c.iter().flatten().any(|&v| v != 0), "FullRes offsets all zero");
801 assert_eq!(fnv1a_i32s(&c), 17244099396955043453, "FullRes golden");
802
803 // Round 10 decision-identity: the POLY arm (poly log2 in the offs
804 // loop + ties-even final round) must yield the IDENTICAL i32 offsets
805 // in every mode — asserted, not argued.
806 crate::fastmath::TEST_POLYTIER.with(|c| c.set(Some(true)));
807 cfg.mbtree_lookahead = LookaheadMode::HalfRes;
808 assert_eq!(gop_qp_offsets(&cfg, &frames, 0.9), a, "poly arm HalfRes differs");
809 cfg.mbtree_lookahead = LookaheadMode::Hybrid;
810 assert_eq!(gop_qp_offsets(&cfg, &frames, 0.9), b, "poly arm Hybrid differs");
811 cfg.mbtree_lookahead = LookaheadMode::FullRes;
812 assert_eq!(gop_qp_offsets(&cfg, &frames, 2.0), c, "poly arm FullRes differs");
813 crate::fastmath::TEST_POLYTIER.with(|c| c.set(None));
814 }
815}