Compare commits

..
Author SHA1 Message Date
Dylan M. TaylorandClaude Sonnet 5.5 3aa190edf4 refactor(av1): split inverse transform into stages and passes (S3776)
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
2026-10-24 08:51:53 +00:00
2 changed files with 199 additions and 202 deletions
+47 -3
View File
@@ -165,6 +165,15 @@ pub(crate) fn inverse_dct<T: Lane>(t: &mut [T], n: u32, r: u32) {
for i in 0..len {
t[i] = copy[brev(n, i)];
}
inverse_dct_stage1(t, n, r);
inverse_dct_stage2(t, n, r);
inverse_dct_stage3(t, n, r);
inverse_dct_stage4(t, n, r);
inverse_dct_stage5(t, n, r);
inverse_dct_stage6(t, n, r);
}
fn inverse_dct_stage1<T: Lane>(t: &mut [T], n: u32, r: u32) {
if n == 6 {
for i in 0..16 {
bf(t, 32 + i, 63 - i, 63 - 4 * brev(4, i) as i32, false);
@@ -190,6 +199,9 @@ pub(crate) fn inverse_dct<T: Lane>(t: &mut [T], n: u32, r: u32) {
hd(t, 16 + 2 * i, 17 + 2 * i, i & 1 == 1, r);
}
}
}
fn inverse_dct_stage2<T: Lane>(t: &mut [T], n: u32, r: u32) {
if n == 6 {
for i in 0..4 {
for j in 0..2 {
@@ -221,6 +233,9 @@ pub(crate) fn inverse_dct<T: Lane>(t: &mut [T], n: u32, r: u32) {
}
}
}
}
fn inverse_dct_stage3<T: Lane>(t: &mut [T], n: u32, r: u32) {
for i in 0..2 {
bf(t, 2 * i, 2 * i + 1, 32 + 16 * i as i32, i == 0);
}
@@ -248,6 +263,9 @@ pub(crate) fn inverse_dct<T: Lane>(t: &mut [T], n: u32, r: u32) {
}
}
}
}
fn inverse_dct_stage4<T: Lane>(t: &mut [T], n: u32, r: u32) {
for i in 0..2 {
hd(t, i, 3 - i, false, r);
}
@@ -273,6 +291,9 @@ pub(crate) fn inverse_dct<T: Lane>(t: &mut [T], n: u32, r: u32) {
}
}
}
}
fn inverse_dct_stage5<T: Lane>(t: &mut [T], n: u32, r: u32) {
if n >= 3 {
for i in 0..4 {
hd(t, i, 7 - i, false, r);
@@ -300,6 +321,9 @@ pub(crate) fn inverse_dct<T: Lane>(t: &mut [T], n: u32, r: u32) {
hd(t, i, 15 - i, false, r);
}
}
}
fn inverse_dct_stage6<T: Lane>(t: &mut [T], n: u32, r: u32) {
if n >= 5 {
for i in 0..4 {
bf(t, 27 - i, 20 + i, 32, true);
@@ -531,10 +555,9 @@ pub(crate) fn inverse_transform_2d(dequant: &[i32], residual: &mut [i32], tx_sz:
let col_clamp = (bit_depth + 6).max(16);
let rk = row_kind(tx_type);
let ck = col_kind(tx_type);
let mut t = [0i32; 64];
let rect2 = log2w.abs_diff(log2h) == 1;
let p = Tx2d { w, h, log2w, log2h, row_shift, col_shift, row_clamp, col_clamp, rk, ck, rect2 };
if !lossless && bit_depth <= 10 {
let p = Tx2d { w, h, log2w, log2h, row_shift, col_shift, row_clamp, col_clamp, rk, ck, rect2 };
if w >= 8 && h >= 8 {
transform_lanes::<8>(dequant, residual, &p);
} else {
@@ -542,6 +565,13 @@ pub(crate) fn inverse_transform_2d(dequant: &[i32], residual: &mut [i32], tx_sz:
}
return;
}
scalar_row_pass(dequant, residual, &p, lossless);
scalar_col_pass(residual, &p, lossless);
}
fn scalar_row_pass(dequant: &[i32], residual: &mut [i32], p: &Tx2d, lossless: bool) {
let (w, h, log2w, row_shift, row_clamp, col_clamp, rk, rect2) = (p.w, p.h, p.log2w, p.row_shift, p.row_clamp, p.col_clamp, p.rk, p.rect2);
let mut t = [0i32; 64];
let in_w = w.min(32);
for i in 0..h {
// Every 1D transform maps an all-zero row to zeros (each step is linear with
@@ -572,6 +602,11 @@ pub(crate) fn inverse_transform_2d(dequant: &[i32], residual: &mut [i32], tx_sz:
residual[i * w + j] = round2(t[j], row_shift).clamp(lo, hi);
}
}
}
fn scalar_col_pass(residual: &mut [i32], p: &Tx2d, lossless: bool) {
let (w, h, log2h, col_shift, col_clamp, ck) = (p.w, p.h, p.log2h, p.col_shift, p.col_clamp, p.ck);
let mut t = [0i32; 64];
for j in 0..w {
for i in 0..h {
t[i] = residual[i * w + j];
@@ -617,11 +652,16 @@ fn transform_1d<T: Lane>(t: &mut [T], kind: Kind1d, n: u32, r: u32) {
/// The 2D inverse transform with `L` rows (row pass) or columns (column pass) at a time;
/// bit-identical to the one-at-a-time process for bit depths up to 10.
fn transform_lanes<const L: usize>(dequant: &[i32], residual: &mut [i32], p: &Tx2d) {
let mut t = [[0i32; L]; 64];
lanes_row_pass(dequant, residual, p, &mut t);
lanes_col_pass(residual, p, &mut t);
}
fn lanes_row_pass<const L: usize>(dequant: &[i32], residual: &mut [i32], p: &Tx2d, t: &mut [[i32; L]; 64]) {
let (w, h) = (p.w, p.h);
let in_w = w.min(32);
let lo = -(1i32 << (p.col_clamp - 1));
let hi = (1i32 << (p.col_clamp - 1)) - 1;
let mut t = [[0i32; L]; 64];
let mut i0 = 0;
while i0 < h {
// Rows without coefficients (and rows 32..63) transform to zeros.
@@ -649,6 +689,10 @@ fn transform_lanes<const L: usize>(dequant: &[i32], residual: &mut [i32], p: &Tx
}
i0 += L;
}
}
fn lanes_col_pass<const L: usize>(residual: &mut [i32], p: &Tx2d, t: &mut [[i32; L]; 64]) {
let (w, h) = (p.w, p.h);
let mut j0 = 0;
while j0 < w {
for i in 0..h {
+152 -199
View File
@@ -3,7 +3,7 @@
use crate::error::{Result, ensure};
use crate::picture::{FrameRef, RefPic};
use crate::slice::{Mmco, RefPicListMod, SliceHeader, SliceType};
use crate::slice::{Mmco, SliceHeader, SliceType};
use std::sync::Arc;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
@@ -174,103 +174,6 @@ impl Dpb {
}
}
/// Initial reference picture lists (8.2.4.2) as DPB entries.
fn initial_entry_lists<'a>(
sh: &SliceHeader,
cur_poc: i32,
short: &[&'a DpbEntry],
long: &[&'a DpbEntry],
pic_num: impl Fn(&DpbEntry) -> i32,
) -> [Vec<&'a DpbEntry>; 2] {
let mut entry_lists: [Vec<&'a DpbEntry>; 2] = [Vec::new(), Vec::new()];
if sh.slice_type != SliceType::B {
let mut s = short.to_vec();
s.sort_by_key(|e| std::cmp::Reverse(pic_num(e)));
entry_lists[0] = s.into_iter().chain(long.iter().copied()).collect();
} else {
let mut before: Vec<&'a DpbEntry> = short.iter().copied().filter(|e| e.poc < cur_poc).collect();
let mut after: Vec<&'a DpbEntry> = short.iter().copied().filter(|e| e.poc > cur_poc).collect();
before.sort_by_key(|e| std::cmp::Reverse(e.poc));
after.sort_by_key(|e| e.poc);
entry_lists[0] = before.iter().chain(after.iter()).chain(long.iter()).copied().collect();
entry_lists[1] = after.iter().chain(before.iter()).chain(long.iter()).copied().collect();
if entry_lists[1].len() > 1 && entry_lists[0].len() == entry_lists[1].len() {
let same = entry_lists[0].iter().zip(entry_lists[1].iter()).all(|(a, b)| std::ptr::eq(*a, *b));
if same {
entry_lists[1].swap(0, 1);
}
}
}
entry_lists
}
/// Reference picture list modification (8.2.4.3) of one list; `n` is num_ref_idx_active.
fn modify_ref_list<'a>(
list: &mut Vec<Option<&'a DpbEntry>>,
mods: &[RefPicListMod],
n: usize,
cur_fn: u32,
max_frame_num: u32,
short: &[&'a DpbEntry],
long: &[&'a DpbEntry],
) -> Result<()> {
let max_pic_num = max_frame_num as i32;
let curr_pic_num = cur_fn as i32;
let mut pred = curr_pic_num;
let mut ref_idx = 0usize;
for m in mods {
ensure!(ref_idx < n, "too many ref_pic_list_modification entries");
list.push(None); // temporarily one longer
let (target, is_long): (Option<&DpbEntry>, bool) = match m.idc {
0 | 1 => {
let d = m.value as i32 + 1;
let no_wrap = if m.idc == 0 {
if pred - d < 0 { pred - d + max_pic_num } else { pred - d }
} else if pred + d >= max_pic_num {
pred + d - max_pic_num
} else {
pred + d
};
pred = no_wrap;
let pn = if no_wrap > curr_pic_num { no_wrap - max_pic_num } else { no_wrap };
(short.iter().copied().find(|e| Self::frame_num_wrap(e, cur_fn, max_frame_num) == pn), false)
}
2 => (long.iter().copied().find(|e| e.long_term_frame_idx == m.value), true),
_ => continue,
};
let Some(target) = target else {
// Missing picture: keep the list unchanged (error concealment).
list.pop();
ref_idx += 1;
continue;
};
for c in (ref_idx + 1..=n).rev() {
list[c] = list[c - 1];
}
list[ref_idx] = Some(target);
ref_idx += 1;
let mut nidx = ref_idx;
for c in ref_idx..=n {
let keep = match list[c] {
None => true,
Some(e) => {
if is_long {
!(e.mark == RefMark::Long && e.long_term_frame_idx == target.long_term_frame_idx)
} else {
!(e.mark == RefMark::Short && std::ptr::eq(e, target))
}
}
};
if keep {
list[nidx] = list[c];
nidx += 1;
}
}
list.truncate(n);
}
Ok(())
}
/// Initial reference picture lists (8.2.4.2) + modification (8.2.4.3) for a frame slice.
pub fn build_ref_lists(&self, sh: &SliceHeader, cur_poc: i32, max_frame_num: u32) -> Result<[Vec<RefPic>; 2]> {
let mut lists: [Vec<RefPic>; 2] = [Vec::new(), Vec::new()];
@@ -283,7 +186,25 @@ impl Dpb {
let mut long: Vec<&DpbEntry> = self.entries.iter().filter(|e| e.mark == RefMark::Long).collect();
long.sort_by_key(|e| e.long_term_frame_idx);
let to_ref = |e: &&DpbEntry| RefPic { frame: e.frame.clone(), long_term: e.mark == RefMark::Long };
let entry_lists = Self::initial_entry_lists(sh, cur_poc, &short, &long, pic_num);
let mut entry_lists: [Vec<&DpbEntry>; 2] = [Vec::new(), Vec::new()];
if sh.slice_type != SliceType::B {
let mut s = short.clone();
s.sort_by_key(|e| std::cmp::Reverse(pic_num(e)));
entry_lists[0] = s.into_iter().chain(long.iter().copied()).collect();
} else {
let mut before: Vec<&DpbEntry> = short.iter().copied().filter(|e| e.poc < cur_poc).collect();
let mut after: Vec<&DpbEntry> = short.iter().copied().filter(|e| e.poc > cur_poc).collect();
before.sort_by_key(|e| std::cmp::Reverse(e.poc));
after.sort_by_key(|e| e.poc);
entry_lists[0] = before.iter().chain(after.iter()).chain(long.iter()).copied().collect();
entry_lists[1] = after.iter().chain(before.iter()).chain(long.iter()).copied().collect();
if entry_lists[1].len() > 1 && entry_lists[0].len() == entry_lists[1].len() {
let same = entry_lists[0].iter().zip(entry_lists[1].iter()).all(|(a, b)| std::ptr::eq(*a, *b));
if same {
entry_lists[1].swap(0, 1);
}
}
}
let nlists = if sh.slice_type == SliceType::B { 2 } else { 1 };
for l in 0..nlists {
let n = sh.num_ref_idx_active[l] as usize;
@@ -291,7 +212,60 @@ impl Dpb {
list.truncate(n);
list.resize(n, None);
if !sh.ref_pic_list_mod[l].is_empty() {
Self::modify_ref_list(&mut list, &sh.ref_pic_list_mod[l], n, cur_fn, max_frame_num, &short, &long)?;
let max_pic_num = max_frame_num as i32;
let curr_pic_num = cur_fn as i32;
let mut pred = curr_pic_num;
let mut ref_idx = 0usize;
for m in &sh.ref_pic_list_mod[l] {
ensure!(ref_idx < n, "too many ref_pic_list_modification entries");
list.push(None); // temporarily one longer
let (target, is_long): (Option<&DpbEntry>, bool) = match m.idc {
0 | 1 => {
let d = m.value as i32 + 1;
let no_wrap = if m.idc == 0 {
if pred - d < 0 { pred - d + max_pic_num } else { pred - d }
} else if pred + d >= max_pic_num {
pred + d - max_pic_num
} else {
pred + d
};
pred = no_wrap;
let pn = if no_wrap > curr_pic_num { no_wrap - max_pic_num } else { no_wrap };
(short.iter().copied().find(|e| pic_num(e) == pn), false)
}
2 => (long.iter().copied().find(|e| e.long_term_frame_idx == m.value), true),
_ => continue,
};
let Some(target) = target else {
// Missing picture: keep the list unchanged (error concealment).
list.pop();
ref_idx += 1;
continue;
};
for c in (ref_idx + 1..=n).rev() {
list[c] = list[c - 1];
}
list[ref_idx] = Some(target);
ref_idx += 1;
let mut nidx = ref_idx;
for c in ref_idx..=n {
let keep = match list[c] {
None => true,
Some(e) => {
if is_long {
!(e.mark == RefMark::Long && e.long_term_frame_idx == target.long_term_frame_idx)
} else {
!(e.mark == RefMark::Short && std::ptr::eq(e, target))
}
}
};
if keep {
list[nidx] = list[c];
nidx += 1;
}
}
list.truncate(n);
}
}
// Fill missing entries (should not happen in conforming streams) with the first valid entry.
let fallback = list.iter().flatten().next().copied().or_else(|| self.entries.last());
@@ -301,81 +275,85 @@ impl Dpb {
Ok(lists)
}
/// One memory management control operation; returns true for MMCO 5.
fn apply_mmco(&mut self, m: &Mmco, cur_frame_num: u32, max_frame_num: u32, cur_long: &mut Option<u32>) -> bool {
let mut had5 = false;
let curr_pic_num = cur_frame_num as i32;
match m.op {
1 => {
let pn = curr_pic_num - (m.difference_of_pic_nums_minus1 as i32 + 1);
for e in &mut self.entries {
if e.mark == RefMark::Short && Self::frame_num_wrap(e, cur_frame_num, max_frame_num) == pn {
e.mark = RefMark::Unused;
}
}
}
2 => {
for e in &mut self.entries {
if e.mark == RefMark::Long && e.long_term_frame_idx == m.long_term_pic_num {
e.mark = RefMark::Unused;
}
}
}
3 => {
let pn = curr_pic_num - (m.difference_of_pic_nums_minus1 as i32 + 1);
for e in &mut self.entries {
if e.mark == RefMark::Long && e.long_term_frame_idx == m.long_term_frame_idx {
e.mark = RefMark::Unused;
}
}
for e in &mut self.entries {
if e.mark == RefMark::Short && Self::frame_num_wrap(e, cur_frame_num, max_frame_num) == pn {
e.mark = RefMark::Long;
e.long_term_frame_idx = m.long_term_frame_idx;
}
}
}
4 => {
let max = m.max_long_term_frame_idx_plus1;
for e in &mut self.entries {
if e.mark == RefMark::Long && (max == 0 || e.long_term_frame_idx > max - 1) {
e.mark = RefMark::Unused;
}
}
self.max_long_term_frame_idx = if max == 0 { None } else { Some(max - 1) };
}
5 => {
for e in &mut self.entries {
e.mark = RefMark::Unused;
}
self.max_long_term_frame_idx = None;
had5 = true;
}
6 => {
for e in &mut self.entries {
if e.mark == RefMark::Long && e.long_term_frame_idx == m.long_term_frame_idx {
e.mark = RefMark::Unused;
}
}
*cur_long = Some(m.long_term_frame_idx);
}
_ => {}
}
had5
}
/// Adaptive memory control (8.2.5.4). Returns true if MMCO 5 was executed.
fn apply_mmcos(&mut self, mmcos: &[Mmco], cur_frame_num: u32, max_frame_num: u32, cur_long: &mut Option<u32>) -> bool {
let mut had5 = false;
let curr_pic_num = cur_frame_num as i32;
for m in mmcos {
had5 |= self.apply_mmco(m, cur_frame_num, max_frame_num, cur_long);
match m.op {
1 => {
let pn = curr_pic_num - (m.difference_of_pic_nums_minus1 as i32 + 1);
for e in &mut self.entries {
if e.mark == RefMark::Short && Self::frame_num_wrap(e, cur_frame_num, max_frame_num) == pn {
e.mark = RefMark::Unused;
}
}
}
2 => {
for e in &mut self.entries {
if e.mark == RefMark::Long && e.long_term_frame_idx == m.long_term_pic_num {
e.mark = RefMark::Unused;
}
}
}
3 => {
let pn = curr_pic_num - (m.difference_of_pic_nums_minus1 as i32 + 1);
for e in &mut self.entries {
if e.mark == RefMark::Long && e.long_term_frame_idx == m.long_term_frame_idx {
e.mark = RefMark::Unused;
}
}
for e in &mut self.entries {
if e.mark == RefMark::Short && Self::frame_num_wrap(e, cur_frame_num, max_frame_num) == pn {
e.mark = RefMark::Long;
e.long_term_frame_idx = m.long_term_frame_idx;
}
}
}
4 => {
let max = m.max_long_term_frame_idx_plus1;
for e in &mut self.entries {
if e.mark == RefMark::Long && (max == 0 || e.long_term_frame_idx > max - 1) {
e.mark = RefMark::Unused;
}
}
self.max_long_term_frame_idx = if max == 0 { None } else { Some(max - 1) };
}
5 => {
for e in &mut self.entries {
e.mark = RefMark::Unused;
}
self.max_long_term_frame_idx = None;
had5 = true;
}
6 => {
for e in &mut self.entries {
if e.mark == RefMark::Long && e.long_term_frame_idx == m.long_term_frame_idx {
e.mark = RefMark::Unused;
}
}
*cur_long = Some(m.long_term_frame_idx);
}
_ => {}
}
}
self.prune();
had5
}
/// Reference marking of the current picture (8.2.5); returns (mark, long-term idx, MMCO 5 executed).
fn mark_current(&mut self, sh: &SliceHeader, max_frame_num: u32, max_num_ref_frames: usize) -> (RefMark, u32, bool) {
/// Reference marking for the just-decoded picture and insertion into the DPB with output bumping.
/// Returns pictures ready for output.
#[allow(clippy::too_many_arguments)]
pub fn store_picture(
&mut self,
sh: &SliceHeader,
frame: FrameRef,
poc: i32,
max_frame_num: u32,
max_num_ref_frames: usize,
meta: Arc<OutputMeta>,
out: &mut Vec<Output>,
) {
let mut mark = RefMark::Unused;
let mut long_idx = 0;
let mut had5 = false;
@@ -407,12 +385,12 @@ impl Dpb {
}
}
}
(mark, long_idx, had5)
}
/// Bumping until there is space (C.4.5.1 / C.4.5.2). Returns true when the (non-reference) current
/// picture has the smallest POC and is to be output directly instead of being stored.
fn bump_for_space(&mut self, mark: RefMark, poc: i32, out: &mut Vec<Output>) -> bool {
if had5 {
// All earlier pictures are output before the current one (C.4.4).
self.flush(out);
}
let frame_num = if had5 { 0 } else { sh.frame_num };
// Bumping until there is space (C.4.5.1 / C.4.5.2).
if mark == RefMark::Unused {
loop {
let waiting = self.entries.len();
@@ -422,7 +400,8 @@ impl Dpb {
// non-reference picture with the smallest POC is output directly
let min_poc = self.entries.iter().filter(|e| e.needed_for_output).map(|e| e.poc).min();
if min_poc.is_none_or(|m| poc < m) {
return true;
out.push(Output { frame, poc, meta });
return;
}
if !self.bump(out) {
break;
@@ -435,32 +414,6 @@ impl Dpb {
}
}
}
false
}
/// Reference marking for the just-decoded picture and insertion into the DPB with output bumping.
/// Returns pictures ready for output.
#[allow(clippy::too_many_arguments)]
pub fn store_picture(
&mut self,
sh: &SliceHeader,
frame: FrameRef,
poc: i32,
max_frame_num: u32,
max_num_ref_frames: usize,
meta: Arc<OutputMeta>,
out: &mut Vec<Output>,
) {
let (mark, long_idx, had5) = self.mark_current(sh, max_frame_num, max_num_ref_frames);
if had5 {
// All earlier pictures are output before the current one (C.4.4).
self.flush(out);
}
let frame_num = if had5 { 0 } else { sh.frame_num };
if self.bump_for_space(mark, poc, out) {
out.push(Output { frame, poc, meta });
return;
}
self.entries.push(DpbEntry {
frame,
poc,