TurboQuant again! (#7829)

## Summary

Tracking issue: #7830

Moves TurboQuant out of `vortex-tensor` into a new `vortex-turboquant`
crate.

The first commit was mostly copying and pasting a bunch of code, as well
as adding the `unpack` method to replace canonicalization. The second
commit was cleaning up everything holistically.

A lot of the code in `vortex-tensor` was reviewed pretty lightly because
we knew that it was unstable, but now that we are more certain about the
implementation (not necessarily about the exact design, but the actual
implementation of the TQ algorithms), I think it is worth reviewing
everything as a whole.

## Testing

These tests were mostly there before, but now there are more!

---------

Signed-off-by: Connor Tsui <connor.tsui20@gmail.com>

Author: connortsui20

Cargo.lock

@@ -13,6 +13,7 @@ members = [
     "vortex-proto",
     "vortex-array",
     "vortex-tensor",
+    "vortex-turboquant",
     "vortex-compressor",
     "vortex-btrblocks",
     "vortex-layout",
@@ -296,6 +297,7 @@ vortex-sequence = { version = "0.1.0", path = "encodings/sequence", default-feat
 vortex-session = { version = "0.1.0", path = "./vortex-session", default-features = false }
 vortex-sparse = { version = "0.1.0", path = "./encodings/sparse", default-features = false }
 vortex-tensor = { version = "0.1.0", path = "./vortex-tensor", default-features = false }
+vortex-turboquant = { version = "0.1.0", path = "./vortex-turboquant", default-features = false }
 vortex-utils = { version = "0.1.0", path = "./vortex-utils", default-features = false }
 vortex-zigzag = { version = "0.1.0", path = "./encodings/zigzag", default-features = false }
 vortex-zstd = { version = "0.1.0", path = "./encodings/zstd", default-features = false }

Cargo.toml

@@ -528,6 +528,10 @@ pub fn vortex_tensor::vector::AnyVector::try_match<'a>(&'a vortex_array::dtype::
 
 pub struct vortex_tensor::vector::Vector
 
+impl vortex_tensor::vector::Vector
+
+pub fn vortex_tensor::vector::Vector::try_new_vector_array(vortex_array::array::erased::ArrayRef) -> vortex_error::VortexResult<vortex_array::array::erased::ArrayRef>
+
 impl core::clone::Clone for vortex_tensor::vector::Vector
 
 pub fn vortex_tensor::vector::Vector::clone(&self) -> vortex_tensor::vector::Vector

vortex-tensor/public-api.lock

@@ -205,8 +205,8 @@ fn turboquant_quantize_core(
     let dimension = fsl.list_size() as usize;
     let num_rows = fsl.len();
 
-    let rotation = SorfMatrix::try_new(seed, dimension, num_rounds as usize)?;
-    let padded_dim = rotation.padded_dim();
+    let padded_dim = dimension.next_power_of_two();
+    let rotation = SorfMatrix::try_new_padded(padded_dim, num_rounds as usize, seed)?;
     let padded_dim_u32 =
         u32::try_from(padded_dim).vortex_expect("padded_dim stays representable as u32");
 

vortex-tensor/src/encodings/turboquant/compress.rs

@@ -259,8 +259,8 @@ fn sorf_transform_roundtrip_isolation() -> VortexResult<()> {
     }
 
     // Forward transform + quantize (mimicking what turboquant_quantize_core does).
-    let rotation = SorfMatrix::try_new(seed, dim, num_rounds as usize)?;
-    let padded_dim = rotation.padded_dim();
+    let padded_dim = dim.next_power_of_two();
+    let rotation = SorfMatrix::try_new_padded(padded_dim, num_rounds as usize, seed)?;
     let centroids = compute_or_get_centroids(padded_dim as u32, 8)?;
     let boundaries = compute_centroid_boundaries(&centroids);
 

vortex-tensor/src/encodings/turboquant/tests/structural.rs

@@ -366,7 +366,7 @@ impl InnerProduct {
         let mut padded_query = vec![0.0f32; padded_dim];
         padded_query[..dim].copy_from_slice(flat.as_slice::<f32>());
 
-        let rotation = SorfMatrix::try_new(seed, dim, num_rounds)?;
+        let rotation = SorfMatrix::try_new_padded(padded_dim, num_rounds, seed)?;
         let mut rotated_query = vec![0.0f32; padded_dim];
         rotation.rotate(&padded_query, &mut rotated_query);
 
@@ -930,7 +930,7 @@ mod tests {
             seed: u64,
             num_rounds: u8,
         ) -> VortexResult<Vec<f32>> {
-            let rotation = SorfMatrix::try_new(seed, dim, num_rounds as usize)?;
+            let rotation = SorfMatrix::try_new_padded(padded_dim, num_rounds as usize, seed)?;
             let mut padded = vec![0.0f32; padded_dim];
             let mut rotated = vec![0.0f32; padded_dim];
             let mut out = Vec::with_capacity(num_rows * dim);

vortex-tensor/src/scalar_fns/inner_product.rs

@@ -77,9 +77,25 @@ impl SorfMatrix {
     /// round-major, block-major order, with each `u64` contributing 64 sign bits in
     /// least-significant-bit-first order.
     pub fn try_new(seed: u64, dimensions: usize, num_rounds: usize) -> VortexResult<Self> {
+        Self::try_new_padded(dimensions.next_power_of_two(), num_rounds, seed)
+    }
+
+    /// Create a new structured Walsh-Hadamard-based orthogonal transform for a padded dimension.
+    ///
+    /// `padded_dimensions` must already be a power of two. Callers that start from an unpadded
+    /// logical dimension should call [`Self::try_new`] instead.
+    pub(crate) fn try_new_padded(
+        padded_dimensions: usize,
+        num_rounds: usize,
+        seed: u64,
+    ) -> VortexResult<Self> {
         vortex_ensure!(num_rounds >= 1, "num_rounds must be >= 1, got {num_rounds}");
+        vortex_ensure!(
+            padded_dimensions.is_power_of_two(),
+            "padded_dimensions must be a power of two, got {padded_dimensions}"
+        );
 
-        let padded_dim = dimensions.next_power_of_two();
+        let padded_dim = padded_dimensions;
         let sign_masks = gen_sign_masks_from_seed(seed, padded_dim, num_rounds);
 
         // Compute in f64 for precision, then store as f32 since the WHT operates on f32 buffers.
@@ -132,8 +148,7 @@ impl SorfMatrix {
     /// Apply the forward structured transform: `norm · H · D_k · ... · H · D₁ · x`.
     fn apply_srht(&self, buf: &mut [f32]) {
         for round in 0..self.num_rounds {
-            let offset = round * self.padded_dim;
-            apply_signs_xor(buf, &self.sign_masks[offset..offset + self.padded_dim]);
+            self.apply_signs_xor(buf, round);
             walsh_hadamard_transform(buf);
         }
 
@@ -148,14 +163,24 @@ impl SorfMatrix {
     fn apply_inverse_srht(&self, buf: &mut [f32]) {
         for round in (0..self.num_rounds).rev() {
             walsh_hadamard_transform(buf);
-            let offset = round * self.padded_dim;
-            apply_signs_xor(buf, &self.sign_masks[offset..offset + self.padded_dim]);
+            self.apply_signs_xor(buf, round);
         }
 
         let norm = self.norm_factor;
         buf.iter_mut().for_each(|val| *val *= norm);
     }
 
+    /// Apply one round's sign masks via XOR on the IEEE 754 sign bit.
+    ///
+    /// This is branchless and auto-vectorizes into `vpxor` (x86) / `veor` (ARM). Equivalent to
+    /// multiplying each element by +/-1.0, but avoids FP dependency chains.
+    fn apply_signs_xor(&self, buf: &mut [f32], round: usize) {
+        let masks = &self.sign_masks[round * self.padded_dim..][..self.padded_dim];
+        for (val, &mask) in buf.iter_mut().zip(masks.iter()) {
+            *val = f32::from_bits(val.to_bits() ^ mask);
+        }
+    }
+
     /// Export the sign vectors as a flat `Vec<u8>` of 0/1 values in inverse application order
     /// `[D_k | ... | D₁]`.
     ///
@@ -263,16 +288,6 @@ fn sign_mask_from_word(word: u64, bit_idx: usize) -> u32 {
     }
 }
 
-/// Apply sign masks via XOR on the IEEE 754 sign bit.
-///
-/// This is branchless and auto-vectorizes into `vpxor` (x86) / `veor` (ARM). Equivalent to
-/// multiplying each element by +/-1.0, but avoids FP dependency chains.
-fn apply_signs_xor(buf: &mut [f32], masks: &[u32]) {
-    for (val, &mask) in buf.iter_mut().zip(masks.iter()) {
-        *val = f32::from_bits(val.to_bits() ^ mask);
-    }
-}
-
 /// In-place Fast Walsh-Hadamard Transform (FWHT), unnormalized and iterative.
 ///
 /// Input length must be a power of 2. Runs in O(n log n) via `log2(n)` stages of `n / 2`
@@ -327,14 +342,24 @@ mod tests {
             .collect()
     }
 
+    fn dim_to_usize(dim: u32) -> usize {
+        usize::try_from(dim).unwrap()
+    }
+
+    fn rounds_to_usize(num_rounds: u8) -> usize {
+        usize::from(num_rounds)
+    }
+
     #[test]
     fn deterministic_from_seed() -> VortexResult<()> {
-        let r1 = SorfMatrix::try_new(42, 64, 3)?;
-        let r2 = SorfMatrix::try_new(42, 64, 3)?;
+        let dim = dim_to_usize(64u32);
+        let num_rounds = rounds_to_usize(3u8);
+        let r1 = SorfMatrix::try_new(42u64, dim, num_rounds)?;
+        let r2 = SorfMatrix::try_new(42u64, dim, num_rounds)?;
         let pd = r1.padded_dim();
 
         let mut input = vec![0.0f32; pd];
-        for i in 0..64 {
+        for i in 0..dim {
             input[i] = i as f32;
         }
         let mut out1 = vec![0.0f32; pd];
@@ -349,41 +374,58 @@ mod tests {
 
     #[test]
     fn export_inverse_signs_matches_golden_words() -> VortexResult<()> {
-        let rot = SorfMatrix::try_new(42, 64, 2)?;
+        let dim = dim_to_usize(64u32);
+        let num_rounds = rounds_to_usize(2u8);
+        let seed = 42u64;
+        let rot = SorfMatrix::try_new(seed, dim, num_rounds)?;
+        let padded_dim = rot.padded_dim();
         let actual = rot.export_inverse_signs_u8();
-        let mut rng = SplitMix64::new(42);
+        let mut rng = SplitMix64::new(seed);
         let round0_word = rng.next_u64();
         let round1_word = rng.next_u64();
 
-        let mut expected = Vec::with_capacity(128);
-        expected.extend(unpack_sign_bits(round1_word, 64));
-        expected.extend(unpack_sign_bits(round0_word, 64));
+        let mut expected = Vec::with_capacity(num_rounds * padded_dim);
+        expected.extend(unpack_sign_bits(round1_word, padded_dim));
+        expected.extend(unpack_sign_bits(round0_word, padded_dim));
 
         assert_eq!(actual, expected);
         Ok(())
     }
 
     #[test]
     fn one_word_generates_64_signs_lsb_first() {
-        let masks = gen_sign_masks_from_seed(42, 64, 1);
-        assert_eq!(masks.len(), 64);
+        let seed = 42u64;
+        let padded_dim = dim_to_usize(64u32);
+        let num_rounds = rounds_to_usize(1u8);
+        let masks = gen_sign_masks_from_seed(seed, padded_dim, num_rounds);
+        assert_eq!(masks.len(), padded_dim);
 
-        let mut rng = SplitMix64::new(42);
+        let mut rng = SplitMix64::new(seed);
         let word = rng.next_u64();
-        let expected: Vec<_> = (0..64)
+        let expected: Vec<_> = (0..padded_dim)
             .map(|bit_idx| sign_mask_from_word(word, bit_idx))
             .collect();
         assert_eq!(masks, expected);
     }
 
+    #[test]
+    fn accepts_non_power_of_two_dimensions() -> VortexResult<()> {
+        let rot = SorfMatrix::try_new(42u64, dim_to_usize(100u32), rounds_to_usize(3u8))?;
+        assert_eq!(rot.padded_dim(), 128);
+        Ok(())
+    }
+
     #[test]
     fn tail_block_uses_only_required_bits() {
-        let masks = gen_sign_masks_from_seed(42, 32, 1);
-        assert_eq!(masks.len(), 32);
+        let seed = 42u64;
+        let padded_dim = dim_to_usize(32u32);
+        let num_rounds = rounds_to_usize(1u8);
+        let masks = gen_sign_masks_from_seed(seed, padded_dim, num_rounds);
+        assert_eq!(masks.len(), padded_dim);
 
-        let mut rng = SplitMix64::new(42);
+        let mut rng = SplitMix64::new(seed);
         let word = rng.next_u64();
-        let expected: Vec<_> = (0..32)
+        let expected: Vec<_> = (0..padded_dim)
             .map(|bit_idx| sign_mask_from_word(word, bit_idx))
             .collect();
         assert_eq!(masks, expected);
@@ -392,19 +434,21 @@ mod tests {
     /// Verify roundtrip is exact to f32 precision across many dimensions and round counts,
     /// including non-power-of-two dimensions that require padding.
     #[rstest]
-    #[case(32, 3)]
-    #[case(64, 3)]
-    #[case(100, 3)]
-    #[case(128, 1)]
-    #[case(128, 2)]
-    #[case(128, 3)]
-    #[case(128, 5)]
-    #[case(256, 3)]
-    #[case(512, 3)]
-    #[case(768, 3)]
-    #[case(1024, 3)]
-    fn roundtrip_exact(#[case] dim: usize, #[case] num_rounds: usize) -> VortexResult<()> {
-        let rot = SorfMatrix::try_new(42, dim, num_rounds)?;
+    #[case(32u32, 3u8)]
+    #[case(64u32, 3u8)]
+    #[case(100u32, 3u8)]
+    #[case(128u32, 1u8)]
+    #[case(128u32, 2u8)]
+    #[case(128u32, 3u8)]
+    #[case(128u32, 5u8)]
+    #[case(256u32, 3u8)]
+    #[case(512u32, 3u8)]
+    #[case(768u32, 3u8)]
+    #[case(1024u32, 3u8)]
+    fn roundtrip_exact(#[case] dim: u32, #[case] num_rounds: u8) -> VortexResult<()> {
+        let dim = dim_to_usize(dim);
+        let num_rounds = rounds_to_usize(num_rounds);
+        let rot = SorfMatrix::try_new(42u64, dim, num_rounds)?;
         let padded_dim = rot.padded_dim();
 
         let mut input = vec![0.0f32; padded_dim];
@@ -435,12 +479,14 @@ mod tests {
 
     /// Verify norm preservation across dimensions and round counts.
     #[rstest]
-    #[case(128, 1)]
-    #[case(128, 3)]
-    #[case(128, 5)]
-    #[case(768, 3)]
-    fn preserves_norm(#[case] dim: usize, #[case] num_rounds: usize) -> VortexResult<()> {
-        let rot = SorfMatrix::try_new(7, dim, num_rounds)?;
+    #[case(128u32, 1u8)]
+    #[case(128u32, 3u8)]
+    #[case(128u32, 5u8)]
+    #[case(768u32, 3u8)]
+    fn preserves_norm(#[case] dim: u32, #[case] num_rounds: u8) -> VortexResult<()> {
+        let dim = dim_to_usize(dim);
+        let num_rounds = rounds_to_usize(num_rounds);
+        let rot = SorfMatrix::try_new(42u64, dim, num_rounds)?;
         let padded_dim = rot.padded_dim();
 
         let mut input = vec![0.0f32; padded_dim];
@@ -465,16 +511,15 @@ mod tests {
 
     /// Verify that export -> [`from_u8_slice`] produces identical transform output.
     #[rstest]
-    #[case(64, 3)]
-    #[case(128, 1)]
-    #[case(128, 3)]
-    #[case(128, 5)]
-    #[case(768, 3)]
-    fn sign_export_import_roundtrip(
-        #[case] dim: usize,
-        #[case] num_rounds: usize,
-    ) -> VortexResult<()> {
-        let rot = SorfMatrix::try_new(42, dim, num_rounds)?;
+    #[case(64u32, 3u8)]
+    #[case(128u32, 1u8)]
+    #[case(128u32, 3u8)]
+    #[case(128u32, 5u8)]
+    #[case(768u32, 3u8)]
+    fn sign_export_import_roundtrip(#[case] dim: u32, #[case] num_rounds: u8) -> VortexResult<()> {
+        let dim = dim_to_usize(dim);
+        let num_rounds = rounds_to_usize(num_rounds);
+        let rot = SorfMatrix::try_new(42u64, dim, num_rounds)?;
         let padded_dim = rot.padded_dim();
 
         let signs_u8 = rot.export_inverse_signs_u8();

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

@@ -65,8 +65,8 @@ fn forward_rotate_and_quantize(
         }
     }
 
-    let rotation = SorfMatrix::try_new(seed, dim, num_rounds)?;
-    let padded_dim = rotation.padded_dim();
+    let padded_dim = dim.next_power_of_two();
+    let rotation = SorfMatrix::try_new_padded(padded_dim, num_rounds, seed)?;
     let centroids = compute_or_get_centroids(padded_dim as u32, bit_width)?;
     let boundaries = compute_centroid_boundaries(&centroids);
 

vortex-tensor/src/scalar_fns/sorf_transform/tests.rs

@@ -164,7 +164,8 @@ impl ScalarFnVTable for SorfTransform {
         let f32_elements = elements_prim.into_buffer::<f32>();
 
         // Reconstruct the orthogonal transform matrix from the seed.
-        let rotation = SorfMatrix::try_new(options.seed, dim, options.num_rounds as usize)?;
+        let rotation =
+            SorfMatrix::try_new_padded(padded_dim, options.num_rounds as usize, options.seed)?;
 
         // Inverse transform each row, truncate to original dimension, cast to target type.
         match_each_float_ptype!(options.element_ptype, |T| {

vortex-tensor/src/scalar_fns/sorf_transform/vtable.rs

@@ -49,7 +49,7 @@ impl Vector {
     /// # Errors
     ///
     /// Returns an error if the [`Vector`] extension dtype rejects the storage array.
-    pub(crate) fn try_new_vector_array(storage: ArrayRef) -> VortexResult<ArrayRef> {
+    pub fn try_new_vector_array(storage: ArrayRef) -> VortexResult<ArrayRef> {
         ExtensionArray::try_new_from_vtable(Vector, EmptyMetadata, storage)
             .map(|ext| ext.into_array())
     }