aleleba/effectcraft
1
0
Fork 0
mirror of https://github.com/storytold/effectcraft.git synced 2026-10-29 20:47:18 -06:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #2402 from vangel76/opusenc-pvq-search

Browse Source 
opusenc: vectorisable candidate filter for the greedy PVQ pulse search
This commit is contained in:
BFlat
2026-10-30 01:17:19 +07:00
committed by GitHub
parent fc4f9579b2 63c9ac8941
commit 4a060dbb5a
1 changed files with 264 additions and 35 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/opusenc/src/bands.rs
+264 -35
View File
@@ -97,12 +97,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;
@@ -114,30 +126,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 {
@@ -147,25 +198,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 {
@@ -174,10 +245,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;
@@ -194,7 +266,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>,
}
@@ -288,7 +359,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;
}
@@ -310,7 +381,7 @@ pub fn quant_all_bands(m: &Mode, frame: CeltFrame, x: &[f32], n_total: usize, al
let BandCoding { spread, total_bits } = coding;
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;
@@ -460,12 +531,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}");
@@ -474,6 +544,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();