diff --git a/crates/aac/src/encoder.rs b/crates/aac/src/encoder.rs --- a/crates/aac/src/encoder.rs +++ b/crates/aac/src/encoder.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, + /// Rate-loop memo per channel, filled while the frame is quantised. + caches: Vec, ms_mask: Option<(u32, Vec)>, } @@ -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 = 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> { + fn quantise_all(frames: &mut [ElementFrame], lambda: f32) -> Vec> { #[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; diff --git a/crates/aac/src/encoder/quant.rs b/crates/aac/src/encoder/quant.rs --- a/crates/aac/src/encoder/quant.rs +++ b/crates/aac/src/encoder/quant.rs @@ -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, + bands: Vec, + /// Quantised values of every band at its memoised scalefactor. + q: Vec, + /// Cost row of a band that quantises to zeros. + zero_cost: Vec>, +} + +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 = (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}"); + } + } } diff --git a/crates/aac/src/huffman.rs b/crates/aac/src/huffman.rs --- a/crates/aac/src/huffman.rs +++ b/crates/aac/src/huffman.rs @@ -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(); diff --git a/crates/aac/src/tables.rs b/crates/aac/src/tables.rs --- a/crates/aac/src/tables.rs +++ b/crates/aac/src/tables.rs @@ -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> = 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> = 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, }