diff --git a/demos/benches/encode.rs b/demos/benches/encode.rs index 0915217..6f5e1e2 100644 --- a/demos/benches/encode.rs +++ b/demos/benches/encode.rs @@ -5,13 +5,12 @@ use quantization::encoded_vectors_u8::EncodedVectorsU8; use rand::Rng; #[cfg(target_arch = "x86_64")] -use demos::metrics::utils_avx2::dot_avx; - +use demos::metrics::utils_avx2::{dot_avx, l1_avx}; #[cfg(target_arch = "x86_64")] -use demos::metrics::utils_sse::dot_sse; +use demos::metrics::utils_sse::{dot_sse, l1_sse}; -fn encode_bench(c: &mut Criterion) { - let mut group = c.benchmark_group("encode"); +fn encode_dot_bench(c: &mut Criterion) { + let mut group = c.benchmark_group("encode dot"); let vectors_count = 100_000; let vector_dim = 1024; @@ -145,10 +144,145 @@ fn encode_bench(c: &mut Criterion) { }); } +fn encode_l1_bench(c: &mut Criterion) { + let mut group = c.benchmark_group("encode l1"); + + let vectors_count = 100_000; + let vector_dim = 1024; + let mut rng = rand::thread_rng(); + let mut list: Vec = Vec::new(); + for _ in 0..vectors_count { + let vector: Vec = (0..vector_dim).map(|_| rng.gen()).collect(); + list.extend_from_slice(&vector); + } + + let i8_encoded = EncodedVectorsU8::encode( + (0..vectors_count).map(|i| &list[i * vector_dim..(i + 1) * vector_dim]), + Vec::::new(), + &VectorParameters { + dim: vector_dim, + count: vectors_count, + distance_type: DistanceType::L1, + invert: true, + }, + None, + || false, + ) + .unwrap(); + + let query: Vec = (0..vector_dim).map(|_| rng.gen()).collect(); + let encoded_query = i8_encoded.encode_query(&query); + + #[cfg(target_arch = "x86_64")] + group.bench_function("score all u8 avx", |b| { + b.iter(|| { + let mut _s = 0.0; + for i in 0..vectors_count as u32 { + _s = i8_encoded.score_point_avx(&encoded_query, i); + } + }); + }); + + #[cfg(target_arch = "x86_64")] + group.bench_function("score all u8 sse", |b| { + b.iter(|| { + let mut _s = 0.0; + for i in 0..vectors_count as u32 { + _s = i8_encoded.score_point_sse(&encoded_query, i); + } + }); + }); + + #[cfg(target_arch = "aarch64")] + group.bench_function("score all u8 neon", |b| { + b.iter(|| { + let mut _s = 0.0; + for i in 0..vectors_count as u32 { + _s = i8_encoded.score_point_neon(&encoded_query, i); + } + }); + }); + + #[cfg(target_arch = "x86_64")] + group.bench_function("score all avx", |b| { + b.iter(|| unsafe { + let mut _s = 0.0; + for i in 0..vectors_count { + _s = l1_avx(&query, &list[i * vector_dim..(i + 1) * vector_dim]); + } + }); + }); + + #[cfg(target_arch = "x86_64")] + group.bench_function("score all sse", |b| { + b.iter(|| unsafe { + let mut _s = 0.0; + for i in 0..vectors_count { + _s = l1_sse(&query, &list[i * vector_dim..(i + 1) * vector_dim]); + } + }); + }); + + let permutor = Permutor::new(vectors_count as u64); + let permutation: Vec = permutor.map(|i| i as u32).collect(); + + #[cfg(target_arch = "x86_64")] + group.bench_function("score random access u8 avx", |b| { + b.iter(|| { + let mut _s = 0.0; + for &i in &permutation { + _s = i8_encoded.score_point_avx(&encoded_query, i); + } + }); + }); + + #[cfg(target_arch = "x86_64")] + group.bench_function("score random access u8 sse", |b| { + let mut _s = 0.0; + b.iter(|| { + for &i in &permutation { + _s = i8_encoded.score_point_sse(&encoded_query, i); + } + }); + }); + + #[cfg(target_arch = "aarch64")] + group.bench_function("score random access u8 neon", |b| { + let mut _s = 0.0; + b.iter(|| { + for &i in &permutation { + _s = i8_encoded.score_point_neon(&encoded_query, i); + } + }); + }); + + #[cfg(target_arch = "x86_64")] + group.bench_function("score random access avx", |b| { + b.iter(|| unsafe { + let mut _s = 0.0; + for &i in &permutation { + let i = i as usize; + _s = l1_avx(&query, &list[i * vector_dim..(i + 1) * vector_dim]); + } + }); + }); + + #[cfg(target_arch = "x86_64")] + group.bench_function("score random access sse", |b| { + let mut _s = 0.0; + b.iter(|| unsafe { + for &i in &permutation { + let i = i as usize; + _s = l1_sse(&query, &list[i * vector_dim..(i + 1) * vector_dim]); + } + }); + }); +} + criterion_group! { name = benches; config = Criterion::default().sample_size(10); - targets = encode_bench + targets = encode_dot_bench, encode_l1_bench } criterion_main!(benches); diff --git a/demos/src/ann_benchmark.rs b/demos/src/ann_benchmark.rs index 39101d1..d74d6cf 100644 --- a/demos/src/ann_benchmark.rs +++ b/demos/src/ann_benchmark.rs @@ -7,7 +7,6 @@ use quantization::{EncodedVectorsU8, VectorParameters}; #[cfg(target_arch = "x86_64")] use crate::metrics::utils_avx2::dot_avx; - #[cfg(target_arch = "x86_64")] use demos::metrics::utils_sse::dot_sse; diff --git a/demos/src/metrics/utils_avx2.rs b/demos/src/metrics/utils_avx2.rs index 32eae73..d3bb9de 100644 --- a/demos/src/metrics/utils_avx2.rs +++ b/demos/src/metrics/utils_avx2.rs @@ -56,3 +56,49 @@ pub unsafe fn dot_avx(v1: &[f32], v2: &[f32]) -> f32 { } result } + +#[target_feature(enable = "avx2")] +#[target_feature(enable = "fma")] +#[allow(clippy::missing_safety_doc, dead_code)] +pub unsafe fn l1_avx(v1: &[f32], v2: &[f32]) -> f32 { + let mask: __m256 = _mm256_set1_ps(-0.0f32); // 1 << 31 used to clear sign bit to mimic abs + + let n = v1.len(); + let m = n - (n % 32); + let mut ptr1: *const f32 = v1.as_ptr(); + let mut ptr2: *const f32 = v2.as_ptr(); + let mut sum256_1: __m256 = _mm256_setzero_ps(); + let mut sum256_2: __m256 = _mm256_setzero_ps(); + let mut sum256_3: __m256 = _mm256_setzero_ps(); + let mut sum256_4: __m256 = _mm256_setzero_ps(); + let mut i: usize = 0; + while i < m { + let sub256_1: __m256 = _mm256_sub_ps(_mm256_loadu_ps(ptr1), _mm256_loadu_ps(ptr2)); + sum256_1 = _mm256_add_ps(_mm256_andnot_ps(mask, sub256_1), sum256_1); + + let sub256_2: __m256 = + _mm256_sub_ps(_mm256_loadu_ps(ptr1.add(8)), _mm256_loadu_ps(ptr2.add(8))); + sum256_2 = _mm256_add_ps(_mm256_andnot_ps(mask, sub256_2), sum256_2); + + let sub256_3: __m256 = + _mm256_sub_ps(_mm256_loadu_ps(ptr1.add(16)), _mm256_loadu_ps(ptr2.add(16))); + sum256_3 = _mm256_add_ps(_mm256_andnot_ps(mask, sub256_3), sum256_3); + + let sub256_4: __m256 = + _mm256_sub_ps(_mm256_loadu_ps(ptr1.add(24)), _mm256_loadu_ps(ptr2.add(24))); + sum256_4 = _mm256_add_ps(_mm256_andnot_ps(mask, sub256_4), sum256_4); + + ptr1 = ptr1.add(32); + ptr2 = ptr2.add(32); + i += 32; + } + + let mut result = hsum256_ps_avx(sum256_1) + + hsum256_ps_avx(sum256_2) + + hsum256_ps_avx(sum256_3) + + hsum256_ps_avx(sum256_4); + for i in 0..n - m { + result += (*ptr1.add(i) - *ptr2.add(i)).abs(); + } + -result +} diff --git a/demos/src/metrics/utils_sse.rs b/demos/src/metrics/utils_sse.rs index 55589a8..9aa4898 100644 --- a/demos/src/metrics/utils_sse.rs +++ b/demos/src/metrics/utils_sse.rs @@ -53,3 +53,45 @@ pub unsafe fn dot_sse(v1: &[f32], v2: &[f32]) -> f32 { } result } + +#[target_feature(enable = "sse4.1")] +#[allow(clippy::missing_safety_doc, dead_code)] +pub unsafe fn l1_sse(v1: &[f32], v2: &[f32]) -> f32 { + let mask: __m128 = _mm_set1_ps(-0.0f32); // 1 << 31 used to clear sign bit to mimic abs + + let n = v1.len(); + let m = n - (n % 16); + let mut ptr1: *const f32 = v1.as_ptr(); + let mut ptr2: *const f32 = v2.as_ptr(); + let mut sum128_1: __m128 = _mm_setzero_ps(); + let mut sum128_2: __m128 = _mm_setzero_ps(); + let mut sum128_3: __m128 = _mm_setzero_ps(); + let mut sum128_4: __m128 = _mm_setzero_ps(); + let mut i: usize = 0; + while i < m { + let sub128_1 = _mm_sub_ps(_mm_loadu_ps(ptr1), _mm_loadu_ps(ptr2)); + sum128_1 = _mm_add_ps(_mm_andnot_ps(mask, sub128_1), sum128_1); + + let sub128_2 = _mm_sub_ps(_mm_loadu_ps(ptr1.add(4)), _mm_loadu_ps(ptr2.add(4))); + sum128_2 = _mm_add_ps(_mm_andnot_ps(mask, sub128_2), sum128_2); + + let sub128_3 = _mm_sub_ps(_mm_loadu_ps(ptr1.add(8)), _mm_loadu_ps(ptr2.add(8))); + sum128_3 = _mm_add_ps(_mm_andnot_ps(mask, sub128_3), sum128_3); + + let sub128_4 = _mm_sub_ps(_mm_loadu_ps(ptr1.add(12)), _mm_loadu_ps(ptr2.add(12))); + sum128_4 = _mm_add_ps(_mm_andnot_ps(mask, sub128_4), sum128_4); + + ptr1 = ptr1.add(16); + ptr2 = ptr2.add(16); + i += 16; + } + + let mut result = hsum128_ps_sse(sum128_1) + + hsum128_ps_sse(sum128_2) + + hsum128_ps_sse(sum128_3) + + hsum128_ps_sse(sum128_4); + for i in 0..n - m { + result += (*ptr1.add(i) - *ptr2.add(i)).abs(); + } + -result +} diff --git a/quantization/cpp/avx2.c b/quantization/cpp/avx2.c index f401944..87dcc85 100644 --- a/quantization/cpp/avx2.c +++ b/quantization/cpp/avx2.c @@ -1,3 +1,4 @@ +#include #include #include @@ -12,6 +13,15 @@ R = _mm_cvtss_f32(x32); \ } +#define HSUM256_EPI32(X, R) \ + int R = 0; \ + { \ + __m128i x128 = _mm_add_epi32(_mm256_extractf128_si256(X, 1), _mm256_castsi256_si128(X)); \ + __m128i x64 = _mm_add_epi32(x128, _mm_srli_si128(x128, 8)); \ + __m128i x32 = _mm_add_epi32(x64, _mm_srli_si128(x64, 4)); \ + R = _mm_cvtsi128_si32(x32); \ + } + EXPORT float impl_score_dot_avx( const uint8_t* query_ptr, const uint8_t* vector_ptr, @@ -34,6 +44,8 @@ EXPORT float impl_score_dot_avx( mul1 = _mm256_add_epi32(mul1, s_low); mul1 = _mm256_add_epi32(mul1, s_high); } + + // the vector sizes are assumed to be multiples of 16, check if one last 16-element part remaining if (dim % 32 != 0) { __m128i v_short = _mm_loadu_si128((const __m128i*)v_ptr); __m128i q_short = _mm_loadu_si128((const __m128i*)q_ptr); @@ -49,3 +61,62 @@ EXPORT float impl_score_dot_avx( HSUM256_PS(mul_ps, mul_scalar); return mul_scalar; } + +EXPORT float impl_score_l1_avx( + const uint8_t* query_ptr, + const uint8_t* vector_ptr, + uint32_t dim +) { + const __m256i* v_ptr = (const __m256i*)vector_ptr; + const __m256i* q_ptr = (const __m256i*)query_ptr; + + uint32_t m = dim - (dim % 32); + __m256i sum256 = _mm256_setzero_si256(); + + for (uint32_t i = 0; i < m; i += 32) { + __m256i v = _mm256_loadu_si256(v_ptr); + __m256i q = _mm256_loadu_si256(q_ptr); + v_ptr++; + q_ptr++; + + // Compute the difference in both directions and take the maximum for abs + __m256i diff1 = _mm256_subs_epu8(v, q); + __m256i diff2 = _mm256_subs_epu8(q, v); + + __m256i abs_diff = _mm256_max_epu8(diff1, diff2); + + __m256i abs_diff16_lo = _mm256_unpacklo_epi8(abs_diff, _mm256_setzero_si256()); + __m256i abs_diff16_hi = _mm256_unpackhi_epi8(abs_diff, _mm256_setzero_si256()); + + sum256 = _mm256_add_epi16(sum256, abs_diff16_lo); + sum256 = _mm256_add_epi16(sum256, abs_diff16_hi); + } + + // the vector sizes are assumed to be multiples of 16, check if one last 16-element part remaining + if (m < dim) { + __m128i v_short = _mm_loadu_si128((const __m128i * ) v_ptr); + __m128i q_short = _mm_loadu_si128((const __m128i * ) q_ptr); + + __m128i diff1 = _mm_subs_epu8(v_short, q_short); + __m128i diff2 = _mm_subs_epu8(q_short, v_short); + + __m128i abs_diff = _mm_max_epu8(diff1, diff2); + + __m128i abs_diff16_lo_128 = _mm_unpacklo_epi8(abs_diff, _mm_setzero_si128()); + __m128i abs_diff16_hi_128 = _mm_unpackhi_epi8(abs_diff, _mm_setzero_si128()); + + __m256i abs_diff16_lo = _mm256_cvtepu16_epi32(abs_diff16_lo_128); + __m256i abs_diff16_hi = _mm256_cvtepu16_epi32(abs_diff16_hi_128); + + sum256 = _mm256_add_epi16(sum256, abs_diff16_lo); + sum256 = _mm256_add_epi16(sum256, abs_diff16_hi); + } + + __m256i sum_epi32 = _mm256_add_epi32( + _mm256_unpacklo_epi16(sum256, _mm256_setzero_si256()), + _mm256_unpackhi_epi16(sum256, _mm256_setzero_si256())); + + HSUM256_EPI32(sum_epi32, sum); + + return (float) sum; +} diff --git a/quantization/cpp/neon.c b/quantization/cpp/neon.c index e20348b..44d8a04 100644 --- a/quantization/cpp/neon.c +++ b/quantization/cpp/neon.c @@ -1,3 +1,4 @@ +#include #include #include "export_macro.h" @@ -63,3 +64,43 @@ EXPORT uint64_t impl_xor_popcnt_neon( return (uint64_t)vaddvq_u32(result); } + +EXPORT float impl_score_l1_neon( + const uint8_t * query_ptr, + const uint8_t * vector_ptr, + uint32_t dim +) { + const uint8_t* v_ptr = (const uint8_t*)vector_ptr; + const uint8_t* q_ptr = (const uint8_t*)query_ptr; + + uint32_t m = dim - (dim % 16); + uint16x8_t sum16_low = vdupq_n_u16(0); + uint16x8_t sum16_high = vdupq_n_u16(0); + + // the vector sizes are assumed to be multiples of 16, no remaining part here + for (uint32_t i = 0; i < m; i += 16) { + uint8x16_t vec1 = vld1q_u8(v_ptr); + uint8x16_t vec2 = vld1q_u8(q_ptr); + + uint8x16_t abs_diff = vabdq_u8(vec1, vec2); + uint16x8_t abs_diff16_low = vmovl_u8(vget_low_u8(abs_diff)); + uint16x8_t abs_diff16_high = vmovl_u8(vget_high_u8(abs_diff)); + + sum16_low = vaddq_u16(sum16_low, abs_diff16_low); + sum16_high = vaddq_u16(sum16_high, abs_diff16_high); + + v_ptr += 16; + q_ptr += 16; + } + + // Horizontal sum of 16-bit integers + uint32x4_t sum32_low = vpaddlq_u16(sum16_low); + uint32x4_t sum32_high = vpaddlq_u16(sum16_high); + uint32x4_t sum32 = vaddq_u32(sum32_low, sum32_high); + + uint32x2_t sum64_low = vadd_u32(vget_low_u32(sum32), vget_high_u32(sum32)); + uint32x2_t sum64_high = vpadd_u32(sum64_low, sum64_low); + uint32_t sum = vget_lane_u32(sum64_high, 0); + + return (float) sum; +} diff --git a/quantization/cpp/sse.c b/quantization/cpp/sse.c index d055199..f3dbd6f 100644 --- a/quantization/cpp/sse.c +++ b/quantization/cpp/sse.c @@ -1,3 +1,4 @@ +#include #include #include @@ -11,6 +12,14 @@ R = _mm_cvtss_f32(x32); \ } +#define HSUM128_EPI16(X, R) \ + int R = 0; \ + { \ + __m128i x64 = _mm_add_epi16(X, _mm_srli_si128(X, 8)); \ + __m128i x32 = _mm_add_epi16(x64, _mm_srli_si128(x64, 4)); \ + R = _mm_extract_epi16(x32, 0) + _mm_extract_epi16(x32, 1); \ + } + EXPORT float impl_score_dot_sse( const uint8_t* query_ptr, const uint8_t* vector_ptr, @@ -53,3 +62,46 @@ EXPORT uint64_t impl_xor_popcnt_sse( } return (uint32_t)result; } + +EXPORT float impl_score_l1_sse( + const uint8_t* query_ptr, + const uint8_t* vector_ptr, + uint32_t dim +) { + const __m128i* v_ptr = (const __m128i*)vector_ptr; + const __m128i* q_ptr = (const __m128i*)query_ptr; + + uint32_t m = dim - (dim % 16); + __m128i sum128 = _mm_setzero_si128(); + + // the vector sizes are assumed to be multiples of 16, no remaining part here + for (uint32_t i = 0; i < m; i += 16) { + __m128i vec2 = _mm_loadu_si128(v_ptr); + __m128i vec1 = _mm_loadu_si128(q_ptr); + v_ptr++; + q_ptr++; + + // Compute the difference in both directions + __m128i diff1 = _mm_subs_epu8(vec1, vec2); + __m128i diff2 = _mm_subs_epu8(vec2, vec1); + + // Take the maximum + __m128i abs_diff = _mm_max_epu8(diff1, diff2); + + __m128i abs_diff16_low = _mm_unpacklo_epi8(abs_diff, _mm_setzero_si128()); + __m128i abs_diff16_high = _mm_unpackhi_epi8(abs_diff, _mm_setzero_si128()); + + sum128 = _mm_add_epi16(sum128, abs_diff16_low); + sum128 = _mm_add_epi16(sum128, abs_diff16_high); + } + + // Convert 16-bit sums to 32-bit and sum them up + __m128i sum_epi32 = _mm_add_epi32( + _mm_unpacklo_epi16(sum128, _mm_setzero_si128()), + _mm_unpackhi_epi16(sum128, _mm_setzero_si128())); + + // Horizontal sum using the macro + HSUM128_EPI16(sum_epi32, sum); + + return (float) sum; +} diff --git a/quantization/src/encoded_vectors.rs b/quantization/src/encoded_vectors.rs index 63ad2fb..f5bfe01 100644 --- a/quantization/src/encoded_vectors.rs +++ b/quantization/src/encoded_vectors.rs @@ -6,6 +6,7 @@ use crate::EncodingError; #[derive(Serialize, Deserialize, Debug, Clone, Copy, PartialEq, Eq)] pub enum DistanceType { Dot, + L1, L2, } @@ -36,8 +37,9 @@ pub trait EncodedVectors: Sized { impl DistanceType { pub fn distance(&self, a: &[f32], b: &[f32]) -> f32 { match self { - DistanceType::Dot => a.iter().zip(b.iter()).map(|(a, b)| a * b).sum(), - DistanceType::L2 => a.iter().zip(b.iter()).map(|(a, b)| (a - b) * (a - b)).sum(), + DistanceType::Dot => a.iter().zip(b).map(|(a, b)| a * b).sum(), + DistanceType::L1 => a.iter().zip(b).map(|(a, b)| (a - b).abs()).sum(), + DistanceType::L2 => a.iter().zip(b).map(|(a, b)| (a - b) * (a - b)).sum(), } } } diff --git a/quantization/src/encoded_vectors_binary.rs b/quantization/src/encoded_vectors_binary.rs index dc7ded4..9d13467 100644 --- a/quantization/src/encoded_vectors_binary.rs +++ b/quantization/src/encoded_vectors_binary.rs @@ -1,7 +1,8 @@ use crate::encoded_vectors::validate_vector_parameters; use crate::utils::{transmute_from_u8_to_slice, transmute_to_u8_slice}; use crate::{ - EncodedStorage, EncodedStorageBuilder, EncodedVectors, EncodingError, VectorParameters, + DistanceType, EncodedStorage, EncodedStorageBuilder, EncodedVectors, EncodingError, + VectorParameters, }; use serde::{Deserialize, Serialize}; use std::path::Path; @@ -120,30 +121,38 @@ impl EncodedVectorsBin { } fn calculate_metric(&self, v1: &[BitsStoreType], v2: &[BitsStoreType]) -> f32 { - let xor_product = Self::xor_product(v1, v2); - // Dot product in a range [-1; 1] is approximated by NXOR in a range [0; 1] + // L1 distance in range [-1; 1] (alpha=2) is approximated by alpha*XOR in a range [0; 1] + // L2 distance in range [-1; 1] (alpha=2) is approximated by alpha*sqrt(XOR) in a range [0; 1] // For example: - // A | B | Dot product - // -0.5 | -0.5 | 0.25 - // -0.5 | 0.5 | -0.25 - // 0.5 | -0.5 | -0.25 - // 0.5 | 0.5 | 0.25 - - // A | B | NXOR - // 0 | 0 | 1 - // 0 | 1 | 0 - // 1 | 0 | 0 - // 1 | 1 | 1 - - // So is `invert` is true, we return XOR, otherwise we return (dim - XOR) - - let zeros_count = self.metadata.vector_parameters.dim - xor_product; - if self.metadata.vector_parameters.invert { - xor_product as f32 - zeros_count as f32 - } else { - zeros_count as f32 - xor_product as f32 + // | A | B | Dot product | L1 | L2 | + // | -0.5 | -0.5 | 0.25 | 0 | 0 | + // | -0.5 | 0.5 | -0.25 | 1 | 1 | + // | 0.5 | -0.5 | -0.25 | 1 | 1 | + // | 0.5 | 0.5 | 0.25 | 0 | 0 | + + // | A | B | NXOR | XOR + // | 0 | 0 | 1 | 0 + // | 0 | 1 | 0 | 1 + // | 1 | 0 | 0 | 1 + // | 1 | 1 | 1 | 0 + + let xor_product = Self::xor_product(v1, v2) as f32; + + let dim = self.metadata.vector_parameters.dim as f32; + let zeros_count = dim - xor_product; + + match ( + self.metadata.vector_parameters.distance_type, + self.metadata.vector_parameters.invert, + ) { + // So if `invert` is true we return XOR, otherwise we return (dim - XOR) + (DistanceType::Dot, true) => xor_product - zeros_count, + (DistanceType::Dot, false) => zeros_count - xor_product, + // This also results in exact ordering as L1 and L2 but reversed. + (DistanceType::L1 | DistanceType::L2, true) => zeros_count - xor_product, + (DistanceType::L1 | DistanceType::L2, false) => xor_product - zeros_count, } } } diff --git a/quantization/src/encoded_vectors_u8.rs b/quantization/src/encoded_vectors_u8.rs index 6f3b832..eeb0e1b 100644 --- a/quantization/src/encoded_vectors_u8.rs +++ b/quantization/src/encoded_vectors_u8.rs @@ -78,17 +78,17 @@ impl EncodedVectorsU8 { let mut encoded_vector = Vec::with_capacity(actual_dim + std::mem::size_of::()); encoded_vector.extend_from_slice(&f32::default().to_ne_bytes()); for &value in vector { - let endoded = Self::f32_to_u8(value, alpha, offset); - encoded_vector.push(endoded); + let encoded = Self::f32_to_u8(value, alpha, offset); + encoded_vector.push(encoded); } if vector_parameters.dim % ALIGNMENT != 0 { for _ in 0..(ALIGNMENT - vector_parameters.dim % ALIGNMENT) { let placeholder = match vector_parameters.distance_type { DistanceType::Dot => 0.0, - DistanceType::L2 => offset, + DistanceType::L1 | DistanceType::L2 => offset, }; - let endoded = Self::f32_to_u8(placeholder, alpha, offset); - encoded_vector.push(endoded); + let encoded = Self::f32_to_u8(placeholder, alpha, offset); + encoded_vector.push(encoded); } } let vector_offset = match vector_parameters.distance_type { @@ -96,6 +96,7 @@ impl EncodedVectorsU8 { actual_dim as f32 * offset * offset + encoded_vector.iter().map(|&x| x as f32).sum::() * alpha * offset } + DistanceType::L1 => 0.0, DistanceType::L2 => { actual_dim as f32 * offset * offset + encoded_vector @@ -117,6 +118,7 @@ impl EncodedVectorsU8 { } let multiplier = match vector_parameters.distance_type { DistanceType::Dot => alpha * alpha, + DistanceType::L1 => alpha, DistanceType::L2 => -2.0 * alpha * alpha, }; let multiplier = if vector_parameters.invert { @@ -138,25 +140,40 @@ impl EncodedVectorsU8 { } pub fn score_point_simple(&self, query: &EncodedQueryU8, i: u32) -> f32 { - unsafe { - let (vector_offset, v_ptr) = self.get_vec_ptr(i); - let mut mul = 0i32; - for i in 0..self.metadata.actual_dim { - mul += query.encoded_query[i] as i32 * (*v_ptr.add(i)) as i32; - } - self.metadata.multiplier * mul as f32 + query.offset + vector_offset - } + let (vector_offset, v_ptr) = self.get_vec_ptr(i); + + let score = match self.metadata.vector_parameters.distance_type { + DistanceType::Dot | DistanceType::L2 => impl_score_dot( + query.encoded_query.as_ptr(), + v_ptr, + self.metadata.actual_dim, + ), + DistanceType::L1 => impl_score_l1( + query.encoded_query.as_ptr(), + v_ptr, + self.metadata.actual_dim, + ), + }; + + self.metadata.multiplier * score as f32 + query.offset + vector_offset } #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] pub fn score_point_neon(&self, query: &EncodedQueryU8, i: u32) -> f32 { unsafe { let (vector_offset, v_ptr) = self.get_vec_ptr(i); - let score = impl_score_dot_neon( - query.encoded_query.as_ptr() as *const u8, - v_ptr, - self.metadata.actual_dim as u32, - ); + let score = match self.metadata.vector_parameters.distance_type { + DistanceType::Dot | DistanceType::L2 => impl_score_dot_neon( + query.encoded_query.as_ptr(), + v_ptr, + self.metadata.actual_dim as u32, + ), + DistanceType::L1 => impl_score_l1_neon( + query.encoded_query.as_ptr(), + v_ptr, + self.metadata.actual_dim as u32, + ), + }; self.metadata.multiplier * score + query.offset + vector_offset } } @@ -165,11 +182,18 @@ impl EncodedVectorsU8 { pub fn score_point_sse(&self, query: &EncodedQueryU8, i: u32) -> f32 { unsafe { let (vector_offset, v_ptr) = self.get_vec_ptr(i); - let score = impl_score_dot_sse( - query.encoded_query.as_ptr(), - v_ptr, - self.metadata.actual_dim as u32, - ); + let score = match self.metadata.vector_parameters.distance_type { + DistanceType::Dot | DistanceType::L2 => impl_score_dot_sse( + query.encoded_query.as_ptr(), + v_ptr, + self.metadata.actual_dim as u32, + ), + DistanceType::L1 => impl_score_l1_sse( + query.encoded_query.as_ptr(), + v_ptr, + self.metadata.actual_dim as u32, + ), + }; self.metadata.multiplier * score + query.offset + vector_offset } } @@ -178,11 +202,18 @@ impl EncodedVectorsU8 { pub fn score_point_avx(&self, query: &EncodedQueryU8, i: u32) -> f32 { unsafe { let (vector_offset, v_ptr) = self.get_vec_ptr(i); - let score = impl_score_dot_avx( - query.encoded_query.as_ptr(), - v_ptr, - self.metadata.actual_dim as u32, - ); + let score = match self.metadata.vector_parameters.distance_type { + DistanceType::Dot | DistanceType::L2 => impl_score_dot_avx( + query.encoded_query.as_ptr(), + v_ptr, + self.metadata.actual_dim as u32, + ), + DistanceType::L1 => impl_score_l1_avx( + query.encoded_query.as_ptr(), + v_ptr, + self.metadata.actual_dim as u32, + ), + }; self.metadata.multiplier * score + query.offset + vector_offset } } @@ -264,11 +295,11 @@ impl EncodedVectors for EncodedVectors for _ in 0..(ALIGNMENT - dim % ALIGNMENT) { let placeholder = match self.metadata.vector_parameters.distance_type { DistanceType::Dot => 0.0, - DistanceType::L2 => self.metadata.offset, + DistanceType::L1 | DistanceType::L2 => self.metadata.offset, }; - let endoded = + let encoded = Self::f32_to_u8(placeholder, self.metadata.alpha, self.metadata.offset); - query.push(endoded); + query.push(encoded); } } let offset = match self.metadata.vector_parameters.distance_type { @@ -277,6 +308,7 @@ impl EncodedVectors for EncodedVectors * self.metadata.alpha * self.metadata.offset } + DistanceType::L1 => 0.0, DistanceType::L2 => { query.iter().map(|&x| x as f32 * x as f32).sum::() * self.metadata.alpha @@ -300,23 +332,50 @@ impl EncodedVectors for EncodedVectors #[cfg(target_arch = "x86_64")] if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") { - let score = - unsafe { impl_score_dot_avx(q_ptr, v_ptr, self.metadata.actual_dim as u32) }; - return self.metadata.multiplier * score + query.offset + vector_offset; + unsafe { + let score = match self.metadata.vector_parameters.distance_type { + DistanceType::Dot | DistanceType::L2 => { + impl_score_dot_avx(q_ptr, v_ptr, self.metadata.actual_dim as u32) + } + DistanceType::L1 => { + impl_score_l1_avx(q_ptr, v_ptr, self.metadata.actual_dim as u32) + } + }; + + return self.metadata.multiplier * score + query.offset + vector_offset; + } } #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] if is_x86_feature_detected!("sse4.1") { - let score = - unsafe { impl_score_dot_sse(q_ptr, v_ptr, self.metadata.actual_dim as u32) }; - return self.metadata.multiplier * score + query.offset + vector_offset; + unsafe { + let score = match self.metadata.vector_parameters.distance_type { + DistanceType::Dot | DistanceType::L2 => { + impl_score_dot_sse(q_ptr, v_ptr, self.metadata.actual_dim as u32) + } + DistanceType::L1 => { + impl_score_l1_sse(q_ptr, v_ptr, self.metadata.actual_dim as u32) + } + }; + + return self.metadata.multiplier * score + query.offset + vector_offset; + } } #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] if std::arch::is_aarch64_feature_detected!("neon") { - let score = - unsafe { impl_score_dot_neon(q_ptr, v_ptr, self.metadata.actual_dim as u32) }; - return self.metadata.multiplier * score + query.offset + vector_offset; + unsafe { + let score = match self.metadata.vector_parameters.distance_type { + DistanceType::Dot | DistanceType::L2 => { + impl_score_dot_neon(q_ptr, v_ptr, self.metadata.actual_dim as u32) + } + DistanceType::L1 => { + impl_score_l1_neon(q_ptr, v_ptr, self.metadata.actual_dim as u32) + } + }; + + return self.metadata.multiplier * score + query.offset + vector_offset; + } } self.score_point_simple(query, i) @@ -335,43 +394,94 @@ impl EncodedVectors for EncodedVectors #[cfg(target_arch = "x86_64")] if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") { - let score = - unsafe { impl_score_dot_avx(q_ptr, v_ptr, self.metadata.actual_dim as u32) }; - return self.metadata.multiplier * score + offset; + unsafe { + let score = match self.metadata.vector_parameters.distance_type { + DistanceType::Dot | DistanceType::L2 => { + impl_score_dot_avx(q_ptr, v_ptr, self.metadata.actual_dim as u32) + } + DistanceType::L1 => { + impl_score_l1_avx(q_ptr, v_ptr, self.metadata.actual_dim as u32) + } + }; + + return self.metadata.multiplier * score + offset; + } } #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] if is_x86_feature_detected!("sse4.1") { - let score = - unsafe { impl_score_dot_sse(q_ptr, v_ptr, self.metadata.actual_dim as u32) }; - return self.metadata.multiplier * score + offset; + unsafe { + let score = match self.metadata.vector_parameters.distance_type { + DistanceType::Dot | DistanceType::L2 => { + impl_score_dot_sse(q_ptr, v_ptr, self.metadata.actual_dim as u32) + } + DistanceType::L1 => { + impl_score_l1_sse(q_ptr, v_ptr, self.metadata.actual_dim as u32) + } + }; + + return self.metadata.multiplier * score + offset; + } } #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] if std::arch::is_aarch64_feature_detected!("neon") { - let score = - unsafe { impl_score_dot_neon(q_ptr, v_ptr, self.metadata.actual_dim as u32) }; - return self.metadata.multiplier * score + offset; + unsafe { + let score = match self.metadata.vector_parameters.distance_type { + DistanceType::Dot | DistanceType::L2 => { + impl_score_dot_neon(q_ptr, v_ptr, self.metadata.actual_dim as u32) + } + DistanceType::L1 => { + impl_score_l1_neon(q_ptr, v_ptr, self.metadata.actual_dim as u32) + } + }; + + return self.metadata.multiplier * score + offset; + } } - unsafe { - let mut mul = 0i32; - for i in 0..self.metadata.actual_dim { - mul += (*q_ptr.add(i)) as i32 * (*v_ptr.add(i)) as i32; + let score = match self.metadata.vector_parameters.distance_type { + DistanceType::Dot | DistanceType::L2 => { + impl_score_dot(q_ptr, v_ptr, self.metadata.actual_dim) } - self.metadata.multiplier * mul as f32 + offset + DistanceType::L1 => impl_score_l1(q_ptr, v_ptr, self.metadata.actual_dim), + }; + + self.metadata.multiplier * score as f32 + offset + } +} + +fn impl_score_dot(q_ptr: *const u8, v_ptr: *const u8, actual_dim: usize) -> i32 { + unsafe { + let mut score = 0i32; + for i in 0..actual_dim { + score += (*q_ptr.add(i)) as i32 * (*v_ptr.add(i)) as i32; + } + score + } +} + +fn impl_score_l1(q_ptr: *const u8, v_ptr: *const u8, actual_dim: usize) -> i32 { + unsafe { + let mut score = 0i32; + for i in 0..actual_dim { + score += (*q_ptr.add(i) as i32).abs_diff(*v_ptr.add(i) as i32) as i32; } + score } } #[cfg(target_arch = "x86_64")] extern "C" { fn impl_score_dot_avx(query_ptr: *const u8, vector_ptr: *const u8, dim: u32) -> f32; + fn impl_score_l1_avx(query_ptr: *const u8, vector_ptr: *const u8, dim: u32) -> f32; fn impl_score_dot_sse(query_ptr: *const u8, vector_ptr: *const u8, dim: u32) -> f32; + fn impl_score_l1_sse(query_ptr: *const u8, vector_ptr: *const u8, dim: u32) -> f32; } #[cfg(all(target_arch = "aarch64", target_feature = "neon"))] extern "C" { fn impl_score_dot_neon(query_ptr: *const u8, vector_ptr: *const u8, dim: u32) -> f32; + fn impl_score_l1_neon(query_ptr: *const u8, vector_ptr: *const u8, dim: u32) -> f32; } diff --git a/quantization/tests/metrics.rs b/quantization/tests/metrics.rs index 0ce58c4..c6ff1e5 100644 --- a/quantization/tests/metrics.rs +++ b/quantization/tests/metrics.rs @@ -3,9 +3,9 @@ pub fn dot_similarity(v1: &[f32], v2: &[f32]) -> f32 { } pub fn l2_similarity(v1: &[f32], v2: &[f32]) -> f32 { - v1.iter() - .copied() - .zip(v2.iter().copied()) - .map(|(a, b)| (a - b).powi(2)) - .sum() + v1.iter().zip(v2).map(|(a, b)| (a - b).powi(2)).sum() +} + +pub fn l1_similarity(v1: &[f32], v2: &[f32]) -> f32 { + v1.iter().zip(v2).map(|(a, b)| (a - b).abs()).sum() } diff --git a/quantization/tests/test_avx2.rs b/quantization/tests/test_avx2.rs index abe51a1..18d1338 100644 --- a/quantization/tests/test_avx2.rs +++ b/quantization/tests/test_avx2.rs @@ -4,7 +4,7 @@ mod metrics; #[cfg(test)] #[cfg(target_arch = "x86_64")] mod tests { - use crate::metrics::{dot_similarity, l2_similarity}; + use crate::metrics::{dot_similarity, l1_similarity, l2_similarity}; use quantization::{ encoded_vectors::{DistanceType, EncodedVectors, VectorParameters}, encoded_vectors_u8::EncodedVectorsU8, @@ -84,4 +84,41 @@ mod tests { assert!((score - orginal_score).abs() < error); } } + + #[test] + fn test_l1_avx() { + let vectors_count = 129; + let vector_dim = 65; + let error = vector_dim as f32 * 0.1; + + //let mut rng = rand::thread_rng(); + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + let mut vector_data: Vec> = Vec::new(); + for _ in 0..vectors_count { + let vector: Vec = (0..vector_dim).map(|_| rng.gen_range(-1.0..=1.0)).collect(); + vector_data.push(vector); + } + let query: Vec = (0..vector_dim).map(|_| rng.gen_range(-1.0..=1.0)).collect(); + + let encoded = EncodedVectorsU8::encode( + vector_data.iter().map(|v| v.as_slice()), + Vec::::new(), + &VectorParameters { + dim: vector_dim, + count: vectors_count, + distance_type: DistanceType::L1, + invert: false, + }, + None, + || false, + ) + .unwrap(); + let query_u8 = encoded.encode_query(&query); + + for (index, vector) in vector_data.iter().enumerate() { + let score = encoded.score_point_avx(&query_u8, index as u32); + let orginal_score = l1_similarity(&query, vector); + assert!((score - orginal_score).abs() < error); + } + } } diff --git a/quantization/tests/test_binary.rs b/quantization/tests/test_binary.rs index 2bffd80..4912741 100644 --- a/quantization/tests/test_binary.rs +++ b/quantization/tests/test_binary.rs @@ -9,7 +9,7 @@ mod tests { }; use rand::{Rng, SeedableRng}; - use crate::metrics::dot_similarity; + use crate::metrics::{dot_similarity, l1_similarity, l2_similarity}; fn generate_number(rng: &mut rand::rngs::StdRng) -> f32 { let n = f32::signum(rng.gen_range(-1.0..1.0)); @@ -161,4 +161,400 @@ mod tests { assert!((score - orginal_score).abs() < error); } } + + #[test] + fn test_binary_l1() { + let vectors_count = 128; + let vector_dim = 3 * 128; + + //let mut rng = rand::thread_rng(); + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + let mut vector_data: Vec> = Vec::new(); + for _ in 0..vectors_count { + vector_data.push(generate_vector(vector_dim, &mut rng)); + } + + let encoded = EncodedVectorsBin::encode( + vector_data.iter().map(|v| v.as_slice()), + Vec::::new(), + &VectorParameters { + dim: vector_dim, + count: vectors_count, + distance_type: DistanceType::L1, + invert: false, + }, + || false, + ) + .unwrap(); + + let query: Vec = generate_vector(vector_dim, &mut rng); + let query_b = encoded.encode_query(&query); + + let mut scores: Vec<_> = vector_data + .iter() + .enumerate() + .map(|(i, _)| (encoded.score_point(&query_b, i as u32), i)) + .collect(); + + scores.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + + let sorted_indices: Vec<_> = scores.into_iter().map(|(_, i)| i).collect(); + + let mut original_scores: Vec<_> = vector_data + .iter() + .enumerate() + .map(|(i, v)| (l1_similarity(&query, v), i)) + .collect(); + + original_scores.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + + let sorted_original_indices: Vec<_> = original_scores.into_iter().map(|(_, i)| i).collect(); + + assert_eq!(sorted_original_indices, sorted_indices); + } + + #[test] + fn test_binary_l1_inverted() { + let vectors_count = 128; + let vector_dim = 128; + + //let mut rng = rand::thread_rng(); + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + let mut vector_data: Vec> = Vec::new(); + for _ in 0..vectors_count { + vector_data.push(generate_vector(vector_dim, &mut rng)); + } + + let encoded = EncodedVectorsBin::encode( + vector_data.iter().map(|v| v.as_slice()), + Vec::::new(), + &VectorParameters { + dim: vector_dim, + count: vectors_count, + distance_type: DistanceType::L1, + invert: true, + }, + || false, + ) + .unwrap(); + + let query: Vec = generate_vector(vector_dim, &mut rng); + let query_b = encoded.encode_query(&query); + + let mut scores: Vec<_> = vector_data + .iter() + .enumerate() + .map(|(i, _)| (encoded.score_point(&query_b, i as u32), i)) + .collect(); + + scores.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + + let sorted_indices: Vec<_> = scores.into_iter().map(|(_, i)| i).collect(); + + let mut original_scores: Vec<_> = vector_data + .iter() + .enumerate() + .map(|(i, v)| (l1_similarity(&query, v), i)) + .collect(); + + original_scores.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap()); + + let sorted_original_indices: Vec<_> = original_scores.into_iter().map(|(_, i)| i).collect(); + + assert_eq!(sorted_original_indices, sorted_indices); + } + + #[test] + fn test_binary_l1_internal() { + let vectors_count = 128; + let vector_dim = 128; + + //let mut rng = rand::thread_rng(); + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + let mut vector_data: Vec> = Vec::new(); + for _ in 0..vectors_count { + vector_data.push(generate_vector(vector_dim, &mut rng)); + } + + let encoded = EncodedVectorsBin::encode( + vector_data.iter().map(|v| v.as_slice()), + Vec::::new(), + &VectorParameters { + dim: vector_dim, + count: vectors_count, + distance_type: DistanceType::L1, + invert: false, + }, + || false, + ) + .unwrap(); + + let mut scores: Vec<_> = vector_data + .iter() + .enumerate() + .map(|(i, _)| (encoded.score_internal(0, i as u32), i)) + .collect(); + + scores.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + + let sorted_indices: Vec<_> = scores.into_iter().map(|(_, i)| i).collect(); + + let mut original_scores: Vec<_> = vector_data + .iter() + .enumerate() + .map(|(i, v)| (l1_similarity(&vector_data[0], v), i)) + .collect(); + + original_scores.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + + let sorted_original_indices: Vec<_> = original_scores.into_iter().map(|(_, i)| i).collect(); + + assert_eq!(sorted_original_indices, sorted_indices); + } + + #[test] + fn test_binary_l1_inverted_internal() { + let vectors_count = 128; + let vector_dim = 128; + + //let mut rng = rand::thread_rng(); + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + let mut vector_data: Vec> = Vec::new(); + for _ in 0..vectors_count { + vector_data.push(generate_vector(vector_dim, &mut rng)); + } + + let encoded = EncodedVectorsBin::encode( + vector_data.iter().map(|v| v.as_slice()), + Vec::::new(), + &VectorParameters { + dim: vector_dim, + count: vectors_count, + distance_type: DistanceType::L1, + invert: true, + }, + || false, + ) + .unwrap(); + + let mut scores: Vec<_> = vector_data + .iter() + .enumerate() + .map(|(i, _)| (encoded.score_internal(0, i as u32), i)) + .collect(); + + scores.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + + let sorted_indices: Vec<_> = scores.into_iter().map(|(_, i)| i).collect(); + + let mut original_scores: Vec<_> = vector_data + .iter() + .enumerate() + .map(|(i, v)| (l1_similarity(&vector_data[0], v), i)) + .collect(); + + original_scores.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap()); + + let sorted_original_indices: Vec<_> = original_scores.into_iter().map(|(_, i)| i).collect(); + + assert_eq!(sorted_original_indices, sorted_indices); + } + + #[test] + fn test_binary_l2() { + let vectors_count = 128; + let vector_dim = 3 * 128; + + //let mut rng = rand::thread_rng(); + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + let mut vector_data: Vec> = Vec::new(); + for _ in 0..vectors_count { + vector_data.push(generate_vector(vector_dim, &mut rng)); + } + + let encoded = EncodedVectorsBin::encode( + vector_data.iter().map(|v| v.as_slice()), + Vec::::new(), + &VectorParameters { + dim: vector_dim, + count: vectors_count, + distance_type: DistanceType::L2, + invert: false, + }, + || false, + ) + .unwrap(); + + let query: Vec = generate_vector(vector_dim, &mut rng); + let query_b = encoded.encode_query(&query); + + let mut scores: Vec<_> = vector_data + .iter() + .enumerate() + .map(|(i, _)| (encoded.score_point(&query_b, i as u32), i)) + .collect(); + + scores.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + + let sorted_indices: Vec<_> = scores.into_iter().map(|(_, i)| i).collect(); + + let mut original_scores: Vec<_> = vector_data + .iter() + .enumerate() + .map(|(i, v)| (l2_similarity(&query, v), i)) + .collect(); + + original_scores.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + + let sorted_original_indices: Vec<_> = original_scores.into_iter().map(|(_, i)| i).collect(); + + assert_eq!(sorted_original_indices, sorted_indices); + } + + #[test] + fn test_binary_l2_inverted() { + let vectors_count = 128; + let vector_dim = 128; + + //let mut rng = rand::thread_rng(); + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + let mut vector_data: Vec> = Vec::new(); + for _ in 0..vectors_count { + vector_data.push(generate_vector(vector_dim, &mut rng)); + } + + let encoded = EncodedVectorsBin::encode( + vector_data.iter().map(|v| v.as_slice()), + Vec::::new(), + &VectorParameters { + dim: vector_dim, + count: vectors_count, + distance_type: DistanceType::L2, + invert: true, + }, + || false, + ) + .unwrap(); + + let query: Vec = generate_vector(vector_dim, &mut rng); + let query_b = encoded.encode_query(&query); + + let mut scores: Vec<_> = vector_data + .iter() + .enumerate() + .map(|(i, _)| (encoded.score_point(&query_b, i as u32), i)) + .collect(); + + scores.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + + let sorted_indices: Vec<_> = scores.into_iter().map(|(_, i)| i).collect(); + + let mut original_scores: Vec<_> = vector_data + .iter() + .enumerate() + .map(|(i, v)| (l2_similarity(&query, v), i)) + .collect(); + + original_scores.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap()); + + let sorted_original_indices: Vec<_> = original_scores.into_iter().map(|(_, i)| i).collect(); + + assert_eq!(sorted_original_indices, sorted_indices); + } + + #[test] + fn test_binary_l2_internal() { + let vectors_count = 128; + let vector_dim = 128; + + //let mut rng = rand::thread_rng(); + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + let mut vector_data: Vec> = Vec::new(); + for _ in 0..vectors_count { + vector_data.push(generate_vector(vector_dim, &mut rng)); + } + + let encoded = EncodedVectorsBin::encode( + vector_data.iter().map(|v| v.as_slice()), + Vec::::new(), + &VectorParameters { + dim: vector_dim, + count: vectors_count, + distance_type: DistanceType::L2, + invert: false, + }, + || false, + ) + .unwrap(); + + let mut scores: Vec<_> = vector_data + .iter() + .enumerate() + .map(|(i, _)| (encoded.score_internal(0, i as u32), i)) + .collect(); + + scores.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + + let sorted_indices: Vec<_> = scores.into_iter().map(|(_, i)| i).collect(); + + let mut original_scores: Vec<_> = vector_data + .iter() + .enumerate() + .map(|(i, v)| (l2_similarity(&vector_data[0], v), i)) + .collect(); + + original_scores.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + + let sorted_original_indices: Vec<_> = original_scores.into_iter().map(|(_, i)| i).collect(); + + assert_eq!(sorted_original_indices, sorted_indices); + } + + #[test] + fn test_binary_l2_inverted_internal() { + let vectors_count = 128; + let vector_dim = 128; + + //let mut rng = rand::thread_rng(); + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + let mut vector_data: Vec> = Vec::new(); + for _ in 0..vectors_count { + vector_data.push(generate_vector(vector_dim, &mut rng)); + } + + let encoded = EncodedVectorsBin::encode( + vector_data.iter().map(|v| v.as_slice()), + Vec::::new(), + &VectorParameters { + dim: vector_dim, + count: vectors_count, + distance_type: DistanceType::L2, + invert: true, + }, + || false, + ) + .unwrap(); + + let mut scores: Vec<_> = vector_data + .iter() + .enumerate() + .map(|(i, _)| (encoded.score_internal(0, i as u32), i)) + .collect(); + + scores.sort_by(|a, b| a.0.partial_cmp(&b.0).unwrap()); + + let sorted_indices: Vec<_> = scores.into_iter().map(|(_, i)| i).collect(); + + let mut original_scores: Vec<_> = vector_data + .iter() + .enumerate() + .map(|(i, v)| (l1_similarity(&vector_data[0], v), i)) + .collect(); + + original_scores.sort_by(|a, b| b.0.partial_cmp(&a.0).unwrap()); + + let sorted_original_indices: Vec<_> = original_scores.into_iter().map(|(_, i)| i).collect(); + + assert_eq!(sorted_original_indices, sorted_indices); + } } diff --git a/quantization/tests/test_neon.rs b/quantization/tests/test_neon.rs index 3f75927..318537c 100644 --- a/quantization/tests/test_neon.rs +++ b/quantization/tests/test_neon.rs @@ -10,7 +10,7 @@ mod tests { }; use rand::{Rng, SeedableRng}; - use crate::metrics::{dot_similarity, l2_similarity}; + use crate::metrics::{dot_similarity, l1_similarity, l2_similarity}; #[test] fn test_dot_neon() { @@ -83,4 +83,40 @@ mod tests { assert!((score - orginal_score).abs() < error); } } + + #[test] + fn test_l1_neon() { + let vectors_count = 129; + let vector_dim = 65; + let error = vector_dim as f32 * 0.1; + + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + let mut vector_data: Vec> = Vec::new(); + for _ in 0..vectors_count { + let vector: Vec = (0..vector_dim).map(|_| rng.gen()).collect(); + vector_data.push(vector); + } + let query: Vec = (0..vector_dim).map(|_| rng.gen()).collect(); + + let encoded = EncodedVectorsU8::encode( + vector_data.iter().map(|v| v.as_slice()), + Vec::::new(), + &VectorParameters { + dim: vector_dim, + count: vectors_count, + distance_type: DistanceType::L1, + invert: false, + }, + None, + || false, + ) + .unwrap(); + let query_u8 = encoded.encode_query(&query); + + for (index, vector) in vector_data.iter().enumerate() { + let score = encoded.score_point_neon(&query_u8, index as u32); + let orginal_score = l1_similarity(&query, vector); + assert!((score - orginal_score).abs() < error); + } + } } diff --git a/quantization/tests/test_pq.rs b/quantization/tests/test_pq.rs index ef4836a..0b8c1ff 100644 --- a/quantization/tests/test_pq.rs +++ b/quantization/tests/test_pq.rs @@ -11,7 +11,7 @@ mod tests { }; use rand::{Rng, SeedableRng}; - use crate::metrics::{dot_similarity, l2_similarity}; + use crate::metrics::{dot_similarity, l1_similarity, l2_similarity}; const VECTORS_COUNT: usize = 513; const VECTOR_DIM: usize = 65; @@ -81,6 +81,38 @@ mod tests { } } + #[test] + fn test_pq_l1() { + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + let mut vector_data: Vec> = vec![]; + for _ in 0..VECTORS_COUNT { + vector_data.push((0..VECTOR_DIM).map(|_| rng.gen()).collect()); + } + let query: Vec<_> = (0..VECTOR_DIM).map(|_| rng.gen()).collect(); + + let encoded = EncodedVectorsPQ::encode( + vector_data.iter().map(Vec::as_slice), + vec![], + &VectorParameters { + dim: VECTOR_DIM, + count: VECTORS_COUNT, + distance_type: DistanceType::L1, + invert: false, + }, + 1, + 1, + || false, + ) + .unwrap(); + let query_u8 = encoded.encode_query(&query); + + for (index, vector) in vector_data.iter().enumerate() { + let score = encoded.score_point(&query_u8, index as u32); + let orginal_score = l1_similarity(&query, vector); + assert!((score - orginal_score).abs() < ERROR); + } + } + #[test] fn test_pq_dot_inverted() { let mut rng = rand::rngs::StdRng::seed_from_u64(42); @@ -145,6 +177,38 @@ mod tests { } } + #[test] + fn test_pq_l1_inverted() { + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + let mut vector_data: Vec> = vec![]; + for _ in 0..VECTORS_COUNT { + vector_data.push((0..VECTOR_DIM).map(|_| rng.gen()).collect()); + } + let query: Vec<_> = (0..VECTOR_DIM).map(|_| rng.gen()).collect(); + + let encoded = EncodedVectorsPQ::encode( + vector_data.iter().map(Vec::as_slice), + vec![], + &VectorParameters { + dim: VECTOR_DIM, + count: VECTORS_COUNT, + distance_type: DistanceType::L1, + invert: true, + }, + 1, + 1, + || false, + ) + .unwrap(); + let query_u8 = encoded.encode_query(&query); + + for (index, vector) in vector_data.iter().enumerate() { + let score = encoded.score_point(&query_u8, index as u32); + let orginal_score = -l1_similarity(&query, vector); + assert!((score - orginal_score).abs() < ERROR); + } + } + #[test] fn test_pq_dot_internal() { let mut rng = rand::rngs::StdRng::seed_from_u64(42); diff --git a/quantization/tests/test_simple.rs b/quantization/tests/test_simple.rs index 0b34229..e032aba 100644 --- a/quantization/tests/test_simple.rs +++ b/quantization/tests/test_simple.rs @@ -9,7 +9,7 @@ mod tests { }; use rand::{Rng, SeedableRng}; - use crate::metrics::{dot_similarity, l2_similarity}; + use crate::metrics::{dot_similarity, l1_similarity, l2_similarity}; #[test] fn test_dot_simple() { @@ -85,6 +85,43 @@ mod tests { } } + #[test] + fn test_l1_simple() { + let vectors_count = 129; + let vector_dim = 65; + let error = vector_dim as f32 * 0.1; + + //let mut rng = rand::thread_rng(); + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + let mut vector_data: Vec> = Vec::new(); + for _ in 0..vectors_count { + let vector: Vec = (0..vector_dim).map(|_| rng.gen_range(-1.0..=1.0)).collect(); + vector_data.push(vector); + } + let query: Vec = (0..vector_dim).map(|_| rng.gen_range(-1.0..=1.0)).collect(); + + let encoded = EncodedVectorsU8::encode( + vector_data.iter().map(|v| v.as_slice()), + Vec::::new(), + &VectorParameters { + dim: vector_dim, + count: vectors_count, + distance_type: DistanceType::L1, + invert: false, + }, + None, + || false, + ) + .unwrap(); + let query_u8 = encoded.encode_query(&query); + + for (index, vector) in vector_data.iter().enumerate() { + let score = encoded.score_point_simple(&query_u8, index as u32); + let orginal_score = l1_similarity(&query, vector); + assert!((score - orginal_score).abs() < error); + } + } + #[test] fn test_dot_inverted_simple() { let vectors_count = 129; @@ -159,6 +196,43 @@ mod tests { } } + #[test] + fn test_l1_inverted_simple() { + let vectors_count = 129; + let vector_dim = 65; + let error = vector_dim as f32 * 0.1; + + //let mut rng = rand::thread_rng(); + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + let mut vector_data: Vec> = Vec::new(); + for _ in 0..vectors_count { + let vector: Vec = (0..vector_dim).map(|_| rng.gen_range(-1.0..=1.0)).collect(); + vector_data.push(vector); + } + let query: Vec = (0..vector_dim).map(|_| rng.gen_range(-1.0..=1.0)).collect(); + + let encoded = EncodedVectorsU8::encode( + vector_data.iter().map(|v| v.as_slice()), + Vec::::new(), + &VectorParameters { + dim: vector_dim, + count: vectors_count, + distance_type: DistanceType::L1, + invert: true, + }, + None, + || false, + ) + .unwrap(); + let query_u8 = encoded.encode_query(&query); + + for (index, vector) in vector_data.iter().enumerate() { + let score = encoded.score_point_simple(&query_u8, index as u32); + let orginal_score = -l1_similarity(&query, vector); + assert!((score - orginal_score).abs() < error); + } + } + #[test] fn test_dot_internal_simple() { let vectors_count: usize = 129; diff --git a/quantization/tests/test_sse.rs b/quantization/tests/test_sse.rs index f95c5ab..7730190 100644 --- a/quantization/tests/test_sse.rs +++ b/quantization/tests/test_sse.rs @@ -4,7 +4,7 @@ mod metrics; #[cfg(test)] #[cfg(any(target_arch = "x86", target_arch = "x86_64"))] mod tests { - use crate::metrics::{dot_similarity, l2_similarity}; + use crate::metrics::{dot_similarity, l1_similarity, l2_similarity}; use quantization::{ encoded_vectors::{DistanceType, EncodedVectors, VectorParameters}, encoded_vectors_u8::EncodedVectorsU8, @@ -84,4 +84,41 @@ mod tests { assert!((score - orginal_score).abs() < error); } } + + #[test] + fn test_l1_sse() { + let vectors_count = 129; + let vector_dim = 65; + let error = vector_dim as f32 * 0.1; + + //let mut rng = rand::thread_rng(); + let mut rng = rand::rngs::StdRng::seed_from_u64(42); + let mut vector_data: Vec> = Vec::new(); + for _ in 0..vectors_count { + let vector: Vec = (0..vector_dim).map(|_| rng.gen()).collect(); + vector_data.push(vector); + } + let query: Vec = (0..vector_dim).map(|_| rng.gen()).collect(); + + let encoded = EncodedVectorsU8::encode( + vector_data.iter().map(|v| v.as_slice()), + Vec::::new(), + &VectorParameters { + dim: vector_dim, + count: vectors_count, + distance_type: DistanceType::L1, + invert: false, + }, + None, + || false, + ) + .unwrap(); + let query_u8 = encoded.encode_query(&query); + + for (index, vector) in vector_data.iter().enumerate() { + let score = encoded.score_point_sse(&query_u8, index as u32); + let orginal_score = l1_similarity(&query, vector); + assert!((score - orginal_score).abs() < error); + } + } }