diff --git a/crates/opusenc/src/bands.rs b/crates/opusenc/src/bands.rs --- a/crates/opusenc/src/bands.rs +++ b/crates/opusenc/src/bands.rs @@ -96,12 +96,24 @@ pub fn exp_rotation_inv(x: &mut [f32], len: usize, k: usize, spread: usize) { } } +/// Relative slack for the pulse-candidate filter in [`pvq_search`]: far above the f32 rounding of +/// the filter score (a few 1e-7), far below the spread of real scores. +const PVQ_SCORE_SLACK: f32 = 1e-5; + /// Greedy PVQ search: the integer vector `y` with `sum |y| = k` closest in direction to `x`. pub fn pvq_search(x: &[f32], k: usize, y: &mut [i32]) { let n = x.len(); y[..n].fill(0); - let abs: Vec = x.iter().map(|v| v.abs()).collect(); - let sum: f32 = abs.iter().sum(); + // `abs`, `dy` (2 y + 1, infinite in the padding) and the per-pass scores `score`, each padded to + // whole 8-lane chunks. + let np = n.next_multiple_of(8); + let mut buf = vec![0f32; 3 * np]; + let (abs, rest) = buf.split_at_mut(np); + let (dy, score) = rest.split_at_mut(np); + for (a, v) in abs.iter_mut().zip(x) { + *a = v.abs(); + } + let sum: f32 = abs[..n].iter().sum(); if !sum.is_finite() || sum <= 1e-15 { y[0] = k as i32; return; @@ -113,30 +125,69 @@ pub fn pvq_search(x: &[f32], k: usize, y: &mut [i32]) { // Pre-project onto the pyramid (rounding down) and place the rest greedily. let r = (k as f32 - 1.0) / sum; for j in 0..n { - let p = (abs[j] * r).floor() as i32; + // Non-negative and at most k, so the truncating cast is the floor. + let p = (abs[j] * r) as i32; y[j] = p; placed += p as usize; xy += abs[j] * p as f32; yy += (p * p) as f32; } } + for j in 0..np { + dy[j] = if j < n { 2.0 * y[j] as f32 + 1.0 } else { f32::INFINITY }; + } while placed < k { + // Vectorisable scores (xy + |x_j|)^2 / (yy + 2 y_j + 1) and their maximum. + for ((sc, &a), &d) in score.iter_mut().zip(abs.iter()).zip(dy.iter()) { + let t = xy + a; + *sc = t * t / (yy + d); + } + let mut lane_max = [0f32; 8]; + for c in score.as_chunks::<8>().0 { + for i in 0..8 { + lane_max[i] = if c[i] > lane_max[i] { c[i] } else { lane_max[i] }; + } + } + let top = lane_max.iter().fold(0f32, |m, &v| if v > m { v } else { m }); + // The exact cross-multiplied test below, in index order with the first maximum winning, only + // matters for entries within rounding distance of the top score: every other entry is beaten + // by the top one whenever it is reached and never beats it afterwards. Running the test on + // those alone therefore picks the same pulse as running it on all of them. + let filter = top > 1e-30 && top.is_finite(); + let lo = top * (1.0 - PVQ_SCORE_SLACK); + // Usually exactly one entry qualifies; its index falls out of a lane-wise count. + let mut count = [0u32; 8]; + let mut index = [0u32; 8]; + for (ci, c) in score.as_chunks::<8>().0.iter().enumerate() { + for i in 0..8 { + let hit = u32::from(c[i] >= lo); + count[i] += hit; + index[i] += hit * (ci * 8 + i) as u32; + } + } let mut best = 0usize; - let mut best_num = -1f32; - let mut best_den = 1f32; - for j in 0..n { - let num = xy + abs[j]; - let num = num * num; - let den = yy + 2.0 * y[j] as f32 + 1.0; - if num * best_den > best_num * den { - best_num = num; - best_den = den; - best = j; + if filter && count.iter().sum::() == 1 { + best = index.iter().sum::() as usize; + } else { + let mut best_num = -1f32; + let mut best_den = 1f32; + for j in 0..n { + if !filter || score[j] >= lo { + let num = xy + abs[j]; + let num = num * num; + let den = yy + 2.0 * y[j] as f32 + 1.0; + if num * best_den > best_num * den { + best_num = num; + best_den = den; + best = j; + } + } } } xy += abs[best]; yy += 2.0 * y[best] as f32 + 1.0; y[best] += 1; + dy[best] = 2.0 * y[best] as f32 + 1.0; placed += 1; } for j in 0..n { @@ -146,25 +197,45 @@ pub fn pvq_search(x: &[f32], k: usize, y: &mut [i32]) { } } -/// Encodes the PVQ codeword `y` (`sum |y| = k`) as the index the RFC 6716 §4.3.4.2 decoder maps -/// back to `y`. -pub fn encode_pulses(y: &[i32], k: usize, enc: &mut RangeEncoder, v: &mut Vec) { - let n = y.len(); - let w = k + 1; - v.clear(); - v.resize((n + 1) * w, 0); +/// Largest band width and pulse count the encoder codes (176 coefficients, `get_pulses(40)` = 128). +const PVQ_MAX_N: usize = 176; +const PVQ_MAX_K: usize = 128; + +/// Row-major `V(n, k)` for `n <= max_n`, `k <= max_k` with row stride `max_k + 1`, saturated at +/// `u32::MAX` (a codebook that large is never coded, see `rate::fits_in32`). +fn build_pvq_table(max_n: usize, max_k: usize) -> Vec { + let w = max_k + 1; + let mut v = vec![0u32; (max_n + 1) * w]; v[0] = 1; - for nn in 1..=n { + for nn in 1..=max_n { v[nn * w] = 1; - for kk in 1..=k { - let a = v[(nn - 1) * w + kk]; - let b = v[nn * w + kk - 1]; - let c = v[(nn - 1) * w + kk - 1]; - v[nn * w + kk] = a.saturating_add(b).saturating_add(c); + for kk in 1..=max_k { + let sum = v[(nn - 1) * w + kk] as u64 + v[nn * w + kk - 1] as u64 + v[(nn - 1) * w + kk - 1] as u64; + v[nn * w + kk] = sum.min(u32::MAX as u64) as u32; } } - let total = v[n * w + k]; - let vv = |nn: usize, kx: usize| v[nn * w + kx]; + v +} + +fn pvq_table() -> &'static [u32] { + static TABLE: std::sync::OnceLock> = std::sync::OnceLock::new(); + TABLE.get_or_init(|| build_pvq_table(PVQ_MAX_N, PVQ_MAX_K)) +} + +/// Encodes the PVQ codeword `y` (`sum |y| = k`) as the index the RFC 6716 §4.3.4.2 decoder maps +/// back to `y`. +pub fn encode_pulses(y: &[i32], k: usize, enc: &mut RangeEncoder) { + let n = y.len(); + // Shapes outside the coded range (never produced by the band layout) get a local table. + let local; + let (v, w) = if n <= PVQ_MAX_N && k <= PVQ_MAX_K { + (pvq_table(), PVQ_MAX_K + 1) + } else { + local = build_pvq_table(n, k); + (&local[..], k + 1) + }; + let row = |nn: usize| &v[nn * w..(nn + 1) * w]; + let total = row(n)[k] as u64; let mut index = 0u64; let mut kk = k; for j in 0..n { @@ -173,10 +244,11 @@ pub fn encode_pulses(y: &[i32], k: usize, enc: &mut RangeEncoder, v: &mut Vec { pub i: usize, pub spread: usize, pub remaining_bits: i32, - pvq_v: Vec, iy: Vec, tmp: Vec, } @@ -287,7 +358,7 @@ impl BandEnc<'_, '_> { iy.clear(); iy.resize(n, 0); pvq_search(&tmp, k, &mut iy); - encode_pulses(&iy, k, self.enc, &mut self.pvq_v); + encode_pulses(&iy, k, self.enc); self.iy = iy; self.tmp = tmp; } @@ -316,7 +387,7 @@ pub fn quant_all_bands( ) { let mm = 1usize << lm; debug_assert!(c == 1 || dual_stereo || intensity <= start, "stereo frames are coded as dual stereo"); - let mut ctx = BandEnc { m, enc, i: 0, spread, remaining_bits: 0, pvq_v: Vec::new(), iy: Vec::new(), tmp: Vec::new() }; + let mut ctx = BandEnc { m, enc, i: 0, spread, remaining_bits: 0, iy: Vec::new(), tmp: Vec::new() }; for i in start..end { ctx.i = i; let xo = mm * EBANDS[i] as usize; @@ -466,12 +537,11 @@ mod tests { #[test] fn pvq_index_roundtrip_exhaustive() { - let mut v = Vec::new(); for n in 1..6usize { for k in 1..6usize { for y in all_codewords(n, k) { let mut e = RangeEncoder::new(); - encode_pulses(&y, k, &mut e, &mut v); + encode_pulses(&y, k, &mut e); let buf = e.done(16); let mut d = RangeDecoder::new(&buf); assert_eq!(decode_pulses(n, k, &mut d), y, "n={n} k={k}"); @@ -480,6 +550,165 @@ mod tests { } } + /// The per-call table encoder this module used before the shared table. + fn encode_pulses_reference(y: &[i32], k: usize, enc: &mut RangeEncoder) { + let n = y.len(); + let w = k + 1; + let mut v = vec![0u64; (n + 1) * w]; + v[0] = 1; + for nn in 1..=n { + v[nn * w] = 1; + for kk in 1..=k { + let a = v[(nn - 1) * w + kk]; + let b = v[nn * w + kk - 1]; + let c = v[(nn - 1) * w + kk - 1]; + v[nn * w + kk] = a.saturating_add(b).saturating_add(c); + } + } + let total = v[n * w + k]; + let vv = |nn: usize, kx: usize| v[nn * w + kx]; + let mut index = 0u64; + let mut kk = k; + for j in 0..n { + if kk == 0 { + break; + } + let rem = n - j; + let t = y[j].unsigned_abs() as usize; + let mut base: u64 = (0..kk - t).map(|m| vv(rem - 1, m)).sum(); + if y[j] < 0 { + base += (vv(rem - 1, kk) + vv(rem, kk)) / 2; + } + index += base; + kk -= t; + } + if total >= 2 { + enc.uint(index as u32, total as u32); + } + } + + #[test] + fn pvq_table_encode_matches_per_call_table() { + let mut seed = 0x9E37_79B9_7F4A_7C15u64; + let mut next = move |m: u64| { + seed ^= seed << 13; + seed ^= seed >> 7; + seed ^= seed << 17; + seed % m + }; + let mut checked = 0; + for n in (1..=PVQ_MAX_N).step_by(5).chain([PVQ_MAX_N]) { + for k in (1..=PVQ_MAX_K).step_by(3).chain([PVQ_MAX_K]) { + if crate::rate::pvq_v(n, k) > u32::MAX as u128 { + continue; + } + for _ in 0..4 { + let mut y = vec![0i32; n]; + for _ in 0..k { + let at = next(n as u64) as usize; + y[at] += if next(2) == 0 { 1 } else { -1 }; + } + // Opposite pulses cancel; top up so that sum |y| = k again. + let mut have: usize = y.iter().map(|v| v.unsigned_abs() as usize).sum(); + while have < k { + let at = next(n as u64) as usize; + y[at] += if y[at] < 0 { -1 } else { 1 }; + have += 1; + } + let (mut a, mut b) = (RangeEncoder::new(), RangeEncoder::new()); + encode_pulses(&y, k, &mut a); + encode_pulses_reference(&y, k, &mut b); + assert_eq!(a.done(16), b.done(16), "n={n} k={k}"); + checked += 1; + } + } + } + assert!(checked > 100, "{checked}"); + } + + /// The scalar greedy search the vectorised one replaced. + fn pvq_search_ref(x: &[f32], k: usize, y: &mut [i32]) { + let n = x.len(); + y[..n].fill(0); + let abs: Vec = x.iter().map(|v| v.abs()).collect(); + let sum: f32 = abs.iter().sum(); + if !sum.is_finite() || sum <= 1e-15 { + y[0] = k as i32; + return; + } + let mut placed = 0usize; + let mut xy = 0f32; + let mut yy = 0f32; + if k > n / 2 { + let r = (k as f32 - 1.0) / sum; + for j in 0..n { + let p = (abs[j] * r).floor() as i32; + y[j] = p; + placed += p as usize; + xy += abs[j] * p as f32; + yy += (p * p) as f32; + } + } + while placed < k { + let mut best = 0usize; + let mut best_num = -1f32; + let mut best_den = 1f32; + for j in 0..n { + let num = xy + abs[j]; + let num = num * num; + let den = yy + 2.0 * y[j] as f32 + 1.0; + if num * best_den > best_num * den { + best_num = num; + best_den = den; + best = j; + } + } + xy += abs[best]; + yy += 2.0 * y[best] as f32 + 1.0; + y[best] += 1; + placed += 1; + } + for j in 0..n { + if x[j] < 0.0 { + y[j] = -y[j]; + } + } + } + + #[test] + fn pvq_search_matches_scalar_reference() { + let mut state = 0x9e3779b97f4a7c15u64; + let mut rnd = move || { + state ^= state << 13; + state ^= state >> 7; + state ^= state << 17; + (state >> 40) as f32 / (1u64 << 24) as f32 * 2.0 - 1.0 + }; + let mut a = vec![0i32; 200]; + let mut b = vec![0i32; 200]; + let mut cases = 0; + for round in 0..6000 { + let n = 1 + (round * 37 + (rnd().abs() * 176.0) as usize) % 176; + let k = 1 + (rnd().abs() * 127.0) as usize; + let scale = [1.0f32, 1e-9, 1e6, 3e-3][round % 4]; + let x: Vec = (0..n) + .map(|_| match round % 6 { + 0 => rnd(), + 1 => if rnd() > 0.8 { rnd() } else { 0.0 }, + 2 => (rnd() * 4.0).round() / 4.0, + 3 => if rnd() > 0.0 { 1.0 } else { -1.0 }, + 4 => rnd().powi(5), + _ => if rnd() > 0.95 { 3.0 * rnd() } else { 1e-4 * rnd() }, + } * scale) + .collect(); + pvq_search(&x, k, &mut a); + pvq_search_ref(&x, k, &mut b); + assert_eq!(a[..n], b[..n], "n={n} k={k} round={round}"); + cases += 1; + } + assert_eq!(cases, 6000); + } + #[test] fn pvq_search_finds_good_direction() { let x: Vec = (0..32).map(|i| ((i as f32) * 0.7).sin()).collect();