collective
Browse / github.com/storytold/filmcraft / patch
aac: memoise the encoder's per-band work across rate-loop iterations (2.4x faster, output identical)
aac: the encoder repeats the per-band search in every rate-loop iteration (2.0-2.5x) ## Problem `crates/aac` encoder: every frame runs the rate loop 5-9 times and redoes the per-band search each time. ## Proposal Cache the per-band work per (band, scalefactor) across the rate-loop iterations, plus two inner-loop fixes. With the `rayon` feature the caches must be `&mut` (per worker), not `RefCell`. ## Measured (prototype) Stereo 128 kbit/s: 17.50 G -> 8.44 G instructions (2.0-2.5x), 1.07 s -> 0.46 s. Output bytes identical in 4 configurations. `cargo test --release -p filmcraft-aac` passes. Not part of this issue: a lambda warm start (estimated another 1.3-1.5x) changes the
- Status
- open
- Score
- 0.00
- Adopters
- 0 other operators
- Author
- vangel76 (imported from pull request #3264, open)
- Base
7428957ae0ece1139ae3af7ea76e97cacdb69570- Size
- +195 −55 in 4 files
- Review
- pending
- Shared
- 2026-10-28
- Id
e0ce8a535768d5ed97ac04e0
How it went
No other operator has reported using it yet.
Diff
raw · crates/aac/src/encoder.rs crates/aac/src/encoder/quant.rs crates/aac/src/huffman.rs crates/aac/src/tables.rs
@@ -25,7 +25,7 @@ use crate::tables::{sample_rate_index, swb_offsets_long, swb_offsets_short, tns_ use crate::tns::{self, TnsData, TnsFilter}; use crate::{Error, FRAME_LEN, Result}; use psy::Psy;-use quant::{ChannelInput, QuantChannel, quantise_channel, write_ics};+use quant::{ChannelInput, QuantCache, QuantChannel, quantise_channel, write_ics}; /// Rate control mode. #[derive(Debug, Clone, Copy, PartialEq)]@@ -110,6 +110,8 @@ struct ElementState { struct ElementFrame { info: IcsInfo, chans: Vec<ChannelInput>,+ /// Rate-loop memo per channel, filled while the frame is quantised.+ caches: Vec<QuantCache>, ms_mask: Option<(u32, Vec<bool>)>, } @@ -575,14 +577,15 @@ impl Encoder { }; ms_mask = Some((present, used)); }- let chans = specs.iter().zip(thrs.iter()).zip(tnss).map(|((s, th), tn)| ChannelInput::new(info, swb, s, th, tn)).collect();+ let chans: Vec<ChannelInput> = specs.iter().zip(thrs.iter()).zip(tnss).map(|((s, th), tn)| ChannelInput::new(info, swb, s, th, tn)).collect();+ let caches = chans.iter().map(QuantCache::new).collect(); let e = &mut self.elements[ei]; e.prev_seq = seq; e.short_next = Some(short_next); if short { self.stats.short_blocks += 1; }- (ElementFrame { info, chans, ms_mask }, pe)+ (ElementFrame { info, chans, caches, ms_mask }, pe) } fn element_overhead(&self, ei: usize, f: &ElementFrame) -> u32 {@@ -599,15 +602,18 @@ impl Encoder { } } - fn quantise_all(frames: &[ElementFrame], lambda: f32) -> Vec<Vec<QuantChannel>> {+ fn quantise_all(frames: &mut [ElementFrame], lambda: f32) -> Vec<Vec<QuantChannel>> { #[cfg(feature = "rayon")] { use rayon::prelude::*;- frames.par_iter().map(|f| f.chans.par_iter().map(|c| quantise_channel(c, lambda)).collect()).collect()+ frames+ .par_iter_mut()+ .map(|f| f.chans.par_iter().zip(f.caches.par_iter_mut()).map(|(c, cache)| quantise_channel(c, cache, lambda)).collect())+ .collect() } #[cfg(not(feature = "rayon"))] {- frames.iter().map(|f| f.chans.iter().map(|c| quantise_channel(c, lambda)).collect()).collect()+ frames.iter_mut().map(|f| f.chans.iter().zip(f.caches.iter_mut()).map(|(c, cache)| quantise_channel(c, cache, lambda)).collect()).collect() } } @@ -660,16 +666,16 @@ impl Encoder { BitrateMode::Vbr(q) => { let d = q.clamp(0.5, 6.0) - 1.0; let mut lambda = 10f32.powf(-0.2 - 1.2 * d + 0.09 * d * d);- let mut res = Self::quantise_all(&frames, lambda);+ let mut res = Self::quantise_all(&mut frames, lambda); while total_bits(&res) > target && lambda < 1e9 { lambda *= 2.0;- res = Self::quantise_all(&frames, lambda);+ res = Self::quantise_all(&mut frames, lambda); } res } BitrateMode::Cbr(_) => { let mut lam = self.lambda.clamp(1e-6, 1e9);- let mut cur = Self::quantise_all(&frames, lam);+ let mut cur = Self::quantise_all(&mut frames, lam); let (mut lo, mut hi, mut best); if total_bits(&cur) <= target { hi = lam;@@ -680,7 +686,7 @@ impl Encoder { if lo < 1e-6 { break; }- cur = Self::quantise_all(&frames, lo);+ cur = Self::quantise_all(&mut frames, lo); if total_bits(&cur) > target { break; }@@ -691,7 +697,7 @@ impl Encoder { lo = lam; loop { lam *= 4.0;- cur = Self::quantise_all(&frames, lam);+ cur = Self::quantise_all(&mut frames, lam); if total_bits(&cur) <= target || lam > 1e9 { break; }@@ -705,7 +711,7 @@ impl Encoder { break; } let mid = (lo * hi).sqrt();- let r = Self::quantise_all(&frames, mid);+ let r = Self::quantise_all(&mut frames, mid); if total_bits(&r) <= target { hi = mid; best = r;@@ -8,7 +8,7 @@ use filmcraft_bitstream::BitWriter; use crate::huffman::{self, ZERO_HCB}; use crate::ics::IcsInfo;-use crate::tables::pow43;+use crate::tables::pow43_table; use crate::tns::TnsData; /// Magic rounding offset for AAC quantisation (`0.5 - 0.0946`), minimising the error for the@@ -67,6 +67,42 @@ impl ChannelInput { } } +/// Scalefactor-keyed results of one band, valid until the band is quantised at another scalefactor.+#[derive(Clone, Copy)]+struct BandMemo {+ /// Scalefactor the cached values belong to (`NO_SF` when empty).+ sf: i32,+ maxq: i32,+ /// Spectral bits per codebook plus one for an all-zero band (see `band_costs`), once computed.+ cost: Option<[u32; 12]>,+}++const NO_SF: i32 = i32::MIN;++/// Work that depends only on (band, scalefactor), not on λ, kept across the rate-loop iterations of+/// one frame. One cache per channel; callers hold it mutably, so parallel channels share nothing.+pub struct QuantCache {+ /// Distortion per (band, scalefactor), NaN while unknown.+ dist: Vec<f32>,+ bands: Vec<BandMemo>,+ /// Quantised values of every band at its memoised scalefactor.+ q: Vec<i32>,+ /// Cost row of a band that quantises to zeros.+ zero_cost: Vec<Option<[u32; 12]>>,+}++impl QuantCache {+ pub fn new(ch: &ChannelInput) -> QuantCache {+ let nb = ch.bands.len();+ QuantCache {+ dist: vec![f32::NAN; nb * 256],+ bands: vec![BandMemo { sf: NO_SF, maxq: 0, cost: None }; nb],+ q: vec![0; ch.x.len()],+ zero_cost: vec![None; nb],+ }+ }+}+ /// A quantised channel ready to be written. #[derive(Clone, Default)] pub struct QuantChannel {@@ -90,46 +126,54 @@ fn gain(sf: i32) -> f32 { (0.25 * (sf as f32 - 100.0)).exp2() } -/// Quantise a slice at `sf`, returning the distortion; writes the (signed) values to `out` if given.-fn quantise(x: &[f32], xabs: &[f32], x34: &[f32], sf: i32, out: Option<&mut [i32]>) -> f32 {+/// Distortion of a slice quantised at `sf`. Kept out of line so the accumulator stays in a register.+#[inline(never)]+fn distortion(xabs: &[f32], x34: &[f32], sf: i32) -> f32 { let g34 = gain34(sf); let g = gain(sf);+ let t = pow43_table(); let mut d = 0f32;- match out {- Some(out) => {- for i in 0..x.len() {- let q = ((x34[i] * g34 + ROUND) as i32).min(MAX_Q);- let e = xabs[i] - pow43(q as u32) * g;- d += e * e;- out[i] = if x[i] < 0.0 { -q } else { q };- }- }- None => {- for i in 0..x.len() {- let q = ((x34[i] * g34 + ROUND) as i32).min(MAX_Q);- let e = xabs[i] - pow43(q as u32) * g;- d += e * e;- }- }+ for (&a, &b) in xabs.iter().zip(x34) {+ let q = ((b * g34 + ROUND) as i32).min(MAX_Q);+ let e = a - t[(q as usize) & 8191] * g;+ d += e * e; } d } +/// Quantise a slice at `sf`, writing the (signed) values to `out`.+fn quantise(x: &[f32], sf: i32, x34: &[f32], out: &mut [i32]) {+ let g34 = gain34(sf);+ for ((o, &xv), &b) in out.iter_mut().zip(x).zip(x34) {+ let q = ((b * g34 + ROUND) as i32).min(MAX_Q);+ *o = if xv < 0.0 { -q } else { q };+ }+}+ /// Largest scalefactor whose distortion stays within `allowed`.-fn search_sf(ch: &ChannelInput, b: &Band, allowed: f32) -> i32 {+fn search_sf(ch: &ChannelInput, bi: usize, dist: &mut [f32], allowed: f32) -> i32 {+ let b = &ch.bands[bi]; let r = b.start..b.end;- let (x, xa, x34) = (&ch.x[r.clone()], &ch.xabs[r.clone()], &ch.x34[r]);+ let (xa, x34) = (&ch.xabs[r.clone()], &ch.x34[r]);+ let dist = &mut dist[bi * 256..bi * 256 + 256];+ let mut dist_at = |sf: i32| {+ let d = &mut dist[sf as usize];+ if d.is_nan() {+ *d = distortion(xa, x34, sf);+ }+ *d+ }; let l2 = b.max34.log2(); let lo = ((100.0 - ((MAX_Q as f32 + 0.5).log2() - l2) / 0.1875).ceil() as i32).clamp(0, 255); let zero = ((100.0 - ((1.0 - ROUND).log2() - l2) / 0.1875).floor() as i32 + 1).clamp(0, 256); let mut a = lo; let mut hi = (zero - 1).clamp(lo, 255);- if quantise(x, xa, x34, a, None) > allowed {+ if dist_at(a) > allowed { return a; } while a < hi { let mid = (a + hi + 1) / 2;- if quantise(x, xa, x34, mid, None) <= allowed {+ if dist_at(mid) <= allowed { a = mid; } else { hi = mid - 1;@@ -151,8 +195,31 @@ fn candidates(m: i32) -> &'static [u8] { } } +/// Bits per candidate codebook for a band holding `vals` (largest magnitude `maxq`); `u32::MAX / 4`+/// marks codebooks that cannot code it. An all-zero band inside a coded section also needs a (zero)+/// scalefactor difference, hence the extra bit.+fn band_costs(vals: &[i32], maxq: i32) -> [u32; 12] {+ const INF: u32 = u32::MAX / 4;+ let mut cost = [INF; 12];+ let extra = (maxq == 0) as u32;+ for &c in candidates(maxq) {+ match c {+ 0 => cost[0] = 0,+ // odd codebook below 11 and its successor share everything but the length table+ c if c % 2 == 1 && c < huffman::ESC_HCB => {+ let (a, b) = huffman::count_bits_pair(c, vals);+ cost[c as usize] = a + extra;+ cost[c as usize + 1] = b + extra;+ }+ c if c % 2 == 0 => {} // filled in by the pair above+ c => cost[c as usize] = huffman::count_bits(c, vals) + extra,+ }+ }+ cost+}+ /// Quantise one channel at rate-distortion multiplier `lambda` and select sections.-pub fn quantise_channel(ch: &ChannelInput, lambda: f32) -> QuantChannel {+pub fn quantise_channel(ch: &ChannelInput, cache: &mut QuantCache, lambda: f32) -> QuantChannel { let nb = ch.bands.len(); let mut sf = vec![-1i32; nb]; for (i, b) in ch.bands.iter().enumerate() {@@ -160,7 +227,7 @@ pub fn quantise_channel(ch: &ChannelInput, lambda: f32) -> QuantChannel { if b.max34 <= 0.0 || b.en <= allowed { continue; }- sf[i] = search_sf(ch, b, allowed);+ sf[i] = search_sf(ch, i, &mut cache.dist, allowed); } // Scalefactor differences between coded bands must stay within ±60: lower (refine) offenders, // then quantise; repeat if a band unexpectedly quantised to all zeros (it leaves the chain).@@ -194,8 +261,13 @@ pub fn quantise_channel(ch: &ChannelInput, lambda: f32) -> QuantChannel { maxq[i] = 0; continue; }- quantise(&ch.x[r.clone()], &ch.xabs[r.clone()], &ch.x34[r.clone()], sf[i], Some(&mut q[r.clone()]));- maxq[i] = q[r].iter().map(|v| v.abs()).max().unwrap_or(0);+ let memo = &mut cache.bands[i];+ if memo.sf != sf[i] {+ quantise(&ch.x[r.clone()], sf[i], &ch.x34[r.clone()], &mut cache.q[r.clone()]);+ *memo = BandMemo { sf: sf[i], maxq: cache.q[r.clone()].iter().map(|v| v.abs()).max().unwrap_or(0), cost: None };+ }+ q[r.clone()].copy_from_slice(&cache.q[r]);+ maxq[i] = memo.maxq; if maxq[i] == 0 { sf[i] = -1; removed = true;@@ -228,14 +300,9 @@ pub fn quantise_channel(ch: &ChannelInput, lambda: f32) -> QuantChannel { let bi = base + i; let b = &ch.bands[bi]; let vals = &q[b.start..b.end];- for &c in candidates(maxq[bi]) {- cost[i][c as usize] = if c == 0 {- 0- } else {- // all-zero band inside a coded section also needs a (zero) scalefactor difference- huffman::count_bits(c, vals) + if maxq[bi] == 0 { 1 } else { 0 }- };- }+ // A band that ends up uncoded (`sf < 0`) holds zeros, whatever its memo says.+ let row = if sf[bi] >= 0 { &mut cache.bands[bi].cost } else { &mut cache.zero_cost[bi] };+ cost[i] = *row.get_or_insert_with(|| band_costs(vals, maxq[bi])); } let new_sec = 4 + sect_bits; let mut best = [INF; 12];@@ -284,8 +351,8 @@ pub fn quantise_channel(ch: &ChannelInput, lambda: f32) -> QuantChannel { _ => secs.push((c, 1)), } if c != ZERO_HCB {- let b = &ch.bands[base + i];- spectral_bits += huffman::count_bits(c, &q[b.start..b.end]);+ // the cost row carries one extra bit for an all-zero band (its scalefactor difference)+ spectral_bits += cost[i][c as usize] - (maxq[base + i] == 0) as u32; } } for &(_, len) in &secs {@@ -373,10 +440,27 @@ mod tests { let thr = vec![vec![1000.0f32; 64]]; let ch = ChannelInput::new(info, swb, &spec, &thr, None); for lambda in [0.01f32, 1.0, 100.0, 1e5] {- let qc = quantise_channel(&ch, lambda);+ let mut cache = QuantCache::new(&ch);+ let qc = quantise_channel(&ch, &mut cache, lambda); let mut bw = BitWriter::new(); write_ics(&mut bw, &ch, &qc, false); assert_eq!(bw.bit_len() as u32, qc.bits, "lambda {lambda}"); } }++ #[test]+ fn reused_cache_gives_the_same_result_as_a_fresh_one() {+ let swb = swb_offsets_long(4);+ let info = IcsInfo { window_sequence: WindowSequence::OnlyLong, window_shape: WindowShape::Sine, max_sfb: 40, ..Default::default() };+ let spec: Vec<f32> = (0..1024).map(|i| ((i * 7919 % 1000) as f32 - 500.0) * (1000.0 / (i as f32 + 10.0))).collect();+ let thr = vec![vec![1000.0f32; 64]];+ let ch = ChannelInput::new(info, swb, &spec, &thr, None);+ let mut shared = QuantCache::new(&ch);+ // up and down, like the rate loop, with repeats+ for lambda in [1.0f32, 4.0, 16.0, 0.25, 1.0, 100.0, 3.0, 3.0, 1e5, 0.01, 1.0] {+ let a = quantise_channel(&ch, &mut shared, lambda);+ let b = quantise_channel(&ch, &mut QuantCache::new(&ch), lambda);+ assert_eq!((&a.q, &a.sf, &a.cb, &a.sections, a.global_gain, a.bits), (&b.q, &b.sf, &b.cb, &b.sections, b.global_gain, b.bits), "lambda {lambda}");+ }+ } }@@ -198,6 +198,27 @@ pub fn count_bits(cb: u8, vals: &[i32]) -> u32 { bits } +/// Bits for codebooks `cb` and `cb + 1` in one pass. `cb` must be odd and below 11: each such pair+/// shares dimension, modulus and signedness, so the index is computed once and only the length table+/// differs.+pub fn count_bits_pair(cb: u8, vals: &[i32]) -> (u32, u32) {+ debug_assert!(cb % 2 == 1 && cb < ESC_HCB);+ let (b, b2) = (book(cb), book(cb + 1));+ let (mut x, mut y) = (0u32, 0u32);+ for chunk in vals.chunks_exact(b.dim) {+ let i = index_of(b, chunk);+ x += b.codes[i].1 as u32;+ y += b2.codes[i].1 as u32;+ if !b.signed {+ // sign bits (no escapes below codebook 11)+ let nz = chunk.iter().filter(|&&v| v != 0).count() as u32;+ x += nz;+ y += nz;+ }+ }+ (x, y)+}+ /// Write `vals` with codebook `cb` (1..=11). pub fn write_spectral(bw: &mut BitWriter, cb: u8, vals: &[i32]) { let b = book(cb);@@ -291,6 +312,25 @@ mod tests { } } + #[test]+ fn pair_count_matches_single_count() {+ let mut seed = 99u32;+ for cb in (1..11u8).step_by(2) {+ let b = book(cb);+ for round in 0..40 {+ let mut vals = Vec::new();+ for _ in 0..(8 * b.dim) {+ seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);+ let m = ((seed >> 8) % (b.lav + 1)) as i32;+ let m = if round % 4 == 0 { 0 } else { m };+ let neg = if b.signed { seed & 0x8000_0000 != 0 } else { seed & 1 == 1 };+ vals.push(if neg { -m } else { m });+ }+ assert_eq!(count_bits_pair(cb, &vals), (count_bits(cb, &vals), count_bits(cb + 1, &vals)), "codebook {cb}");+ }+ }+ }+ #[test] fn scalefactor_roundtrip() { let mut bw = BitWriter::new();@@ -99,12 +99,22 @@ pub fn tns_max_order(short: bool) -> usize { if short { 7 } else { 12 } } -/// `x^(4/3)` for the integer quantised magnitudes 0..=8191.-pub fn pow43(q: u32) -> f32 {+/// The `x^(4/3)` table for the integer quantised magnitudes 0..=8191.+pub fn pow43_table() -> &'static [f32; 8192] { use std::sync::OnceLock;- static TABLE: OnceLock<Vec<f32>> = OnceLock::new();- let t = TABLE.get_or_init(|| (0..8192u32).map(|i| (i as f64).powf(4.0 / 3.0) as f32).collect());- match t.get(q as usize) {+ static TABLE: OnceLock<Box<[f32; 8192]>> = OnceLock::new();+ TABLE.get_or_init(|| {+ let mut t = Box::new([0f32; 8192]);+ for (i, v) in t.iter_mut().enumerate() {+ *v = (i as f64).powf(4.0 / 3.0) as f32;+ }+ t+ })+}++/// `x^(4/3)` for an integer quantised magnitude (table lookup up to 8191).+pub fn pow43(q: u32) -> f32 {+ match pow43_table().get(q as usize) { Some(v) => *v, None => (q as f64).powf(4.0 / 3.0) as f32, }