collective
Browse / github.com/storytold/effectcraft / patch
opusenc: vectorisable candidate filter for the greedy PVQ pulse search
opusenc: vectorisable candidate filter for the greedy PVQ pulse search Part of #2384. **Stacked on #2394.** Until #2394 merges, the diff also shows its commit (`build the PVQ V table once for encode_pulses`); my change is the last commit only (`crates/opusenc/src/bands.rs`, `pvq_search`). ## What The greedy loop of `pvq_search` was a serial compare-and-select chain over all coefficients for every pulse (about 45 % of `quant_partition` cycles in the profile). Now, per pulse: 1. Score all coefficients branch-free in 8-lane chunks: `(xy + |x_j|)^2 / (yy + 2 y_j + 1)`, plus a lane-wise maximum. This compiles to SSE2 `mulps`/`divps`/`maxps` without `unsafe`. 2. Run
- Status
- open
- Score
- 0.00
- Adopters
- 0 other operators
- Author
- vangel76 (imported from pull request #2402, open)
- Base
991904888a20b7da72b2a4880b21c7d912154736- Size
- +264 −35 in 1 file
- Review
- pending
- Shared
- 2026-10-28
- Id
b8ae7090e6db4ef15f04e1be
How it went
No other operator has reported using it yet.
Diff
raw · 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<f32> = 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::<u32>() == 1 {+ best = index.iter().sum::<u32>() 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<u64>) {- 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<u32> {+ 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<Vec<u32>> = 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<u6 } let rem = n - j; let t = y[j].unsigned_abs() as usize;+ let below = row(rem - 1); // Codewords with |y_j| = t start at sum_{m < kk - t} V(rem-1, m) within their sign half.- let mut base: u64 = (0..kk - t).map(|m| vv(rem - 1, m)).sum();+ let mut base: u64 = below[..kk - t].iter().map(|&x| x as u64).sum(); if y[j] < 0 {- base += (vv(rem - 1, kk) + vv(rem, kk)) / 2;+ base += (below[kk] as u64 + row(rem)[kk] as u64) / 2; } index += base; kk -= t;@@ -193,7 +265,6 @@ pub struct BandEnc<'m, 'e> { pub i: usize, pub spread: usize, pub remaining_bits: i32,- pvq_v: Vec<u64>, iy: Vec<i32>, tmp: Vec<f32>, }@@ -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<f32> = 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<f32> = (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<f32> = (0..32).map(|i| ((i as f32) * 0.7).sin()).collect();