From 88ca7b45fc31152d2a7469a49230c45d08e99c2b Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sat, 3 Oct 2026 05:40:03 -0400 Subject: [PATCH] M13.4: av1enc speedups (vectorisable forward transform, separable sub-sample paths, cheaper loop filter search), more oracle cases Co-Authored-By: Claude Opus 5.5 --- crates/av1enc/src/decide.rs | 12 +++++-- crates/av1enc/src/lib.rs | 12 ++++--- crates/av1enc/src/loopfilter.rs | 10 ++++-- crates/av1enc/src/pred.rs | 62 +++++++++++++++++++++++++++------ crates/av1enc/src/transform.rs | 61 +++++++++++++++++++++++--------- crates/av1enc/tests/bench.rs | 30 +++++++++++++++- 6 files changed, 149 insertions(+), 38 deletions(-) diff --git a/crates/av1enc/src/decide.rs b/crates/av1enc/src/decide.rs index b877fc23..4c1bf65d 100644 --- a/crates/av1enc/src/decide.rs +++ b/crates/av1enc/src/decide.rs @@ -514,7 +514,13 @@ impl<'a> Decider<'a> { best = (m, cst); } } - let mut step = 16 * 8; + let mut step = 8 * if n >= 32 { + 16 + } else if n == 16 { + 8 + } else { + 4 + }; while step >= 8 { let mut moved = true; let mut iters = 0; @@ -659,7 +665,7 @@ impl<'a> Decider<'a> { let rr = refp.row((ry + i as i32).clamp(0, ly) as usize); if inside { let rrow = &rr[rx as usize..rx as usize + w]; - sad += s.iter().zip(rrow).map(|(&a, &b)| (a as i32 - b as i32).unsigned_abs() as u64).sum::(); + sad += s.iter().zip(rrow).map(|(&a, &b)| (a as i32 - b as i32).unsigned_abs()).sum::() as u64; } else { for j in 0..w { let b = rr[(rx + j as i32).clamp(0, lx) as usize]; @@ -679,7 +685,7 @@ impl<'a> Decider<'a> { let mut s = 0u64; for i in 0..h { let row = &src.row(y + i)[x..x + w]; - s += row.iter().zip(&pred[i * n..i * n + w]).map(|(&a, &b)| (a as i32 - b as i32).unsigned_abs() as u64).sum::(); + s += row.iter().zip(&pred[i * n..i * n + w]).map(|(&a, &b)| (a as i32 - b as i32).unsigned_abs()).sum::() as u64; } s } diff --git a/crates/av1enc/src/lib.rs b/crates/av1enc/src/lib.rs index f1ab74a4..0eeded9c 100644 --- a/crates/av1enc/src/lib.rs +++ b/crates/av1enc/src/lib.rs @@ -373,7 +373,7 @@ impl Encoder { } } let lf_level = self.pick_filter_level(&dec.rec, &tw.lf_info(), src, qidx); - loop_filter(&mut dec.rec, &tw.lf_info(), [lf_level; 4], self.cfg.bit_depth as u32, false); + loop_filter(&mut dec.rec, &tw.lf_info(), [lf_level; 4], self.cfg.bit_depth as u32, false, 1); let (tile, cdfs) = tw.finish(); let mut w = BitWriter::new(); self.frame_header(&mut w, key, qidx, lf_level); @@ -388,12 +388,16 @@ impl Encoder { } /// The loop filter level minimising the luma error (searched around a guess from the - /// quantizer; luma only). + /// quantizer, on the luma of a subset of the superblock rows). fn pick_filter_level(&self, rec: &[Plane; 3], info: &LfInfo, src: &[Plane; 3], qidx: u8) -> u32 { let (w, h) = (self.geom.width, self.geom.height); + let sb_row_step = if self.geom.sb_rows >= 9 { 3 } else { 1 }; let sse = |pl: &Plane| -> u64 { let mut s = 0u64; for y in 0..h { + if !(y >> 6).is_multiple_of(sb_row_step) { + continue; + } for (a, b) in pl.row(y)[..w].iter().zip(&src[0].row(y)[..w]) { let d = *a as i64 - *b as i64; s += (d * d) as u64; @@ -408,7 +412,7 @@ impl Encoder { return e; } let mut planes = [rec[0].clone(), Plane::new(0, 0), Plane::new(0, 0)]; - loop_filter(&mut planes, info, [lvl; 4], bd, true); + loop_filter(&mut planes, info, [lvl; 4], bd, true, sb_row_step); let e = sse(&planes[0]); cache.push((lvl, e)); e @@ -421,7 +425,7 @@ impl Encoder { if e < best.1 { best = (center, e); } - while step >= 1 { + while step >= 2 { for lvl in [center.saturating_sub(step), (center + step).min(63)] { let e = eval(lvl); if e < best.1 { diff --git a/crates/av1enc/src/loopfilter.rs b/crates/av1enc/src/loopfilter.rs index 6e74e411..2add41ed 100644 --- a/crates/av1enc/src/loopfilter.rs +++ b/crates/av1enc/src/loopfilter.rs @@ -18,9 +18,9 @@ pub(crate) struct LfInfo<'a> { pub tx: [&'a [u8]; 3], } -/// Applies the loop filter with `loop_filter_level[ 0..4 ]` to `planes` (only plane 0 when -/// `luma_only`). -pub(crate) fn loop_filter(planes: &mut [Plane; 3], info: &LfInfo, level: [u32; 4], bit_depth: u32, luma_only: bool) { +/// Applies the loop filter with `loop_filter_level[ 0..4 ]` to `planes`. For level search +/// (`luma_only`), only plane 0 and only every `sb_row_step`-th superblock row are filtered. +pub(crate) fn loop_filter(planes: &mut [Plane; 3], info: &LfInfo, level: [u32; 4], bit_depth: u32, luma_only: bool, sb_row_step: usize) { if level[0] == 0 && level[1] == 0 { return; } @@ -33,6 +33,10 @@ pub(crate) fn loop_filter(planes: &mut [Plane; 3], info: &LfInfo, level: [u32; 4 let step = if plane == 0 { 1 } else { 2 }; let mut row = 0; while row < info.mi_rows { + if !(row >> 4).is_multiple_of(sb_row_step) { + row += step; + continue; + } let mut col = 0; while col < info.mi_cols { edge(pl, info, level, bit_depth, plane, pass, row, col); diff --git a/crates/av1enc/src/pred.rs b/crates/av1enc/src/pred.rs index 48d48ad8..d61a6fa9 100644 --- a/crates/av1enc/src/pred.rs +++ b/crates/av1enc/src/pred.rs @@ -224,19 +224,61 @@ pub(crate) fn predict_inter( let filt_v = if h <= 4 { 4 } else { EIGHTTAP }; let kh = &SUBPEL_FILTERS[filt_h][fx]; let kv = &SUBPEL_FILTERS[filt_v][fy]; + let bx = ix - 3; + let x_inside = bx >= 0 && bx + w as i32 + 7 <= last_x + 1; + // horizontal pass of one reference row into `out` (Round2( sum, InterRound0 = 3 )) + let hfilter = |ry: usize, out: &mut [i32]| { + let row = refp.row(ry); + if x_inside { + let seg = &row[bx as usize..bx as usize + w + 7]; + for (c, o) in out.iter_mut().enumerate() { + let mut s = 0i32; + for t in 0..8 { + s += kh[t] as i32 * seg[c + t] as i32; + } + *o = round2(s, 3); + } + } else { + for (c, o) in out.iter_mut().enumerate() { + let mut s = 0i32; + for t in 0..8 { + s += kh[t] as i32 * row[(bx + c as i32 + t as i32).clamp(0, last_x) as usize] as i32; + } + *o = round2(s, 3); + } + } + }; + let mut inter = [0i32; (64 + 7) * 64]; + if fy == 0 { + // the vertical kernel is 128 at the centre: Round2( 128 * v, 11 ) = Round2( v, 4 ) + for r in 0..h { + let ry = (iy + r as i32).clamp(0, last_y) as usize; + hfilter(ry, &mut inter[..w]); + for (p, &v) in pred[r * w..r * w + w].iter_mut().zip(&inter[..w]) { + *p = round2(v, 4).clamp(0, max) as u16; + } + } + return; + } + if fx == 0 { + // the horizontal pass is exactly 16 * sample: Round2( 16 * sum, 11 ) = Round2( sum, 7 ) + for r in 0..h { + let rows: [&[u16]; 8] = std::array::from_fn(|t| refp.row((iy + r as i32 + t as i32 - 3).clamp(0, last_y) as usize)); + for c in 0..w { + let xx = if x_inside { (ix + c as i32) as usize } else { (ix + c as i32).clamp(0, last_x) as usize }; + let mut s = 0i32; + for t in 0..8 { + s += kv[t] as i32 * rows[t][xx] as i32; + } + pred[r * w + c] = round2(s, 7).clamp(0, max) as u16; + } + } + return; + } let inter_h = h + 7; - let mut inter = vec![0i32; inter_h * w]; for r in 0..inter_h { let ry = (iy + r as i32 - 3).clamp(0, last_y) as usize; - let row = refp.row(ry); - for c in 0..w { - let bx = ix + c as i32 - 3; - let mut s = 0i32; - for t in 0..8 { - s += kh[t] as i32 * row[(bx + t as i32).clamp(0, last_x) as usize] as i32; - } - inter[r * w + c] = round2(s, 3); - } + hfilter(ry, &mut inter[r * w..r * w + w]); } for r in 0..h { for c in 0..w { diff --git a/crates/av1enc/src/transform.rs b/crates/av1enc/src/transform.rs index a2b2d657..386133b5 100644 --- a/crates/av1enc/src/transform.rs +++ b/crates/av1enc/src/transform.rs @@ -344,6 +344,20 @@ pub(crate) fn inverse_2d(dequant: &[i32], residual: &mut [i32], tx_sz: usize, tx struct Bases { dct: [Vec; 4], adst: [Vec; 3], + /// The same matrices transposed (`[i][k]`). + dct_t: [Vec; 4], + adst_t: [Vec; 3], +} + +fn transpose(m: &[f32]) -> Vec { + let n = (m.len() as f64).sqrt() as usize; + let mut t = vec![0f32; n * n]; + for k in 0..n { + for i in 0..n { + t[i * n + k] = m[k * n + i]; + } + } + t } fn bases() -> &'static Bases { @@ -376,10 +390,20 @@ fn bases() -> &'static Bases { } m }); - Bases { dct, adst } + let dct_t = std::array::from_fn(|l| transpose(&dct[l])); + let adst_t = std::array::from_fn(|l| transpose(&adst[l])); + Bases { dct, adst, dct_t, adst_t } }) } +fn basis_t(kind: Kind, n: usize) -> &'static [f32] { + let l = n.trailing_zeros() as usize - 2; + match kind { + Kind::Dct => &bases().dct_t[l], + Kind::Adst => &bases().adst_t[l], + } +} + fn basis(kind: Kind, n: usize) -> &'static [f32] { let l = n.trailing_zeros() as usize - 2; match kind { @@ -393,29 +417,32 @@ fn basis(kind: Kind, n: usize) -> &'static [f32] { pub(crate) fn forward_2d(residual: &[i32], out: &mut [f32], n: usize, tx_type: usize) { let (ck, rk) = kinds(tx_type); let bc = basis(ck, n); - let br = basis(rk, n); - // tmp[i][k] = sum_j residual[i][j] * br[k][j] (horizontal) - let mut tmp = vec![0f32; n * n]; + let brt = basis_t(rk, n); + // tmp[i][k] = sum_j residual[i][j] * br[k][j] (horizontal; written as row updates so that + // the inner loops vectorise) + let mut tmp = [0f32; 32 * 32]; for i in 0..n { - let row = &residual[i * n..i * n + n]; - for k in 0..n { - let b = &br[k * n..k * n + n]; - let mut s = 0f32; - for j in 0..n { - s += row[j] as f32 * b[j]; + let t = &mut tmp[i * n..i * n + n]; + for j in 0..n { + let r = residual[i * n + j]; + if r == 0 { + continue; + } + let r = r as f32; + for (o, &b) in t.iter_mut().zip(&brt[j * n..j * n + n]) { + *o += r * b; } - tmp[i * n + k] = s; } } // out[k][j] = 8 * sum_i bc[k][i] * tmp[i][j] (vertical) + out[..n * n].iter_mut().for_each(|v| *v = 0.0); for k in 0..n { - let b = &bc[k * n..k * n + n]; - for j in 0..n { - let mut s = 0f32; - for i in 0..n { - s += b[i] * tmp[i * n + j]; + let o = &mut out[k * n..k * n + n]; + for i in 0..n { + let b = 8.0 * bc[k * n + i]; + for (v, &t) in o.iter_mut().zip(&tmp[i * n..i * n + n]) { + *v += b * t; } - out[k * n + j] = 8.0 * s; } } } diff --git a/crates/av1enc/tests/bench.rs b/crates/av1enc/tests/bench.rs index 164c6f66..8203d505 100644 --- a/crates/av1enc/tests/bench.rs +++ b/crates/av1enc/tests/bench.rs @@ -11,7 +11,7 @@ use effectcraft_av1enc::{Encoder, EncoderConfig, RateControl}; /// Raw 8-bit 4:2:0 frames of an ffmpeg lavfi source, or None without ffmpeg. fn lavfi(src: &str, w: usize, h: usize, frames: usize) -> Option> { - let out = std::env::temp_dir().join(format!("effectcraft-av1enc-bench-{}-{}.yuv", std::process::id(), src.split('=').next().unwrap_or("src"))); + let out = std::env::temp_dir().join(format!("effectcraft-av1enc-bench-{}-{src}-{w}x{h}.yuv", std::process::id())); let st = Command::new("ffmpeg") .args(["-v", "error", "-y", "-f", "lavfi", "-i"]) .arg(format!("{src}=size={w}x{h}:rate=30")) @@ -93,3 +93,31 @@ fn bench_lavfi() { run("mandelbrot", &pics, RateControl::Bitrate { kbps: 1000 }, 30); } } + +#[test] +#[ignore] +fn bench_1080p_loop() { + let Some(pics) = lavfi("mandelbrot", 1920, 1080, 4) else { + return; + }; + run("mandelbrot", &pics, RateControl::ConstantQ(120), 30); + run("mandelbrot", &pics, RateControl::ConstantQ(120), 30); +} + +#[test] +#[ignore] +fn bench_rate_control() { + let Some(pics) = lavfi("mandelbrot", 480, 272, 60) else { + eprintln!("ffmpeg not available"); + return; + }; + for kbps in [200, 600, 2000] { + run("mandelbrot", &pics, RateControl::Bitrate { kbps }, 30); + } + let Some(pics) = lavfi("testsrc2", 480, 272, 60) else { + return; + }; + for kbps in [300, 3000] { + run("testsrc2", &pics, RateControl::Bitrate { kbps }, 60); + } +}