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
7 changed files with 123 additions and 62 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 {
+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};
@@ -669,7 +668,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(());
}
@@ -698,7 +697,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;
@@ -715,8 +714,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;
}
@@ -730,7 +728,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);
@@ -900,7 +898,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
@@ -925,10 +923,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 =
@@ -937,18 +935,10 @@ 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)
}
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];
@@ -1274,7 +1264,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> {
@@ -1293,7 +1283,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];
@@ -1345,8 +1335,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 {
@@ -1361,7 +1362,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 {
@@ -1405,7 +1406,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];
}
@@ -1503,11 +1504,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);
}
}
}