From 5c889bab7b81bbccc6c1335db7e0aa74888493d9 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 10:04:07 +0800 Subject: [PATCH 1/5] refactor(simd): share uint8 and uint16 quantizer bodies per backend The AVX2, AVX-512 and NEON uint8 and uint16 scalar quantizers repeated the zero-range fill, reciprocal setup and scalar tail, differing only in the final narrowing and store. Each backend now has one template over the code type. --- src/simd/query_neon.cpp | 45 +++++++++++++++------------------- src/simd/space_avx2.cpp | 51 ++++++++++++++------------------------- src/simd/space_avx512.cpp | 51 ++++++++++++++++----------------------- 3 files changed, 58 insertions(+), 89 deletions(-) diff --git a/src/simd/query_neon.cpp b/src/simd/query_neon.cpp index 67940232..a264e4b8 100644 --- a/src/simd/query_neon.cpp +++ b/src/simd/query_neon.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include "rabitqlib/core/detail/simd/space_dispatch.hpp" @@ -25,13 +26,12 @@ inline uint32x4_t quantize4(const float* data, float32x4_t lo, float32x4_t recip // Callers provide finite values whose rounded codes fit the output type. return vcvtaq_u32_f32(vmulq_f32(vsubq_f32(vld1q_f32(data), lo), reciprocal)); } -} // namespace -void scalar_quantize_uint8_neon( - uint8_t* result, const float* data, size_t dim, float lo, float delta -) { +// Code is uint8_t or uint16_t; only the final narrowing and store differ. +template +void quantize(Code* result, const float* data, size_t dim, float lo, float delta) { if (delta == 0.0F) { - std::fill_n(result, dim, uint8_t{0}); + std::fill_n(result, dim, Code{0}); return; } const float reciprocal = 1.0F / delta; @@ -42,35 +42,28 @@ void scalar_quantize_uint8_neon( vmovn_u32(quantize4(data + i, low, scale)), vmovn_u32(quantize4(data + i + 4, low, scale)) ); - vst1_u8(result + i, vmovn_u16(codes)); + if constexpr (std::is_same_v) { + vst1_u8(result + i, vmovn_u16(codes)); + } else { + vst1q_u16(result + i, codes); + } } for (; i < dim; ++i) { - result[i] = static_cast(std::round((data[i] - lo) * reciprocal)); + result[i] = static_cast(std::round((data[i] - lo) * reciprocal)); } } +} // namespace + +void scalar_quantize_uint8_neon( + uint8_t* result, const float* data, size_t dim, float lo, float delta +) { + quantize(result, data, dim, lo, delta); +} void scalar_quantize_uint16_neon( uint16_t* result, const float* data, size_t dim, float lo, float delta ) { - if (delta == 0.0F) { - std::fill_n(result, dim, uint16_t{0}); - return; - } - const float reciprocal = 1.0F / delta; - const float32x4_t low = vdupq_n_f32(lo), scale = vdupq_n_f32(reciprocal); - size_t i = 0; - for (; i + 8 <= dim; i += 8) { - vst1q_u16( - result + i, - vcombine_u16( - vmovn_u32(quantize4(data + i, low, scale)), - vmovn_u32(quantize4(data + i + 4, low, scale)) - ) - ); - } - for (; i < dim; ++i) { - result[i] = static_cast(std::round((data[i] - lo) * reciprocal)); - } + quantize(result, data, dim, lo, delta); } void new_transpose_bin_neon( diff --git a/src/simd/space_avx2.cpp b/src/simd/space_avx2.cpp index e7d8add1..0a0fbcaa 100644 --- a/src/simd/space_avx2.cpp +++ b/src/simd/space_avx2.cpp @@ -5,6 +5,7 @@ #include #include #include +#include #include "packed_tail_avx2.hpp" #include "rabitqlib/core/detail/simd/space_dispatch.hpp" @@ -43,13 +44,14 @@ float l2norm_sqr_avx2(const float* a, size_t dim) { } namespace { -// NonNegative requires (value - lo) / delta >= 0, as for FastScan LUT entries at -// least their minimum; truncation plus a half test then rounds half away from zero. -// Signed inputs keep round_away_from_zero. Do not replace either with adding 0.5f. -template -void quantize_uint8(uint8_t* result, const float* vec0, size_t dim, float lo, float delta) { +// Code is uint8_t or uint16_t. NonNegative requires (value - lo) / delta >= 0, as for +// FastScan LUT entries at least their minimum; truncation plus a half test then +// rounds half away from zero. Signed inputs keep round_away_from_zero. Do not replace +// either with adding 0.5f. +template +void quantize(Code* result, const float* vec0, size_t dim, float lo, float delta) { if (delta == 0.0F) { - std::fill_n(result, dim, uint8_t{0}); + std::fill_n(result, dim, Code{0}); return; } size_t mul8 = dim - (dim & 0b111); @@ -57,7 +59,6 @@ void quantize_uint8(uint8_t* result, const float* vec0, size_t dim, float lo, fl float one_over_delta = 1.0F / delta; __m256 lo256 = _mm256_set1_ps(lo); __m256 od256 = _mm256_set1_ps(one_over_delta); - __m128i zero = _mm_setzero_si128(); for (; i < mul8; i += 8) { __m256 cur = _mm256_loadu_ps(&vec0[i]); @@ -78,11 +79,15 @@ void quantize_uint8(uint8_t* result, const float* vec0, size_t dim, float lo, fl __m128i lo32 = _mm256_castsi256_si128(i32); __m128i hi32 = _mm256_extracti128_si256(i32, 1); __m128i i16 = _mm_packus_epi32(lo32, hi32); - __m128i i8 = _mm_packus_epi16(i16, zero); - _mm_storel_epi64(reinterpret_cast<__m128i*>(&result[i]), i8); + if constexpr (std::is_same_v) { + __m128i i8 = _mm_packus_epi16(i16, _mm_setzero_si128()); + _mm_storel_epi64(reinterpret_cast<__m128i*>(&result[i]), i8); + } else { + _mm_storeu_si128(reinterpret_cast<__m128i*>(result + i), i16); + } } for (; i < dim; ++i) { - result[i] = static_cast(std::round((vec0[i] - lo) * one_over_delta)); + result[i] = static_cast(std::round((vec0[i] - lo) * one_over_delta)); } } } // namespace @@ -90,39 +95,19 @@ void quantize_uint8(uint8_t* result, const float* vec0, size_t dim, float lo, fl void scalar_quantize_uint8_avx2( uint8_t* result, const float* vec0, size_t dim, float lo, float delta ) { - quantize_uint8(result, vec0, dim, lo, delta); + quantize(result, vec0, dim, lo, delta); } void scalar_quantize_nonnegative_uint8_avx2( uint8_t* result, const float* vec0, size_t dim, float lo, float delta ) { - quantize_uint8(result, vec0, dim, lo, delta); + quantize(result, vec0, dim, lo, delta); } void scalar_quantize_uint16_avx2( uint16_t* result, const float* vec0, size_t dim, float lo, float delta ) { - if (delta == 0.0F) { - std::fill_n(result, dim, uint16_t{0}); - return; - } - size_t mul8 = dim - (dim & 0b111); - size_t i = 0; - float one_over_delta = 1.0F / delta; - __m256 lo256 = _mm256_set1_ps(lo); - __m256 ow256 = _mm256_set1_ps(one_over_delta); - for (; i < mul8; i += 8) { - __m256 cur = _mm256_loadu_ps(&vec0[i]); - cur = _mm256_mul_ps(_mm256_sub_ps(cur, lo256), ow256); - __m256i i32 = _mm256_cvttps_epi32(round_away_from_zero(cur)); - __m128i lo32 = _mm256_castsi256_si128(i32); - __m128i hi32 = _mm256_extracti128_si256(i32, 1); - __m128i i16 = _mm_packus_epi32(lo32, hi32); - _mm_storeu_si128(reinterpret_cast<__m128i*>(result + i), i16); - } - for (; i < dim; ++i) { - result[i] = static_cast(std::round((vec0[i] - lo) * one_over_delta)); - } + quantize(result, vec0, dim, lo, delta); } void new_transpose_bin_avx2( diff --git a/src/simd/space_avx512.cpp b/src/simd/space_avx512.cpp index 98efea36..0594e4ef 100644 --- a/src/simd/space_avx512.cpp +++ b/src/simd/space_avx512.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include "rabitqlib/core/detail/simd/space_dispatch.hpp" #include "space_float_kernels.hpp" @@ -43,13 +44,14 @@ float l2norm_sqr_avx512(const float* a, size_t dim) { } namespace { -// NonNegative requires (value - lo) / delta >= 0, as for FastScan LUT entries at -// least their minimum; truncation plus a half test then rounds half away from zero. -// Signed inputs keep round_away_from_zero. Do not replace either with adding 0.5f. -template -void quantize_uint8(uint8_t* result, const float* vec0, size_t dim, float lo, float delta) { +// Code is uint8_t or uint16_t. NonNegative requires (value - lo) / delta >= 0, as for +// FastScan LUT entries at least their minimum; truncation plus a half test then +// rounds half away from zero. Signed inputs keep round_away_from_zero. Do not replace +// either with adding 0.5f. +template +void quantize(Code* result, const float* vec0, size_t dim, float lo, float delta) { if (delta == 0.0F) { - std::fill_n(result, dim, uint8_t{0}); + std::fill_n(result, dim, Code{0}); return; } size_t mul16 = dim - (dim & 0b1111); @@ -72,11 +74,18 @@ void quantize_uint8(uint8_t* result, const float* vec0, size_t dim, float lo, fl } else { i32 = _mm512_cvttps_epi32(round_away_from_zero(cur)); } - __m128i i8 = _mm512_cvtusepi32_epi8(i32); - _mm_storeu_si128(reinterpret_cast<__m128i*>(&result[i]), i8); + if constexpr (std::is_same_v) { + _mm_storeu_si128( + reinterpret_cast<__m128i*>(&result[i]), _mm512_cvtusepi32_epi8(i32) + ); + } else { + _mm256_storeu_si256( + reinterpret_cast<__m256i*>(&result[i]), _mm512_cvtusepi32_epi16(i32) + ); + } } for (; i < dim; ++i) { - result[i] = static_cast(std::round((vec0[i] - lo) * one_over_delta)); + result[i] = static_cast(std::round((vec0[i] - lo) * one_over_delta)); } } } // namespace @@ -84,37 +93,19 @@ void quantize_uint8(uint8_t* result, const float* vec0, size_t dim, float lo, fl void scalar_quantize_uint8_avx512( uint8_t* result, const float* vec0, size_t dim, float lo, float delta ) { - quantize_uint8(result, vec0, dim, lo, delta); + quantize(result, vec0, dim, lo, delta); } void scalar_quantize_nonnegative_uint8_avx512( uint8_t* result, const float* vec0, size_t dim, float lo, float delta ) { - quantize_uint8(result, vec0, dim, lo, delta); + quantize(result, vec0, dim, lo, delta); } void scalar_quantize_uint16_avx512( uint16_t* result, const float* vec0, size_t dim, float lo, float delta ) { - if (delta == 0.0F) { - std::fill_n(result, dim, uint16_t{0}); - return; - } - size_t mul16 = dim - (dim & 0b1111); - size_t i = 0; - float one_over_delta = 1.0F / delta; - __m512 lo512 = _mm512_set1_ps(lo); - __m512 ow512 = _mm512_set1_ps(one_over_delta); - for (; i < mul16; i += 16) { - __m512 cur = _mm512_loadu_ps(&vec0[i]); - cur = _mm512_mul_ps(_mm512_sub_ps(cur, lo512), ow512); - __m256i i16 = - _mm512_cvtusepi32_epi16(_mm512_cvttps_epi32(round_away_from_zero(cur))); - _mm256_storeu_si256(reinterpret_cast<__m256i*>(&result[i]), i16); - } - for (; i < dim; ++i) { - result[i] = static_cast(std::round((vec0[i] - lo) * one_over_delta)); - } + quantize(result, vec0, dim, lo, delta); } void new_transpose_bin_avx512( From 531b1ad9680154fdebb3db56295c77ad38b0f3b9 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 10:04:07 +0800 Subject: [PATCH 2/5] refactor(simd): share the AVX2 horizontal float sum Three translation-unit copies of the same tree reduction become one internal-linkage helper. It replaces an external-linkage inline copy in an AVX2 translation unit. The generated kernels are unchanged. --- src/simd/packed_tail_avx2.hpp | 14 ++++++++++---- src/simd/space_excode_avx2.cpp | 24 ++++++++---------------- src/simd/space_float_kernels.hpp | 8 +++----- 3 files changed, 21 insertions(+), 25 deletions(-) diff --git a/src/simd/packed_tail_avx2.hpp b/src/simd/packed_tail_avx2.hpp index dda28af7..267e9dfe 100644 --- a/src/simd/packed_tail_avx2.hpp +++ b/src/simd/packed_tail_avx2.hpp @@ -83,6 +83,13 @@ inline void accumulate_excode_tail_avx2( } } +// Sum the eight lanes as a tree rather than serially. +static inline float reduce_add_avx2(__m256 v) { + __m128 lanes = _mm_add_ps(_mm256_castps256_ps128(v), _mm256_extractf128_ps(v, 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 mask_ip_avx2( const float* query, const uint8_t* data, size_t padded_dim ) { @@ -139,9 +146,8 @@ static inline float mask_ip_avx2( sum2 = accumulate(sum2, it_query + 16, bits, shifts2); sum3 = accumulate(sum3, it_query + 24, bits, shifts3); } - 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))); + return reduce_add_avx2( + _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)) + ); } } // namespace rabitqlib::simd::detail diff --git a/src/simd/space_excode_avx2.cpp b/src/simd/space_excode_avx2.cpp index 8d197f6d..a6a5d20c 100644 --- a/src/simd/space_excode_avx2.cpp +++ b/src/simd/space_excode_avx2.cpp @@ -49,14 +49,6 @@ inline void contribute_ip_signed( 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) { - __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_fxu1 consumes groups of 16 float query values and packed 1-bit codes. // This translation unit uses AVX2/FMA; callers should use runtime dispatch. // Packed extra-bit layouts are owned by core/pack_excode.hpp. @@ -81,7 +73,7 @@ float ip16_fxu1_avx2( query += 8; ++compact_code; } - result = mm256_reduce_add_ps(sum); + result = detail::reduce_add_avx2(sum); return result; } @@ -112,7 +104,7 @@ float ip64_fxu2_avx2( simd::detail::accumulate_excode_tail_avx2<2>( query + dim - dim % 64, compact_code, dim % 64, sum0, sum1 ); - result = mm256_reduce_add_ps( + result = detail::reduce_add_avx2( _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)) ); @@ -162,7 +154,7 @@ float ip64_fxu3_avx2( simd::detail::accumulate_excode_tail_avx2<3>( query + dim - dim % 64, compact_code, dim % 64, sum0, sum1 ); - result = mm256_reduce_add_ps( + result = detail::reduce_add_avx2( _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)) ); @@ -188,7 +180,7 @@ float ip16_fxu4_avx2( compact_code += 8; } - result = mm256_reduce_add_ps(_mm256_add_ps(sum0, sum1)); + result = detail::reduce_add_avx2(_mm256_add_ps(sum0, sum1)); return result; } @@ -239,7 +231,7 @@ float ip64_fxu5_avx2( simd::detail::accumulate_excode_tail_avx2<5>( query + dim - dim % 64, compact_code, dim % 64, sum0, sum1 ); - result = mm256_reduce_add_ps( + result = detail::reduce_add_avx2( _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)) ); @@ -281,7 +273,7 @@ float ip64_fxu6_avx2( simd::detail::accumulate_excode_tail_avx2<6>( query + dim - dim % 64, compact_code, dim % 64, sum0, sum1 ); - result = mm256_reduce_add_ps( + result = detail::reduce_add_avx2( _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)) ); @@ -341,7 +333,7 @@ float ip64_fxu7_avx2( simd::detail::accumulate_excode_tail_avx2<7>( query + dim - dim % 64, compact_code, dim % 64, sum0, sum1 ); - result = mm256_reduce_add_ps( + result = detail::reduce_add_avx2( _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)) ); @@ -358,7 +350,7 @@ float ip16_fxu8_avx2( contribute_ip(c8, &query[i], sum0, sum1); code += 16; } - return mm256_reduce_add_ps(_mm256_add_ps(sum0, sum1)); + return detail::reduce_add_avx2(_mm256_add_ps(sum0, sum1)); } } // namespace rabitqlib::simd::excode_ipimpl diff --git a/src/simd/space_float_kernels.hpp b/src/simd/space_float_kernels.hpp index 8a0a4159..95521486 100644 --- a/src/simd/space_float_kernels.hpp +++ b/src/simd/space_float_kernels.hpp @@ -4,6 +4,8 @@ #include +#include "packed_tail_avx2.hpp" + namespace rabitqlib::simd { namespace { #ifdef __AVX512F__ @@ -23,11 +25,7 @@ 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))); -} +inline float reduce(Vec a) { return detail::reduce_add_avx2(a); } #endif inline Vec tail_load(const float* p, size_t remaining) { #ifdef __AVX512F__ From 4073d3e8b6a9e7f5c1deb4f24bec856d4a7877fa Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 10:04:07 +0800 Subject: [PATCH 3/5] refactor(simd): share the rescale search tail and fallback The AVX2 and AVX-512 evaluators repeated the exact event-scale recovery and scalar tail, and four wrappers repeated the -1 fallback convention. Both now live in rescale_search.hpp with internal linkage. --- src/simd/quantization_avx2.cpp | 69 +++++++------------------------- src/simd/quantization_avx512.cpp | 69 +++++++------------------------- src/simd/rescale_search.hpp | 68 +++++++++++++++++++++++++++++++ 3 files changed, 98 insertions(+), 108 deletions(-) diff --git a/src/simd/quantization_avx2.cpp b/src/simd/quantization_avx2.cpp index c85c925f..7d6f065f 100644 --- a/src/simd/quantization_avx2.cpp +++ b/src/simd/quantization_avx2.cpp @@ -117,33 +117,8 @@ static RescaleSearchState evaluate_scale_state_avx2( : static_cast(next_positions[k]); } size_t fk = std::max_element(fs, fs + 4) - fs, nk = std::min_element(ts, ts + 4) - ts; - // Reciprocal products selected the coordinates; division now recovers the - // exact heap event scales, clipped to the configured search interval. - if (fi[fk] != std::numeric_limits::max()) - s.first = std::max( - s.first, - static_cast( - quantized_level_at_scale(magnitudes[fi[fk]], scale, max_code) - ) / magnitudes[fi[fk]] - ); - if (ni[nk] != std::numeric_limits::max()) - s.next = std::min( - s.next, - (quantized_level_at_scale(magnitudes[ni[nk]], scale, max_code) + 1.0) / - magnitudes[ni[nk]] - ); - // Handle the remaining zero to three coordinates with the same definitions. - for (; i < dim; ++i) { - int c = quantized_level_at_scale(magnitudes[i], scale, max_code); - double q = c + .5; - s.numerator += magnitudes[i] * q; - s.squared_norm += q * q; - if (c > 0) - s.first = std::max(s.first, static_cast(c) / magnitudes[i]); - if (c < max_code && magnitudes[i] > 0) - s.next = std::min(s.next, (c + 1.0) / magnitudes[i]); - } - return s; + // Recover exact boundary events and handle the remaining zero to three coordinates. + return finish_scale_state(s, magnitudes, fi[fk], ni[nk], i, dim, max_code, scale); } } // namespace rabitqlib::simd::detail @@ -152,19 +127,9 @@ namespace rabitqlib::simd { double best_rescale_factor_avx2( const float* magnitudes, size_t dim, int max_code, double start, double end ) { - // A negative result asks the caller to rerun the scalar heap search. - double best_t = 0; - return detail::try_find_best_scale_by_interval_search( - magnitudes, - dim, - max_code, - start, - end, - detail::evaluate_scale_state_avx2, - best_t - ) - ? best_t - : -1; + return detail::best_scale_or_fallback( + magnitudes, dim, max_code, start, end, detail::evaluate_scale_state_avx2 + ); } double best_rescale_factor_with_workspace_avx2( const float* magnitudes, @@ -175,19 +140,15 @@ double best_rescale_factor_with_workspace_avx2( double* promoted, double* reciprocals ) { - double best_t = 0; - return detail::try_find_best_scale_by_interval_search( - magnitudes, - dim, - max_code, - start, - end, - detail::evaluate_scale_state_avx2, - best_t, - promoted, - reciprocals - ) - ? best_t - : -1; + return detail::best_scale_or_fallback( + magnitudes, + dim, + max_code, + start, + end, + detail::evaluate_scale_state_avx2, + promoted, + reciprocals + ); } } // namespace rabitqlib::simd diff --git a/src/simd/quantization_avx512.cpp b/src/simd/quantization_avx512.cpp index bab967b9..67deca86 100644 --- a/src/simd/quantization_avx512.cpp +++ b/src/simd/quantization_avx512.cpp @@ -109,33 +109,8 @@ static RescaleSearchState evaluate_scale_state_avx512( nk = bitops::countr_zero32( static_cast(_mm512_cmp_pd_mask(vnext, _mm512_set1_pd(n), _CMP_EQ_OQ)) ); - // Reciprocal products selected the coordinates; division now recovers the - // exact heap event scales, clipped to the configured search interval. - if (fi[fk] != std::numeric_limits::max()) - s.first = std::max( - s.first, - static_cast( - quantized_level_at_scale(magnitudes[fi[fk]], scale, max_code) - ) / magnitudes[fi[fk]] - ); - if (ni[nk] != std::numeric_limits::max()) - s.next = std::min( - s.next, - (quantized_level_at_scale(magnitudes[ni[nk]], scale, max_code) + 1.0) / - magnitudes[ni[nk]] - ); - // Handle the remaining zero to seven coordinates with the same definitions. - for (; i < dim; ++i) { - int c = quantized_level_at_scale(magnitudes[i], scale, max_code); - double q = c + .5; - s.numerator += magnitudes[i] * q; - s.squared_norm += q * q; - if (c > 0) - s.first = std::max(s.first, static_cast(c) / magnitudes[i]); - if (c < max_code && magnitudes[i] > 0) - s.next = std::min(s.next, (c + 1.0) / magnitudes[i]); - } - return s; + // Recover exact boundary events and handle the remaining zero to seven coordinates. + return finish_scale_state(s, magnitudes, fi[fk], ni[nk], i, dim, max_code, scale); } } // namespace rabitqlib::simd::detail @@ -144,19 +119,9 @@ namespace rabitqlib::simd { double best_rescale_factor_avx512( const float* magnitudes, size_t dim, int max_code, double start, double end ) { - // A negative result asks the caller to rerun the scalar heap search. - double best_t = 0; - return detail::try_find_best_scale_by_interval_search( - magnitudes, - dim, - max_code, - start, - end, - detail::evaluate_scale_state_avx512, - best_t - ) - ? best_t - : -1; + return detail::best_scale_or_fallback( + magnitudes, dim, max_code, start, end, detail::evaluate_scale_state_avx512 + ); } double best_rescale_factor_with_workspace_avx512( const float* magnitudes, @@ -167,19 +132,15 @@ double best_rescale_factor_with_workspace_avx512( double* promoted, double* reciprocals ) { - double best_t = 0; - return detail::try_find_best_scale_by_interval_search( - magnitudes, - dim, - max_code, - start, - end, - detail::evaluate_scale_state_avx512, - best_t, - promoted, - reciprocals - ) - ? best_t - : -1; + return detail::best_scale_or_fallback( + magnitudes, + dim, + max_code, + start, + end, + detail::evaluate_scale_state_avx512, + promoted, + reciprocals + ); } } // namespace rabitqlib::simd diff --git a/src/simd/rescale_search.hpp b/src/simd/rescale_search.hpp index b9b2a999..f800e40e 100644 --- a/src/simd/rescale_search.hpp +++ b/src/simd/rescale_search.hpp @@ -38,6 +38,46 @@ struct RescaleSearchState { // compiled with different ISA flags; merging external inline/template symbols // could make the AVX2 backend call an AVX512-compiled helper. +// Recover exact event scales at the boundary coordinates a vector evaluator selected +// (size_t max for none): it compares reciprocal products, while the heap sweep +// divides. Then score coordinates [tail, dim) with the scalar definitions. +static inline RescaleSearchState finish_scale_state( + RescaleSearchState s, + const double* magnitudes, + size_t first_index, + size_t next_index, + size_t tail, + size_t dim, + int max_code, + double scale +) { + constexpr size_t kNone = std::numeric_limits::max(); + if (first_index != kNone) + s.first = std::max( + s.first, + static_cast( + quantized_level_at_scale(magnitudes[first_index], scale, max_code) + ) / magnitudes[first_index] + ); + if (next_index != kNone) + s.next = std::min( + s.next, + (quantized_level_at_scale(magnitudes[next_index], scale, max_code) + 1.0) / + magnitudes[next_index] + ); + for (size_t i = tail; i < dim; ++i) { + int c = quantized_level_at_scale(magnitudes[i], scale, max_code); + double q = c + .5; + s.numerator += magnitudes[i] * q; + s.squared_norm += q * q; + if (c > 0) + s.first = std::max(s.first, static_cast(c) / magnitudes[i]); + if (c < max_code && magnitudes[i] > 0) + s.next = std::min(s.next, (c + 1.0) / magnitudes[i]); + } + return s; +} + // Finite IEEE-754 binary64 neighbors for outward-rounded bounds. // Adjusting the representation leaves the floating-point rounding mode unchanged. static inline double next_double_up(double value) { @@ -334,4 +374,32 @@ static bool try_find_best_scale_by_interval_search( return true; } +// The searched scale, or -1 to ask the caller to rerun the exact heap search. +template +static double best_scale_or_fallback( + const float* input, + size_t dim, + int max_code, + double t_start, + double t_end, + Evaluator evaluator, + double* magnitudes = nullptr, + double* reciprocals = nullptr +) { + double best_t = 0; + return try_find_best_scale_by_interval_search( + input, + dim, + max_code, + t_start, + t_end, + evaluator, + best_t, + magnitudes, + reciprocals + ) + ? best_t + : -1; +} + } // namespace rabitqlib::simd::detail From dcd0736859ffd18afdddee92c3e9601b5a08c221 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 10:04:07 +0800 Subject: [PATCH 4/5] refactor(fastscan): use the generic HACC LUT transfer for AVX2 transfer_lut_hacc_avx2 used no vector instructions and wrote the same Pair32 layout as transfer_lut_hacc_generic, which the NEON backend already uses. The outputs are byte-identical. --- include/rabitqlib/fastscan/detail/simd.hpp | 1 - src/fastscan/fastscan_avx2.cpp | 27 ------------------- src/fastscan/fastscan_dispatch.cpp | 3 ++- .../unit/rabitqlib/fastscan/fastscan_test.cpp | 6 ++--- 4 files changed, 5 insertions(+), 32 deletions(-) diff --git a/include/rabitqlib/fastscan/detail/simd.hpp b/include/rabitqlib/fastscan/detail/simd.hpp index 0dce7bc7..cedeba0f 100644 --- a/include/rabitqlib/fastscan/detail/simd.hpp +++ b/include/rabitqlib/fastscan/detail/simd.hpp @@ -58,7 +58,6 @@ void accumulate_avx2( size_t dim ); void transfer_lut_hacc_generic(const uint16_t* lut, size_t dim, uint8_t* hc_lut); -void transfer_lut_hacc_avx2(const uint16_t* lut, size_t dim, uint8_t* hc_lut); void accumulate_hacc_generic( const uint8_t* __restrict__ codes, const uint8_t* __restrict__ lut, diff --git a/src/fastscan/fastscan_avx2.cpp b/src/fastscan/fastscan_avx2.cpp index 34b03088..4ddc67b7 100644 --- a/src/fastscan/fastscan_avx2.cpp +++ b/src/fastscan/fastscan_avx2.cpp @@ -195,33 +195,6 @@ void accumulate_avx2( } } -void transfer_lut_hacc_avx2(const uint16_t* lut, size_t dim, uint8_t* hc_lut) { - size_t num_codebook = dim >> 2; - - for (size_t i = 0; i < num_codebook; i++) { - constexpr size_t kRegBits = 256; - constexpr size_t kLaneBits = 128; - constexpr size_t kByteBits = 8; - - constexpr size_t kLutPerIter = kRegBits / kLaneBits; - constexpr size_t kCodePerIter = 2 * kRegBits / kByteBits; - constexpr size_t kCodePerLine = kLaneBits / kByteBits; - - uint8_t* fill_lo = - hc_lut + (i / kLutPerIter * kCodePerIter) + ((i % kLutPerIter) * kCodePerLine); - uint8_t* fill_hi = fill_lo + (kRegBits / kByteBits); - - for (size_t j = 0; j < 16; ++j) { - int tmp = lut[j]; - uint8_t lo = static_cast(tmp); - uint8_t hi = static_cast(tmp >> 8); - fill_lo[j] = lo; - fill_hi[j] = hi; - } - lut += 16; - } -} - namespace { void accumulate_hacc_chunk_avx2( diff --git a/src/fastscan/fastscan_dispatch.cpp b/src/fastscan/fastscan_dispatch.cpp index 6379ab2d..b3ba0886 100644 --- a/src/fastscan/fastscan_dispatch.cpp +++ b/src/fastscan/fastscan_dispatch.cpp @@ -37,7 +37,8 @@ constexpr Backend kAvx2Backend{ simd::pack_lut_avx2, simd::pack_lut_range_avx2, rabitqlib::simd::scalar_quantize_nonnegative_uint8_avx2, - simd::transfer_lut_hacc_avx2, + // AVX2 consumes the Pair32 layout the generic producer writes. + simd::transfer_lut_hacc_generic, simd::accumulate_avx2, simd::accumulate_hacc_avx2, rabitqlib::simd::qg_batch_estdist_avx2, diff --git a/tests/unit/rabitqlib/fastscan/fastscan_test.cpp b/tests/unit/rabitqlib/fastscan/fastscan_test.cpp index 7b8af047..65dccb1c 100644 --- a/tests/unit/rabitqlib/fastscan/fastscan_test.cpp +++ b/tests/unit/rabitqlib/fastscan/fastscan_test.cpp @@ -297,7 +297,7 @@ TEST(FastScanPackingTest, AccumulatesReferenceLutValuesOnEverySupportedBackend) SCOPED_TRACE("AVX2"); check_backend( simd::accumulate_avx2, - simd::transfer_lut_hacc_avx2, + simd::transfer_lut_hacc_generic, simd::accumulate_hacc_avx2 ); } @@ -401,7 +401,7 @@ TEST(FastScanHighAccuracyTest, AccumulatesLargeDimensionsWithoutLaneOverflow) { #if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { SCOPED_TRACE("AVX2"); - check_backend(simd::transfer_lut_hacc_avx2, simd::accumulate_hacc_avx2); + check_backend(simd::transfer_lut_hacc_generic, simd::accumulate_hacc_avx2); } #endif #if defined(__x86_64__) || defined(_M_X64) @@ -437,7 +437,7 @@ TEST(FastScanHighAccuracyTest, RejectsAccumulationsExceedingInt32Range) { #if defined(__x86_64__) || defined(_M_X64) if (cpu::has_avx2()) { SCOPED_TRACE("AVX2"); - check_backend(simd::transfer_lut_hacc_avx2, simd::accumulate_hacc_avx2); + check_backend(simd::transfer_lut_hacc_generic, simd::accumulate_hacc_avx2); } #endif #if defined(__x86_64__) || defined(_M_X64) From 89a1f102ddf211f777aca512416c6629af1204e9 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 10:04:07 +0800 Subject: [PATCH 5/5] refactor(simd): resolve NEON or x86 kernels with one macro Thirteen dispatch wrappers repeated an #if block choosing resolve_optional_kernel on ARM64 or RABITQ_RESOLVE elsewhere, with inconsistent ARM64 guards. RABITQ_RESOLVE_NEON keeps each target referencing only its own ISA's symbols; the x86 dispatch code is unchanged. --- src/simd/dispatch.cpp | 127 +++++++++------------------------ src/simd/dispatch_resolver.hpp | 9 +++ 2 files changed, 41 insertions(+), 95 deletions(-) diff --git a/src/simd/dispatch.cpp b/src/simd/dispatch.cpp index 5a96a5a8..923c3527 100644 --- a/src/simd/dispatch.cpp +++ b/src/simd/dispatch.cpp @@ -90,51 +90,36 @@ void pairwise_distances_lower( } float euclidean_sqr(const float* a, const float* b, size_t dim) { -#if defined(__aarch64__) - static const auto kEuclideanSqrFn = - resolve_optional_kernel(euclidean_sqr_neon, euclidean_sqr_generic, cpu::has_neon()); -#else - static const auto kEuclideanSqrFn = - RABITQ_RESOLVE(euclidean_sqr_avx512, euclidean_sqr_avx2, euclidean_sqr_generic); -#endif + static const auto kEuclideanSqrFn = RABITQ_RESOLVE_NEON( + euclidean_sqr_neon, euclidean_sqr_avx512, euclidean_sqr_avx2, euclidean_sqr_generic + ); return kEuclideanSqrFn(a, b, dim); } float dot_product(const float* a, const float* b, size_t dim) { -#if defined(__aarch64__) - static const auto kDotProductFn = - resolve_optional_kernel(dot_product_neon, dot_product_generic, cpu::has_neon()); -#else - static const auto kDotProductFn = - RABITQ_RESOLVE(dot_product_avx512, dot_product_avx2, dot_product_generic); -#endif + static const auto kDotProductFn = RABITQ_RESOLVE_NEON( + dot_product_neon, dot_product_avx512, dot_product_avx2, dot_product_generic + ); return kDotProductFn(a, b, dim); } float dot_product_dis(const float* a, const float* b, size_t dim) { -#if defined(__aarch64__) - static const auto kDotProductDisFn = resolve_optional_kernel( - dot_product_dis_neon, dot_product_dis_generic, cpu::has_neon() + static const auto kDotProductDisFn = RABITQ_RESOLVE_NEON( + dot_product_dis_neon, + dot_product_dis_avx512, + dot_product_dis_avx2, + dot_product_dis_generic ); -#else - static const auto kDotProductDisFn = RABITQ_RESOLVE( - dot_product_dis_avx512, dot_product_dis_avx2, dot_product_dis_generic - ); -#endif return kDotProductDisFn(a, b, dim); } float l2norm_sqr(const float* a, size_t dim) { -#if defined(__aarch64__) - static const auto kL2normSqrFn = - resolve_optional_kernel(l2norm_sqr_neon, l2norm_sqr_generic, cpu::has_neon()); -#else - static const auto kL2normSqrFn = - RABITQ_RESOLVE(l2norm_sqr_avx512, l2norm_sqr_avx2, l2norm_sqr_generic); -#endif + static const auto kL2normSqrFn = RABITQ_RESOLVE_NEON( + l2norm_sqr_neon, l2norm_sqr_avx512, l2norm_sqr_avx2, l2norm_sqr_generic + ); return kL2normSqrFn(a, dim); } @@ -186,13 +171,9 @@ void fht_rotate( float scale_factor, const uint8_t* flip ) { -#if defined(__aarch64__) || defined(_M_ARM64) - static const auto kFhtRotateFn = - resolve_optional_kernel(fht_rotate_neon, fht_rotate_generic, cpu::has_neon()); -#else - static const auto kFhtRotateFn = - RABITQ_RESOLVE(fht_rotate_avx512, fht_rotate_avx2, fht_rotate_generic); -#endif + static const auto kFhtRotateFn = RABITQ_RESOLVE_NEON( + fht_rotate_neon, fht_rotate_avx512, fht_rotate_avx2, fht_rotate_generic + ); kFhtRotateFn(data, rotated_data, dim, padded_dim, trunc_dim, scale_factor, flip); } @@ -256,42 +237,29 @@ static ExcodeIpTable resolve_excode_ip_table() { } void flip_sign(const uint8_t* flip, float* data, size_t dim) { -#if defined(__aarch64__) || defined(_M_ARM64) - static const auto kFlipSignFn = - resolve_optional_kernel(flip_sign_neon, flip_sign_generic, cpu::has_neon()); -#else - static const auto kFlipSignFn = - RABITQ_RESOLVE(flip_sign_avx512, flip_sign_avx2, flip_sign_generic); -#endif + static const auto kFlipSignFn = RABITQ_RESOLVE_NEON( + flip_sign_neon, flip_sign_avx512, flip_sign_avx2, flip_sign_generic + ); kFlipSignFn(flip, data, dim); } void kacs_walk(float* data, size_t len) { -#if defined(__aarch64__) || defined(_M_ARM64) - static const auto kKacsWalkFn = - resolve_optional_kernel(kacs_walk_neon, kacs_walk_generic, cpu::has_neon()); -#else - static const auto kKacsWalkFn = - RABITQ_RESOLVE(kacs_walk_avx512, kacs_walk_avx2, kacs_walk_generic); -#endif + static const auto kKacsWalkFn = RABITQ_RESOLVE_NEON( + kacs_walk_neon, kacs_walk_avx512, kacs_walk_avx2, kacs_walk_generic + ); kKacsWalkFn(data, len); } void scalar_quantize_uint8( uint8_t* result, const float* vec0, size_t dim, float lo, float delta ) { -#if defined(__aarch64__) || defined(_M_ARM64) - static const auto kScalarQuantizeUint8Fn = resolve_optional_kernel( - scalar_quantize_uint8_neon, scalar_quantize_uint8_generic, cpu::has_neon() - ); -#else - static const auto kScalarQuantizeUint8Fn = RABITQ_RESOLVE( + static const auto kScalarQuantizeUint8Fn = RABITQ_RESOLVE_NEON( + scalar_quantize_uint8_neon, scalar_quantize_uint8_avx512, scalar_quantize_uint8_avx2, scalar_quantize_uint8_generic ); -#endif kScalarQuantizeUint8Fn(result, vec0, dim, lo, delta); } @@ -299,17 +267,12 @@ void scalar_quantize_uint8( void scalar_quantize_uint16( uint16_t* result, const float* vec0, size_t dim, float lo, float delta ) { -#if defined(__aarch64__) || defined(_M_ARM64) - static const auto kScalarQuantizeUint16Fn = resolve_optional_kernel( - scalar_quantize_uint16_neon, scalar_quantize_uint16_generic, cpu::has_neon() - ); -#else - static const auto kScalarQuantizeUint16Fn = RABITQ_RESOLVE( + static const auto kScalarQuantizeUint16Fn = RABITQ_RESOLVE_NEON( + scalar_quantize_uint16_neon, scalar_quantize_uint16_avx512, scalar_quantize_uint16_avx2, scalar_quantize_uint16_generic ); -#endif kScalarQuantizeUint16Fn(result, vec0, dim, lo, delta); } @@ -424,17 +387,12 @@ float excode_ipimpl::ip64_fxu7_avx( } void new_transpose_bin(const uint16_t* q, uint64_t* tq, size_t padded_dim, size_t b_query) { -#if defined(__aarch64__) || defined(_M_ARM64) - static const auto kNewTransposeBinFn = simd::resolve_optional_kernel( - simd::new_transpose_bin_neon, simd::new_transpose_bin_generic, cpu::has_neon() - ); -#else - static const auto kNewTransposeBinFn = RABITQ_RESOLVE( + static const auto kNewTransposeBinFn = RABITQ_RESOLVE_NEON( + simd::new_transpose_bin_neon, simd::new_transpose_bin_avx512, simd::new_transpose_bin_avx2, simd::new_transpose_bin_generic ); -#endif kNewTransposeBinFn(q, tq, padded_dim, b_query); } @@ -442,37 +400,23 @@ void new_transpose_bin(const uint16_t* q, uint64_t* tq, size_t padded_dim, size_ void new_transpose_bin_512( const uint8_t* q, uint64_t* tq, size_t padded_dim, size_t b_query ) { -#if defined(__aarch64__) || defined(_M_ARM64) - static const auto kNewTransposeBin512Fn = simd::resolve_optional_kernel( + static const auto kNewTransposeBin512Fn = RABITQ_RESOLVE_NEON( simd::new_transpose_bin_512_neon, - simd::new_transpose_bin_512_generic, - cpu::has_neon() - ); -#else - static const auto kNewTransposeBin512Fn = RABITQ_RESOLVE( simd::new_transpose_bin_512_avx512, simd::new_transpose_bin_512_avx2, simd::new_transpose_bin_512_generic ); -#endif kNewTransposeBin512Fn(q, tq, padded_dim, b_query); } float mask_ip_x0_q(const float* query, const uint8_t* data, size_t padded_dim) { -#if defined(__aarch64__) || defined(_M_ARM64) - static const MaskIpX0QFn kMaskIpX0QFn = simd::resolve_optional_kernel( + static const MaskIpX0QFn kMaskIpX0QFn = RABITQ_RESOLVE_NEON( static_cast(simd::mask_ip_x0_q_neon), - static_cast(simd::mask_ip_x0_q_generic), - cpu::has_neon() - ); -#else - static const MaskIpX0QFn kMaskIpX0QFn = RABITQ_RESOLVE( static_cast(simd::mask_ip_x0_q_avx512), static_cast(simd::mask_ip_x0_q_avx2), static_cast(simd::mask_ip_x0_q_generic) ); -#endif return kMaskIpX0QFn(query, data, padded_dim); } @@ -496,20 +440,13 @@ float warmup_ip_x0_q_512( size_t padded_dim, size_t b_query ) { -#if defined(__aarch64__) || defined(_M_ARM64) - static const WarmupIpX0Q512Fn kWarmupIpX0Q512Fn = simd::resolve_optional_kernel( + static const WarmupIpX0Q512Fn kWarmupIpX0Q512Fn = RABITQ_RESOLVE_NEON( static_cast(simd::warmup_ip_x0_q_512_neon), - static_cast(simd::warmup_ip_x0_q_512_generic), - cpu::has_neon() - ); -#else - static const WarmupIpX0Q512Fn kWarmupIpX0Q512Fn = RABITQ_RESOLVE( static_cast(rabitqlib::simd::warmup_ip_x0_q_512_avx512), static_cast(rabitqlib::simd::warmup_ip_x0_q_512_avx2), static_cast(rabitqlib::simd::warmup_ip_x0_q_512_generic), cpu::has_avx512_popcnt() ); -#endif return kWarmupIpX0Q512Fn(data, query, delta, vl, padded_dim, b_query); } diff --git a/src/simd/dispatch_resolver.hpp b/src/simd/dispatch_resolver.hpp index b00409d1..2722e69f 100644 --- a/src/simd/dispatch_resolver.hpp +++ b/src/simd/dispatch_resolver.hpp @@ -48,3 +48,12 @@ Function resolve_optional_kernel(Function preferred, Function fallback, bool sup #else #define RABITQ_RESOLVE(avx512, avx2, fallback, ...) fallback #endif + +// NEON when available on ARM64, otherwise RABITQ_RESOLVE over the AVX kernels and an +// optional AVX-512 predicate. Each target references only its own ISA's symbols. +#if defined(__aarch64__) || defined(_M_ARM64) +#define RABITQ_RESOLVE_NEON(neon, avx512, avx2, fallback, ...) \ + rabitqlib::simd::resolve_optional_kernel(neon, fallback, rabitqlib::cpu::has_neon()) +#else +#define RABITQ_RESOLVE_NEON(neon, ...) RABITQ_RESOLVE(__VA_ARGS__) +#endif