collective
Browse / github.com/storytold/effectcraft / patch
av1enc: reuse scratch buffers in the transform path, the coefficient writer and the level search
av1enc: per-block heap allocations and a plane clone in the loop-filter search (-2..3 %) ## Problem Heap allocation in per-block paths: - `decide.rs` `try_tx` / `reconstruct` allocate six Vecs per candidate (res, coef, levels, deq, resid, recon). - `writer.rs:504` allocates `quant` per coded transform block; `writer.rs:184` has a `to_vec` per edge block. - The loop-filter level search (`loopfilter.rs`, about 1.7 % of the time in total) clones the whole luma plane for every candidate level (about 1 %). calloc + free + malloc are about 1.5 % and memset about 2 % of the encode (libc). ## Proposal Scratch buffers owned by the `Decider` and the `TileWriter` (reused Vec fields or fixed
- Status
- open
- Score
- 0.00
- Adopters
- 0 other operators
- Author
- vangel76 (imported from pull request #2405, open)
- Base
991904888a20b7da72b2a4880b21c7d912154736- Size
- +113 −50 in 3 files
- Review
- pending
- Shared
- 2026-10-28
- Id
336f4bedb9dcee600d3d6b82
How it went
No other operator has reported using it yet.
Diff
raw · crates/av1enc/src/decide.rs crates/av1enc/src/lib.rs crates/av1enc/src/writer.rs
@@ -107,6 +107,31 @@ struct TxResult { bits: f64, } +/// Outcome of trying one transform block; the levels and the reconstruction stay in `TxScratch`.+#[derive(Clone, Copy)]+struct TxEval {+ nonzero: bool,+ dist: f64,+ bits: f64,+}++/// Buffers of the transform path, sized for the largest coded transform block (32x32). Only the+/// first `n * n` entries are used; they are reused across candidates instead of being allocated.+struct TxScratch {+ res: [i32; 1024],+ coef: [f32; 1024],+ levels: [i32; 1024],+ deq: [i32; 1024],+ resid: [i32; 1024],+ recon: [u16; 1024],+}++impl Default for TxScratch {+ fn default() -> Self {+ TxScratch { res: [0; 1024], coef: [0.0; 1024], levels: [0; 1024], deq: [0; 1024], resid: [0; 1024], recon: [0; 1024] }+ }+}+ /// Saved encoder state of a block area (reconstruction, BlockDecoded flags). struct Snapshot { planes: [Vec<u16>; 3],@@ -131,6 +156,8 @@ pub(crate) struct Decider<'a> { lambda_sad: f64, /// Half-sample interpolations of the reference luma (visible area): (8, 0), (0, 8), (8, 8). half: [Vec<u16>; 3],+ /// Transform buffers, taken out while a block is being tried (see `code_tx`).+ scratch: Option<Box<TxScratch>>, } const MODE_BITS_INTRA: f64 = 3.0;@@ -166,6 +193,7 @@ impl<'a> Decider<'a> { Some(rf) => [(8, 0), (0, 8), (8, 8)].map(|(fx, fy)| subpel_plane(&rf.planes[0], g.width, g.height, fx, fy, bd)), None => Default::default(), },+ scratch: None, } } @@ -324,8 +352,10 @@ impl<'a> Decider<'a> { let log2 = (n4 * 4).trailing_zeros(); let (sbr, sbc) = ((r & 15) as isize, (c & 15) as isize); let bd = self.bd;- let mut pred = vec![0u16; 32 * 32];- let mut best_pred = vec![0u16; 32 * 32];+ let mut pred = [0u16; 32 * 32];+ let mut best_pred = [0u16; 32 * 32];+ let mut best_levels = [0i32; 32 * 32];+ let mut best_recon = [0u16; 32 * 32]; let n = 1usize << log2; // luma let e = IntraEdge {@@ -343,26 +373,26 @@ impl<'a> Decider<'a> { let ed = edges(&self.rec[0], &e); // modes ranked by SATD, then a rate-distortion choice of mode and transform type among // the best three- let mut ranked: Vec<(f64, usize)> = INTRA_MODES_USED- .iter()- .map(|&mode| {- predict_intra(&ed, &e, mode, &mut pred);- let sad = self.satd(0, c * 4, r * 4, n, &pred[..n * n]) as f64;- (sad + self.lambda_sad * if mode == DC_PRED { 1.0 } else { MODE_BITS_INTRA }, mode)- })- .collect();+ let mut ranked: [(f64, usize); INTRA_MODES_USED.len()] = INTRA_MODES_USED.map(|mode| {+ predict_intra(&ed, &e, mode, &mut pred);+ let sad = self.satd(0, c * 4, r * 4, n, &pred[..n * n]) as f64;+ (sad + self.lambda_sad * if mode == DC_PRED { 1.0 } else { MODE_BITS_INTRA }, mode)+ }); ranked.sort_by(|a, b| a.0.total_cmp(&b.0)); let tx_y = tx_size_for(log2);- let mut best: Option<(f64, usize, usize, TxResult, Option<Vec<u16>>)> = None;+ let mut best: Option<(f64, usize, usize, TxEval, bool)> = None;+ let mut sc = self.scratch.take().unwrap_or_default(); for &(_, mode) in ranked.iter().take(3) { predict_intra(&ed, &e, mode, &mut pred);- let mut types = vec![DCT_DCT];+ let mut types = [DCT_DCT; 2];+ let mut n_types = 1; let alt = MODE_TO_TXFM[mode] as usize; if alt != DCT_DCT && tx_in_set(tx_set(tx_y, false), false, alt) {- types.push(alt);+ types[1] = alt;+ n_types = 2; }- for &t in &types {- let (res, recon) = self.try_tx(0, c * 4, r * 4, log2, t, &pred[..n * n], true);+ for &t in &types[..n_types] {+ let (res, has_recon) = self.try_tx(&mut sc, 0, c * 4, r * 4, log2, t, &pred[..n * n], true); let type_bits = if !res.nonzero || tx_set(tx_y, false) == TX_SET_DCTONLY { 0.0 } else if t == DCT_DCT {@@ -374,18 +404,24 @@ impl<'a> Decider<'a> { let cost = res.dist + self.lambda * (res.bits + type_bits + mode_bits); if best.as_ref().is_none_or(|b| cost < b.0) { best_pred[..n * n].copy_from_slice(&pred[..n * n]);- best = Some((cost, mode, t, res, recon));+ best_levels[..n * n].copy_from_slice(&sc.levels[..n * n]);+ if has_recon {+ best_recon[..n * n].copy_from_slice(&sc.recon[..n * n]);+ }+ best = Some((cost, mode, t, res, has_recon)); } } }+ self.scratch = Some(sc); // Every ranked mode tries DCT_DCT, so `best` is always set; without a candidate the block // can't be coded intra (an infinite cost keeps it from being chosen).- let Some((_, y_mode, tx_type_y, ty, recon)) = best else {+ let Some((_, y_mode, tx_type_y, ty, has_recon)) = best else { let leaf = Leaf { bsize, is_inter: false, y_mode: DC_PRED, uv_mode: DC_PRED, tx_type_y: DCT_DCT, mv: [0, 0], skip: true, coefs: Default::default() }; return (leaf, f64::INFINITY); };- self.put(0, c * 4, r * 4, n, recon.as_deref().unwrap_or(&best_pred[..n * n]));+ self.put(0, c * 4, r * 4, n, if has_recon { &best_recon[..n * n] } else { &best_pred[..n * n] });+ let ty = TxResult { levels: best_levels[..n * n].to_vec(), nonzero: ty.nonzero, dist: ty.dist, bits: ty.bits }; let tx_type_y = if ty.nonzero { tx_type_y } else { DCT_DCT }; self.mark_decoded(0, r, c, n4); // chroma (one mode for both planes)@@ -772,31 +808,35 @@ impl<'a> Decider<'a> { /// Transform, quantise and reconstruct one square transform block of `plane` at (x, y) /// over the prediction `pred`; the reconstruction is written to `rec`. fn code_tx(&mut self, p: usize, x: usize, y: usize, log2: u32, tx_type: usize, pred: &[u16], intra: bool) -> TxResult {- let (res, recon) = self.try_tx(p, x, y, log2, tx_type, pred, intra);+ let mut sc = self.scratch.take().unwrap_or_default();+ let (res, has_recon) = self.try_tx(&mut sc, p, x, y, log2, tx_type, pred, intra); let n = 1usize << log2;- self.put(p, x, y, n, recon.as_deref().unwrap_or(pred));- res+ self.put(p, x, y, n, if has_recon { &sc.recon[..n * n] } else { pred });+ let levels = sc.levels[..n * n].to_vec();+ self.scratch = Some(sc);+ TxResult { levels, nonzero: res.nonzero, dist: res.dist, bits: res.bits } } - /// Transform and quantise one block; returns the result and its reconstruction (None: the- /// prediction) without writing it.- fn try_tx(&self, p: usize, x: usize, y: usize, log2: u32, tx_type: usize, pred: &[u16], intra: bool) -> (TxResult, Option<Vec<u16>>) {+ /// Transform and quantise one block without writing it. The levels are left in `sc.levels`;+ /// the flag tells whether `sc.recon` holds the reconstruction (false: it is the prediction).+ fn try_tx(&self, sc: &mut TxScratch, p: usize, x: usize, y: usize, log2: u32, tx_type: usize, pred: &[u16], intra: bool) -> (TxEval, bool) { let n = 1usize << log2;+ let nn = n * n; let tx_sz = tx_size_for(log2); let bd = self.bd; let src = &self.src[p];- let mut res = vec![0i32; n * n];+ let TxScratch { res, coef, levels, deq, resid, recon } = sc;+ let (res, coef, levels) = (&mut res[..nn], &mut coef[..nn], &mut levels[..nn]); for i in 0..n { let row = &src.row(y + i)[x..x + n]; for j in 0..n { res[i * n + j] = row[j] as i32 - pred[i * n + j] as i32; } }- let mut coef = vec![0f32; n * n];- forward_2d(&res, &mut coef, n, tx_type);+ forward_2d(res, coef, n, tx_type); // dead-zone quantiser: rounding offset below one half let rnd = if intra { 0.38f32 } else { 0.30f32 };- let mut levels = vec![0i32; n * n];+ levels.fill(0); let mut nonzero = false; for (k, (&cf, l)) in coef.iter().zip(levels.iter_mut()).enumerate() { let q = if k == 0 { self.dc_q } else { self.ac_q };@@ -808,23 +848,36 @@ impl<'a> Decider<'a> { } let zero_dist = self.sse(p, x, y, n, pred); if !nonzero {- return (TxResult { levels, nonzero, dist: zero_dist, bits: 1.0 }, None);+ return (TxEval { nonzero, dist: zero_dist, bits: 1.0 }, false); }- let recon = self.reconstruct(&levels, tx_sz, tx_type, pred, n, bd);- let dist = self.sse(p, x, y, n, &recon);- let bits = estimate_bits(&levels, tx_sz);+ let recon = &mut recon[..nn];+ self.reconstruct(levels, &mut deq[..nn], &mut resid[..nn], recon, tx_sz, tx_type, pred, n, bd);+ let dist = self.sse(p, x, y, n, recon);+ let bits = estimate_bits(levels, tx_sz); if dist + self.lambda * bits >= zero_dist + self.lambda {- levels.iter_mut().for_each(|l| *l = 0);- return (TxResult { levels, nonzero: false, dist: zero_dist, bits: 1.0 }, None);- }- (TxResult { levels, nonzero: true, dist, bits }, Some(recon))- }-- /// Dequantisation (7.12.3) and the inverse transform, added to the prediction.- fn reconstruct(&self, levels: &[i32], tx_sz: usize, tx_type: usize, pred: &[u16], n: usize, bd: u32) -> Vec<u16> {+ levels.fill(0);+ return (TxEval { nonzero: false, dist: zero_dist, bits: 1.0 }, false);+ }+ (TxEval { nonzero: true, dist, bits }, true)+ }++ /// Dequantisation (7.12.3) and the inverse transform, added to the prediction; the result is+ /// written to `recon`. `deq` and `resid` are scratch buffers of `n * n` entries.+ fn reconstruct(+ &self,+ levels: &[i32],+ deq: &mut [i32],+ resid: &mut [i32],+ recon: &mut [u16],+ tx_sz: usize,+ tx_type: usize,+ pred: &[u16],+ n: usize,+ bd: u32,+ ) { let denom: i64 = if tx_sz == TX_32X32 { 2 } else { 1 }; let lim = 1i64 << (7 + bd);- let mut deq = vec![0i32; n * n];+ deq.fill(0); for (k, (&l, d)) in levels.iter().zip(deq.iter_mut()).enumerate() { if l == 0 { continue;@@ -834,10 +887,12 @@ impl<'a> Decider<'a> { let mag = (dq.abs() & 0xFF_FFFF) / denom; *d = (if dq < 0 { -mag } else { mag }).clamp(-lim, lim - 1) as i32; }- let mut resid = vec![0i32; n * n];- inverse_2d(&deq, &mut resid, tx_sz, tx_type, bd);+ inverse_2d(deq, resid, tx_sz, tx_type, bd); let max = (1i32 << bd) - 1;- pred.iter().zip(&resid).map(|(&p, &r)| (p as i32 + r).clamp(0, max) as u16).collect()+ for ((o, &p), &r) in recon.iter_mut().zip(pred).zip(resid.iter()) {+ *o = (p as i32 + r).clamp(0, max) as u16;+ }+ debug_assert_eq!(recon.len(), n * n); } } @@ -409,11 +409,13 @@ impl Encoder { }; let bd = self.cfg.bit_depth as u32; let mut cache: Vec<(u32, u64)> = vec![(0, sse(&rec[0]))];+ // one working copy of the luma plane for all candidate levels+ let mut planes = [rec[0].clone(), Plane::new(0, 0), Plane::new(0, 0)]; let mut eval = |lvl: u32| -> u64 { if let Some(&(_, e)) = cache.iter().find(|(l, _)| *l == lvl) { return e; }- let mut planes = [rec[0].clone(), Plane::new(0, 0), Plane::new(0, 0)];+ planes[0].data.copy_from_slice(&rec[0].data); loop_filter(&mut planes, info, [lvl; 4], bd, true, sb_row_step); let e = sse(&planes[0]); cache.push((lvl, e));@@ -62,6 +62,8 @@ pub(crate) struct TileWriter<'a> { left_dc: [Vec<u8>; 3], /// LoopfilterTxSizes per plane (row stride mi_cols + 32). lf_tx: [Vec<u8>; 3],+ /// Scratch for the levels coded so far in `coeffs`.+ quant: Vec<i32>, // current block r: usize, c: usize,@@ -114,6 +116,7 @@ impl<'a> TileWriter<'a> { left_level: std::array::from_fn(|_| vec![0; g.mi_rows + 64]), left_dc: std::array::from_fn(|_| vec![0; g.mi_rows + 64]), lf_tx: std::array::from_fn(|_| vec![0; (g.mi_cols + 32) * (g.mi_rows + 32)]),+ quant: Vec::new(), r: 0, c: 0, bsize: 0,@@ -181,10 +184,10 @@ impl<'a> TileWriter<'a> { self.sw.symbol(partition, cdf); } else if has_cols || has_rows { assert!(split, "blocks crossing the frame edge are split");- let pcdf: Vec<u16> = match bsl {- 2 => self.cdf.partition_w16[ctx].to_vec(),- 3 => self.cdf.partition_w32[ctx].to_vec(),- _ => self.cdf.partition_w64[ctx].to_vec(),+ let pcdf: &[u16] = match bsl {+ 2 => &self.cdf.partition_w16[ctx],+ 3 => &self.cdf.partition_w32[ctx],+ _ => &self.cdf.partition_w64[ctx], }; let pr = |k: usize| pcdf[k] as i32 - if k == 0 { 0 } else { pcdf[k - 1] as i32 }; let psum = if has_cols {@@ -501,7 +504,9 @@ impl<'a> TileWriter<'a> { let height = TX_HEIGHT[tx_sz] as usize; let width = 1usize << bwl; // Quant[] as the decoder sees it while parsing: the levels coded so far (capped)- let mut quant = vec![0i32; width * height];+ let mut quant = std::mem::take(&mut self.quant);+ quant.clear();+ quant.resize(width * height, 0); for c in (0..eob).rev() { let pos = scan[c] as usize; let a = levels[pos].unsigned_abs().min(15) as i32;@@ -526,6 +531,7 @@ impl<'a> TileWriter<'a> { } quant[pos] = a; }+ self.quant = quant; for c in 0..eob { let pos = scan[c] as usize; let v = levels[pos];