Skip to content

Commit 45ee193

Browse files
committed
add frozen prng algo for SORF
Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>
1 parent da2f733 commit 45ee193

6 files changed

Lines changed: 181 additions & 29 deletions

File tree

‎vortex-tensor/Cargo.toml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,8 @@ half = { workspace = true }
2929
itertools = { workspace = true }
3030
num-traits = { workspace = true }
3131
prost = { workspace = true }
32-
rand = { workspace = true }
3332

3433
[dev-dependencies]
34+
rand = { workspace = true }
3535
rand_distr = { workspace = true }
3636
rstest = { workspace = true }

‎vortex-tensor/src/encodings/turboquant/centroids.rs‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -176,7 +176,6 @@ fn mean_between_centroids(lo: f64, hi: f64, exponent: HalfIntExponent) -> f64 {
176176
/// Uses `powi` + `sqrt` instead of `powf` for the half-integer exponents that arise from `(d-3)/2`.
177177
/// This is significantly faster than the general `powf` which goes through
178178
/// `exp(exponent * ln(base))`.
179-
#[inline]
180179
fn pdf_unnormalized(x_val: f64, exponent: HalfIntExponent) -> f64 {
181180
let base = (1.0 - x_val * x_val).max(0.0);
182181

‎vortex-tensor/src/encodings/turboquant/mod.rs‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -42,9 +42,9 @@
4242
//!
4343
//! The TurboQuant paper analyzes a full random orthogonal rotation. The current Vortex
4444
//! implementation instead uses a fixed 3-round Walsh-Hadamard-based structured transform with
45-
//! random sign diagonals. This is a practical approximation chosen for encode/decode efficiency,
46-
//! and should be understood as an implementation choice rather than the exact construction used in
47-
//! the paper's proofs.
45+
//! random sign diagonals generated by Vortex's frozen local SplitMix64 stream. This is a practical
46+
//! approximation chosen for encode/decode efficiency, and should be understood as an
47+
//! implementation choice rather than the exact construction used in the paper's proofs.
4848
//!
4949
//! The current encoding is also intentionally MSE-only. It does not yet implement the paper's QJL
5050
//! residual correction for unbiased inner-product estimation, and it still uses internal

‎vortex-tensor/src/scalar_fns/sorf_transform/mod.rs‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,9 @@
1212
//! transform at execution time, producing a [`Vector`] extension array with the original
1313
//! (pre-padding) dimensionality.
1414
//!
15-
//! The transform parameters are stored as a deterministic PRNG seed in [`SorfOptions`], so the
16-
//! [`SorfMatrix`] is reconstructed cheaply at decode time.
15+
//! The transform parameters are stored as a deterministic seed in [`SorfOptions`], so the
16+
//! [`SorfMatrix`] is reconstructed cheaply at decode time. Sign diagonals are defined by Vortex's
17+
//! frozen local SplitMix64 stream contract rather than by an external RNG crate.
1718
//!
1819
//! **All SORF computation happens in f32.** Input elements of other float types (f16, f64) are cast
1920
//! to f32 before the transform, and the result is cast back to the target type specified by
@@ -33,6 +34,7 @@ use vortex_error::VortexResult;
3334
use vortex_error::vortex_ensure;
3435

3536
mod rotation;
37+
mod splitmix64;
3638
pub use rotation::SorfMatrix;
3739

3840
mod vtable;
@@ -51,7 +53,7 @@ pub struct SorfTransform;
5153
/// [`SorfMatrix`] at decode time.
5254
#[derive(Clone, Debug, PartialEq, Eq, Hash)]
5355
pub struct SorfOptions {
54-
/// PRNG seed used to generate the random sign diagonals.
56+
/// Seed used to generate the structured sign diagonals via Vortex's frozen SplitMix64 stream.
5557
pub seed: u64,
5658
/// Number of sign-diagonal + WHT rounds in the structured orthogonal transform.
5759
pub num_rounds: u8,

‎vortex-tensor/src/scalar_fns/sorf_transform/rotation.rs‎

Lines changed: 99 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,19 @@
1212
//! applies a random sign diagonal `D_i` and then the Hadamard matrix `H`, giving O(d log d) cost
1313
//! per matrix-vector product instead of the O(d^2) cost of a dense orthogonal matrix.
1414
//!
15+
//! Vortex defines those sign diagonals using a frozen local SplitMix64 stream rather than an
16+
//! external RNG crate. The contract is:
17+
//!
18+
//! - state is a single `u64` seed,
19+
//! - each `next_u64()` call uses the SplitMix64 reference algorithm with wrapping `u64`
20+
//! arithmetic,
21+
//! - signs are generated in round-major, block-major order,
22+
//! - each generated `u64` contributes 64 signs in least-significant-bit-first order,
23+
//! - bit `1` means `+1` and bit `0` means `-1`.
24+
//!
25+
//! This makes SORF sign generation stable as a Vortex format contract even if external RNG
26+
//! implementations change.
27+
//!
1528
//! [sorf-paper]: https://proceedings.neurips.cc/paper_files/paper/2016/file/53adaf494dc89ef7196d73636eb2451b-Paper.pdf
1629
//!
1730
//! The FWHT exploits the Kronecker product structure of the Hadamard matrix (`H_n = H_2 (x) H_2
@@ -29,12 +42,11 @@
2942
//! floating-point multiply, which avoids FP dependency chains and auto-vectorizes into
3043
//! `vpxor`/`veor`.
3144
32-
use rand::RngExt;
33-
use rand::SeedableRng;
34-
use rand::rngs::StdRng;
3545
use vortex_error::VortexResult;
3646
use vortex_error::vortex_ensure;
3747

48+
use super::splitmix64::SplitMix64;
49+
3850
/// IEEE 754 sign bit mask for f32.
3951
const F32_SIGN_BIT: u32 = 0x8000_0000;
4052

@@ -58,17 +70,17 @@ pub struct SorfMatrix {
5870
}
5971

6072
impl SorfMatrix {
61-
/// Create a new structured Walsh-Hadamard-based orthogonal transform from a deterministic seed.
73+
/// Create a new structured Walsh-Hadamard-based orthogonal transform from a deterministic
74+
/// seed.
75+
///
76+
/// The seed is expanded using Vortex's frozen local SplitMix64 stream. Signs are generated in
77+
/// round-major, block-major order, with each `u64` contributing 64 sign bits in
78+
/// least-significant-bit-first order.
6279
pub fn try_new(seed: u64, dimension: usize, num_rounds: usize) -> VortexResult<Self> {
6380
vortex_ensure!(num_rounds >= 1, "num_rounds must be >= 1, got {num_rounds}");
6481

6582
let padded_dim = dimension.next_power_of_two();
66-
let mut rng = StdRng::seed_from_u64(seed);
67-
68-
let mut sign_masks = Vec::with_capacity(num_rounds * padded_dim);
69-
for _ in 0..num_rounds {
70-
sign_masks.extend(gen_random_sign_masks(&mut rng, padded_dim));
71-
}
83+
let sign_masks = gen_sign_masks_from_seed(seed, padded_dim, num_rounds);
7284

7385
// Compute in f64 for precision, then store as f32 since the WHT operates on f32 buffers.
7486
// The result is always in (0, 1] for any valid padded_dim >= 2 and num_rounds >= 1, so
@@ -221,17 +233,34 @@ impl SorfMatrix {
221233
}
222234
}
223235

224-
/// Generate a vector of random XOR sign masks.
225-
fn gen_random_sign_masks(rng: &mut StdRng, len: usize) -> Vec<u32> {
226-
(0..len)
227-
.map(|_| {
228-
if rng.random::<bool>() {
229-
0u32 // +1: no-op
230-
} else {
231-
F32_SIGN_BIT // -1: flip sign bit
232-
}
233-
})
234-
.collect()
236+
/// Generate XOR sign masks from the frozen local SplitMix64 stream.
237+
///
238+
/// Signs are produced in round-major, block-major order. For each block we call
239+
/// [`SplitMix64::next_u64`] exactly once and unpack its bits from least significant to most
240+
/// significant. Bit `1` means positive sign / `0x00000000`; bit `0` means negative sign /
241+
/// [`F32_SIGN_BIT`].
242+
fn gen_sign_masks_from_seed(seed: u64, padded_dim: usize, num_rounds: usize) -> Vec<u32> {
243+
let mut rng = SplitMix64::new(seed);
244+
let mut sign_masks = Vec::with_capacity(num_rounds * padded_dim);
245+
246+
for _round in 0..num_rounds {
247+
for base_idx in (0..padded_dim).step_by(64) {
248+
let word = rng.next_u64();
249+
let bits_in_block = (padded_dim - base_idx).min(64);
250+
sign_masks.extend((0..bits_in_block).map(|bit_idx| sign_mask_from_word(word, bit_idx)));
251+
}
252+
}
253+
254+
sign_masks
255+
}
256+
257+
/// Convert one bit from a SplitMix64 output word into an XOR sign mask.
258+
fn sign_mask_from_word(word: u64, bit_idx: usize) -> u32 {
259+
if ((word >> bit_idx) & 1) != 0 {
260+
0u32
261+
} else {
262+
F32_SIGN_BIT
263+
}
235264
}
236265

237266
/// Apply sign masks via XOR on the IEEE 754 sign bit.
@@ -290,6 +319,13 @@ mod tests {
290319
use vortex_error::VortexResult;
291320

292321
use super::*;
322+
use crate::scalar_fns::sorf_transform::splitmix64::SplitMix64;
323+
324+
fn unpack_sign_bits(word: u64, count: usize) -> Vec<u8> {
325+
(0..count)
326+
.map(|bit_idx| u8::from(((word >> bit_idx) & 1) != 0))
327+
.collect()
328+
}
293329

294330
#[test]
295331
fn deterministic_from_seed() -> VortexResult<()> {
@@ -311,6 +347,48 @@ mod tests {
311347
Ok(())
312348
}
313349

350+
#[test]
351+
fn export_inverse_signs_matches_golden_words() -> VortexResult<()> {
352+
let rot = SorfMatrix::try_new(42, 64, 2)?;
353+
let actual = rot.export_inverse_signs_u8();
354+
let mut rng = SplitMix64::new(42);
355+
let round0_word = rng.next_u64();
356+
let round1_word = rng.next_u64();
357+
358+
let mut expected = Vec::with_capacity(128);
359+
expected.extend(unpack_sign_bits(round1_word, 64));
360+
expected.extend(unpack_sign_bits(round0_word, 64));
361+
362+
assert_eq!(actual, expected);
363+
Ok(())
364+
}
365+
366+
#[test]
367+
fn one_word_generates_64_signs_lsb_first() {
368+
let masks = gen_sign_masks_from_seed(42, 64, 1);
369+
assert_eq!(masks.len(), 64);
370+
371+
let mut rng = SplitMix64::new(42);
372+
let word = rng.next_u64();
373+
let expected: Vec<_> = (0..64)
374+
.map(|bit_idx| sign_mask_from_word(word, bit_idx))
375+
.collect();
376+
assert_eq!(masks, expected);
377+
}
378+
379+
#[test]
380+
fn tail_block_uses_only_required_bits() {
381+
let masks = gen_sign_masks_from_seed(42, 32, 1);
382+
assert_eq!(masks.len(), 32);
383+
384+
let mut rng = SplitMix64::new(42);
385+
let word = rng.next_u64();
386+
let expected: Vec<_> = (0..32)
387+
.map(|bit_idx| sign_mask_from_word(word, bit_idx))
388+
.collect();
389+
assert_eq!(masks, expected);
390+
}
391+
314392
/// Verify roundtrip is exact to f32 precision across many dimensions and round counts,
315393
/// including non-power-of-two dimensions that require padding.
316394
#[rstest]
Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,73 @@
1+
// SPDX-License-Identifier: Apache-2.0
2+
// SPDX-FileCopyrightText: Copyright the Vortex contributors
3+
4+
//! Frozen local SplitMix64 stream used to define SORF sign diagonals.
5+
//!
6+
//! This is a direct translation of the `splitmix64.c` reference implementation. The state is a
7+
//! single `u64`, and `next_u64()` first adds [`SPLITMIX64_INCREMENT`] with wrapping arithmetic,
8+
//! then applies the two reference mixing steps and final xor-shift.
9+
10+
/// SplitMix64 additive constant from the reference implementation.
11+
const SPLITMIX64_INCREMENT: u64 = 0x9E37_79B9_7F4A_7C15;
12+
13+
/// First SplitMix64 mixing multiplier from the reference implementation.
14+
const SPLITMIX64_MUL1: u64 = 0xBF58_476D_1CE4_E5B9;
15+
16+
/// Second SplitMix64 mixing multiplier from the reference implementation.
17+
const SPLITMIX64_MUL2: u64 = 0x94D0_49BB_1331_11EB;
18+
19+
/// Frozen local SplitMix64 stream used to define SORF sign diagonals.
20+
pub(crate) struct SplitMix64 {
21+
state: u64,
22+
}
23+
24+
impl SplitMix64 {
25+
pub(crate) fn new(seed: u64) -> Self {
26+
Self { state: seed }
27+
}
28+
29+
pub(crate) fn next_u64(&mut self) -> u64 {
30+
self.state = self.state.wrapping_add(SPLITMIX64_INCREMENT);
31+
let mut z = self.state;
32+
z = (z ^ (z >> 30)).wrapping_mul(SPLITMIX64_MUL1);
33+
z = (z ^ (z >> 27)).wrapping_mul(SPLITMIX64_MUL2);
34+
z ^ (z >> 31)
35+
}
36+
}
37+
38+
#[cfg(test)]
39+
mod tests {
40+
use super::SplitMix64;
41+
42+
const SPLITMIX64_SEED0_GOLDEN: [u64; 4] = [
43+
0xE220_A839_7B1D_CDAF,
44+
0x6E78_9E6A_A1B9_65F4,
45+
0x06C4_5D18_8009_454F,
46+
0xF88B_B8A8_724C_81EC,
47+
];
48+
49+
const SPLITMIX64_SEED42_GOLDEN: [u64; 4] = [
50+
0xBDD7_3226_2FEB_6E95,
51+
0x28EF_E333_B266_F103,
52+
0x4752_6757_130F_9F52,
53+
0x581C_E1FF_0E4A_E394,
54+
];
55+
56+
#[test]
57+
fn splitmix64_seed0_matches_golden_outputs() {
58+
let mut rng = SplitMix64::new(0);
59+
let actual: Vec<_> = (0..SPLITMIX64_SEED0_GOLDEN.len())
60+
.map(|_| rng.next_u64())
61+
.collect();
62+
assert_eq!(actual, SPLITMIX64_SEED0_GOLDEN);
63+
}
64+
65+
#[test]
66+
fn splitmix64_seed42_matches_golden_outputs() {
67+
let mut rng = SplitMix64::new(42);
68+
let actual: Vec<_> = (0..SPLITMIX64_SEED42_GOLDEN.len())
69+
.map(|_| rng.next_u64())
70+
.collect();
71+
assert_eq!(actual, SPLITMIX64_SEED42_GOLDEN);
72+
}
73+
}

0 commit comments

Comments
 (0)