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
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 ;
3545use vortex_error:: VortexResult ;
3646use vortex_error:: vortex_ensure;
3747
48+ use super :: splitmix64:: SplitMix64 ;
49+
3850/// IEEE 754 sign bit mask for f32.
3951const F32_SIGN_BIT : u32 = 0x8000_0000 ;
4052
@@ -58,17 +70,17 @@ pub struct SorfMatrix {
5870}
5971
6072impl 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]
0 commit comments