diff --git a/CMakeLists.txt b/CMakeLists.txt index 0c84ac3..d260908 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -47,6 +47,7 @@ endfunction() set(RABITQ_COMMON_SOURCES src/utils/cpu_features.cpp src/simd/dispatch.cpp + src/simd/space_float.cpp ) set(RABITQ_AVX2_SOURCES diff --git a/include/rabitqlib/simd/space_dispatch.hpp b/include/rabitqlib/simd/space_dispatch.hpp index 5cf08eb..67d5e64 100644 --- a/include/rabitqlib/simd/space_dispatch.hpp +++ b/include/rabitqlib/simd/space_dispatch.hpp @@ -4,6 +4,27 @@ #include namespace rabitqlib::simd { +// Raw float32 distances accept arbitrary dimensions and unaligned inputs. +float euclidean_sqr(const float* a, const float* b, size_t dim); +float dot_product(const float* a, const float* b, size_t dim); +float dot_product_dis(const float* a, const float* b, size_t dim); +float l2norm_sqr(const float* a, size_t dim); + +float euclidean_sqr_generic(const float* a, const float* b, size_t dim); +float dot_product_generic(const float* a, const float* b, size_t dim); +float dot_product_dis_generic(const float* a, const float* b, size_t dim); +float l2norm_sqr_generic(const float* a, size_t dim); + +float euclidean_sqr_avx2(const float* a, const float* b, size_t dim); +float dot_product_avx2(const float* a, const float* b, size_t dim); +float dot_product_dis_avx2(const float* a, const float* b, size_t dim); +float l2norm_sqr_avx2(const float* a, size_t dim); + +float euclidean_sqr_avx512(const float* a, const float* b, size_t dim); +float dot_product_avx512(const float* a, const float* b, size_t dim); +float dot_product_dis_avx512(const float* a, const float* b, size_t dim); +float l2norm_sqr_avx512(const float* a, size_t dim); + namespace excode_ipimpl { float ip16_fxu1_avx2( diff --git a/include/rabitqlib/utils/space.hpp b/include/rabitqlib/utils/space.hpp index 4ffe035..8a47d4a 100644 --- a/include/rabitqlib/utils/space.hpp +++ b/include/rabitqlib/utils/space.hpp @@ -83,31 +83,47 @@ inline void vec_rescale(T* data, size_t dim, T val) { template inline T euclidean_sqr(const T* __restrict__ vec0, const T* __restrict__ vec1, size_t dim) { - ConstVectorMap v0(vec0, dim); - ConstVectorMap v1(vec1, dim); - return (v0 - v1).dot(v0 - v1); + if constexpr (std::is_same_v) { + return simd::euclidean_sqr(vec0, vec1, dim); + } else { + ConstVectorMap v0(vec0, dim); + ConstVectorMap v1(vec1, dim); + return (v0 - v1).dot(v0 - v1); + } } template inline T dot_product_dis( const T* __restrict__ vec0, const T* __restrict__ vec1, size_t dim ) { - ConstVectorMap v0(vec0, dim); - ConstVectorMap v1(vec1, dim); - return 1 - v0.dot(v1); + if constexpr (std::is_same_v) { + return simd::dot_product_dis(vec0, vec1, dim); + } else { + ConstVectorMap v0(vec0, dim); + ConstVectorMap v1(vec1, dim); + return 1 - v0.dot(v1); + } } template inline T l2norm_sqr(const T* __restrict__ vec0, size_t dim) { - ConstVectorMap v0(vec0, dim); - return v0.dot(v0); + if constexpr (std::is_same_v) { + return simd::l2norm_sqr(vec0, dim); + } else { + ConstVectorMap v0(vec0, dim); + return v0.dot(v0); + } } template inline T dot_product(const T* __restrict__ vec0, const T* __restrict__ vec1, size_t dim) { - ConstVectorMap v0(vec0, dim); - ConstVectorMap v1(vec1, dim); - return v0.dot(v1); + if constexpr (std::is_same_v) { + return simd::dot_product(vec0, vec1, dim); + } else { + ConstVectorMap v0(vec0, dim); + ConstVectorMap v1(vec1, dim); + return v0.dot(v1); + } } template diff --git a/src/index/hnsw_search_avx2_kernels.hpp b/src/index/hnsw_search_avx2_kernels.hpp index bf20fc0..4206857 100644 --- a/src/index/hnsw_search_avx2_kernels.hpp +++ b/src/index/hnsw_search_avx2_kernels.hpp @@ -4,6 +4,7 @@ #include #include +#include #include "rabitqlib/index/query.hpp" #include "rabitqlib/simd/space_dispatch.hpp" @@ -68,40 +69,45 @@ static inline float hnsw_mask_ip_x0_q_avx2( const float* query, const uint64_t* data, size_t padded_dim ) { const size_t num_blk = padded_dim / 64; - const uint8_t* it_data = reinterpret_cast(data); + const auto* it_data = reinterpret_cast(data); const float* it_query = query; - - __m256 sum = _mm256_setzero_ps(); - __m256i bit_checker = _mm256_set_epi32(0x80, 0x40, 0x20, 0x10, 0x08, 0x04, 0x02, 0x01); + const __m256i shifts0 = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7); + const __m256i shifts1 = _mm256_setr_epi32(8, 9, 10, 11, 12, 13, 14, 15); + const __m256i shifts2 = _mm256_setr_epi32(16, 17, 18, 19, 20, 21, 22, 23); + const __m256i shifts3 = _mm256_setr_epi32(24, 25, 26, 27, 28, 29, 30, 31); + __m256 sum0 = _mm256_setzero_ps(); + __m256 sum1 = _mm256_setzero_ps(); + __m256 sum2 = _mm256_setzero_ps(); + __m256 sum3 = _mm256_setzero_ps(); for (size_t i = 0; i < num_blk; ++i) { - uint64_t bits = rabitqlib::reverse_bits_u64(rabitqlib::load_unaligned_u64(it_data)); - - // 64 bits / 8 floats = 8 iterations - for (int j = 0; j < 8; ++j) { - uint8_t current_byte = static_cast(bits >> (j * 8)); - __m256i v_byte = _mm256_set1_epi32(current_byte); - __m256i masked_bits = _mm256_and_si256(v_byte, bit_checker); - __m256i mask = _mm256_cmpgt_epi32(masked_bits, _mm256_setzero_si256()); - - __m256 q_vals = _mm256_loadu_ps(it_query); - __m256 masked = _mm256_and_ps(q_vals, _mm256_castsi256_ps(mask)); - - sum = _mm256_add_ps(sum, masked); - - it_query += 8; + // Stored coordinates run from bit 63 to bit 0: high word first, + // with each selected bit shifted into a maskload lane's sign bit. + for (size_t half = 0; half < 2; ++half) { + int32_t word; + std::memcpy(&word, it_data + (1 - half) * sizeof(word), sizeof(word)); + const __m256i bits = _mm256_set1_epi32(word); + sum0 = _mm256_add_ps( + sum0, _mm256_maskload_ps(it_query, _mm256_sllv_epi32(bits, shifts0)) + ); + sum1 = _mm256_add_ps( + sum1, _mm256_maskload_ps(it_query + 8, _mm256_sllv_epi32(bits, shifts1)) + ); + sum2 = _mm256_add_ps( + sum2, _mm256_maskload_ps(it_query + 16, _mm256_sllv_epi32(bits, shifts2)) + ); + sum3 = _mm256_add_ps( + sum3, _mm256_maskload_ps(it_query + 24, _mm256_sllv_epi32(bits, shifts3)) + ); + it_query += 32; } it_data += sizeof(uint64_t); } - alignas(32) float lanes[8]; - _mm256_store_ps(lanes, sum); - - float result = 0.0f; - for (float lane : lanes) { - result += lane; - } - return result; + const __m256 sum = _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)); + __m128 lanes = _mm_add_ps(_mm256_castps256_ps128(sum), _mm256_extractf128_ps(sum, 1)); + lanes = _mm_add_ps(lanes, _mm_movehl_ps(lanes, lanes)); + return _mm_cvtss_f32(_mm_add_ss(lanes, _mm_movehdup_ps(lanes))); } static inline float hnsw_warmup_ip_x0_q_512_avx2( diff --git a/src/simd/dispatch.cpp b/src/simd/dispatch.cpp index 56f379f..5e17b8b 100644 --- a/src/simd/dispatch.cpp +++ b/src/simd/dispatch.cpp @@ -16,6 +16,33 @@ namespace rabitqlib::simd { +const auto kEuclideanSqrFn = cpu::has_avx512_core() ? euclidean_sqr_avx512 + : cpu::has_avx2() ? euclidean_sqr_avx2 + : euclidean_sqr_generic; +const auto kDotProductFn = cpu::has_avx512_core() ? dot_product_avx512 + : cpu::has_avx2() ? dot_product_avx2 + : dot_product_generic; +const auto kDotProductDisFn = cpu::has_avx512_core() ? dot_product_dis_avx512 + : cpu::has_avx2() ? dot_product_dis_avx2 + : dot_product_dis_generic; +const auto kL2normSqrFn = cpu::has_avx512_core() ? l2norm_sqr_avx512 + : cpu::has_avx2() ? l2norm_sqr_avx2 + : l2norm_sqr_generic; + +float euclidean_sqr(const float* a, const float* b, size_t dim) { + return kEuclideanSqrFn(a, b, dim); +} + +float dot_product(const float* a, const float* b, size_t dim) { + return kDotProductFn(a, b, dim); +} + +float dot_product_dis(const float* a, const float* b, size_t dim) { + return kDotProductDisFn(a, b, dim); +} + +float l2norm_sqr(const float* a, size_t dim) { return kL2normSqrFn(a, dim); } + namespace detail { RescaleScratch& get_thread_local_rescale_scratch(size_t dim) { diff --git a/src/simd/space_avx2.cpp b/src/simd/space_avx2.cpp index a618569..be19cb2 100644 --- a/src/simd/space_avx2.cpp +++ b/src/simd/space_avx2.cpp @@ -2,11 +2,29 @@ #include #include +#include #include "rabitqlib/utils/space.hpp" +#include "space_float_kernels.hpp" namespace rabitqlib::simd { +float euclidean_sqr_avx2(const float* a, const float* b, size_t dim) { + return raw_float(a, b, dim); +} + +float dot_product_avx2(const float* a, const float* b, size_t dim) { + return raw_float(a, b, dim); +} + +float dot_product_dis_avx2(const float* a, const float* b, size_t dim) { + return raw_float(a, b, dim); +} + +float l2norm_sqr_avx2(const float* a, size_t dim) { + return raw_float(a, a, dim); +} + void scalar_quantize_uint8_avx2( uint8_t* result, const float* vec0, size_t dim, float lo, float delta ) { @@ -153,38 +171,45 @@ void new_transpose_bin_512_avx2( float mask_ip_x0_q_avx2(const float* query, const uint64_t* data, size_t padded_dim) { const size_t num_blk = padded_dim / 64; - const uint8_t* it_data = reinterpret_cast(data); + const auto* it_data = reinterpret_cast(data); const float* it_query = query; - - __m256 sum = _mm256_setzero_ps(); - - __m256i bit_checker = _mm256_set_epi32(0x80, 0x40, 0x20, 0x10, 0x08, 0x04, 0x02, 0x01); + const __m256i shifts0 = _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7); + const __m256i shifts1 = _mm256_setr_epi32(8, 9, 10, 11, 12, 13, 14, 15); + const __m256i shifts2 = _mm256_setr_epi32(16, 17, 18, 19, 20, 21, 22, 23); + const __m256i shifts3 = _mm256_setr_epi32(24, 25, 26, 27, 28, 29, 30, 31); + __m256 sum0 = _mm256_setzero_ps(); + __m256 sum1 = _mm256_setzero_ps(); + __m256 sum2 = _mm256_setzero_ps(); + __m256 sum3 = _mm256_setzero_ps(); for (size_t i = 0; i < num_blk; ++i) { - uint64_t bits = reverse_bits_u64(load_unaligned_u64(it_data)); - - // 64 bits / 8 floats = 8 iterations - for (int j = 0; j < 8; ++j) { - uint8_t current_byte = static_cast(bits >> (j * 8)); - __m256i v_byte = _mm256_set1_epi32(current_byte); - __m256i masked_bits = _mm256_and_si256(v_byte, bit_checker); - __m256i mask = _mm256_cmpgt_epi32(masked_bits, _mm256_setzero_si256()); - - __m256 q_vals = _mm256_loadu_ps(it_query); - __m256 masked = _mm256_and_ps(q_vals, _mm256_castsi256_ps(mask)); - - sum = _mm256_add_ps(sum, masked); - - it_query += 8; + // Stored coordinates run from bit 63 to bit 0: high word first, + // with each selected bit shifted into a maskload lane's sign bit. + for (size_t half = 0; half < 2; ++half) { + int32_t word; + std::memcpy(&word, it_data + (1 - half) * sizeof(word), sizeof(word)); + const __m256i bits = _mm256_set1_epi32(word); + sum0 = _mm256_add_ps( + sum0, _mm256_maskload_ps(it_query, _mm256_sllv_epi32(bits, shifts0)) + ); + sum1 = _mm256_add_ps( + sum1, _mm256_maskload_ps(it_query + 8, _mm256_sllv_epi32(bits, shifts1)) + ); + sum2 = _mm256_add_ps( + sum2, _mm256_maskload_ps(it_query + 16, _mm256_sllv_epi32(bits, shifts2)) + ); + sum3 = _mm256_add_ps( + sum3, _mm256_maskload_ps(it_query + 24, _mm256_sllv_epi32(bits, shifts3)) + ); + it_query += 32; } it_data += sizeof(uint64_t); } - float result = 0.0f; - for (int i = 0; i < 8; ++i) { - result += reinterpret_cast(&sum)[i]; - } - return result; + const __m256 sum = _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)); + __m128 lanes = _mm_add_ps(_mm256_castps256_ps128(sum), _mm256_extractf128_ps(sum, 1)); + lanes = _mm_add_ps(lanes, _mm_movehl_ps(lanes, lanes)); + return _mm_cvtss_f32(_mm_add_ss(lanes, _mm_movehdup_ps(lanes))); } } // namespace rabitqlib::simd diff --git a/src/simd/space_avx512.cpp b/src/simd/space_avx512.cpp index 349cc66..ac812e6 100644 --- a/src/simd/space_avx512.cpp +++ b/src/simd/space_avx512.cpp @@ -4,9 +4,26 @@ #include #include "rabitqlib/utils/space.hpp" +#include "space_float_kernels.hpp" namespace rabitqlib::simd { +float euclidean_sqr_avx512(const float* a, const float* b, size_t dim) { + return raw_float(a, b, dim); +} + +float dot_product_avx512(const float* a, const float* b, size_t dim) { + return raw_float(a, b, dim); +} + +float dot_product_dis_avx512(const float* a, const float* b, size_t dim) { + return raw_float(a, b, dim); +} + +float l2norm_sqr_avx512(const float* a, size_t dim) { + return raw_float(a, a, dim); +} + void scalar_quantize_uint8_avx512( uint8_t* result, const float* vec0, size_t dim, float lo, float delta ) { diff --git a/src/simd/space_excode_avx2.cpp b/src/simd/space_excode_avx2.cpp index 6a77fbf..3d478c6 100644 --- a/src/simd/space_excode_avx2.cpp +++ b/src/simd/space_excode_avx2.cpp @@ -1,6 +1,5 @@ #include -#include #include #include @@ -24,19 +23,21 @@ namespace { } } // namespace -// helper function for AVX2 inner product -inline void contribute_ip(__m128i vec, const float* __restrict__ query, __m256& sum) { +// Accumulate the two halves independently to shorten the FMA dependency chain. +inline void contribute_ip( + __m128i vec, const float* __restrict__ query, __m256& sum, __m256& sum_hi +) { __m256 q = _mm256_loadu_ps(query); __m256 cf = _mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(vec)); sum = _mm256_fmadd_ps(q, cf, sum); q = _mm256_loadu_ps(query + 8); cf = _mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(_mm_srli_si128(vec, 8))); - sum = _mm256_fmadd_ps(q, cf, sum); + sum_hi = _mm256_fmadd_ps(q, cf, sum_hi); }; inline void contribute_ip_signed( - __m128i vec, const float* __restrict__ query, __m256& sum + __m128i vec, const float* __restrict__ query, __m256& sum, __m256& sum_hi ) { __m256 q = _mm256_loadu_ps(query); __m256 cf = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(vec)); @@ -44,17 +45,15 @@ inline void contribute_ip_signed( q = _mm256_loadu_ps(query + 8); cf = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128(vec, 8))); - sum = _mm256_fmadd_ps(cf, q, sum); + sum_hi = _mm256_fmadd_ps(cf, q, sum_hi); }; +// Reduce in a tree rather than serially adding eight scalar lanes. inline float mm256_reduce_add_ps(__m256 v) { - std::array accumulator{}; - _mm256_storeu_ps(accumulator.data(), v); - float result = 0.0F; - for (const auto& i : accumulator) { - result += i; - } - return result; + __m128 h = _mm_add_ps(_mm256_castps256_ps128(v), _mm256_extractf128_ps(v, 1)); + h = _mm_add_ps(h, _mm_movehl_ps(h, h)); + h = _mm_add_ss(h, _mm_movehdup_ps(h)); + return _mm_cvtss_f32(h); } // ip16: this function is used to compute inner product of @@ -90,7 +89,7 @@ float ip16_fxu1_avx2( float ip64_fxu2_avx2( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - __m256 sum = _mm256_setzero_ps(); + __m256 sum0 = _mm256_setzero_ps(), sum1 = sum0, sum2 = sum0, sum3 = sum0; float result = 0; const __m128i mask = _mm_set1_epi8(0b00000011); @@ -102,15 +101,17 @@ float ip64_fxu2_avx2( __m128i vec_16_to_31 = _mm_and_si128(_mm_srli_epi16(compact, 2), mask); __m128i vec_32_to_47 = _mm_and_si128(_mm_srli_epi16(compact, 4), mask); __m128i vec_48_to_63 = _mm_and_si128(_mm_srli_epi16(compact, 6), mask); - contribute_ip(vec_00_to_15, &query[i], sum); - contribute_ip(vec_16_to_31, &query[i + 16], sum); - contribute_ip(vec_32_to_47, &query[i + 32], sum); - contribute_ip(vec_48_to_63, &query[i + 48], sum); + contribute_ip(vec_00_to_15, &query[i], sum0, sum1); + contribute_ip(vec_16_to_31, &query[i + 16], sum2, sum3); + contribute_ip(vec_32_to_47, &query[i + 32], sum0, sum1); + contribute_ip(vec_48_to_63, &query[i + 48], sum2, sum3); compact_code += 16; } - result = mm256_reduce_add_ps(sum); + result = mm256_reduce_add_ps( + _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)) + ); return result; } @@ -118,7 +119,7 @@ float ip64_fxu2_avx2( float ip64_fxu3_avx2( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - __m256 sum = _mm256_setzero_ps(); + __m256 sum0 = _mm256_setzero_ps(), sum1 = sum0, sum2 = sum0, sum3 = sum0; float result = 0; const __m128i mask = _mm_set1_epi8(0b11); @@ -149,13 +150,15 @@ float ip64_fxu3_avx2( vec_16_to_31 = _mm_or_si128(top_16_to_31, vec_16_to_31); vec_32_to_47 = _mm_or_si128(top_32_to_47, vec_32_to_47); vec_48_to_63 = _mm_or_si128(top_48_to_63, vec_48_to_63); - contribute_ip(vec_00_to_15, &query[i], sum); - contribute_ip(vec_16_to_31, &query[i + 16], sum); - contribute_ip(vec_32_to_47, &query[i + 32], sum); - contribute_ip(vec_48_to_63, &query[i + 48], sum); + contribute_ip(vec_00_to_15, &query[i], sum0, sum1); + contribute_ip(vec_16_to_31, &query[i + 16], sum2, sum3); + contribute_ip(vec_32_to_47, &query[i + 32], sum0, sum1); + contribute_ip(vec_48_to_63, &query[i + 48], sum2, sum3); } - result = mm256_reduce_add_ps(sum); + result = mm256_reduce_add_ps( + _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)) + ); return result; } @@ -163,21 +166,23 @@ float ip64_fxu3_avx2( float ip16_fxu4_avx2( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - __m256 sum = _mm256_setzero_ps(); + __m256 sum0 = _mm256_setzero_ps(); + __m256 sum1 = _mm256_setzero_ps(); float result = 0.0F; constexpr int64_t kMask = 0x0f0f0f0f0f0f0f0f; for (size_t i = 0; i < dim; i += 16) { - int64_t compact = *reinterpret_cast(compact_code); + int64_t compact; + std::memcpy(&compact, compact_code, sizeof(compact)); int64_t code0 = compact & kMask; int64_t code1 = (compact >> 4) & kMask; __m128i c8 = _mm_set_epi64x(code1, code0); - contribute_ip_signed(c8, &query[i], sum); + contribute_ip_signed(c8, &query[i], sum0, sum1); compact_code += 8; } - result = mm256_reduce_add_ps(sum); + result = mm256_reduce_add_ps(_mm256_add_ps(sum0, sum1)); return result; } @@ -185,7 +190,7 @@ float ip16_fxu4_avx2( float ip64_fxu5_avx2( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - __m256 sum = _mm256_setzero_ps(); + __m256 sum0 = _mm256_setzero_ps(), sum1 = sum0, sum2 = sum0, sum3 = sum0; float result = 0.0F; const __m128i mask = _mm_set1_epi8(0b1111); @@ -220,12 +225,14 @@ float ip64_fxu5_avx2( vec_32_to_47 = _mm_or_si128(top_32_to_47, vec_32_to_47); vec_48_to_63 = _mm_or_si128(top_48_to_63, vec_48_to_63); - contribute_ip(vec_00_to_15, &query[i], sum); - contribute_ip(vec_16_to_31, &query[i + 16], sum); - contribute_ip(vec_32_to_47, &query[i + 32], sum); - contribute_ip(vec_48_to_63, &query[i + 48], sum); + contribute_ip(vec_00_to_15, &query[i], sum0, sum1); + contribute_ip(vec_16_to_31, &query[i + 16], sum2, sum3); + contribute_ip(vec_32_to_47, &query[i + 32], sum0, sum1); + contribute_ip(vec_48_to_63, &query[i + 48], sum2, sum3); } - result = mm256_reduce_add_ps(sum); + result = mm256_reduce_add_ps( + _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)) + ); return result; } @@ -233,7 +240,7 @@ float ip64_fxu5_avx2( float ip64_fxu6_avx2( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - __m256 sum = _mm256_setzero_ps(); + __m256 sum0 = _mm256_setzero_ps(), sum1 = sum0, sum2 = sum0, sum3 = sum0; float result = 0.0F; const __m128i mask6 = _mm_set1_epi8(0b00111111); @@ -257,12 +264,14 @@ float ip64_fxu6_avx2( _mm_srli_epi16(_mm_and_si128(cpt3, mask2), 2) ); - contribute_ip(vec_00_to_15, &query[i], sum); - contribute_ip(vec_16_to_31, &query[i + 16], sum); - contribute_ip(vec_32_to_47, &query[i + 32], sum); - contribute_ip(vec_48_to_63, &query[i + 48], sum); + contribute_ip(vec_00_to_15, &query[i], sum0, sum1); + contribute_ip(vec_16_to_31, &query[i + 16], sum2, sum3); + contribute_ip(vec_32_to_47, &query[i + 32], sum0, sum1); + contribute_ip(vec_48_to_63, &query[i + 48], sum2, sum3); } - result = mm256_reduce_add_ps(sum); + result = mm256_reduce_add_ps( + _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)) + ); return result; } @@ -270,7 +279,7 @@ float ip64_fxu6_avx2( float ip64_fxu7_avx2( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - __m256 sum = _mm256_setzero_ps(); + __m256 sum0 = _mm256_setzero_ps(), sum1 = sum0, sum2 = sum0, sum3 = sum0; float result = 0.0F; const __m128i mask6 = _mm_set1_epi8(0b00111111); @@ -311,13 +320,15 @@ float ip64_fxu7_avx2( vec_32_to_47 = _mm_or_si128(top_32_to_47, vec_32_to_47); vec_48_to_63 = _mm_or_si128(top_48_to_63, vec_48_to_63); - contribute_ip(vec_00_to_15, &query[i], sum); - contribute_ip(vec_16_to_31, &query[i + 16], sum); - contribute_ip(vec_32_to_47, &query[i + 32], sum); - contribute_ip(vec_48_to_63, &query[i + 48], sum); + contribute_ip(vec_00_to_15, &query[i], sum0, sum1); + contribute_ip(vec_16_to_31, &query[i + 16], sum2, sum3); + contribute_ip(vec_32_to_47, &query[i + 32], sum0, sum1); + contribute_ip(vec_48_to_63, &query[i + 48], sum2, sum3); } - result = mm256_reduce_add_ps(sum); + result = mm256_reduce_add_ps( + _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)) + ); return result; } @@ -325,13 +336,14 @@ float ip64_fxu7_avx2( float ip16_fxu8_avx2( const float* __restrict__ query, const uint8_t* __restrict__ code, size_t dim ) { - __m256 sum = _mm256_setzero_ps(); + __m256 sum0 = _mm256_setzero_ps(); + __m256 sum1 = _mm256_setzero_ps(); for (size_t i = 0; i < dim; i += 16) { __m128i c8 = _mm_loadu_si128(reinterpret_cast(code)); - contribute_ip(c8, &query[i], sum); + contribute_ip(c8, &query[i], sum0, sum1); code += 16; } - return mm256_reduce_add_ps(sum); + return mm256_reduce_add_ps(_mm256_add_ps(sum0, sum1)); } } // namespace rabitqlib::simd::excode_ipimpl diff --git a/src/simd/space_excode_avx512.cpp b/src/simd/space_excode_avx512.cpp index a317cc7..e1bc210 100644 --- a/src/simd/space_excode_avx512.cpp +++ b/src/simd/space_excode_avx512.cpp @@ -31,27 +31,31 @@ namespace { float ip16_fxu1_avx512( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - float result = 0; - __m512 sum = _mm512_setzero_ps(); - - for (size_t i = 0; i < dim; i += 16) { - __mmask16 mask = *reinterpret_cast(compact_code); - __m512 q = _mm512_loadu_ps(query); - - sum = _mm512_add_ps(_mm512_maskz_mov_ps(mask, q), sum); - - compact_code += 2; - query += 16; + __m512 sum0 = _mm512_setzero_ps(), sum1 = sum0, sum2 = sum0, sum3 = sum0; + const auto load_query = [&](size_t k) { + __mmask16 mask = 0; + std::memcpy(&mask, compact_code + k / 8, sizeof(mask)); + return _mm512_maskz_loadu_ps(mask, query + k); + }; + size_t i = 0; + for (; i + 64 <= dim; i += 64) { + sum0 = _mm512_add_ps(sum0, load_query(i)); + sum1 = _mm512_add_ps(sum1, load_query(i + 16)); + sum2 = _mm512_add_ps(sum2, load_query(i + 32)); + sum3 = _mm512_add_ps(sum3, load_query(i + 48)); } - result = _mm512_reduce_add_ps(sum); - - return result; + for (; i < dim; i += 16) { + sum0 = _mm512_add_ps(sum0, load_query(i)); + } + return _mm512_reduce_add_ps( + _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)) + ); } float ip64_fxu2_avx512( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - __m512 sum = _mm512_setzero_ps(); + __m512 sum0 = _mm512_setzero_ps(), sum1 = sum0, sum2 = sum0, sum3 = sum0; float result = 0; const __m128i mask = _mm_set1_epi8(0b00000011); @@ -68,24 +72,26 @@ float ip64_fxu2_avx512( q = _mm512_loadu_ps(&query[i]); cf = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(vec_00_to_15)); - sum = _mm512_fmadd_ps(q, cf, sum); + sum0 = _mm512_fmadd_ps(q, cf, sum0); q = _mm512_loadu_ps(&query[i + 16]); cf = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(vec_16_to_31)); - sum = _mm512_fmadd_ps(q, cf, sum); + sum1 = _mm512_fmadd_ps(q, cf, sum1); q = _mm512_loadu_ps(&query[i + 32]); cf = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(vec_32_to_47)); - sum = _mm512_fmadd_ps(q, cf, sum); + sum2 = _mm512_fmadd_ps(q, cf, sum2); q = _mm512_loadu_ps(&query[i + 48]); cf = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(vec_48_to_63)); - sum = _mm512_fmadd_ps(q, cf, sum); + sum3 = _mm512_fmadd_ps(q, cf, sum3); compact_code += 16; } - result = _mm512_reduce_add_ps(sum); + result = _mm512_reduce_add_ps( + _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)) + ); return result; } @@ -152,25 +158,29 @@ float ip64_fxu3_avx512( float ip16_fxu4_avx512( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - __m512 sum = _mm512_setzero_ps(); - - float result = 0.0F; - constexpr int64_t kMask = 0x0f0f0f0f0f0f0f0f; - for (size_t i = 0; i < dim; i += 16) { - int64_t compact = *reinterpret_cast(compact_code); - int64_t code0 = compact & kMask; - int64_t code1 = (compact >> 4) & kMask; - - __m128i c8 = _mm_set_epi64x(code1, code0); - __m512 q = _mm512_loadu_ps(&query[i]); - __m512 cf = _mm512_cvtepi32_ps(_mm512_cvtepi8_epi32(c8)); - sum = _mm512_fmadd_ps(cf, q, sum); - - compact_code += 8; + __m512 sum0 = _mm512_setzero_ps(), sum1 = sum0, sum2 = sum0, sum3 = sum0; + // Each eight-byte block stores dimensions 0-7 in low nibbles and 8-15 in high nibbles. + const auto unpack_code = [&](size_t k) { + __m128i bytes = + _mm_loadl_epi64(reinterpret_cast(compact_code + k / 2)); + __m128i lo = _mm_and_si128(bytes, _mm_set1_epi8(15)); + __m128i hi = _mm_and_si128(_mm_srli_epi16(bytes, 4), _mm_set1_epi8(15)); + __m128i expanded = _mm_unpacklo_epi64(lo, hi); + return _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(expanded)); + }; + size_t i = 0; + for (; i + 64 <= dim; i += 64) { + sum0 = _mm512_fmadd_ps(unpack_code(i), _mm512_loadu_ps(query + i), sum0); + sum1 = _mm512_fmadd_ps(unpack_code(i + 16), _mm512_loadu_ps(query + i + 16), sum1); + sum2 = _mm512_fmadd_ps(unpack_code(i + 32), _mm512_loadu_ps(query + i + 32), sum2); + sum3 = _mm512_fmadd_ps(unpack_code(i + 48), _mm512_loadu_ps(query + i + 48), sum3); } - result = _mm512_reduce_add_ps(sum); - - return result; + for (; i < dim; i += 16) { + sum0 = _mm512_fmadd_ps(unpack_code(i), _mm512_loadu_ps(query + i), sum0); + } + return _mm512_reduce_add_ps( + _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)) + ); } float ip64_fxu5_avx512( @@ -238,7 +248,7 @@ float ip64_fxu5_avx512( float ip64_fxu6_avx512( const float* __restrict__ query, const uint8_t* __restrict__ compact_code, size_t dim ) { - __m512 sum = _mm512_setzero_ps(); + __m512 sum0 = _mm512_setzero_ps(), sum1 = sum0, sum2 = sum0, sum3 = sum0; float result = 0.0F; const __m128i mask6 = _mm_set1_epi8(0b00111111); @@ -267,21 +277,23 @@ float ip64_fxu6_avx512( q = _mm512_loadu_ps(&query[i]); cf = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(vec_00_to_15)); - sum = _mm512_fmadd_ps(q, cf, sum); + sum0 = _mm512_fmadd_ps(q, cf, sum0); q = _mm512_loadu_ps(&query[i + 16]); cf = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(vec_16_to_31)); - sum = _mm512_fmadd_ps(q, cf, sum); + sum1 = _mm512_fmadd_ps(q, cf, sum1); q = _mm512_loadu_ps(&query[i + 32]); cf = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(vec_32_to_47)); - sum = _mm512_fmadd_ps(q, cf, sum); + sum2 = _mm512_fmadd_ps(q, cf, sum2); q = _mm512_loadu_ps(&query[i + 48]); cf = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(vec_48_to_63)); - sum = _mm512_fmadd_ps(q, cf, sum); + sum3 = _mm512_fmadd_ps(q, cf, sum3); } - result = _mm512_reduce_add_ps(sum); + result = _mm512_reduce_add_ps( + _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)) + ); return result; } @@ -358,15 +370,25 @@ float ip64_fxu7_avx512( float ip16_fxu8_avx512( const float* __restrict__ query, const uint8_t* __restrict__ code, size_t dim ) { - __m512 sum = _mm512_setzero_ps(); - for (size_t i = 0; i < dim; i += 16) { - __m128i c8 = _mm_loadu_si128(reinterpret_cast(code)); - __m512 q = _mm512_loadu_ps(&query[i]); - __m512 cf = _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32(c8)); - sum = _mm512_fmadd_ps(cf, q, sum); - code += 16; + __m512 sum0 = _mm512_setzero_ps(), sum1 = sum0, sum2 = sum0, sum3 = sum0; + const auto unpack_code = [&](size_t k) { + return _mm512_cvtepi32_ps(_mm512_cvtepu8_epi32( + _mm_loadu_si128(reinterpret_cast(code + k)) + )); + }; + size_t i = 0; + for (; i + 64 <= dim; i += 64) { + sum0 = _mm512_fmadd_ps(unpack_code(i), _mm512_loadu_ps(query + i), sum0); + sum1 = _mm512_fmadd_ps(unpack_code(i + 16), _mm512_loadu_ps(query + i + 16), sum1); + sum2 = _mm512_fmadd_ps(unpack_code(i + 32), _mm512_loadu_ps(query + i + 32), sum2); + sum3 = _mm512_fmadd_ps(unpack_code(i + 48), _mm512_loadu_ps(query + i + 48), sum3); + } + for (; i < dim; i += 16) { + sum0 = _mm512_fmadd_ps(unpack_code(i), _mm512_loadu_ps(query + i), sum0); } - return _mm512_reduce_add_ps(sum); + return _mm512_reduce_add_ps( + _mm512_add_ps(_mm512_add_ps(sum0, sum1), _mm512_add_ps(sum2, sum3)) + ); } } // namespace rabitqlib::simd::excode_ipimpl diff --git a/src/simd/space_float.cpp b/src/simd/space_float.cpp new file mode 100644 index 0000000..129e94a --- /dev/null +++ b/src/simd/space_float.cpp @@ -0,0 +1,35 @@ +#include "rabitqlib/simd/space_dispatch.hpp" + +namespace rabitqlib::simd { + +// Use wider scalar accumulation to limit rounding error on long inputs. +float euclidean_sqr_generic(const float* a, const float* b, size_t dim) { + double sum = 0; + for (size_t i = 0; i < dim; ++i) { + const double delta = static_cast(a[i]) - b[i]; + sum += delta * delta; + } + return static_cast(sum); +} + +float dot_product_generic(const float* a, const float* b, size_t dim) { + double sum = 0; + for (size_t i = 0; i < dim; ++i) { + sum += static_cast(a[i]) * b[i]; + } + return static_cast(sum); +} + +float dot_product_dis_generic(const float* a, const float* b, size_t dim) { + return 1.0F - dot_product_generic(a, b, dim); +} + +float l2norm_sqr_generic(const float* a, size_t dim) { + double sum = 0; + for (size_t i = 0; i < dim; ++i) { + sum += static_cast(a[i]) * a[i]; + } + return static_cast(sum); +} + +} // namespace rabitqlib::simd diff --git a/src/simd/space_float_kernels.hpp b/src/simd/space_float_kernels.hpp new file mode 100644 index 0000000..8a0a415 --- /dev/null +++ b/src/simd/space_float_kernels.hpp @@ -0,0 +1,89 @@ +#pragma once + +#include + +#include + +namespace rabitqlib::simd { +namespace { +#ifdef __AVX512F__ +using Vec = __m512; +constexpr size_t kWidth = 16; +inline Vec zero() { return _mm512_setzero_ps(); } +inline Vec load(const float* p) { return _mm512_loadu_ps(p); } +inline Vec add(Vec a, Vec b) { return _mm512_add_ps(a, b); } +inline Vec sub(Vec a, Vec b) { return _mm512_sub_ps(a, b); } +inline Vec fma(Vec a, Vec b, Vec c) { return _mm512_fmadd_ps(a, b, c); } +inline float reduce(Vec a) { return _mm512_reduce_add_ps(a); } +#else +using Vec = __m256; +constexpr size_t kWidth = 8; +inline Vec zero() { return _mm256_setzero_ps(); } +inline Vec load(const float* p) { return _mm256_loadu_ps(p); } +inline Vec add(Vec a, Vec b) { return _mm256_add_ps(a, b); } +inline Vec sub(Vec a, Vec b) { return _mm256_sub_ps(a, b); } +inline Vec fma(Vec a, Vec b, Vec c) { return _mm256_fmadd_ps(a, b, c); } +inline float reduce(Vec a) { + __m128 h = _mm_add_ps(_mm256_castps256_ps128(a), _mm256_extractf128_ps(a, 1)); + h = _mm_add_ps(h, _mm_movehl_ps(h, h)); + return _mm_cvtss_f32(_mm_add_ss(h, _mm_movehdup_ps(h))); +} +#endif +inline Vec tail_load(const float* p, size_t remaining) { +#ifdef __AVX512F__ + return _mm512_maskz_loadu_ps(static_cast<__mmask16>((1U << remaining) - 1), p); +#else + __m256i mask = _mm256_cmpgt_epi32( + _mm256_set1_epi32(static_cast(remaining)), + _mm256_setr_epi32(0, 1, 2, 3, 4, 5, 6, 7) + ); + return _mm256_maskload_ps(p, mask); +#endif +} +// Four independent accumulation chains avoid a single FMA dependency bottleneck. +// Partial loads preserve the unpadded-input contract. +enum class FloatOperation { SquaredL2, Dot, InnerProductDistance, SquaredNorm }; + +template +inline float raw_float(const float* a, const float* b, size_t n) { + Vec s0 = zero(), s1 = zero(), s2 = zero(), s3 = zero(); + size_t i = 0; + auto accumulate = [](Vec sum, Vec x, Vec y) { + if constexpr (Op == FloatOperation::SquaredL2) { + Vec d = sub(x, y); + return fma(d, d, sum); + } else if constexpr (Op == FloatOperation::SquaredNorm) + return fma(x, x, sum); + else + return fma(x, y, sum); + }; + auto step = [&](Vec sum, size_t at) { + Vec x = load(a + at); + if constexpr (Op == FloatOperation::SquaredNorm) + return accumulate(sum, x, x); + else + return accumulate(sum, x, load(b + at)); + }; + for (; n - i >= kWidth * 4; i += kWidth * 4) { + s0 = step(s0, i); + s1 = step(s1, i + kWidth); + s2 = step(s2, i + kWidth * 2); + s3 = step(s3, i + kWidth * 3); + } + Vec sum = add(add(s0, s1), add(s2, s3)); + for (; n - i >= kWidth; i += kWidth) + sum = step(sum, i); + if (i < n) { + Vec x = tail_load(a + i, n - i); + if constexpr (Op == FloatOperation::SquaredNorm) + sum = accumulate(sum, x, x); + else + sum = accumulate(sum, x, tail_load(b + i, n - i)); + } + float result = reduce(sum); + if constexpr (Op == FloatOperation::InnerProductDistance) + return 1.0F - result; + return result; +} +} // namespace +} // namespace rabitqlib::simd diff --git a/tests/python/test_ivf.py b/tests/python/test_ivf.py index 583b260..0922d74 100644 --- a/tests/python/test_ivf.py +++ b/tests/python/test_ivf.py @@ -230,19 +230,20 @@ def test_save_load_roundtrip(built_ivf, query_data, tmp_path): np.testing.assert_allclose(dists_orig, dists_load, rtol=1e-5) +@pytest.mark.parametrize("dim", [65, 420, 960]) @pytest.mark.parametrize("metric", ["l2", "ip"]) @pytest.mark.parametrize("high_accuracy", [False, True]) @pytest.mark.parametrize("fast_quantization", [False, True]) -def test_raw_reranking(metric, high_accuracy, fast_quantization, tmp_path): +def test_raw_reranking(dim, metric, high_accuracy, fast_quantization, tmp_path): rng = np.random.default_rng(71) - data = rng.standard_normal((65, 65)).astype(np.float32) + data = rng.standard_normal((65, dim)).astype(np.float32) data[0] = 0 # zero residual, padded dimension, and tail batches - queries = rng.standard_normal((3, 65)).astype(np.float32) - idx = IvfIndex(65, len(data), 3, 32, metric) + queries = rng.standard_normal((3, dim)).astype(np.float32) + idx = IvfIndex(dim, len(data), 3, 32, metric) cluster_ids = np.arange(len(data), dtype=np.uint32) % 2 # empty third cluster idx.build( data, - np.zeros((3, 65), dtype=np.float32), + np.zeros((3, dim), dtype=np.float32), cluster_ids, fast_quantization=fast_quantization, num_threads=2, @@ -258,7 +259,7 @@ def test_raw_reranking(metric, high_accuracy, fast_quantization, tmp_path): idx.save(str(path)) loaded = IvfIndex.load(str(path)) assert idx.nbits == loaded.nbits == 32 - assert loaded.dim == 65 + assert loaded.dim == dim assert loaded.metric == metric for k in (5, len(data)): ids, distances = idx.search(queries, k, 3, high_accuracy, 2) @@ -282,8 +283,8 @@ def test_raw_reranking(metric, high_accuracy, fast_quantization, tmp_path): ) order = np.argsort(cluster_ids, kind="stable") np.testing.assert_array_equal(stored.reshape(original.shape), original[order]) - one_bit = IvfIndex(65, len(data), 3, 1, metric) - one_bit.build(original, np.zeros((3, 65), dtype=np.float32), cluster_ids) + one_bit = IvfIndex(dim, len(data), 3, 1, metric) + one_bit.build(original, np.zeros((3, dim), dtype=np.float32), cluster_ids) one_bit_path = tmp_path / "one-bit.index" one_bit.save(str(one_bit_path)) assert len(payload) - one_bit_path.stat().st_size == raw_bytes + 12 diff --git a/tests/unit/rabitqlib/utils/excode_ip_test.cpp b/tests/unit/rabitqlib/utils/excode_ip_test.cpp new file mode 100644 index 0000000..bf4da7f --- /dev/null +++ b/tests/unit/rabitqlib/utils/excode_ip_test.cpp @@ -0,0 +1,113 @@ +#include + +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/simd/pack_excode_dispatch.hpp" +#include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/utils/cpu_features.hpp" +#include "rabitqlib/utils/space.hpp" + +TEST(ExcodeIp, BackendsMatchScalarAcrossWidthsAndBlockBoundaries) { + using namespace rabitqlib; + if (!cpu::has_avx2()) { + GTEST_SKIP() << "Packed-code tests require AVX2/FMA"; + } + + const std::array avx2_functions{ + simd::excode_ipimpl::ip16_fxu1_avx2, + simd::excode_ipimpl::ip64_fxu2_avx2, + simd::excode_ipimpl::ip64_fxu3_avx2, + simd::excode_ipimpl::ip16_fxu4_avx2, + simd::excode_ipimpl::ip64_fxu5_avx2, + simd::excode_ipimpl::ip64_fxu6_avx2, + simd::excode_ipimpl::ip64_fxu7_avx2, + simd::excode_ipimpl::ip16_fxu8_avx2, + }; + const std::array avx512_functions{ + simd::excode_ipimpl::ip16_fxu1_avx512, + simd::excode_ipimpl::ip64_fxu2_avx512, + simd::excode_ipimpl::ip64_fxu3_avx512, + simd::excode_ipimpl::ip16_fxu4_avx512, + simd::excode_ipimpl::ip64_fxu5_avx512, + simd::excode_ipimpl::ip64_fxu6_avx512, + simd::excode_ipimpl::ip64_fxu7_avx512, + simd::excode_ipimpl::ip16_fxu8_avx512, + }; + using PackFunction = void (*)(const uint8_t*, uint8_t*, size_t); + const std::array pack_functions{ + simd::packing_2bit_excode_avx2, + simd::packing_3bit_excode_avx2, + simd::packing_4bit_excode_avx2, + simd::packing_5bit_excode_avx2, + simd::packing_6bit_excode_avx2, + simd::packing_7bit_excode_avx2, + }; + + for (size_t bits = 1; bits <= 8; ++bits) { + SCOPED_TRACE(bits); + const size_t block_dim = (bits == 1 || bits == 4 || bits == 8) ? 16 : 64; + const auto max_code = static_cast((1U << bits) - 1); + for (size_t dim : std::array{ + 0, 16, 32, 48, 64, 80, 96, 112, 128, 192, 256, 576, 960, 1024, 4096}) { + if (dim % block_dim != 0) { + continue; + } + SCOPED_TRACE(dim); + for (size_t pattern = 0; pattern < 4; ++pattern) { + SCOPED_TRACE(pattern); + std::vector query_storage(dim + 1); + auto* query = query_storage.data() + 1; + std::vector codes(dim); + std::vector storage((dim * bits / 8) + 1, 0); + auto* compact = storage.data() + 1; + double expected = 0; + double sum_abs = 0; + for (size_t i = 0; i < dim; ++i) { + query[i] = + pattern == 0 + ? 0.0F + : static_cast(static_cast(i % 23) - 11) / 7.0F; + if (pattern == 3) { + query[i] = (i % 2 == 0) ? 1.0F : -1.0F; + } + codes[i] = pattern >= 2 + ? max_code + : static_cast((i * 37U + 19U) & max_code); + const double product = static_cast(query[i]) * codes[i]; + expected += product; + sum_abs += std::abs(product); + } + if (bits == 1) { + for (size_t i = 0; i < dim; ++i) { + compact[i / 8] |= static_cast(codes[i] << (i % 8)); + } + } else if (bits == 8) { + std::copy(codes.begin(), codes.end(), compact); + } else if (dim != 0) { + pack_functions[bits - 2](codes.data(), compact, dim); + } + + // Scale by product magnitudes, since cancellation can make the dot + // product itself near zero. Different reduction trees may round + // differently. + const double tolerance = 2e-6 * std::max(1.0, sum_abs); + EXPECT_NEAR( + avx2_functions[bits - 1](query, compact, dim), expected, tolerance + ); + if (cpu::has_avx512_core()) { + EXPECT_NEAR( + avx512_functions[bits - 1](query, compact, dim), expected, tolerance + ); + } + EXPECT_NEAR( + select_excode_ipfunc(bits)(query, compact, dim), expected, tolerance + ); + } + } + } +} diff --git a/tests/unit/rabitqlib/utils/float_distance_test.cpp b/tests/unit/rabitqlib/utils/float_distance_test.cpp new file mode 100644 index 0000000..7db1dd3 --- /dev/null +++ b/tests/unit/rabitqlib/utils/float_distance_test.cpp @@ -0,0 +1,193 @@ +#include + +#include +#include +#include +#include +#include + +#include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/utils/cpu_features.hpp" +#include "rabitqlib/utils/space.hpp" + +#if defined(__linux__) +#include +#include +#endif + +namespace { +using BinaryFn = float (*)(const float*, const float*, size_t); +using NormFn = float (*)(const float*, size_t); +struct Backend { + const char* name; + BinaryFn l2; + BinaryFn dot; + BinaryFn ip; + NormFn norm; +}; + +std::vector backends() { + using namespace rabitqlib; + std::vector result{ + {"public", + euclidean_sqr, + dot_product, + dot_product_dis, + l2norm_sqr}, + {"generic", + simd::euclidean_sqr_generic, + simd::dot_product_generic, + simd::dot_product_dis_generic, + simd::l2norm_sqr_generic}}; + if (cpu::has_avx2()) { + result.push_back( + {"avx2", + simd::euclidean_sqr_avx2, + simd::dot_product_avx2, + simd::dot_product_dis_avx2, + simd::l2norm_sqr_avx2} + ); + } + if (cpu::has_avx512_core()) { + result.push_back( + {"avx512", + simd::euclidean_sqr_avx512, + simd::dot_product_avx512, + simd::dot_product_dis_avx512, + simd::l2norm_sqr_avx512} + ); + } + return result; +} + +void check_reference(const float* a, const float* b, size_t dim) { + long double l2 = 0, dot = 0, norm = 0, absolute_products = 0; + for (size_t i = 0; i < dim; ++i) { + const long double x = a[i], y = b[i]; + l2 += (x - y) * (x - y); + dot += x * y; + norm += x * x; + absolute_products += std::abs(x * y); + } + // Longer float reductions accumulate more rounding error. + const long double tolerance = 2e-6L * std::max(1.0L, dim / 4096.0L); + for (const auto& backend : backends()) { + SCOPED_TRACE(backend.name); + EXPECT_NEAR(backend.l2(a, b, dim), l2, tolerance * std::max(1.0L, l2)); + EXPECT_NEAR( + backend.dot(a, b, dim), dot, tolerance * std::max(1.0L, absolute_products) + ); + EXPECT_NEAR(backend.ip(a, b, dim), 1 - dot, tolerance * (1 + absolute_products)); + EXPECT_NEAR(backend.norm(a, dim), norm, tolerance * std::max(1.0L, norm)); + } +} +} // namespace + +TEST(FloatDistance, BackendsMatchReferenceAcrossDimensionsAndAlignments) { + std::vector dimensions; + for (size_t dim = 0; dim <= 80; ++dim) { + dimensions.push_back(dim); + } + for (size_t dim : + {127, + 128, + 129, + 255, + 256, + 257, + 419, + 420, + 421, + 959, + 960, + 961, + 1023, + 1024, + 1025, + 4096, + 65536}) { + dimensions.push_back(dim); + } + for (size_t dim : dimensions) { + SCOPED_TRACE(dim); + for (size_t offset : {0, 1, 7, 15}) { + std::vector a(std::max(size_t{1}, dim + offset)); + std::vector b(std::max(size_t{1}, dim + offset)); + for (int pattern = 0; pattern < 6; ++pattern) { + SCOPED_TRACE(pattern); + for (size_t i = 0; i < dim; ++i) { + float x = static_cast(static_cast(i % 23) - 11) / 7; + float y = static_cast(static_cast(i % 17) - 8) / 9; + if (pattern == 1) { + y = x; + } else if (pattern == 2) { + y = -x; + } else if (pattern == 3) { + x = (i % 4 < 2) ? 4096.0F : -4096.0F; + y = 1; + } else if (pattern == 4) { + x = std::ldexp(x, 30); + y = std::ldexp(y, 30); + } else if (pattern == 5) { + y = std::nextafter(x, 2.0F); + } + a[i + offset] = x; + b[i + offset] = y; + } + check_reference(a.data() + offset, b.data() + offset, dim); + } + } + } +} + +TEST(FloatDistance, EmptyAndAliasedInputsPreserveDistanceConventions) { + const std::array a{1, -2, 3, -4, 5}; + for (const auto& backend : backends()) { + SCOPED_TRACE(backend.name); + EXPECT_EQ(backend.l2(nullptr, nullptr, 0), 0); + EXPECT_EQ(backend.dot(nullptr, nullptr, 0), 0); + EXPECT_EQ(backend.ip(nullptr, nullptr, 0), 1); + EXPECT_EQ(backend.norm(nullptr, 0), 0); + EXPECT_EQ(backend.l2(a.data(), a.data(), a.size()), 0); + EXPECT_EQ(backend.dot(a.data(), a.data(), a.size()), 55); + EXPECT_EQ(backend.ip(a.data(), a.data(), a.size()), -54); + EXPECT_EQ(backend.norm(a.data(), a.size()), 55); + } +} + +TEST(FloatDistance, DoubleTemplatesRetainDoublePrecision) { + const std::array a{1.0000000001, 2.0000000001, 3.0000000001}; + const std::array b{1, 2, 3}; + EXPECT_NEAR(rabitqlib::euclidean_sqr(a.data(), b.data(), a.size()), 3e-20, 1e-26); + EXPECT_DOUBLE_EQ(rabitqlib::dot_product(a.data(), b.data(), a.size()), 14.0000000006); + EXPECT_DOUBLE_EQ( + rabitqlib::dot_product_dis(a.data(), b.data(), a.size()), -13.0000000006 + ); + EXPECT_DOUBLE_EQ(rabitqlib::l2norm_sqr(a.data(), a.size()), 14.0000000012); +} + +#if defined(__linux__) +TEST(FloatDistance, TailsDoNotReadPastGuardPage) { + const size_t page = static_cast(sysconf(_SC_PAGESIZE)); + void* memory = + mmap(nullptr, page * 2, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0); + ASSERT_NE(memory, MAP_FAILED); + struct Mapping { + void* address; + size_t bytes; + ~Mapping() { munmap(address, bytes); } + } mapping{memory, page * 2}; + auto* end = static_cast(memory) + page; + ASSERT_EQ(mprotect(end, page, PROT_NONE), 0); + for (size_t dim : + {0, 1, 7, 8, 9, 15, 16, 17, 31, 32, 33, 63, 64, 65, 127, 128, 129, 420, 960}) { + SCOPED_TRACE(dim); + ASSERT_LE(dim * sizeof(float), page); + auto* data = reinterpret_cast(end) - dim; + for (size_t i = 0; i < dim; ++i) { + data[i] = static_cast(i % 7); + } + check_reference(data, data, dim); + } +} +#endif diff --git a/tests/unit/rabitqlib/utils/mask_ip_test.cpp b/tests/unit/rabitqlib/utils/mask_ip_test.cpp new file mode 100644 index 0000000..0ca5869 --- /dev/null +++ b/tests/unit/rabitqlib/utils/mask_ip_test.cpp @@ -0,0 +1,83 @@ +#include + +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/simd/space_dispatch.hpp" +#include "rabitqlib/utils/cpu_features.hpp" +#include "rabitqlib/utils/space.hpp" + +TEST(MaskIpX0Q, BackendsMatchScalarAcrossBlocksAndAlignments) { + using namespace rabitqlib; + if (!cpu::has_avx2()) { + GTEST_SKIP() << "Binary dot product tests require AVX2/FMA"; + } + using Function = float (*)(const float*, const uint64_t*, size_t); + std::vector functions{simd::mask_ip_x0_q_avx2, mask_ip_x0_q}; + if (cpu::has_avx512_core()) { + functions.push_back(simd::mask_ip_x0_q_avx512); + } + + for (size_t dim : {0, 64, 128, 192, 256, 448, 576, 1024, 4096}) { + SCOPED_TRACE(dim); + for (size_t pattern = 0; pattern < 5; ++pattern) { + SCOPED_TRACE(pattern); + std::vector query(dim + 1); + std::vector words(dim / 64, 0); + double expected = 0; + double sum_abs = 0; + for (size_t i = 0; i < dim; ++i) { + query[i + 1] = static_cast(static_cast(i % 23) - 11) / 7.0F; + if (pattern == 4) { + query[i + 1] = (i % 2 == 0) ? 65536.0F : -65535.0F; + } + const bool selected = pattern == 1 || pattern == 4 || + (pattern == 2 && i % 3 == 0) || + (pattern == 3 && (i % 32 == 0 || i % 32 == 31)); + if (selected) { + words[i / 64] |= uint64_t{1} << (63 - i % 64); + expected += query[i + 1]; + sum_abs += std::abs(static_cast(query[i + 1])); + } + } + for (size_t offset : {0, 1, 3, 7}) { + SCOPED_TRACE(offset); + std::vector storage(dim / 8 + 8, 0); + if (dim != 0) { + std::memcpy(storage.data() + offset, words.data(), dim / 8); + } + const auto* codes = + reinterpret_cast(storage.data() + offset); + for (auto function : functions) { + const float result = function(query.data() + 1, codes, dim); + EXPECT_NEAR(result, expected, 2e-6 * std::max(1.0, sum_abs)); + if (pattern == 0 || pattern == 4) { + EXPECT_EQ(result, expected); + } + } + } + } + } +} + +TEST(MaskIpX0Q, Avx2PreservesEveryStoredBitPosition) { + using namespace rabitqlib; + if (!cpu::has_avx2()) { + GTEST_SKIP() << "Binary dot product tests require AVX2/FMA"; + } + constexpr size_t dim = 192; + std::vector query(dim); + for (size_t i = 0; i < dim; ++i) { + query[i] = static_cast(i + 1); + } + for (size_t bit = 0; bit < dim; ++bit) { + SCOPED_TRACE(bit); + std::vector words(dim / 64, 0); + words[bit / 64] = uint64_t{1} << (63 - bit % 64); + EXPECT_EQ(simd::mask_ip_x0_q_avx2(query.data(), words.data(), dim), query[bit]); + } +}