Compare commits

..
Author SHA1 Message Date
Dylan M. Taylor c32cdf7eca refactor(av1): split OBU parsing in decoder (S3776) 2026-10-25 19:05:31 -04:00
7 changed files with 154 additions and 114 deletions
+78 -55
View File
@@ -405,49 +405,35 @@ impl Decoder {
r
}
/// Whether an OBU belongs to the selected operating point (spec 7.5: OBUs of other layers are dropped).
fn in_operating_point(&self, obu: &ObuUnit) -> bool {
if obu.obu_type != OBU_SEQUENCE_HEADER_T
&& obu.obu_type != OBU_TEMPORAL_DELIMITER_T
&& obu.ext
&& let Some(seq) = &self.seq
&& seq.op_idc != 0
{
let idc = seq.op_idc;
let in_t = (idc >> obu.temporal_id) & 1;
let in_s = (idc >> (obu.spatial_id + 8)) & 1;
if in_t == 0 || in_s == 0 {
return false;
}
}
true
}
fn parse_obus(&mut self, data: &[u8], out: &mut Vec<(Output, usize, u32)>) -> Result<()> {
let mut tu = 0usize;
let mut pos = 0;
while pos < data.len() {
let h = data[pos];
if h & 0x80 != 0 {
return Err(Error::Invalid("obu_forbidden_bit"));
let obu = parse_obu_unit(data, pos)?;
pos = obu.end;
if !self.in_operating_point(&obu) {
continue;
}
let obu_type = ((h >> 3) & 0xf) as u32;
let ext = h & 4 != 0;
let has_size = h & 2 != 0;
let mut p = pos + 1;
let (mut temporal_id, mut spatial_id) = (0, 0);
if ext {
let e = *data.get(p).ok_or(Error::Truncated)?;
temporal_id = (e >> 5) as u32;
spatial_id = ((e >> 3) & 3) as u32;
p += 1;
}
let size = if has_size {
let (v, n) = leb128(data.get(p..).ok_or(Error::Truncated)?)?;
p += n;
v as usize
} else {
data.len().checked_sub(p).ok_or(Error::Truncated)?
};
let end = p.checked_add(size).filter(|&e| e <= data.len()).ok_or(Error::Truncated)?;
let payload = &data[p..end];
pos = end;
if obu_type != OBU_SEQUENCE_HEADER_T
&& obu_type != OBU_TEMPORAL_DELIMITER_T
&& ext
&& let Some(seq) = &self.seq
&& seq.op_idc != 0
{
let idc = seq.op_idc;
let in_t = (idc >> temporal_id) & 1;
let in_s = (idc >> (spatial_id + 8)) & 1;
if in_t == 0 || in_s == 0 {
continue;
}
}
match obu_type {
let (payload, spatial_id) = (obu.payload, obu.spatial_id);
match obu.obu_type {
OBU_SEQUENCE_HEADER_T => {
let s = SequenceHeader::parse(payload)?;
self.seq = Some(s);
@@ -457,27 +443,11 @@ impl Decoder {
tu += 1;
}
OBU_FRAME_HEADER_T | OBU_REDUNDANT_FRAME_HEADER_T | OBU_FRAME_T => {
if self.seen_frame_header {
// frame_header_copy(): identical to the active header
if obu_type != OBU_FRAME_T {
continue;
}
}
let mut r = BitReader::new(payload);
let shown = self.frame_header(&mut r, temporal_id, spatial_id)?;
if let Some(o) = shown {
if let Some(o) = self.frame_obu(&obu)? {
if out.iter().filter(|(o, _, _)| matches!(o, Output::Existing { .. })).count() >= 32 {
return Err(Error::Invalid("too many frame outputs in chunk"));
}
out.push((o, tu, spatial_id));
continue;
}
if obu_type == OBU_FRAME_T {
r.byte_align();
let rest = payload.get(r.byte_pos()..).ok_or(Error::Truncated)?;
if let Some(o) = self.tile_group(rest)? {
out.push((o, tu, spatial_id));
}
}
}
OBU_TILE_GROUP_T => {
@@ -492,6 +462,22 @@ impl Decoder {
Ok(())
}
/// A frame header or frame OBU: returns the picture it completes, if any.
fn frame_obu(&mut self, obu: &ObuUnit) -> Result<Option<Output>> {
// frame_header_copy(): identical to the active header
if self.seen_frame_header && obu.obu_type != OBU_FRAME_T {
return Ok(None);
}
let mut r = BitReader::new(obu.payload);
let shown = self.frame_header(&mut r, obu.temporal_id, obu.spatial_id)?;
if shown.is_some() || obu.obu_type != OBU_FRAME_T {
return Ok(shown);
}
r.byte_align();
let rest = obu.payload.get(r.byte_pos()..).ok_or(Error::Truncated)?;
self.tile_group(rest)
}
/// frame_header_obu( ): returns the output for show_existing_frame.
fn frame_header(&mut self, r: &mut BitReader, temporal_id: u32, spatial_id: u32) -> Result<Option<Output>> {
let seq = self.seq.clone().ok_or(Error::Invalid("frame before sequence header"))?;
@@ -1017,6 +1003,43 @@ mod tests {
}
}
/// One OBU: its header fields, payload and the position just past it.
struct ObuUnit<'a> {
obu_type: u32,
ext: bool,
temporal_id: u32,
spatial_id: u32,
payload: &'a [u8],
end: usize,
}
fn parse_obu_unit(data: &[u8], pos: usize) -> Result<ObuUnit<'_>> {
let h = data[pos];
if h & 0x80 != 0 {
return Err(Error::Invalid("obu_forbidden_bit"));
}
let obu_type = ((h >> 3) & 0xf) as u32;
let ext = h & 4 != 0;
let has_size = h & 2 != 0;
let mut p = pos + 1;
let (mut temporal_id, mut spatial_id) = (0, 0);
if ext {
let e = *data.get(p).ok_or(Error::Truncated)?;
temporal_id = (e >> 5) as u32;
spatial_id = ((e >> 3) & 3) as u32;
p += 1;
}
let size = if has_size {
let (v, n) = leb128(data.get(p..).ok_or(Error::Truncated)?)?;
p += n;
v as usize
} else {
data.len().checked_sub(p).ok_or(Error::Truncated)?
};
let end = p.checked_add(size).filter(|&e| e <= data.len()).ok_or(Error::Truncated)?;
Ok(ObuUnit { obu_type, ext, temporal_id, spatial_id, payload: &data[p..end], end })
}
#[cfg(test)]
mod guard_tests {
use super::*;
+1 -1
View File
@@ -250,7 +250,7 @@ impl Frame {
/// Copy the w x h window of chroma component `c` at (x0, y0) (clamped) into `out`.
#[inline]
pub fn chroma_window(&self, c: usize, (x0, y0): (i32, i32), (w, h): (usize, usize), out: &mut [u8], os: usize) {
pub fn chroma_window(&self, c: usize, x0: i32, y0: i32, w: usize, h: usize, out: &mut [u8], os: usize) {
let (cw, ch) = (self.cwidth, self.cheight);
for r in 0..h {
let yy = (y0 + r as i32).clamp(0, ch as i32 - 1) as usize;
+1 -1
View File
@@ -232,7 +232,7 @@ fn predict_block(f: &Frame, mv: [i16; 2], px: usize, py: usize, bw: usize, bh: u
let (fx, fy) = ((mv[0] & 7) as u32, (mv[1] & 7) as u32);
for (comp, out) in c.into_iter().enumerate() {
let mut w = [0u8; 16 * 9];
f.chroma_window(comp, (cx, cy), (cw + 1, ch + 1), &mut w, 16);
f.chroma_window(comp, cx, cy, cw + 1, ch + 1, &mut w, 16);
inter::mc_chroma_win(&w, 16, fx, fy, cw, ch, out, cs);
}
}
+9 -4
View File
@@ -128,7 +128,7 @@ pub fn mc_chroma(f: &Frame, c: usize, x: i32, y: i32, mv: [i16; 2], w: usize, h:
let yi = y + (mv[1] as i32 >> 3);
let shift1 = bd.min(12) - 8;
let shift3 = 14 - bd;
f.chroma_window(c, (xi - 1, yi - 1), (w + 3, h + 3), &mut sc.win, WIN);
f.chroma_window(c, xi - 1, yi - 1, w + 3, h + 3, &mut sc.win, WIN);
interpolate::<4>(&sc.win, &mut sc.tmp, fx, fy, &CHROMA_FILTER[fx], &CHROMA_FILTER[fy], w, h, shift1, shift3, out);
}
@@ -158,7 +158,7 @@ pub fn put_bi(p0: &[i16], p1: &[i16], w: usize, h: usize, bd: u32, dst: &mut [u1
}
/// Explicit weighted prediction, one list (8-265). `o` is the offset already scaled to the bit depth.
pub fn put_weighted_uni(p: &[i16], (w, h): (usize, usize), bd: u32, (log2wd, wt, o): (u32, i32, i32), dst: &mut [u16], ds: usize) {
pub fn put_weighted_uni(p: &[i16], w: usize, h: usize, bd: u32, log2wd: u32, wt: i32, o: i32, dst: &mut [u16], ds: usize) {
let max = (1i32 << bd) - 1;
for r in 0..h {
for c in 0..w {
@@ -173,9 +173,14 @@ pub fn put_weighted_uni(p: &[i16], (w, h): (usize, usize), bd: u32, (log2wd, wt,
pub fn put_weighted_bi(
p0: &[i16],
p1: &[i16],
(w, h): (usize, usize),
w: usize,
h: usize,
bd: u32,
(log2wd, w0, w1, o0, o1): (u32, i32, i32, i32, i32),
log2wd: u32,
w0: i32,
w1: i32,
o0: i32,
o1: i32,
dst: &mut [u16],
ds: usize,
) {
+34 -23
View File
@@ -36,23 +36,9 @@ fn same_motion(a: &MvField, b: &MvField) -> bool {
true
}
/// Coding block and prediction block geometry (6.4.2): coding block (xc, yc, ncb), prediction block (xp, yp, w, h, part_idx).
#[derive(Clone, Copy)]
pub struct PbGeom {
pub xc: i32,
pub yc: i32,
pub ncb: i32,
pub xp: i32,
pub yp: i32,
pub w: i32,
pub h: i32,
pub part_idx: usize,
}
impl<'a> SliceDecoder<'a> {
/// Prediction block availability (6.4.2).
fn pb_avail(&self, g: &PbGeom, xn: i32, yn: i32) -> bool {
let PbGeom { xc, yc, ncb, xp, yp, w, h, part_idx } = *g;
fn pb_avail(&self, xc: i32, yc: i32, ncb: i32, xp: i32, yp: i32, w: i32, h: i32, part_idx: usize, xn: i32, yn: i32) -> bool {
let same_cb = xc <= xn && yc <= yn && xc + ncb > xn && yc + ncb > yn;
let avail = if !same_cb {
self.pic.zavail(xp, yp, xn, yn, self.slice_addr)
@@ -73,13 +59,23 @@ impl<'a> SliceDecoder<'a> {
}
/// Merge mode motion (8.5.3.2.2).
pub fn derive_merge(&self, g: PbGeom, part: PartMode, merge_idx: usize) -> MvField {
let PbGeom { xc, yc, ncb, xp: xp0, yp: yp0, w: w0, h: h0, part_idx: part_idx0 } = g;
pub fn derive_merge(
&self,
xc: i32,
yc: i32,
ncb: i32,
xp0: i32,
yp0: i32,
w0: i32,
h0: i32,
part_idx0: usize,
part: PartMode,
merge_idx: usize,
) -> MvField {
let pml = self.pps.log2_parallel_merge_level;
let (xp, yp, w, h, part_idx) = if pml > 2 && ncb == 8 { (xc, yc, ncb, ncb, 0) } else { (xp0, yp0, w0, h0, part_idx0) };
let g = PbGeom { xc, yc, ncb, xp, yp, w, h, part_idx };
let same_mer = |xn: i32, yn: i32| (xp >> pml) == (xn >> pml) && (yp >> pml) == (yn >> pml);
let avail = |xn: i32, yn: i32| self.pb_avail(&g, xn, yn) && !same_mer(xn, yn);
let avail = |xn: i32, yn: i32| self.pb_avail(xc, yc, ncb, xp, yp, w, h, part_idx, xn, yn) && !same_mer(xn, yn);
let mut list: [MvField; 5] = [MvField::default(); 5];
let mut n = 0;
let max = self.sh.max_num_merge_cand as usize;
@@ -241,8 +237,20 @@ impl<'a> SliceDecoder<'a> {
}
/// AMVP predictor (8.5.3.2.6, 8.5.3.2.7).
pub fn derive_amvp(&self, g: PbGeom, l: usize, ref_idx: usize, mvp_flag: usize) -> [i16; 2] {
let PbGeom { xp, yp, w, h, .. } = g;
pub fn derive_amvp(
&self,
xc: i32,
yc: i32,
ncb: i32,
xp: i32,
yp: i32,
w: i32,
h: i32,
part_idx: usize,
l: usize,
ref_idx: usize,
mvp_flag: usize,
) -> [i16; 2] {
let target = &self.refs[l][ref_idx];
let (tpoc, tlt) = (target.poc, target.long_term);
let cur_poc = self.pic.poc;
@@ -269,7 +277,10 @@ impl<'a> SliceDecoder<'a> {
pick.map(|(mv, r)| if r.poc != tpoc && !r.long_term && !tlt { scale_mv(mv, cur_poc - r.poc, cur_poc - tpoc) } else { mv })
};
let pa = [(xp - 1, yp + h), (xp - 1, yp + h - 1)];
let av_a = [self.pb_avail(&g, pa[0].0, pa[0].1), self.pb_avail(&g, pa[1].0, pa[1].1)];
let av_a = [
self.pb_avail(xc, yc, ncb, xp, yp, w, h, part_idx, pa[0].0, pa[0].1),
self.pb_avail(xc, yc, ncb, xp, yp, w, h, part_idx, pa[1].0, pa[1].1),
];
let is_scaled = av_a[0] || av_a[1];
let mut a = None;
for k in 0..2 {
@@ -285,7 +296,7 @@ impl<'a> SliceDecoder<'a> {
}
}
let pb = [(xp + w, yp - 1), (xp + w - 1, yp - 1), (xp - 1, yp - 1)];
let av_b = [0, 1, 2].map(|k| self.pb_avail(&g, pb[k].0, pb[k].1));
let av_b = [0, 1, 2].map(|k| self.pb_avail(xc, yc, ncb, xp, yp, w, h, part_idx, pb[k].0, pb[k].1));
let mut b = None;
for k in 0..3 {
if av_b[k] && b.is_none() {
+1 -1
View File
@@ -125,7 +125,7 @@ impl Frame {
}
/// Copy a w x h window of chroma component `c` (0 = Cb, 1 = Cr) with edge replication.
pub fn chroma_window(&self, c: usize, (x0, y0): (i32, i32), (w, h): (usize, usize), out: &mut [i16], os: usize) {
pub fn chroma_window(&self, c: usize, x0: i32, y0: i32, w: usize, h: usize, out: &mut [i16], os: usize) {
let (cw, ch) = (self.cwidth as i32, self.cheight as i32);
let sh = self.log2_ctb - 1;
let mask = (1i32 << sh) - 1;
+30 -29
View File
@@ -7,7 +7,6 @@ use crate::cabac::{Cabac, Contexts, init_contexts};
use crate::error::{Result, ensure, invalid};
use crate::inter;
use crate::intra::{self, Refs};
use crate::mvpred::PbGeom;
use crate::params::{Layout, Pps, Sps};
use crate::picture::{FrameRef, RefPic};
use crate::slice::{SliceHeader, SliceType};
@@ -679,7 +678,7 @@ impl<'a> SliceDecoder<'a> {
self.mark_edges(x0, y0, n, n, E_TU_V | E_PU_V, E_TU_H | E_PU_H);
if skip {
self.cu_intra = false;
self.prediction_unit(PbGeom { xc: x0, yc: y0, ncb: n, xp: x0, yp: y0, w: n, h: n, part_idx: 0 }, PartMode::P2Nx2N, true)?;
self.prediction_unit(x0, y0, n, x0, y0, n, n, 0, PartMode::P2Nx2N, true)?;
self.finish_cu(x0, y0, n);
return Ok(());
}
@@ -708,7 +707,7 @@ impl<'a> SliceDecoder<'a> {
}
self.parse_intra_modes(x0, y0, n, part)?;
self.max_trafo_depth = self.sps.max_th_depth_intra + self.intra_split as u32;
self.transform_tree((x0, y0), (x0, y0), log2, 0, 0, [true, true])?;
self.transform_tree(x0, y0, x0, y0, log2, 0, 0, [true, true])?;
} else {
self.intra_split = false;
let mut merge2nx2n = false;
@@ -725,8 +724,7 @@ impl<'a> SliceDecoder<'a> {
PartMode::PNxN => &[(0, 0, h, h), (h, 0, h, h), (0, h, h, h), (h, h, h, h)],
};
for (i, &(px, py, pw, ph)) in parts.iter().enumerate() {
let merge =
self.prediction_unit(PbGeom { xc: x0, yc: y0, ncb: n, xp: x0 + px, yp: y0 + py, w: pw, h: ph, part_idx: i }, part, false)?;
let merge = self.prediction_unit(x0, y0, n, x0 + px, y0 + py, pw, ph, i, part, false)?;
if part == PartMode::P2Nx2N {
merge2nx2n = merge;
}
@@ -740,7 +738,7 @@ impl<'a> SliceDecoder<'a> {
let root_cbf = if !merge2nx2n { self.dd(RQT_ROOT_CBF) == 1 } else { true };
if root_cbf {
self.max_trafo_depth = self.sps.max_th_depth_inter;
self.transform_tree((x0, y0), (x0, y0), log2, 0, 0, [true, true])?;
self.transform_tree(x0, y0, x0, y0, log2, 0, 0, [true, true])?;
}
}
self.finish_cu(x0, y0, n);
@@ -914,7 +912,7 @@ impl<'a> SliceDecoder<'a> {
Ok(())
}
fn transform_tree(&mut self, (x0, y0): (i32, i32), base: (i32, i32), log2: u32, depth: u32, blk_idx: u32, parent_cbf_c: [bool; 2]) -> Result<()> {
fn transform_tree(&mut self, x0: i32, y0: i32, xb: i32, yb: i32, log2: u32, depth: u32, blk_idx: u32, parent_cbf_c: [bool; 2]) -> Result<()> {
let split =
if log2 <= self.sps.log2_max_tb && log2 > self.sps.log2_min_tb && depth < self.max_trafo_depth && !(self.intra_split && depth == 0) {
self.dd(SPLIT_TRANSFORM + (5 - log2 as usize)) == 1
@@ -939,10 +937,10 @@ impl<'a> SliceDecoder<'a> {
}
if split {
let h = 1 << (log2 - 1);
self.transform_tree((x0, y0), (x0, y0), log2 - 1, depth + 1, 0, cbf_c)?;
self.transform_tree((x0 + h, y0), (x0, y0), log2 - 1, depth + 1, 1, cbf_c)?;
self.transform_tree((x0, y0 + h), (x0, y0), log2 - 1, depth + 1, 2, cbf_c)?;
self.transform_tree((x0 + h, y0 + h), (x0, y0), log2 - 1, depth + 1, 3, cbf_c)?;
self.transform_tree(x0, y0, x0, y0, log2 - 1, depth + 1, 0, cbf_c)?;
self.transform_tree(x0 + h, y0, x0, y0, log2 - 1, depth + 1, 1, cbf_c)?;
self.transform_tree(x0, y0 + h, x0, y0, log2 - 1, depth + 1, 2, cbf_c)?;
self.transform_tree(x0 + h, y0 + h, x0, y0, log2 - 1, depth + 1, 3, cbf_c)?;
return Ok(());
}
let cbf_luma =
@@ -951,7 +949,7 @@ impl<'a> SliceDecoder<'a> {
if depth > 0 {
self.mark_edges(x0, y0, n, n, E_TU_V, E_TU_H);
}
self.transform_unit((x0, y0), base, log2, blk_idx, cbf_luma, cbf_c)
self.transform_unit(x0, y0, xb, yb, log2, blk_idx, cbf_luma, cbf_c)
}
/// cu_qp_delta_abs / sign and QP update (7.3.8.14).
@@ -981,15 +979,7 @@ impl<'a> SliceDecoder<'a> {
Ok(())
}
fn transform_unit(
&mut self,
(x0, y0): (i32, i32),
(xb, yb): (i32, i32),
log2: u32,
blk_idx: u32,
cbf_luma: bool,
cbf_c: [bool; 2],
) -> Result<()> {
fn transform_unit(&mut self, x0: i32, y0: i32, xb: i32, yb: i32, log2: u32, blk_idx: u32, cbf_luma: bool, cbf_c: [bool; 2]) -> Result<()> {
let n = 1i32 << log2;
let chroma_here = log2 > 2 || blk_idx == 3;
let cbf_chroma = cbf_c[0] || cbf_c[1];
@@ -1294,7 +1284,7 @@ impl<'a> SliceDecoder<'a> {
max_y = max_y.max(yc);
}
}
self.reconstruct((x0, y0), log2, c, qp, ts, (max_x, max_y))
self.reconstruct(x0, y0, log2, c, qp, ts, max_x, max_y)
}
fn coeff_abs_level_remaining(&mut self, rice: u32) -> Result<i32> {
@@ -1313,7 +1303,7 @@ impl<'a> SliceDecoder<'a> {
}
/// Scaling, transform and addition to the prediction.
fn reconstruct(&mut self, (x0, y0): (i32, i32), log2: u32, c: usize, qp: i32, ts: bool, (max_x, max_y): (usize, usize)) -> Result<()> {
fn reconstruct(&mut self, x0: i32, y0: i32, log2: u32, c: usize, qp: i32, ts: bool, max_x: usize, max_y: usize) -> Result<()> {
let n = 1usize << log2;
let bd = if c == 0 { self.pic.bd_y } else { self.pic.bd_c };
let coeffs = &mut self.coeffs[..n * n];
@@ -1365,8 +1355,19 @@ impl<'a> SliceDecoder<'a> {
}
/// prediction_unit(): parse, derive motion, store it and predict. Returns merge_flag.
fn prediction_unit(&mut self, g: PbGeom, part: PartMode, skip: bool) -> Result<bool> {
let PbGeom { xc, yc, ncb, xp, yp, w, h, part_idx } = g;
fn prediction_unit(
&mut self,
xc: i32,
yc: i32,
ncb: i32,
xp: i32,
yp: i32,
w: i32,
h: i32,
part_idx: usize,
part: PartMode,
skip: bool,
) -> Result<bool> {
let merge = if skip { true } else { self.dd(MERGE_FLAG) == 1 };
let mvf = if merge {
let idx = if self.sh.max_num_merge_cand > 1 {
@@ -1381,7 +1382,7 @@ impl<'a> SliceDecoder<'a> {
} else {
0
};
self.derive_merge(PbGeom { xc, yc, ncb, xp, yp, w, h, part_idx }, part, idx as usize)
self.derive_merge(xc, yc, ncb, xp, yp, w, h, part_idx, part, idx as usize)
} else {
let pred_idc = if self.sh.is_b() {
if w + h != 12 {
@@ -1425,7 +1426,7 @@ impl<'a> SliceDecoder<'a> {
if ref_idx[l] < 0 {
continue;
}
let mvp = self.derive_amvp(PbGeom { xc, yc, ncb, xp, yp, w, h, part_idx }, l, ref_idx[l] as usize, mvp_flag[l] as usize);
let mvp = self.derive_amvp(xc, yc, ncb, xp, yp, w, h, part_idx, l, ref_idx[l] as usize, mvp_flag[l] as usize);
// mvLX = (mvp + mvd) wrapped to 16 bits (8-94)
f.mv[l] = [(mvp[0] as i32 + mvd[l][0]) as i16, (mvp[1] as i32 + mvd[l][1]) as i16];
}
@@ -1523,11 +1524,11 @@ impl<'a> SliceDecoder<'a> {
if bi {
let (w0, o0) = get(0);
let (w1, o1) = get(1);
inter::put_weighted_bi(&self.pred0, &self.pred1, (bw, bh), bd, (log2wd, w0, w1, o0, o1), dst, stride);
inter::put_weighted_bi(&self.pred0, &self.pred1, bw, bh, bd, log2wd, w0, w1, o0, o1, dst, stride);
} else {
let l = if lists[0] { 0 } else { 1 };
let (wt, o) = get(l);
inter::put_weighted_uni(&self.pred0, (bw, bh), bd, (log2wd, wt, o), dst, stride);
inter::put_weighted_uni(&self.pred0, bw, bh, bd, log2wd, wt, o, dst, stride);
}
}
}