From 4e87f2b850064881f5c868181db95100e4f3997f Mon Sep 17 00:00:00 2001 From: "Dylan M. Taylor" Date: Wed, 11 Nov 2026 03:17:58 -0500 Subject: [PATCH] H.264: bound decoder work from overlapping slices, gaps and queued output (#1706) * H.264: bound decoder work from overlapping slices, gaps and queued output - A macroblock that an earlier slice of the same picture already decoded is refused, so overlapping slices cannot redo the picture's reconstruction once per slice. - The memory budget check counts pictures waiting in the output queue that have left the DPB (one decode call may carry many access units) and the DPB's actual size, counting shared buffers once. - frame_num gap concealment advances the POC state across the skipped run in constant time instead of once per missing frame_num. - Storing a reference picture evicts until the DPB is under capacity, rather than evicting one picture and inserting regardless. * address review: charge pictures already held at their own size in the decoder budget * address review: keep the h264 output-queue byte total incrementally --- crates/h264/src/cabac_mb.rs | 2 +- crates/h264/src/decoder.rs | 53 +++++++++++++++++++++------ crates/h264/src/dpb.rs | 63 ++++++++++++++++++++++++++------ crates/h264/src/picture.rs | 8 +++++ crates/h264/src/slicedec.rs | 12 ++++--- crates/h264/src/synth_tests.rs | 65 +++++++++++++++++++++++++++++++++- 6 files changed, 177 insertions(+), 26 deletions(-) diff --git a/crates/h264/src/cabac_mb.rs b/crates/h264/src/cabac_mb.rs index 78c257e9b..368750292 100644 --- a/crates/h264/src/cabac_mb.rs +++ b/crates/h264/src/cabac_mb.rs @@ -48,7 +48,7 @@ impl SliceDecoder<'_> { let intra_slice = self.sh.slice_type.is_intra(); loop { ensure!(addr < total, "macroblock address {addr} out of range"); - self.start_mb(addr); + self.start_mb(addr)?; let skip = !intra_slice && self.decode_skip_flag(&mut c); if skip { self.decode_skip()?; diff --git a/crates/h264/src/decoder.rs b/crates/h264/src/decoder.rs index a3244d473..455701f92 100644 --- a/crates/h264/src/decoder.rs +++ b/crates/h264/src/decoder.rs @@ -15,7 +15,7 @@ use crate::slicedec::{PicState, SliceDecoder}; use crate::transform::LevelScale; use crate::{ColorInfo, Picture}; use deckcraft_bitstream::{BitReader, annexb_nals, length_prefixed_nals, unescape_rbsp}; -use std::collections::VecDeque; +use std::collections::{HashMap, VecDeque}; use std::sync::{Arc, Mutex}; /// One slice of a picture job. @@ -90,6 +90,10 @@ pub struct Decoder { ls_cache: Vec<(Arc, Arc)>, /// Pictures leaving the DPB, converted once their decoding job has finished. out_queue: VecDeque, + /// Queued entries per frame id and their summed bytes, kept in step with `out_queue` so the + /// budget check does not rescan it. + queued_ids: HashMap, + queued_bytes: u64, shared: Arc, #[cfg(feature = "threads")] pool: Option, @@ -175,6 +179,8 @@ impl Decoder { active_sps: None, ls_cache: Vec::new(), out_queue: VecDeque::new(), + queued_ids: HashMap::new(), + queued_bytes: 0, shared: Arc::new(Shared::default()), #[cfg(feature = "threads")] pool, @@ -283,7 +289,7 @@ impl Decoder { if !o.frame.is_published(o.frame.mb_h() - 1) && self.out_queue.len() <= max_queue { break; } - let Some(o) = self.out_queue.pop_front() else { break }; + let Some(o) = self.pop_out() else { break }; pics.push(make_picture(&o)); } Ok(pics) @@ -295,7 +301,7 @@ impl Decoder { let mut outs = Vec::new(); self.dpb.flush(&mut outs); self.emit(outs); - let pics = self.out_queue.drain(..).map(|o| make_picture(&o)).collect(); + let pics = std::iter::from_fn(|| self.pop_out()).map(|o| make_picture(&o)).collect(); while let Some(f) = self.in_flight.pop_front() { f.wait_complete(); } @@ -451,13 +457,24 @@ impl Decoder { chroma_planes: if chroma == 0 { 0 } else { 2 }, bytes_per_sample: 1, }; - let pictures = (sps.max_dpb_frames() + 1 + self.max_in_flight) as u64; - ensure!( - shape.fits(pictures, self.memory_budget), - "{}x{} pictures with {pictures} in the decoder exceed the memory budget", - sps.width(), - sps.height() - ); + // Pictures already held count at their own size, not the incoming SPS's: after a + // resolution change, larger pictures of the old size can still be in the DPB or waiting + // for output (one decode call may carry many access units). A buffer shared with a DPB + // entry is counted once. + let pic = shape.bytes().unwrap_or(u64::MAX); + let held = |frames: &mut dyn Iterator>| frames.fold(0u64, |n, f| n.saturating_add(f.bytes())); + // Queued bytes minus the queue entries that share a DPB frame (each distinct DPB frame once). + let mut ids: Vec<(u32, u64)> = self.dpb.entries.iter().map(|e| (e.frame.id, e.frame.bytes())).collect(); + ids.sort_unstable(); + ids.dedup(); + let shared = ids.iter().fold(0u64, |n, (id, b)| n.saturating_add(self.queued_ids.get(id).copied().unwrap_or(0).saturating_mul(*b))); + let queued = self.queued_bytes.saturating_sub(shared); + let in_dpb = held(&mut self.dpb.entries.iter().map(|e| &e.frame)); + let dpb = sps.max_dpb_frames().max(self.dpb.entries.len()) as u64; + // The DPB (the larger of its new slots and what it holds now), the picture being decoded, + // and what is in flight or queued. + let bytes = pic.saturating_mul(dpb).max(in_dpb).saturating_add(pic).saturating_add(pic.saturating_mul(self.max_in_flight as u64).max(queued)); + ensure!(bytes <= self.memory_budget, "{}x{} pictures with {bytes} bytes in the decoder exceed the memory budget", sps.width(), sps.height()); Ok(()) } @@ -562,8 +579,24 @@ impl Decoder { } fn emit(&mut self, outs: Vec) { + for o in &outs { + *self.queued_ids.entry(o.frame.id).or_insert(0) += 1; + self.queued_bytes = self.queued_bytes.saturating_add(o.frame.bytes()); + } self.out_queue.extend(outs); } + + fn pop_out(&mut self) -> Option { + let o = self.out_queue.pop_front()?; + if let Some(n) = self.queued_ids.get_mut(&o.frame.id) { + *n -= 1; + if *n == 0 { + self.queued_ids.remove(&o.frame.id); + } + } + self.queued_bytes = self.queued_bytes.saturating_sub(o.frame.bytes()); + Some(o) + } } /// Decode all slices of one picture and publish its rows. diff --git a/crates/h264/src/dpb.rs b/crates/h264/src/dpb.rs index 46e8f060a..940b1a4c5 100644 --- a/crates/h264/src/dpb.rs +++ b/crates/h264/src/dpb.rs @@ -139,8 +139,9 @@ impl Dpb { let mut unused = (prev_ref_frame_num + 1) % max_frame_num; // Only the last `max_num_ref_frames` "non-existing" frames can survive the sliding window, and // inserting more than that evicts every earlier short-term reference. Skip straight to that tail - // (keeping the POC state walk for every missing frame_num) so a hostile gap of up to 65535 - // frames costs at most `max_num_ref_frames` frame allocations instead of one per missing frame. + // (advancing the POC state across the skipped run in constant time) so a hostile gap of up to + // 65535 frames costs at most `max_num_ref_frames` frame allocations and POC updates instead of + // one per missing frame. let gap = frame_num.wrapping_add(max_frame_num).wrapping_sub(unused) % max_frame_num; let keep = u32::try_from(max_num_ref_frames.max(1)).unwrap_or(u32::MAX); let skip = gap.saturating_sub(keep); @@ -149,10 +150,17 @@ impl Dpb { e.mark = RefMark::Unused; } self.prune(); - for _ in 0..skip { - poc_state(unused); - unused = (unused + 1) % max_frame_num; + // The POC state only advances FrameNumOffset where frame_num decreases, so walking the + // skipped run is the same as visiting its first value, the wrap (if the run crosses + // it) and its last value: constant work however long the gap. + let last = (unused + (skip - 1) % max_frame_num) % max_frame_num; + poc_state(unused); + if last < unused { + poc_state(max_frame_num - 1); + poc_state(0); } + poc_state(last); + unused = (last + 1) % max_frame_num; } let mut guard = 0; while unused != frame_num && guard < max_frame_num { @@ -438,14 +446,15 @@ impl Dpb { } } } else { - while self.entries.len() >= self.capacity { + // Evict until there is room (frame_num gap concealment can leave the DPB over-full), so + // the store never holds more pictures than the memory budget was checked for. + while self.entries.len() >= self.capacity.max(1) { if !self.bump(out) { if let Some(pos) = self.entries.iter().position(|e| !e.needed_for_output) { self.entries.remove(pos); } else { self.entries.remove(0); } - break; } } } @@ -546,6 +555,37 @@ mod tests { (made, walk, left) } + /// (FrameNumOffset, prev frame_num) after `walk`, by the rule of `PocState::update_gap_frame`. + fn poc_after(walk: &[u32], prev: u32, max: u32) -> (u64, u32) { + walk.iter().fold((0, prev), |(off, prev), &f| (if prev > f { off + u64::from(max) } else { off }, f)) + } + + #[test] + fn gap_walk_matches_every_frame_num_for_any_skip() { + for (prev, cur) in [(1, 9), (10, 3), (14, 0), (15, 14), (0, 15), (7, 6)] { + for refs in 1..4 { + let mut dpb = Dpb::new(); + let (_, walk, _) = fill(&mut dpb, prev, cur, 16, refs); + let every: Vec = (1..16).map(|i| (prev + i) % 16).take_while(|&f| f != cur).collect(); + assert_eq!(poc_after(&walk, prev, 16), poc_after(&every, prev, 16), "{prev} -> {cur}, {refs} refs"); + assert_eq!(walk.last(), every.last(), "{prev} -> {cur}, {refs} refs"); + } + } + } + + #[test] + fn reference_store_evicts_down_to_capacity() { + let mut dpb = Dpb::new(); + dpb.capacity = 2; + // Gap concealment can leave more non-output references than the capacity. + for i in 0..5 { + short(&mut dpb, i, i); + } + let mut out = Vec::new(); + assert!(!dpb.bump_for_space(RefMark::Short, 0, &mut out)); + assert!(dpb.entries.len() < 2, "{} pictures left before storing into a 2-picture DPB", dpb.entries.len()); + } + #[test] fn huge_gap_makes_only_the_frames_that_stay_referenced() { let mut dpb = Dpb::new(); @@ -553,8 +593,9 @@ mod tests { short(&mut dpb, 2, 1); let (made, walk, left) = fill(&mut dpb, 1, 65535, 65536, 4); assert_eq!(made, 4); - // The POC state still walks every missing frame_num, in order. - assert_eq!(walk.len(), 65533); + // The POC state ends where walking every missing frame_num would leave it, in a few steps. + assert_eq!(poc_after(&walk, 1, 65536), poc_after(&(2..65535).collect::>(), 1, 65536)); + assert!(walk.len() <= 8, "{} POC updates", walk.len()); assert_eq!(walk.first(), Some(&2)); assert_eq!(walk.last(), Some(&65534)); assert_eq!(left, vec![(65531, true), (65532, true), (65533, true), (65534, true)]); @@ -566,7 +607,9 @@ mod tests { short(&mut dpb, 1, 65000); let (made, walk, left) = fill(&mut dpb, 65000, 3, 65536, 2); assert_eq!(made, 2); - assert_eq!(walk.len(), 538); + let every: Vec = (65001..65536).chain(0..3).collect(); + assert_eq!(poc_after(&walk, 65000, 65536), poc_after(&every, 65000, 65536)); + assert!(walk.len() <= 6, "{} POC updates", walk.len()); assert_eq!(walk.first(), Some(&65001)); assert_eq!(walk.last(), Some(&2)); assert_eq!(left, vec![(1, true), (2, true)]); diff --git a/crates/h264/src/picture.rs b/crates/h264/src/picture.rs index 39add6a08..83331cb95 100644 --- a/crates/h264/src/picture.rs +++ b/crates/h264/src/picture.rs @@ -160,6 +160,14 @@ impl Frame { self.rows.len() } + /// Sample bytes this picture holds (luma and both chroma planes), as the decoder's memory + /// budget counts a picture. + pub fn bytes(&self) -> u64 { + (self.width as u64) + .saturating_mul(self.height as u64) + .saturating_add((self.cwidth as u64).saturating_mul(self.cheight as u64).saturating_mul(2)) + } + /// Publish MB row `r` (ignored if already published). pub fn publish(&self, r: usize, row: FrameRow) { let _ = self.rows[r].set(row); diff --git a/crates/h264/src/slicedec.rs b/crates/h264/src/slicedec.rs index 058ccfe16..649739986 100644 --- a/crates/h264/src/slicedec.rs +++ b/crates/h264/src/slicedec.rs @@ -346,8 +346,11 @@ impl<'a> SliceDecoder<'a> { // Neighbour helpers // ------------------------------------------------------------------------------------------ - /// Set up the current macroblock address and neighbour availability. - pub fn start_mb(&mut self, addr: usize) { + /// Set up the current macroblock address and neighbour availability. A macroblock an earlier + /// slice of this picture already decoded is refused: overlapping slices would otherwise redo + /// the picture's reconstruction once per slice. + pub fn start_mb(&mut self, addr: usize) -> Result<()> { + ensure!(self.pic.mbs[addr].slice_num == u32::MAX, "macroblock {addr} overlaps an earlier slice"); self.mb_addr = addr; self.mb_x = addr % self.mb_w; self.mb_y = addr / self.mb_w; @@ -364,6 +367,7 @@ impl<'a> SliceDecoder<'a> { ]; self.mv_done = 0; self.pic.mbs[addr] = MbState { slice_num: sn, ..MbState::default() }; + Ok(()) } #[inline] @@ -1467,7 +1471,7 @@ impl<'a> SliceDecoder<'a> { let run = r.read_ue()? as usize; ensure!(addr + run <= total, "mb_skip_run out of range"); for _ in 0..run { - self.start_mb(addr); + self.start_mb(addr)?; self.decode_skip()?; self.pic.mb_done(addr); addr += 1; @@ -1477,7 +1481,7 @@ impl<'a> SliceDecoder<'a> { } ensure!(addr < total, "macroblock address {addr} out of range"); } - self.start_mb(addr); + self.start_mb(addr)?; self.decode_mb_cavlc(r)?; self.pic.mb_done(addr); addr += 1; diff --git a/crates/h264/src/synth_tests.rs b/crates/h264/src/synth_tests.rs index 76609dbde..199f5e69d 100644 --- a/crates/h264/src/synth_tests.rs +++ b/crates/h264/src/synth_tests.rs @@ -113,6 +113,7 @@ fn avcc_rejects_a_valid_picture_slice_without_starting_a_picture() { poc_lsb: Some(0), poc_delta: None, wt: None, + first_mb: 0, }, false, ); @@ -153,6 +154,8 @@ struct Pic { /// Explicit weights (luma, cb, cr) as (weight, offset) for refIdx 0 of list 0 and list 1; /// log2 denominators are 5 (luma) and 3 (chroma). wt: Option<[[(i32, i32); 3]; 2]>, + /// first_mb_in_slice (`mbs` then lists the slice's macroblocks from there). + first_mb: u32, } impl Pic { @@ -165,7 +168,7 @@ impl Pic { } fn slice_header(w: &mut BitWriter, p: &Pic, cabac: bool) { - w.write_ue(0); // first_mb + w.write_ue(p.first_mb); w.write_ue(if p.b { 6 } else if p.is_p() { @@ -660,3 +663,63 @@ fn frame_num_gap_is_filled_when_the_sps_allows_gaps() { assert_eq!(out.len(), 2); assert_eq!(dec.stats().frame_num_gaps, 1); } + +/// An IDR picture followed by `pictures` reference P pictures of P_Skip macroblocks. +fn skip_stream(pictures: u32) -> (Vec, Vec>) { + let mut seed = 11u32; + let a: Vec>> = (0..NMB).map(|_| Some(random_mb(&mut seed))).collect(); + let mut head = sps(66, 2, false, PocType::T2); + head.extend(pps(false, false)); + head.extend(slice_nal(&Pic { idr: true, frame_num: 0, ref_idc: 3, mbs: a, ..Default::default() }, false)); + let ps = (1..=pictures).map(|k| slice_nal(&Pic { frame_num: k % 16, ref_idc: 2, mbs: vec![None; NMB], ..Default::default() }, false)).collect(); + (head, ps) +} + +#[test] +fn overlapping_slices_are_not_reconstructed_twice() { + let (head, ps) = skip_stream(1); + let mut once = Decoder::with_threads(1); + let mut out = once.decode(&[head.clone(), ps[0].clone()].concat(), 0).unwrap(); + out.extend(once.flush()); + assert!(once.take_error().is_none()); + assert_eq!(out.len(), 2); + // A second slice of the same picture starting at macroblock 1 covers those macroblocks again. + let again = slice_nal(&Pic { frame_num: 1, ref_idc: 2, mbs: vec![None; NMB - 1], first_mb: 1, ..Default::default() }, false); + let mut twice = Decoder::with_threads(1); + let result = twice.decode(&[head, ps[0].clone(), again].concat(), 0); + twice.flush(); + let err = result.err().or_else(|| twice.take_error()).expect("overlapping slice refused"); + assert!(err.to_string().contains("overlaps"), "{err}"); +} + +#[test] +fn output_queue_counts_against_the_memory_budget() { + let (head, ps) = skip_stream(40); + let sps_nal = sps(66, 2, false, PocType::T2); + let parsed = crate::params::Sps::parse(&deckcraft_bitstream::unescape_rbsp(&sps_nal[5..])).unwrap(); + let one = (NMB * 256 + 2 * NMB * 64) as u64; + let budget = (parsed.max_dpb_frames() + 1 + 3) as u64 * one; + // One access unit per call: output drains as it goes, and the stream fits. + let mut dec = Decoder::with_threads(1); + dec.set_memory_budget(budget); + let mut out = dec.decode(&head, 0).unwrap(); + for (i, p) in ps.iter().enumerate() { + out.extend(dec.decode(p, i as i64 + 1).unwrap()); + } + out.extend(dec.flush()); + assert!(dec.take_error().is_none()); + assert_eq!(out.len(), 41); + // Every access unit in one call: the pictures held for output stop fitting the budget. + let mut dec = Decoder::with_threads(1); + dec.set_memory_budget(budget); + let err = dec.decode(&[head, ps.concat()].concat(), 0).expect_err("output queue is over budget"); + assert!(err.to_string().contains("memory budget"), "{err}"); +} + +/// A held picture is charged at its own size (luma and both 4:2:0 chroma planes), the same +/// count the budget uses for the incoming SPS's pictures. +#[test] +fn a_held_frame_is_charged_at_its_own_size() { + let f = crate::picture::Frame::new(0, 0, 2, 3); + assert_eq!(f.bytes(), (2 * 3 * 256 + 2 * 2 * 3 * 64) as u64); +}