diff --git a/include/rabitqlib/clustering/qgkmeans.hpp b/include/rabitqlib/clustering/qgkmeans.hpp index 58b6400f..bc6a97a9 100644 --- a/include/rabitqlib/clustering/qgkmeans.hpp +++ b/include/rabitqlib/clustering/qgkmeans.hpp @@ -74,7 +74,8 @@ class QuerySumCache { public: QuerySumCache(size_t padded_dim, const QGKMeansParameters& parameters) { - if (padded_dim >= 1024 && + constexpr size_t kMinCachedDim = 512; + if (padded_dim >= kMinCachedDim && (parameters.niter > 1 || parameters.final_assignment != FinalAssignmentMode::Exact)) { state_ = State::Pending; diff --git a/include/rabitqlib/core/detail/simd/space_dispatch.hpp b/include/rabitqlib/core/detail/simd/space_dispatch.hpp index d79f01c1..c3a0dc86 100644 --- a/include/rabitqlib/core/detail/simd/space_dispatch.hpp +++ b/include/rabitqlib/core/detail/simd/space_dispatch.hpp @@ -154,6 +154,11 @@ void scalar_quantize_uint8_generic( void scalar_quantize_uint8_avx2( uint8_t* result, const float* vec0, size_t dim, float lo, float delta ); +// Only for (value - lo) / delta >= 0, such as FastScan LUT entries at least their +// minimum; there it matches scalar_quantize_uint8_avx2 with fewer operations. +void scalar_quantize_nonnegative_uint8_avx2( + uint8_t* result, const float* vec0, size_t dim, float lo, float delta +); void scalar_quantize_uint16_generic( uint16_t* result, const float* data, size_t dim, float lo, float delta ); @@ -172,6 +177,10 @@ float mask_ip_x0_q_avx512(const float* query, const uint64_t* data, size_t padde void scalar_quantize_uint8_avx512( uint8_t* result, const float* vec0, size_t dim, float lo, float delta ); +// Only for (value - lo) / delta >= 0; see scalar_quantize_nonnegative_uint8_avx2. +void scalar_quantize_nonnegative_uint8_avx512( + uint8_t* result, const float* vec0, size_t dim, float lo, float delta +); void scalar_quantize_uint16_avx512( uint16_t* result, const float* vec0, size_t dim, float lo, float delta ); diff --git a/include/rabitqlib/fastscan/detail/lut.hpp b/include/rabitqlib/fastscan/detail/lut.hpp index 8a876908..32bf4be8 100644 --- a/include/rabitqlib/fastscan/detail/lut.hpp +++ b/include/rabitqlib/fastscan/detail/lut.hpp @@ -9,6 +9,7 @@ #include #include +#include "rabitqlib/fastscan/detail/lut_prepare.hpp" #include "rabitqlib/fastscan/fastscan.hpp" #include "rabitqlib/fastscan/highacc_fastscan.hpp" #include "rabitqlib/utils/space.hpp" @@ -95,10 +96,17 @@ class Lut { uint16_t* lut_u16_scratch, bool use_hacc ) { - fastscan::pack_lut(padded_dim, rotated_query, lut_scratch); T vl_lut; T vr_lut; - data_range(lut_scratch, table_length_, vl_lut, vr_lut); + bool finite_lut = false; + if constexpr (std::is_same_v) { + finite_lut = fastscan::detail::pack_lut_with_range( + padded_dim, rotated_query, lut_scratch, vl_lut, vr_lut + ); + } else { + fastscan::pack_lut(padded_dim, rotated_query, lut_scratch); + data_range(lut_scratch, table_length_, vl_lut, vr_lut); + } if (use_hacc) { delta_ = (vr_lut - vl_lut) / ((1 << kNumBitsHacc) - 1); @@ -106,7 +114,17 @@ class Lut { fastscan::transfer_lut_hacc(lut_u16_scratch, padded_dim, lut_.data()); } else { delta_ = (vr_lut - vl_lut) / ((1 << kNumBits) - 1); - quantize_entries(lut_.data(), lut_scratch, vl_lut); + // finite_lut is only set for float, the only type the uint8 kernels take. + if constexpr (std::is_same_v) { + if (finite_lut) { + fastscan::detail::quantize_lut_uint8( + lut_.data(), lut_scratch, table_length_, vl_lut, delta_ + ); + } + } + if (!finite_lut) { + quantize_entries(lut_.data(), lut_scratch, vl_lut); + } } const size_t num_tables = table_length_ / 16; diff --git a/include/rabitqlib/fastscan/detail/lut_prepare.hpp b/include/rabitqlib/fastscan/detail/lut_prepare.hpp new file mode 100644 index 00000000..5f8df013 --- /dev/null +++ b/include/rabitqlib/fastscan/detail/lut_prepare.hpp @@ -0,0 +1,19 @@ +#pragma once + +#include +#include + +namespace rabitqlib::fastscan::detail { +// Caller validates a positive dimension divisible by 16 and non-overlapping scratch. +// Return whether every generated LUT entry is finite. Exceptional and zero bounds +// retain the original range reduction, including its signed-zero behavior. +bool pack_lut_with_range( + size_t dim, const float* query, float* lut, float& minimum, float& maximum +); + +// Use only after pack_lut_with_range returns true, with its minimum and range delta. +// General quantization and high-accuracy uint16 quantization remain unchanged. +void quantize_lut_uint8( + uint8_t* result, const float* values, size_t count, float lo, float delta +); +} // namespace rabitqlib::fastscan::detail diff --git a/include/rabitqlib/fastscan/detail/simd.hpp b/include/rabitqlib/fastscan/detail/simd.hpp index eaa6ffbf..0dce7bc7 100644 --- a/include/rabitqlib/fastscan/detail/simd.hpp +++ b/include/rabitqlib/fastscan/detail/simd.hpp @@ -24,6 +24,20 @@ void pack_lut_generic(size_t dim, const float* query, float* lut); void pack_lut_avx2(size_t dim, const float* query, float* lut); void pack_lut_avx512(size_t dim, const float* query, float* lut); +// Raw fused kernels; the wrapper preserves exceptional/zero-bound reduction behavior. +void pack_lut_range_generic( + size_t dim, const float* query, float* lut, float& minimum, float& maximum +); +void pack_lut_range_avx2( + size_t dim, const float* query, float* lut, float& minimum, float& maximum +); +void pack_lut_range_avx512( + size_t dim, const float* query, float* lut, float& minimum, float& maximum +); +void pack_lut_range_neon( + size_t dim, const float* query, float* lut, float& minimum, float& maximum +); + // Scalar reference for backend correctness tests; never selected by runtime dispatch. void accumulate_generic( const uint8_t* __restrict__ codes, diff --git a/include/rabitqlib/index/detail/visited_set_hash.hpp b/include/rabitqlib/index/detail/visited_set_hash.hpp index 7fdb39f2..caec2c92 100644 --- a/include/rabitqlib/index/detail/visited_set_hash.hpp +++ b/include/rabitqlib/index/detail/visited_set_hash.hpp @@ -77,7 +77,8 @@ class HashBasedVisitedSet { bit_size++; bit >>= 1; } - size_t bucket_size = static_cast(0x1) << ((bit_size + 4) / 2 + 3); + // Limit per-visit spillover allocations while keeping the direct table sublinear. + size_t bucket_size = static_cast(0x1) << ((bit_size + 4) / 2 + 4); initialize_table(bucket_size); } diff --git a/src/core/api.cpp b/src/core/api.cpp index c3497eba..22afd727 100644 --- a/src/core/api.cpp +++ b/src/core/api.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include #include #include @@ -89,10 +90,19 @@ void check_input(const float* input, size_t dim) { check_buffer( {input, multiply(dim, sizeof(float))}, multiply(dim, sizeof(float)) ); + static_assert( + sizeof(float) == sizeof(uint32_t) && std::numeric_limits::is_iec559 + ); + // An integer reduction lets portable builds check several float32 values + // per instruction without relaxing floating-point compiler options. + uint32_t nonfinite = 0; for (size_t i = 0; i < dim; ++i) { - if (!std::isfinite(input[i])) { - throw std::invalid_argument("Input vectors must contain only finite values"); - } + uint32_t bits; + std::memcpy(&bits, input + i, sizeof(bits)); + nonfinite |= static_cast((bits & 0x7f800000U) == 0x7f800000U); + } + if (nonfinite != 0) { + throw std::invalid_argument("Input vectors must contain only finite values"); } } diff --git a/src/fastscan/fastscan_avx2.cpp b/src/fastscan/fastscan_avx2.cpp index f187e0e9..34b03088 100644 --- a/src/fastscan/fastscan_avx2.cpp +++ b/src/fastscan/fastscan_avx2.cpp @@ -11,7 +11,26 @@ namespace rabitqlib::fastscan::simd { -void pack_lut_avx2(size_t dim, const float* query, float* lut) { +namespace { +// Evaluated at compile time: unoptimized builds would otherwise export these library +// functions compiled with this file's ISA flags for callers on other CPUs. +constexpr float kInfinity = std::numeric_limits::infinity(); +constexpr float kQuietNaN = std::numeric_limits::quiet_NaN(); + +template +void pack_lut_impl( + size_t dim, const float* query, float* lut, float& minimum, float& maximum +) { + auto min_values = _mm256_set1_ps(kInfinity), max_values = _mm256_set1_ps(-kInfinity); + unsigned nan_mask = 0; + const auto update_range = [&](__m256 values) { + min_values = _mm256_min_ps(min_values, values); + max_values = _mm256_max_ps(max_values, values); + nan_mask |= static_cast( + _mm256_movemask_ps(_mm256_cmp_ps(values, values, _CMP_UNORD_Q)) + ); + }; + // Keep the initial +0 addition observable for signed-zero inputs. volatile float zero = 0; for (size_t group = 0; group < dim / 4; ++group) { @@ -27,11 +46,43 @@ void pack_lut_avx2(size_t dim, const float* query, float* lut) { hi = _mm256_blend_ps(hi, _mm256_add_ps(hi, q2), 0xCC); lo = _mm256_blend_ps(lo, _mm256_add_ps(lo, q3), 0xAA); hi = _mm256_blend_ps(hi, _mm256_add_ps(hi, q3), 0xAA); + if constexpr (ComputeRange) { + update_range(lo); + update_range(hi); + } _mm256_storeu_ps(lut, lo); _mm256_storeu_ps(lut + 8, hi); query += 4; lut += 16; } + if constexpr (ComputeRange) { + float min_lanes[8], max_lanes[8]; + _mm256_storeu_ps(min_lanes, min_values); + _mm256_storeu_ps(max_lanes, max_values); + // Plain comparisons avoid std algorithms, for the reason given at kInfinity. + minimum = min_lanes[0]; + maximum = max_lanes[0]; + for (size_t lane = 1; lane < 8; ++lane) { + minimum = min_lanes[lane] < minimum ? min_lanes[lane] : minimum; + maximum = max_lanes[lane] > maximum ? max_lanes[lane] : maximum; + } + if (nan_mask) + minimum = kQuietNaN; + } +} +} // namespace + +void pack_lut_avx2(size_t dim, const float* query, float* lut) { + float unused_minimum = 0, unused_maximum = 0; + pack_lut_impl(dim, query, lut, unused_minimum, unused_maximum); +} + +// Preserve the existing ordered subset additions; reduce their values in registers. +// A NaN minimum signals that the wrapper must use the original reduction. +void pack_lut_range_avx2( + size_t dim, const float* query, float* lut, float& minimum, float& maximum +) { + pack_lut_impl(dim, query, lut, minimum, maximum); } namespace { diff --git a/src/fastscan/fastscan_avx512.cpp b/src/fastscan/fastscan_avx512.cpp index e4e48f7f..39117296 100644 --- a/src/fastscan/fastscan_avx512.cpp +++ b/src/fastscan/fastscan_avx512.cpp @@ -11,7 +11,19 @@ namespace rabitqlib::fastscan::simd { -void pack_lut_avx512(size_t dim, const float* query, float* lut) { +namespace { +// Evaluated at compile time: unoptimized builds would otherwise export these library +// functions compiled with this file's ISA flags for callers on other CPUs. +constexpr float kInfinity = std::numeric_limits::infinity(); +constexpr float kQuietNaN = std::numeric_limits::quiet_NaN(); + +template +void pack_lut_impl( + size_t dim, const float* query, float* lut, float& minimum, float& maximum +) { + auto min_values = _mm512_set1_ps(kInfinity), max_values = _mm512_set1_ps(-kInfinity); + unsigned nan_mask = 0; + // Keep the initial +0 addition observable for signed-zero inputs. volatile float zero = 0; for (size_t group = 0; group < dim / 4; ++group) { @@ -22,10 +34,33 @@ void pack_lut_avx512(size_t dim, const float* query, float* lut) { values = _mm512_mask_add_ps(values, 0xF0F0, values, _mm512_set1_ps(query[1])); values = _mm512_mask_add_ps(values, 0xCCCC, values, _mm512_set1_ps(query[2])); values = _mm512_mask_add_ps(values, 0xAAAA, values, _mm512_set1_ps(query[3])); + if constexpr (ComputeRange) { + min_values = _mm512_min_ps(min_values, values); + max_values = _mm512_max_ps(max_values, values); + nan_mask |= _mm512_cmp_ps_mask(values, values, _CMP_UNORD_Q); + } _mm512_storeu_ps(lut, values); query += 4; lut += 16; } + if constexpr (ComputeRange) { + minimum = nan_mask ? kQuietNaN : _mm512_reduce_min_ps(min_values); + maximum = _mm512_reduce_max_ps(max_values); + } +} +} // namespace + +void pack_lut_avx512(size_t dim, const float* query, float* lut) { + float unused_minimum = 0, unused_maximum = 0; + pack_lut_impl(dim, query, lut, unused_minimum, unused_maximum); +} + +// Preserve the existing ordered subset additions; reduce their values in registers. +// A NaN minimum signals that the wrapper must use the original reduction. +void pack_lut_range_avx512( + size_t dim, const float* query, float* lut, float& minimum, float& maximum +) { + pack_lut_impl(dim, query, lut, minimum, maximum); } namespace { diff --git a/src/fastscan/fastscan_backend.hpp b/src/fastscan/fastscan_backend.hpp index e71a3b06..46b06667 100644 --- a/src/fastscan/fastscan_backend.hpp +++ b/src/fastscan/fastscan_backend.hpp @@ -1,5 +1,6 @@ #pragma once +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" #include "rabitqlib/fastscan/detail/estimator_dispatch.hpp" #include "rabitqlib/fastscan/detail/simd.hpp" @@ -16,6 +17,8 @@ struct Backend { const char* name; HaccLutLayout hacc_layout; decltype(&simd::pack_lut_generic) pack_lut; + decltype(&simd::pack_lut_range_generic) pack_lut_range; + decltype(&rabitqlib::simd::scalar_quantize_uint8_generic) quantize_lut; decltype(&simd::transfer_lut_hacc_generic) transfer_hacc; decltype(&simd::accumulate_generic) accumulate; decltype(&simd::accumulate_hacc_generic) accumulate_hacc; diff --git a/src/fastscan/fastscan_dispatch.cpp b/src/fastscan/fastscan_dispatch.cpp index 46634c9a..6379ab2d 100644 --- a/src/fastscan/fastscan_dispatch.cpp +++ b/src/fastscan/fastscan_dispatch.cpp @@ -1,3 +1,4 @@ +#include #include #include #include @@ -5,9 +6,11 @@ #include "../simd/dispatch_resolver.hpp" #include "fastscan_backend.hpp" #include "rabitqlib/fastscan/detail/estimator_dispatch.hpp" +#include "rabitqlib/fastscan/detail/lut_prepare.hpp" #include "rabitqlib/fastscan/detail/simd.hpp" #include "rabitqlib/fastscan/fastscan.hpp" #include "rabitqlib/fastscan/highacc_fastscan.hpp" +#include "rabitqlib/utils/space.hpp" #if defined(__aarch64__) || defined(_M_ARM64) #include "rabitqlib/utils/cpu_features.hpp" #endif @@ -18,6 +21,8 @@ constexpr Backend kGenericBackend{ "generic", HaccLutLayout::Pair32, simd::pack_lut_generic, + simd::pack_lut_range_generic, + rabitqlib::simd::scalar_quantize_uint8_generic, simd::transfer_lut_hacc_generic, simd::accumulate_unsupported, simd::accumulate_hacc_generic, @@ -30,6 +35,8 @@ constexpr Backend kAvx2Backend{ "avx2", HaccLutLayout::Pair32, simd::pack_lut_avx2, + simd::pack_lut_range_avx2, + rabitqlib::simd::scalar_quantize_nonnegative_uint8_avx2, simd::transfer_lut_hacc_avx2, simd::accumulate_avx2, simd::accumulate_hacc_avx2, @@ -40,6 +47,8 @@ constexpr Backend kAvx512Backend{ "avx512", HaccLutLayout::Quad64, simd::pack_lut_avx512, + simd::pack_lut_range_avx512, + rabitqlib::simd::scalar_quantize_nonnegative_uint8_avx512, simd::transfer_lut_hacc_avx512, simd::accumulate_avx512, simd::accumulate_hacc_avx512, @@ -52,6 +61,8 @@ constexpr Backend kNeonBackend{ "neon", HaccLutLayout::Pair32, simd::pack_lut_neon, + simd::pack_lut_range_neon, + rabitqlib::simd::scalar_quantize_uint8_neon, simd::transfer_lut_hacc_generic, simd::accumulate_neon, simd::accumulate_hacc_neon, @@ -72,6 +83,26 @@ RABITQ_COLD_INIT const Backend& selected_backend() { #endif return *backend; } + +bool pack_lut_with_range( + size_t dim, const float* query, float* lut, float& minimum, float& maximum +) { + static const auto kPackLutRangeFn = selected_backend().pack_lut_range; + kPackLutRangeFn(dim, query, lut, minimum, maximum); + // Raw kernels signal any NaN via minimum; infinities remain in the extrema. + const bool finite = std::isfinite(minimum) && std::isfinite(maximum); + if (!finite || minimum == 0.0F || maximum == 0.0F) { + data_range(lut, dim * 4, minimum, maximum); + } + return finite; +} + +void quantize_lut_uint8( + uint8_t* result, const float* values, size_t count, float lo, float delta +) { + static const auto kQuantizeLutFn = selected_backend().quantize_lut; + kQuantizeLutFn(result, values, count, lo, delta); +} } // namespace rabitqlib::fastscan::detail namespace rabitqlib::simd { diff --git a/src/fastscan/fastscan_generic.cpp b/src/fastscan/fastscan_generic.cpp index a093f5ea..d39908e6 100644 --- a/src/fastscan/fastscan_generic.cpp +++ b/src/fastscan/fastscan_generic.cpp @@ -1,3 +1,5 @@ +#include +#include #include #include #include @@ -8,7 +10,15 @@ #include "rabitqlib/fastscan/pack_codes.hpp" namespace rabitqlib::fastscan::simd { -void pack_lut_generic(size_t dim, const float* query, float* lut) { +namespace { +template +void pack_lut_impl( + size_t dim, const float* query, float* lut, float& minimum, float& maximum +) { + float min_values = std::numeric_limits::infinity(), + max_values = -std::numeric_limits::infinity(); + bool nan_mask = false; + // Keep the initial +0 addition observable; MSVC otherwise turns +0 + -0 into -0. volatile float zero = 0; for (size_t group = 0; group < dim / 4; ++group) { @@ -17,9 +27,34 @@ void pack_lut_generic(size_t dim, const float* query, float* lut) { const size_t lowbit = j & (size_t{0} - j); lut[j] = lut[j - lowbit] + query[kPos[j]]; } + if constexpr (ComputeRange) { + for (size_t j = 0; j < 16; ++j) { + min_values = std::min(min_values, lut[j]); + max_values = std::max(max_values, lut[j]); + nan_mask |= std::isnan(lut[j]); + } + } query += 4; lut += 16; } + if constexpr (ComputeRange) { + minimum = nan_mask ? std::numeric_limits::quiet_NaN() : min_values; + maximum = max_values; + } +} +} // namespace + +void pack_lut_generic(size_t dim, const float* query, float* lut) { + float unused_minimum = 0, unused_maximum = 0; + pack_lut_impl(dim, query, lut, unused_minimum, unused_maximum); +} + +// Preserve the existing ordered subset additions; reduce their values in registers. +// A NaN minimum signals that the wrapper must use the original reduction. +void pack_lut_range_generic( + size_t dim, const float* query, float* lut, float& minimum, float& maximum +) { + pack_lut_impl(dim, query, lut, minimum, maximum); } } // namespace rabitqlib::fastscan::simd diff --git a/src/fastscan/fastscan_neon.cpp b/src/fastscan/fastscan_neon.cpp index 0cc93cbd..ba943aba 100644 --- a/src/fastscan/fastscan_neon.cpp +++ b/src/fastscan/fastscan_neon.cpp @@ -85,7 +85,15 @@ void accumulate_hacc_neon( } // namespace rabitqlib::fastscan::simd namespace rabitqlib::fastscan::simd { -void pack_lut_neon(size_t dim, const float* query, float* lut) { +namespace { +template +void pack_lut_impl( + size_t dim, const float* query, float* lut, float& minimum, float& maximum +) { + auto min_values = vdupq_n_f32(std::numeric_limits::infinity()), + max_values = vdupq_n_f32(-std::numeric_limits::infinity()); + auto nan_mask = vdupq_n_u32(0); + // Each table enumerates q0..q3 using bits 3..0. Add in that order, // including the initial +0, to preserve cancellation and signed zeros. volatile float zero = 0; @@ -101,10 +109,34 @@ void pack_lut_neon(size_t dim, const float* query, float* lut) { auto table = tables[i]; table = vbslq_f32(select2, vaddq_f32(table, q2), table); table = vbslq_f32(select3, vaddq_f32(table, q3), table); + if constexpr (ComputeRange) { + min_values = vminq_f32(min_values, table); + max_values = vmaxq_f32(max_values, table); + nan_mask = vorrq_u32(nan_mask, vmvnq_u32(vceqq_f32(table, table))); + } vst1q_f32(lut + i * 4, table); } query += 4; lut += 16; } + if constexpr (ComputeRange) { + minimum = vmaxvq_u32(nan_mask) ? std::numeric_limits::quiet_NaN() + : vminvq_f32(min_values); + maximum = vmaxvq_f32(max_values); + } +} +} // namespace + +void pack_lut_neon(size_t dim, const float* query, float* lut) { + float unused_minimum = 0, unused_maximum = 0; + pack_lut_impl(dim, query, lut, unused_minimum, unused_maximum); +} + +// Preserve the existing ordered subset additions; reduce their values in registers. +// A NaN minimum signals that the wrapper must use the original reduction. +void pack_lut_range_neon( + size_t dim, const float* query, float* lut, float& minimum, float& maximum +) { + pack_lut_impl(dim, query, lut, minimum, maximum); } } // namespace rabitqlib::fastscan::simd diff --git a/src/simd/space_avx2.cpp b/src/simd/space_avx2.cpp index e7f28639..e7d8add1 100644 --- a/src/simd/space_avx2.cpp +++ b/src/simd/space_avx2.cpp @@ -42,9 +42,12 @@ 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 -) { +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) { if (delta == 0.0F) { std::fill_n(result, dim, uint8_t{0}); return; @@ -59,7 +62,19 @@ void scalar_quantize_uint8_avx2( for (; i < mul8; i += 8) { __m256 cur = _mm256_loadu_ps(&vec0[i]); cur = _mm256_mul_ps(_mm256_sub_ps(cur, lo256), od256); - __m256i i32 = _mm256_cvttps_epi32(round_away_from_zero(cur)); + __m256i i32; + if constexpr (NonNegative) { + const __m256 truncated = + _mm256_round_ps(cur, _MM_FROUND_TO_ZERO | _MM_FROUND_NO_EXC); + const __m256 increment = _mm256_cmp_ps( + _mm256_sub_ps(cur, truncated), _mm256_set1_ps(0.5F), _CMP_GE_OQ + ); + i32 = _mm256_sub_epi32( + _mm256_cvttps_epi32(truncated), _mm256_castps_si256(increment) + ); + } else { + 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); @@ -70,6 +85,19 @@ void scalar_quantize_uint8_avx2( result[i] = static_cast(std::round((vec0[i] - lo) * one_over_delta)); } } +} // namespace + +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); +} + +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); +} void scalar_quantize_uint16_avx2( uint16_t* result, const float* vec0, size_t dim, float lo, float delta diff --git a/src/simd/space_avx512.cpp b/src/simd/space_avx512.cpp index 5945e8c6..98efea36 100644 --- a/src/simd/space_avx512.cpp +++ b/src/simd/space_avx512.cpp @@ -42,9 +42,12 @@ 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 -) { +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) { if (delta == 0.0F) { std::fill_n(result, dim, uint8_t{0}); return; @@ -57,13 +60,38 @@ void scalar_quantize_uint8_avx512( for (; i < mul16; i += 16) { __m512 cur = _mm512_loadu_ps(&vec0[i]); cur = _mm512_mul_ps(_mm512_sub_ps(cur, lo512), od512); - __m128i i8 = _mm512_cvtusepi32_epi8(_mm512_cvttps_epi32(round_away_from_zero(cur))); + __m512i i32; + if constexpr (NonNegative) { + const __m512 truncated = + _mm512_roundscale_ps(cur, _MM_FROUND_TO_ZERO | _MM_FROUND_NO_EXC); + const __mmask16 increment = _mm512_cmp_ps_mask( + _mm512_sub_ps(cur, truncated), _mm512_set1_ps(0.5F), _CMP_GE_OQ + ); + const __m512i integer = _mm512_cvttps_epi32(truncated); + i32 = _mm512_mask_add_epi32(integer, increment, integer, _mm512_set1_epi32(1)); + } 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); } for (; i < dim; ++i) { result[i] = static_cast(std::round((vec0[i] - lo) * one_over_delta)); } } +} // namespace + +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); +} + +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); +} void scalar_quantize_uint16_avx512( uint16_t* result, const float* vec0, size_t dim, float lo, float delta diff --git a/tests/python/test_qgkmeans.py b/tests/python/test_qgkmeans.py index af7deaf4..25980235 100644 --- a/tests/python/test_qgkmeans.py +++ b/tests/python/test_qgkmeans.py @@ -474,15 +474,16 @@ def test_grouped_training_preserves_centroid_means_with_dimension_tail(bits, thr @pytest.mark.parametrize("bits", [0, 4, 8]) @pytest.mark.parametrize("spherical", [False, True]) -def test_refit_recomputes_sums_for_mutated_storage_seed_and_count(bits, spherical): +@pytest.mark.parametrize("dim", [511, 960, 1024]) +def test_refit_recomputes_sums_for_mutated_storage_seed_and_count(bits, spherical, dim): rng = np.random.default_rng(823) - storage = np.zeros((257, 1024), dtype=np.float32) + storage = np.zeros((257, dim), dtype=np.float32) storage[:, :8] = rng.standard_normal((257, 8)).astype(np.float32) original_pointer = storage.ctypes.data def make_model(seed): return QGKMeans( - 1024, + dim, 65, niter=3, num_threads=1, diff --git a/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp b/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp index a600ed28..2d2be63b 100644 --- a/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp +++ b/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp @@ -36,7 +36,10 @@ TEST(QGKMeansQuerySumTest, CachePolicyAndSuccessfulFillControlValidity) { bool enabled; }; for (const auto configuration : - {Case{960, 25, FinalAssignmentMode::Approximate, false}, + {Case{480, 25, FinalAssignmentMode::Approximate, false}, + Case{512, 1, FinalAssignmentMode::Exact, false}, + Case{512, 2, FinalAssignmentMode::Exact, true}, + Case{960, 25, FinalAssignmentMode::Approximate, true}, Case{1024, 1, FinalAssignmentMode::Exact, false}, Case{1024, 1, FinalAssignmentMode::Approximate, true}, Case{1024, 2, FinalAssignmentMode::Exact, true}}) { @@ -77,7 +80,7 @@ TEST(QGKMeansQuerySumTest, CachePolicyAndSuccessfulFillControlValidity) { TEST(QGAssignerTest, CallerOwnedQuerySumsMatchUncachedAcrossRebuilds) { constexpr size_t kClusters = 33, kPoints = 64; - for (size_t dim : {64U, 1024U}) { + for (size_t dim : {64U, 511U, 512U, 960U, 1024U}) { std::vector data(kPoints * dim), centroids(kClusters * dim); for (size_t i = 0; i < data.size(); ++i) { data[i] = std::sin(static_cast(i) * 0.13F); diff --git a/tests/unit/rabitqlib/fastscan/lut_prepare_test.cpp b/tests/unit/rabitqlib/fastscan/lut_prepare_test.cpp new file mode 100644 index 00000000..c7bbcf39 --- /dev/null +++ b/tests/unit/rabitqlib/fastscan/lut_prepare_test.cpp @@ -0,0 +1,262 @@ +#include "rabitqlib/fastscan/detail/lut_prepare.hpp" + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include "rabitqlib/core/detail/simd/space_dispatch.hpp" +#include "rabitqlib/fastscan/detail/lut.hpp" +#include "rabitqlib/fastscan/detail/simd.hpp" +#include "rabitqlib/utils/cpu_features.hpp" + +namespace rabitqlib::fastscan { +namespace { +using Pack = void (*)(size_t, const float*, float*); +using PackRange = void (*)(size_t, const float*, float*, float&, float&); +using Quantize = void (*)(uint8_t*, const float*, size_t, float, float); + +struct BackendCase { + const char* name; + Pack pack; + PackRange pack_range; + Quantize quantize; + Quantize reference_quantize; +}; + +std::vector available_backends() { + std::vector backends{ + {"generic", + simd::pack_lut_generic, + simd::pack_lut_range_generic, + rabitqlib::simd::scalar_quantize_uint8_generic, + rabitqlib::simd::scalar_quantize_uint8_generic}}; +#if defined(__x86_64__) || defined(_M_X64) + if (cpu::has_avx2()) { + backends.push_back( + {"avx2", + simd::pack_lut_avx2, + simd::pack_lut_range_avx2, + rabitqlib::simd::scalar_quantize_nonnegative_uint8_avx2, + rabitqlib::simd::scalar_quantize_uint8_avx2} + ); + } + if (cpu::has_avx512_core()) { + backends.push_back( + {"avx512", + simd::pack_lut_avx512, + simd::pack_lut_range_avx512, + rabitqlib::simd::scalar_quantize_nonnegative_uint8_avx512, + rabitqlib::simd::scalar_quantize_uint8_avx512} + ); + } +#endif +#if defined(__aarch64__) || defined(_M_ARM64) + if (cpu::has_neon()) { + backends.push_back( + {"neon", + simd::pack_lut_neon, + simd::pack_lut_range_neon, + rabitqlib::simd::scalar_quantize_uint8_neon, + rabitqlib::simd::scalar_quantize_uint8_neon} + ); + } +#endif + return backends; +} + +struct SavedEnvironment { + std::fenv_t environment{}; + SavedEnvironment() { std::fegetenv(&environment); } + ~SavedEnvironment() { std::fesetenv(&environment); } +}; + +template +void expect_bits_equal(T a, T b) { + EXPECT_EQ(std::memcmp(&a, &b, sizeof(T)), 0); +} + +TEST(LutPreparationTest, FusedBackendsPreserveOrderedTablesAndFiniteExtrema) { + const float tiny = std::numeric_limits::denorm_min(); + const std::array, 6> patterns{ + {{1, 2, 4, 8}, + {-1, -2, -4, -8}, + {-0.0F, -0.0F, 0, 0}, + {1e20F, tiny, -1e20F, 1}, + {tiny, tiny, -tiny, -tiny}, + {-3.5F, .25F, 2, -.125F}}}; + for (const auto& backend : available_backends()) { + SCOPED_TRACE(backend.name); + for (size_t dim : {4U, 12U, 16U, 64U, 960U, 1536U}) { + SCOPED_TRACE(dim); + for (const auto& pattern : patterns) { + std::vector query(dim + 1); + for (size_t i = 0; i < dim; ++i) + query[i + 1] = pattern[i % 4]; + const auto original = query; + std::vector expected(dim * 4 + 2, 123), actual = expected; + backend.pack(dim, query.data() + 1, expected.data() + 1); + float lo, hi, expected_lo, expected_hi; + data_range(expected.data() + 1, dim * 4, expected_lo, expected_hi); + backend.pack_range(dim, query.data() + 1, actual.data() + 1, lo, hi); + EXPECT_EQ( + std::memcmp( + actual.data(), expected.data(), actual.size() * sizeof(float) + ), + 0 + ); + EXPECT_EQ( + std::memcmp( + query.data(), original.data(), query.size() * sizeof(float) + ), + 0 + ); + // Signed-zero selection is handled by the wrapper, not the raw kernel. + EXPECT_EQ(lo, expected_lo); + EXPECT_EQ(hi, expected_hi); + } + } + } +} + +TEST(LutPreparationTest, WrapperPreservesExceptionalAndSignedZeroBounds) { + const SavedEnvironment saved; + for (int rounding : {FE_TONEAREST, FE_DOWNWARD, FE_UPWARD, FE_TOWARDZERO}) { + ASSERT_EQ(std::fesetround(rounding), 0); + for (float value : + {0.0F, + -0.0F, + std::numeric_limits::infinity(), + -std::numeric_limits::infinity(), + std::numeric_limits::quiet_NaN()}) { + std::array query{}; + query.fill(value); + std::array expected{}, actual{}; + expected.front() = actual.front() = 19; + expected.back() = actual.back() = 23; + pack_lut(query.size(), query.data(), expected.data() + 1); + float lo, hi, expected_lo, expected_hi; + data_range(expected.data() + 1, size_t{64}, expected_lo, expected_hi); + const bool finite = detail::pack_lut_with_range( + query.size(), query.data(), actual.data() + 1, lo, hi + ); + EXPECT_EQ( + finite, std::all_of(expected.begin() + 1, expected.end() - 1, [](float x) { + return std::isfinite(x); + }) + ); + EXPECT_EQ(std::memcmp(actual.data(), expected.data(), sizeof(actual)), 0); + expect_bits_equal(lo, expected_lo); + expect_bits_equal(hi, expected_hi); + } + } +} + +TEST(LutPreparationTest, Uint8PathPreservesHalfBoundariesFlagsAndZeroRange) { + const SavedEnvironment saved; + std::vector boundaries{0, -0.0F, 255}; + for (int value = 0; value < 255; ++value) { + const float half = static_cast(value) + .5F; + boundaries.push_back(std::nextafter(half, -std::numeric_limits::infinity())); + boundaries.push_back(half); + boundaries.push_back(std::nextafter(half, std::numeric_limits::infinity())); + } + for (const auto& backend : available_backends()) { + SCOPED_TRACE(backend.name); + for (int rounding : {FE_TONEAREST, FE_DOWNWARD, FE_UPWARD, FE_TOWARDZERO}) { + ASSERT_EQ(std::fesetround(rounding), 0); + for (size_t count : {0U, 1U, 7U, 8U, 15U, 16U, 17U, 768U}) { + std::vector values(count + 1); + for (size_t i = 0; i < count; ++i) + values[i + 1] = boundaries[i % boundaries.size()]; + std::vector expected(count + 2, 71), actual = expected; + std::feclearexcept(FE_ALL_EXCEPT); + backend.reference_quantize( + expected.data() + 1, values.data() + 1, count, 0, 1 + ); + const int expected_flags = std::fetestexcept(FE_ALL_EXCEPT); + std::feclearexcept(FE_ALL_EXCEPT); + backend.quantize(actual.data() + 1, values.data() + 1, count, 0, 1); + const int actual_flags = std::fetestexcept(FE_ALL_EXCEPT); + EXPECT_EQ(actual, expected); + EXPECT_EQ(actual_flags, expected_flags); + // Round-half-away results are promised only in the documented + // round-to-nearest environment: Clang inlines the scalar-tail + // std::round in SIMD translation units with a mode-dependent addition. + if (rounding == FE_TONEAREST) { + for (size_t i = 0; i < count; ++i) { + EXPECT_EQ( + actual[i + 1], static_cast(std::round(values[i + 1])) + ); + } + } + for (float delta : {0.0F, -0.0F}) { + backend.quantize( + actual.data() + 1, values.data() + 1, count, -0.0F, delta + ); + EXPECT_EQ(actual.front(), 71); + EXPECT_EQ(actual.back(), 71); + EXPECT_TRUE( + std::all_of(actual.begin() + 1, actual.end() - 1, [](uint8_t x) { + return x == 0; + }) + ); + } + } + } + } +} + +template +void check_prepared_lut() { + for (size_t dim : {16U, 128U, 960U, 1536U, 4096U}) { + for (int pattern = 0; pattern < 4; ++pattern) { + std::vector query(dim), table(dim * 4); + for (size_t i = 0; i < dim; ++i) { + query[i] = pattern == 0 ? T{0} + : pattern == 1 ? -T{0} + : pattern == 2 + ? static_cast(static_cast(i % 19) - 9) / T{7} + : std::numeric_limits::denorm_min() * + static_cast(static_cast(i % 3) - 1); + } + pack_lut(dim, query.data(), table.data()); + T lo, hi; + data_range(table.data(), table.size(), lo, hi); + for (bool hacc : {false, true}) { + const T delta = (hi - lo) / static_cast(hacc ? 65535 : 255); + const T reciprocal = delta == 0 ? T{0} : T{1} / delta; + std::vector expected(table.size() * (hacc ? 2 : 1)); + std::vector codes(table.size()); + for (size_t i = 0; i < table.size(); ++i) { + const T value = + delta == 0 ? T{0} : std::round((table[i] - lo) * reciprocal); + codes[i] = static_cast(value); + if (!hacc) + expected[i] = static_cast(value); + } + if (hacc) + transfer_lut_hacc(codes.data(), dim, expected.data()); + Lut actual(query.data(), dim, hacc); + expect_bits_equal(actual.delta(), delta); + expect_bits_equal(actual.sum_vl(), lo * static_cast(dim / 4)); + EXPECT_EQ(std::memcmp(actual.lut(), expected.data(), expected.size()), 0); + } + } + } +} + +TEST(LutPreparationTest, FloatLutPreservesStandardAndHaccCodes) { + check_prepared_lut(); +} +TEST(LutPreparationTest, DoubleLutRetainsOriginalPreparation) { + check_prepared_lut(); +} +} // namespace +} // namespace rabitqlib::fastscan diff --git a/tests/unit/rabitqlib/quantization/core_api_test.cpp b/tests/unit/rabitqlib/quantization/core_api_test.cpp index e7a9ae89..42980bcc 100644 --- a/tests/unit/rabitqlib/quantization/core_api_test.cpp +++ b/tests/unit/rabitqlib/quantization/core_api_test.cpp @@ -42,6 +42,56 @@ api::CodeView view(api::CodeOutput output) { {output.extra.data, output.extra.bytes}}; } +TEST(CoreApi, RejectsNonFiniteValuesAtEveryInputPosition) { + for (size_t dim : {2U, 3U, 15U, 16U, 17U, 31U, 32U, 33U, 65U}) { + const api::Parameters parameters{dim, 1}; + const auto sizes = api::requirements(parameters, api::CodeLayout::UnpackedU16); + Workspace workspace(sizes.workspace_bytes); + std::vector codes(sizes.code_bytes, 0x5A); + const api::CodeOutput output{ + api::CodeLayout::UnpackedU16, {codes.data(), codes.size()}, {}}; + std::vector values(dim, 0.0F), finite(dim, -0.0F); + for (size_t position = 0; position < dim; ++position) { + for (float bad : + {std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity()}) { + SCOPED_TRACE(::testing::Message() << dim << "/" << position << "/" << bad); + values[position] = bad; + const auto rejects = [](auto operation) { + try { + operation(); + FAIL() << "Non-finite input was accepted"; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ( + error.what(), "Input vectors must contain only finite values" + ); + } + }; + rejects([&] { api::prepare_query(parameters, values.data(), nullptr); }); + rejects([&] { + api::prepare_query(parameters, finite.data(), values.data()); + }); + rejects([&] { + api::encode( + parameters, values.data(), nullptr, output, workspace.view() + ); + }); + rejects([&] { + api::encode( + parameters, finite.data(), values.data(), output, workspace.view() + ); + }); + values[position] = 0.0F; + } + } + EXPECT_TRUE(std::all_of(codes.begin(), codes.end(), [](uint8_t v) { + return v == 0x5A; + })); + EXPECT_NO_THROW(api::prepare_query(parameters, finite.data(), nullptr)); + } +} + TEST(CoreApi, ExactWorkspaceSizeProtectsBoundaryBytes) { for (size_t dim : {2U, 3U, 31U, 33U, 65U, 96U}) { for (size_t bits : {1U, 4U, 9U}) {