From 4201e679efd1a6e80aaa5e7db3da27c08bdd48e5 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 21:33:46 +0800 Subject: [PATCH 1/2] fix(core): validate advanced encoder preconditions code_factor_terms divides by dim - 1, so the advanced encoders returned a NaN error factor for dim == 1 (the stable API and the indexes already require more). rabitq_full_impl accepted a code element type too narrow for total_bits, truncating the sign bit for uint8_t at nine bits. SplitSingleQuery checked bit widths but not that padded_dim is a multiple of 32, although the AVX2 and AVX-512 transposition kernels load 32-byte blocks. All three now throw std::invalid_argument. The non-template helpers in utils/space.hpp gain internal linkage, as the implementation-linkage rule for headers reached from ISA-flagged sources requires. --- docs/docs/rabitq/quantizer.md | 6 ++ include/rabitqlib/core/detail/rabitq_impl.hpp | 13 ++++ include/rabitqlib/core/query.hpp | 6 ++ include/rabitqlib/utils/space.hpp | 12 +-- .../rabitqlib/quantization/rabitq_test.cpp | 78 +++++++++++++++++++ 5 files changed, 110 insertions(+), 5 deletions(-) diff --git a/docs/docs/rabitq/quantizer.md b/docs/docs/rabitq/quantizer.md index e0e31d59..03e7fe9f 100644 --- a/docs/docs/rabitq/quantizer.md +++ b/docs/docs/rabitq/quantizer.md @@ -171,6 +171,12 @@ int main() { } ``` +The advanced encoders validate their preconditions and throw +`std::invalid_argument` otherwise: `dim` must be at least 2 (the error factor +divides by `dim - 1`), the code element type must hold `bits` (nine bits need +`uint16_t`), inputs must be finite, and `SplitSingleQuery` needs a padded +dimension that is a positive multiple of 32. + #### Querying Score a query using the encoded vector and its factors. diff --git a/include/rabitqlib/core/detail/rabitq_impl.hpp b/include/rabitqlib/core/detail/rabitq_impl.hpp index 6fef536a..cae633be 100644 --- a/include/rabitqlib/core/detail/rabitq_impl.hpp +++ b/include/rabitqlib/core/detail/rabitq_impl.hpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -50,6 +51,11 @@ inline CodeFactorTerms code_factor_terms( MetricType metric_type, const char* invalid_metric_message ) { + // The error factor divides by dim - 1; the stable API and the indexes already + // enforce this, the advanced encoders reach here directly. + if (dim < 2) { + throw std::invalid_argument("RaBitQ factors require dim >= 2"); + } // A zero residual needs no centered code or inner-product divisions. if (l2_sqr == 0) { if (metric_type == METRIC_L2) { @@ -863,6 +869,13 @@ static inline void rabitq_full_impl_with_scratch( double t_const, Scratch* scratch ) { + // One sign bit plus at most eight magnitude bits (xy_bits::kMaxCombinedBits). + if (total_bits < 1 || total_bits > 9 || + total_bits > static_cast(std::numeric_limits::digits)) { + throw std::invalid_argument( + "total_bits must be in [1, 9] and fit the code element type" + ); + } const size_t ex_bits = total_bits - 1; if (ex_bits > 0) { RowMajorArray owned_residual; diff --git a/include/rabitqlib/core/query.hpp b/include/rabitqlib/core/query.hpp index 647f6b2a..46a5746d 100644 --- a/include/rabitqlib/core/query.hpp +++ b/include/rabitqlib/core/query.hpp @@ -66,6 +66,12 @@ class SplitSingleQuery { "and base_bits + ex_bits <= 9" ); } + // The query-bit transposition kernels load whole 32-coordinate blocks. + if (padded_dim == 0 || padded_dim % 32 != 0) { + throw std::invalid_argument( + "SplitSingleQuery requires padded_dim to be a positive multiple of 32" + ); + } validate_metric_type(metric_type); diff --git a/include/rabitqlib/utils/space.hpp b/include/rabitqlib/utils/space.hpp index 51c0c7f0..e4398658 100644 --- a/include/rabitqlib/utils/space.hpp +++ b/include/rabitqlib/utils/space.hpp @@ -19,7 +19,7 @@ namespace rabitqlib { -inline uint64_t load_unaligned_u64(const void* data) { +static inline uint64_t load_unaligned_u64(const void* data) { uint64_t value = 0; std::memcpy(&value, data, sizeof(value)); return value; @@ -302,11 +302,13 @@ void new_transpose_bin_512( float mask_ip_x0_q(const float* query, const uint8_t* data, size_t padded_dim); float mask_ip_x0_q(const float* query, const uint64_t* data, size_t padded_dim); -inline float mask_ip_x0_q_old(const float* query, const uint64_t* data, size_t padded_dim) { +static inline float mask_ip_x0_q_old( + const float* query, const uint64_t* data, size_t padded_dim +) { return mask_ip_x0_q(query, data, padded_dim); } -inline float ip_x0_q( +static inline float ip_x0_q( const uint64_t* data, const uint64_t* query, float delta, @@ -337,7 +339,7 @@ static inline uint32_t ip_bin_bin(const uint64_t* q, const uint64_t* d, size_t p return ret; } -inline uint32_t ip_byte_bin( +static inline uint32_t ip_byte_bin( const uint64_t* q, const uint64_t* d, size_t padded_dim, size_t b_query ) { uint32_t ret = 0; @@ -349,7 +351,7 @@ inline uint32_t ip_byte_bin( return ret; } -inline size_t popcount(const uint64_t* __restrict__ d, size_t length) { +static inline size_t popcount(const uint64_t* __restrict__ d, size_t length) { size_t ret = 0; const size_t words = length / 64; const auto* data = reinterpret_cast(d); diff --git a/tests/unit/rabitqlib/quantization/rabitq_test.cpp b/tests/unit/rabitqlib/quantization/rabitq_test.cpp index 598e944c..bb91b69e 100644 --- a/tests/unit/rabitqlib/quantization/rabitq_test.cpp +++ b/tests/unit/rabitqlib/quantization/rabitq_test.cpp @@ -6,6 +6,7 @@ #include #include #include +#include #include #include #include @@ -675,6 +676,83 @@ TEST(RabitqSignConventionTest, ZeroResidualDoesNotInflateReconstructionNorm) { << std::to_string(tan_sq); } +TEST(RabitqPreconditionTest, EncodersRejectSingleDimensionInput) { + const std::array data{0.5F}; + const std::array centroid{0.0F}; + std::array code{}; + float f_add = 0; + float f_rescale = 0; + float f_error = 0; + // code_factor_terms divides by dim - 1; every advanced encoder reaches it. + EXPECT_THROW( + quantize_full_single( + data.data(), centroid.data(), 1, 4, code.data(), f_add, f_rescale, f_error + ), + std::invalid_argument + ); + EXPECT_THROW( + quantize_full_single(data.data(), 1, 1, code.data(), f_add, f_rescale, f_error), + std::invalid_argument + ); +} + +TEST(RabitqPreconditionTest, FullCodeRejectsBitWidthsTheElementTypeCannotHold) { + constexpr size_t kDim = 64; + std::array data; + data.fill(0.25F); + const std::array centroid{}; + std::array narrow{}; + std::array wide{}; + float f_add = 0; + float f_rescale = 0; + float f_error = 0; + // Nine bits (sign plus eight magnitude bits) need more than eight bits per code. + EXPECT_THROW( + quantize_full_single( + data.data(), centroid.data(), kDim, 9, narrow.data(), f_add, f_rescale, f_error + ), + std::invalid_argument + ); + EXPECT_NO_THROW(quantize_full_single( + data.data(), centroid.data(), kDim, 8, narrow.data(), f_add, f_rescale, f_error + )); + EXPECT_NO_THROW(quantize_full_single( + data.data(), centroid.data(), kDim, 9, wide.data(), f_add, f_rescale, f_error + )); + for (size_t bits : {0U, 10U}) { + EXPECT_THROW( + quantize_full_single( + data.data(), + centroid.data(), + kDim, + bits, + wide.data(), + f_add, + f_rescale, + f_error + ), + std::invalid_argument + ); + } +} + +TEST(RabitqPreconditionTest, SplitSingleQueryRequiresPaddedDimensionMultipleOf32) { + std::vector query(96, 0.5F); + for (size_t padded_dim : {0U, 16U, 48U, 95U}) { + SCOPED_TRACE(padded_dim); + EXPECT_THROW( + (SplitSingleQuery(query.data(), padded_dim, 3, RabitqConfig{})), + std::invalid_argument + ); + } + for (size_t padded_dim : {32U, 64U, 96U}) { + SCOPED_TRACE(padded_dim); + EXPECT_NO_THROW( + (SplitSingleQuery(query.data(), padded_dim, 3, RabitqConfig{})) + ); + } +} + } // namespace } // namespace rabitqlib::quant From ffe231d8cfd04799d9603dcf6e54d5dbd7dca1b9 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 21:33:46 +0800 Subject: [PATCH 2/2] fix(hnsw): derive prefetch addresses only from live link slots and pool entries search_layer prefetched the first two link slots and the base-layer search the first four regardless of the link count, and the base-layer search prefetched level0_links(search_pool.next_id()) after every insert even when the pool had no next entry, in which case next_id() reads the spare slot that a full insert fills with the displaced entry, possibly carrying the checked bit. Those reads are stale or unwritten IDs, and the resulting pointer arithmetic leaves the allocation. Bound the look-ahead by the link count and guard the pool prefetch with has_next(). SIFT1M (1M x 128, HNSW M=16, ef_construction=200, 4 bits), Xeon Gold 6418H, GCC 14 -O3 portable build, same saved index for both builds: single-thread QPS medians 9568/6432/3910 before and 9487/6416/3898 after at ef 64/128/256 (within the run-to-run spread), 16-thread 84174/50656 before and 84168/50278 after at ef 128/256; recall@10 identical (0.9007/0.9179/0.9230). Construction with 32 threads: 18.5 s and 18.1 s before, 18.5 s and 18.2 s after. --- src/index/hnsw.cpp | 6 ++++-- src/index/hnsw_search_impl.hpp | 13 +++++++++---- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/src/index/hnsw.cpp b/src/index/hnsw.cpp index 984638ee..16719952 100644 --- a/src/index/hnsw.cpp +++ b/src/index/hnsw.cpp @@ -620,8 +620,10 @@ maxheap> HierarchicalNSW::search_layer( const size_t num_neighbors = link_count(links); const PID* neighbors = links + 1; - prefetch(*neighbors); - prefetch(*(neighbors + 1)); + // Slots beyond the link count hold stale or unwritten IDs. + for (size_t j = 0; j < std::min(2, num_neighbors); ++j) { + prefetch(neighbors[j]); + } for (size_t j = 0; j < num_neighbors; ++j) { PID candidate_id = neighbors[j]; diff --git a/src/index/hnsw_search_impl.hpp b/src/index/hnsw_search_impl.hpp index e271db6a..587c51ae 100644 --- a/src/index/hnsw_search_impl.hpp +++ b/src/index/hnsw_search_impl.hpp @@ -256,7 +256,8 @@ inline void HierarchicalNSW::search_base_layer_direct( const PID* links = level0_links(current_node_id); const size_t num_neighbors = link_count(links); - for (size_t p = 0; p < kPrefetchLookahead; ++p) { + // Slots beyond the link count hold stale IDs; never form addresses from them. + for (size_t p = 0; p < std::min(kPrefetchLookahead, num_neighbors); ++p) { memory::mem_prefetch_l1(bin_data(*(links + 1 + p)), prefetch_size); } // Iterate over neighbors. (List starts at index 1.) @@ -299,9 +300,13 @@ inline void HierarchicalNSW::search_base_layer_direct( search_pool.insert(candidate_id, candidate_estimate.distance); } - memory::mem_prefetch_l2( - reinterpret_cast(level0_links(search_pool.next_id())), 2 - ); + // A full pool parks the displaced entry, possibly with its checked bit + // set, in the spare slot that next_id() would read. + if (search_pool.has_next()) { + memory::mem_prefetch_l2( + reinterpret_cast(level0_links(search_pool.next_id())), 2 + ); + } } } }