diff --git a/diskann-quantization/src/algorithms/transforms/double_hadamard.rs b/diskann-quantization/src/algorithms/transforms/double_hadamard.rs index bb84e10880..50002b33b6 100644 --- a/diskann-quantization/src/algorithms/transforms/double_hadamard.rs +++ b/diskann-quantization/src/algorithms/transforms/double_hadamard.rs @@ -243,9 +243,23 @@ where ) -> Result<(), TransformFailed> { check_dims(dst, src, self.input_dim(), self.output_dim())?; - // Copy and flip signs + if let Some(indices) = self.subsample.as_ref() { + let mut tmp = Poly::broadcast(0.0f32, self.intermediate_dim(), allocator)?; + self.transform_full(&mut tmp, src); + let rescale = ((tmp.len() as f32) / (indices.len() as f32)).sqrt(); + dst.iter_mut() + .zip(indices.iter()) + .for_each(|(d, s)| *d = tmp[*s as usize] * rescale); + } else { + // The destination already has room for the full transform. + dst[src.len()..].fill(0.0); + self.transform_full(dst, src); + } + Ok(()) + } + + fn transform_full(&self, tmp: &mut [f32], src: &[f32]) { let intermediate_dim = self.intermediate_dim(); - let mut tmp = Poly::broadcast(0.0f32, intermediate_dim, allocator)?; std::iter::zip(tmp.iter_mut(), src.iter()) .zip(self.signs0.iter()) @@ -269,21 +283,6 @@ where #[expect(clippy::unwrap_used)] hadamard_transform(&mut tmp[intermediate_dim - split..]).unwrap(); - - match self.subsample.as_ref() { - None => { - dst.copy_from_slice(&tmp); - } - Some(indices) => { - let rescale = ((tmp.len() as f32) / (indices.len() as f32)).sqrt(); - debug_assert_eq!(dst.len(), indices.len()); - dst.iter_mut() - .zip(indices.iter()) - .for_each(|(d, s)| *d = tmp[*s as usize] * rescale); - } - } - - Ok(()) } } @@ -392,7 +391,7 @@ mod tests { use crate::{ algorithms::transforms::{Transform, TransformKind, test_utils}, alloc::GlobalAllocator, - test_util::Check, + test_util::{Check, LimitedAllocator}, }; test_utils::delegate_transformer!(DoubleHadamard); @@ -485,7 +484,7 @@ mod tests { &mut checker, &mut rng, ctx, - ) + ); } // Abstract Transformer @@ -508,6 +507,37 @@ mod tests { ) } } + + let mut allocation_rng = rng.clone(); + let transformer = + DoubleHadamard::new(input_nz, target, &mut allocation_rng, GlobalAllocator) + .unwrap(); + let mut src = vec![1.0; input]; + let mut dst = vec![f32::NAN; output + 2]; + for value in [1.0, 0.0, -0.0] { + src.fill(value); + if !preserves_norms { + let allocator = LimitedAllocator::new(0); + assert_eq!( + transformer.transform_into( + &mut dst[1..=output], + &src, + ScopedAllocator::new(&allocator), + ), + Err(TransformFailed::AllocatorError(AllocatorError)) + ); + } + let allocator = LimitedAllocator::new(usize::from(!preserves_norms)); + transformer + .transform_into(&mut dst[1..=output], &src, ScopedAllocator::new(&allocator)) + .unwrap(); + assert_eq!(allocator.remaining(), 0); + assert!(dst[1..=output].iter().all(|value| value.is_finite())); + if value == 0.0 { + assert!(dst[1..=output].iter().all(|value| *value == 0.0)); + } + assert!(dst[0].is_nan() && dst[output + 1].is_nan()); + } } } diff --git a/diskann-quantization/src/spherical/quantizer.rs b/diskann-quantization/src/spherical/quantizer.rs index dcb1072bce..2a21be38f4 100644 --- a/diskann-quantization/src/spherical/quantizer.rs +++ b/diskann-quantization/src/spherical/quantizer.rs @@ -940,9 +940,9 @@ where // Informally, "what's the smallest scaling factor so `v[position]` gets rounded to // the next value. // -// 2. Imposes a total ordering on `f32` values so it can be used in a `BinaryHeap`. -// Additionally, ordering is reverse so that `BinaryHeap` models a min-heap instead -// of a max-heap. +// 2. Imposes a total ordering on nonnegative, non-NaN critical values using unsigned +// bit order, which matches numeric order for these values. Additionally, ordering +// is reversed so that `BinaryHeap` models a min-heap instead of a max-heap. #[derive(Debug, Clone, Copy)] struct Pair { value: f32, @@ -951,7 +951,7 @@ struct Pair { impl PartialEq for Pair { fn eq(&self, other: &Self) -> bool { - self.value.eq(&other.value) + self.value.to_bits().eq(&other.value.to_bits()) } } @@ -963,10 +963,7 @@ impl PartialOrd for Pair { } impl Ord for Pair { fn cmp(&self, other: &Self) -> std::cmp::Ordering { - other - .value - .partial_cmp(&self.value) - .unwrap_or(std::cmp::Ordering::Equal) + other.value.to_bits().cmp(&self.value.to_bits()) } } @@ -992,6 +989,10 @@ impl Ord for Pair { /// /// Refer to algorithm 1 in . /// +/// # Errors +/// +/// Returns an error if `v` contains NaN or scratch allocation fails. +/// /// # Panics /// /// Panics is `v.is_empty()`. @@ -1005,12 +1006,15 @@ fn maximize_cosine_similarity( v: &[f32], num_bits: NonZeroUsize, allocator: ScopedAllocator<'_>, -) -> Result { +) -> Result { // Initially, the lattice element has the value `0.5` for all dimensions. // This means the initial inner product between `v` and the rounded term is simply // `0.5 * sum(abs.(v))`. The absolute value is used because the latice element is // always in the direction of the components in `v`. let mut current_ip = 0.5 * v.iter().map(|i| i.abs() as f64).sum::(); + if current_ip.is_nan() { + return Err(CompressionError::InputContainsNaN); + } let mut current_square_norm = 0.25 * (v.len() as f64); // Book keeping for the current value of the rounded vector. @@ -1254,7 +1258,40 @@ mod tests { test_util, }; - // Test cosine-similarity maximizer + #[test] + fn cosine_similarity_maximizer_errors() { + let num_bits = NonZeroUsize::new(4).unwrap(); + for nan in [f32::NAN, f32::from_bits(0xffc0_0001)] { + for position in 0..4 { + let mut target = [1.0; 4]; + target[position] = nan; + assert_eq!( + maximize_cosine_similarity( + &target, + num_bits, + ScopedAllocator::new(&test_util::AlwaysFails), + ), + Err(CompressionError::InputContainsNaN), + ); + } + } + for budget in 0..=2 { + let result = maximize_cosine_similarity( + &[1.0, 0.5], + num_bits, + ScopedAllocator::new(&test_util::LimitedAllocator::new(budget)), + ); + if budget < 2 { + assert_eq!( + result, + Err(CompressionError::AllocatorError(AllocatorError)) + ); + } else { + assert!(result.is_ok()); + } + } + } + #[test] fn test_cosine_similarity_maximizer() { let mut rng = StdRng::seed_from_u64(0x070d9ff8cf5e0f8c); @@ -1305,7 +1342,11 @@ mod tests { // We expect to either get the best value found via exhaustive search, or some // scalar multiple of it (since that will have the same cosine similarity). - let passed = if best == clamped { + let passed = if best == clamped + || (best_similarity == clamped_cosine.into_inner() + && std::iter::zip(target, std::iter::zip(best, clamped)) + .all(|(input, (expected, actual))| input == 0.0 || expected == actual)) + { true } else { let ratio: Vec = std::iter::zip(best, clamped) @@ -1337,6 +1378,20 @@ mod tests { } }; + for target in [ + [1.0, 0.0, -0.0, -1.0], + [0.0, 0.5, -0.5, 0.0], + [1.0, 1.0, -1.0, -1.0], + [ + f32::from_bits(1.0f32.to_bits() - 1), + 1.0, + f32::from_bits(1.0f32.to_bits() + 1), + -1.0, + ], + ] { + run_test(target); + } + // Run targeted tests. let min = -3.5; for i0 in (0..8).step_by(2) { @@ -2423,6 +2478,7 @@ mod tests { { let mut query: Vec = quantizer.shift().to_vec(); let mut d = Data::<1, _>::new_boxed(quantizer.output_dim()); + let mut multibit_data = Data::<4, _>::new_boxed(quantizer.output_dim()); let mut q = Query::<4, BitTranspose, _>::new_boxed(quantizer.output_dim()); for i in 0..query.len() { @@ -2436,6 +2492,12 @@ mod tests { assert_eq!(err.to_string(), "input contains NaN", "failed for {}", v); + let err = quantizer + .compress_into_with(&*query, multibit_data.reborrow_mut(), scoped_global) + .unwrap_err(); + + assert_eq!(err, CompressionError::InputContainsNaN); + let err = quantizer .compress_into_with(&*query, q.reborrow_mut(), scoped_global) .unwrap_err(); diff --git a/diskann-quantization/src/test_util.rs b/diskann-quantization/src/test_util.rs index d7f8f95cbb..b8abaa267e 100644 --- a/diskann-quantization/src/test_util.rs +++ b/diskann-quantization/src/test_util.rs @@ -53,6 +53,10 @@ impl LimitedAllocator { remaining: Arc::new(AtomicUsize::new(allocations)), } } + + pub(crate) fn remaining(&self) -> usize { + self.remaining.load(Ordering::Relaxed) + } } #[cfg(not(miri))]