From 3e9c90d9d76e373ada455e28019fa8293612c5cd Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 14:08:45 +0800 Subject: [PATCH 1/5] fix(qg): sum builder degrees in size_t QGBuilder::avg_degree passed 0U to std::accumulate over uint32_t degrees, so the sum was an unsigned int and wrapped once the total degree exceeded 2^32 - 1, for example about 67 million points at degree 64. Start from size_t{0}. The new test sets 65 degrees of 2^31 and fails before the fix. --- src/index/qg_builder.cpp | 3 ++- tests/unit/rabitqlib/index/qg_test.cpp | 12 ++++++++++++ 2 files changed, 14 insertions(+), 1 deletion(-) diff --git a/src/index/qg_builder.cpp b/src/index/qg_builder.cpp index a091269f..2027812e 100644 --- a/src/index/qg_builder.cpp +++ b/src/index/qg_builder.cpp @@ -205,7 +205,8 @@ bool QGBuilder::check_dup() const { } float QGBuilder::avg_degree() const { - size_t degrees = std::accumulate(degrees_.begin(), degrees_.end(), 0U); + // A size_t initial value keeps the sum from wrapping at 2^32 on large graphs. + size_t degrees = std::accumulate(degrees_.begin(), degrees_.end(), size_t{0}); return static_cast(degrees) / static_cast(num_points_); } void QGBuilder::run_passes(size_t num_iter) { diff --git a/tests/unit/rabitqlib/index/qg_test.cpp b/tests/unit/rabitqlib/index/qg_test.cpp index ddf682db..6be998bb 100644 --- a/tests/unit/rabitqlib/index/qg_test.cpp +++ b/tests/unit/rabitqlib/index/qg_test.cpp @@ -242,6 +242,9 @@ struct QGConstructionTestAccess { return builder.new_neighbors_; } static const auto& degrees(const QGBuilder& builder) { return builder.degrees_; } + static void set_degrees(QGBuilder& builder, uint32_t degree) { + std::fill(builder.degrees_.begin(), builder.degrees_.end(), degree); + } static void retain_seed_scores(QGBuilder& builder) { auto& graph = builder.qg_; for (PID i = 0; i < builder.num_points_; ++i) { @@ -1083,6 +1086,15 @@ TEST(QGConstructionTest, DefaultsToPipnnInitialization) { } } +TEST(QGConstructionTest, AverageDegreeDoesNotWrapAboveUint32Range) { + std::vector data(65 * 65, 0.5F); + QuantizedGraph graph(65, 65, 32); + QGBuilder builder(graph, 64, data.data(), 1, QGInitialization::Random); + // 65 degrees of 2^31 sum past 2^32; a 32-bit accumulator would wrap. + QGConstructionTestAccess::set_degrees(builder, uint32_t{1} << 31U); + EXPECT_FLOAT_EQ(builder.avg_degree(), 2147483648.0F); +} + TEST(QGConstructionTest, SupportsExplicitRandomInitialization) { std::vector data(65 * 65, 0.5F); QuantizedGraph graph(65, 65, 32); From efb4c0480070f480b7829b15a1dfe9bf702e3b00 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 14:08:45 +0800 Subject: [PATCH 2/5] fix(ivf): bound centroid search output in every initializer FlatInitializer::centroids_distances partially sorted to nprobe and copied nprobe entries into the caller's vector without checking either bound, so nprobe above the cluster count or an undersized output read and wrote out of range. HNSWInitializer also wrote into the output without sizing it. Both now clamp nprobe to the cluster count and resize the output, as FlatRaBitQInitializer already did, and the Initializer interface documents that contract. IVF always passed a clamped, presized vector, so index search is unchanged. The new test fails before the fix with a segmentation fault. --- include/rabitqlib/index/ivf/initializer.hpp | 29 +++++++++------ .../unit/rabitqlib/index/initializer_test.cpp | 35 +++++++++++++++++++ 2 files changed, 53 insertions(+), 11 deletions(-) diff --git a/include/rabitqlib/index/ivf/initializer.hpp b/include/rabitqlib/index/ivf/initializer.hpp index 0f04cff5..69d6825a 100644 --- a/include/rabitqlib/index/ivf/initializer.hpp +++ b/include/rabitqlib/index/ivf/initializer.hpp @@ -52,6 +52,8 @@ class Initializer { virtual ~Initializer() = 0; [[nodiscard]] virtual const float* centroid(PID) const = 0; virtual void add_vectors(const float*, size_t) = 0; + // Resize the output to min(nprobe, number of clusters) and fill it with the + // nearest centroids in no particular order; the caller need not presize it. virtual void centroids_distances( const float*, size_t, std::vector>& ) const = 0; @@ -96,20 +98,18 @@ class FlatInitializer : public Initializer { ? dot_product_dis(query, centroid(i), dim_) : std::sqrt(euclidean_sqr(query, centroid(i), dim_)); } + const size_t num_probes = std::min(nprobe, num_cluster_); std::partial_sort( centroid_candidates.begin(), - centroid_candidates.begin() + static_cast(nprobe), + centroid_candidates.begin() + static_cast(num_probes), centroid_candidates.end() ); - std::memcpy( - candidates.data(), - centroid_candidates.data(), - sizeof(AnnCandidate) * nprobe - ); + candidates.resize(num_probes); + std::copy_n(centroid_candidates.begin(), num_probes, candidates.begin()); } - // for flat initer, we save & load into the ifstream + // The flat initializer stores its centroids inline in the index file stream. void save(std::ofstream& output, const char*) const override { output.write( reinterpret_cast(centroids_.data()), @@ -258,13 +258,20 @@ class HNSWInitializer : public Initializer { void centroids_distances( const float* query, size_t nprobe, std::vector>& candidates ) const override { + const size_t num_probes = std::min(nprobe, num_cluster_); + candidates.resize(num_probes); + if (num_probes == 0) { + return; + } std::priority_queue> result = - alg_hnsw_->searchKnnWithEf(query, nprobe, std::max(size_t{768}, 2 * nprobe)); - if (result.size() != nprobe) { + alg_hnsw_->searchKnnWithEf( + query, num_probes, std::max(size_t{768}, 2 * num_probes) + ); + if (result.size() != num_probes) { throw std::runtime_error("HNSW centroid graph returned too few candidates"); } - for (size_t i = 0; i < nprobe; ++i) { + for (size_t i = 0; i < num_probes; ++i) { candidates[i].distance = metric_type_ == METRIC_IP ? result.top().first : std::sqrt(result.top().first); @@ -273,7 +280,7 @@ class HNSWInitializer : public Initializer { } } - // for hnsw initer, we save & load into a separate file by hnswlib + // The HNSW initializer stores its graph in a separate hnswlib sidecar file. void save(std::ofstream&, const char* filename) const override { std::string hnsw(filename); hnsw += ".hnsw"; diff --git a/tests/unit/rabitqlib/index/initializer_test.cpp b/tests/unit/rabitqlib/index/initializer_test.cpp index 2114960a..3355beeb 100644 --- a/tests/unit/rabitqlib/index/initializer_test.cpp +++ b/tests/unit/rabitqlib/index/initializer_test.cpp @@ -2,6 +2,7 @@ #include +#include #include #include #include @@ -149,6 +150,40 @@ TEST(HNSWInitializerTest, RoutesByInnerProductWithoutTakingSquareRoot) { EXPECT_FLOAT_EQ(candidates[0].distance, -99.0F); } +TEST(IvfInitializerTest, CentroidSearchSizesOutputAndClampsProbeCount) { + constexpr size_t kDim = 17; + constexpr size_t kCount = 32; + std::vector centroids(kCount * kDim, 0.0F); + for (size_t i = 0; i < kCount; ++i) { + centroids[i * kDim] = static_cast(i) * 2.0F; + } + FlatInitializer flat(kDim, kCount); + HNSWInitializer hnsw(kDim, kCount); + flat.add_vectors(centroids.data(), 1); + hnsw.add_vectors(centroids.data(), 1); + std::array query{}; + query[0] = 20.25F; + for (const Initializer* initializer : + {static_cast(&flat), static_cast(&hnsw)}) { + for (size_t nprobe : {0U, 1U, 7U, 99U}) { + SCOPED_TRACE(nprobe); + // An empty output must be resized, never written past its end. + std::vector> candidates; + initializer->centroids_distances(query.data(), nprobe, candidates); + ASSERT_EQ(candidates.size(), std::min(nprobe, kCount)); + std::vector seen(kCount); + for (const auto& candidate : candidates) { + ASSERT_LT(candidate.id, kCount); + EXPECT_FALSE(seen[candidate.id]); + seen[candidate.id] = true; + } + if (nprobe == 1) { + EXPECT_EQ(candidates[0].id, 10U); + } + } + } +} + TEST(HNSWInitializerTest, ParallelInsertionPreservesRandomLevelDraws) { constexpr size_t kDim = 17; constexpr size_t kCount = 2048; From f53748048604c2c1d0076d9e898e556d06290728 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 14:08:45 +0800 Subject: [PATCH 3/5] fix(hnsw): report allocation failure and zero capacity consistently Out-of-memory in construct and add threw std::runtime_error, while load and resize threw std::bad_alloc, so Python saw RuntimeError or MemoryError for the same failure. Throw std::bad_alloc for the link-list allocations, and drop the null checks after huge_page_allocate, which throws instead of returning null. A zero capacity reached the allocator, which returns null for zero bytes: the constructor reported "Not enough memory" and resize(0) on an empty index threw std::bad_alloc. Reject both with std::invalid_argument (Python ValueError) before allocating, and document the minimum. Also correct the entry-point promotion comment in add_point. --- python_bindings/__init__.pyi | 4 +-- python_bindings/hnsw_bindings.cpp | 2 +- src/index/hnsw.cpp | 33 +++++++----------------- tests/python/test_hnsw.py | 7 +++++ tests/unit/rabitqlib/index/hnsw_test.cpp | 19 ++++++++++++++ 5 files changed, 38 insertions(+), 27 deletions(-) diff --git a/python_bindings/__init__.pyi b/python_bindings/__init__.pyi index 22247a3b..4823f37e 100644 --- a/python_bindings/__init__.pyi +++ b/python_bindings/__init__.pyi @@ -371,7 +371,7 @@ class HnswIndex: dim : int Original vector dimension, from 64 through 65536. max_elements : int - Storage capacity, including removed points. Use resize to grow it. + Storage capacity, at least 1, including removed points. Use resize to grow it. M : int, optional Graph connectivity, at least 2. Default is 16. ef_construction : int, optional @@ -407,7 +407,7 @@ class HnswIndex: dim : int Original vector dimension, from 64 through 65536. max_elements : int - Storage capacity, including removed points. Use resize to grow it. + Storage capacity, at least 1, including removed points. Use resize to grow it. M : int, optional Graph connectivity, at least 2. Default is 16. ef_construction : int, optional diff --git a/python_bindings/hnsw_bindings.cpp b/python_bindings/hnsw_bindings.cpp index 889e356f..a7630ed1 100644 --- a/python_bindings/hnsw_bindings.cpp +++ b/python_bindings/hnsw_bindings.cpp @@ -293,7 +293,7 @@ Parameters dim : int Original vector dimension, from 64 through 65536. max_elements : int - Storage capacity, including removed points. Use resize to grow it. + Storage capacity, at least 1, including removed points. Use resize to grow it. M : int, optional Graph connectivity, at least 2. Default is 16. ef_construction : int, optional diff --git a/src/index/hnsw.cpp b/src/index/hnsw.cpp index 4f6f1d2c..366bb647 100644 --- a/src/index/hnsw.cpp +++ b/src/index/hnsw.cpp @@ -79,6 +79,9 @@ HierarchicalNSW::HierarchicalNSW( if (M < 2) { throw std::invalid_argument("HNSW M must be at least 2"); } + if (max_elements == 0) { + throw std::invalid_argument("HNSW max_elements must be greater than zero"); + } max_elements_ = max_elements; dim_ = dim; rotator_.reset( @@ -117,9 +120,6 @@ HierarchicalNSW::HierarchicalNSW( data_level0_memory_.reset( memory::huge_page_allocate(max_elements_ * level0_record_bytes_) ); - if (data_level0_memory_ == nullptr) { - throw std::runtime_error("Not enough memory"); - } level_generator_.seed(random_seed); @@ -346,15 +346,9 @@ void HierarchicalNSW::load(const char* filename) { if (centroids_bytes != 0) { loaded.centroids_memory_.reset(memory::huge_page_allocate(centroids_bytes)); - if (loaded.centroids_memory_ == nullptr) { - throw std::bad_alloc(); - } read_payload(loaded.centroids_memory_.get(), centroids_bytes); } loaded.data_level0_memory_.reset(memory::huge_page_allocate(data_capacity_bytes)); - if (loaded.data_level0_memory_ == nullptr) { - throw std::bad_alloc(); - } read_payload(loaded.data_level0_memory_.get(), data_bytes); loaded.element_levels_ = std::vector(loaded.max_elements_); @@ -543,11 +537,6 @@ void HierarchicalNSW::construct( num_clusters_ = num_clusters; const size_t centroids_bytes = num_clusters_ * padded_dim_ * sizeof(float); centroids_memory_.reset(memory::huge_page_allocate(centroids_bytes)); - if (centroids_memory_ == nullptr) { - throw std::runtime_error( - "Not enough memory: HNSW failed to allocate centroids" - ); - } for (size_t i = 0; i < num_clusters; ++i) { this->rotator_->rotate( @@ -734,9 +723,7 @@ void HierarchicalNSW::add_point( static_cast(std::malloc((upper_layer_links_bytes_ * node_level) + 1)) ); if (link_lists_[new_node_id] == nullptr) { - throw std::runtime_error( - "Not enough memory: add_point failed to allocate linklist" - ); + throw std::bad_alloc(); } memset( link_lists_[new_node_id].get(), 0, (upper_layer_links_bytes_ * node_level) + 1 @@ -767,7 +754,7 @@ void HierarchicalNSW::add_point( } } - // Releasing lock for the maximum level + // A taller node becomes the entry point; entry_point_lock is still held then. if (node_level > previous_max_level) { entry_point_id_ = new_node_id; max_level_ = node_level; @@ -1170,6 +1157,9 @@ void HierarchicalNSW::resize(size_t new_max_elements) { if (level0_record_bytes_ == 0) { throw std::logic_error("HNSW index must be initialized before resize"); } + if (new_max_elements == 0) { + throw std::invalid_argument("HNSW resize capacity must be greater than zero"); + } if (new_max_elements < num_stored_points_) { throw std::invalid_argument("HNSW resize: capacity below the element count"); } @@ -1188,9 +1178,6 @@ void HierarchicalNSW::resize(size_t new_max_elements) { // comes from aligned_alloc, which realloc does not keep aligned. const size_t level0_bytes = new_max_elements * level0_record_bytes_; AlignedStorage level0(memory::huge_page_allocate(level0_bytes)); - if (level0 == nullptr) { - throw std::bad_alloc(); - } std::memset(level0.get(), 0, level0_bytes); std::memcpy( level0.get(), data_level0_memory_.get(), num_stored_points_ * level0_record_bytes_ @@ -1294,9 +1281,7 @@ std::vector HierarchicalNSW::add( std::calloc(1, (upper_layer_links_bytes_ * levels[i]) + 1) ); if (raw == nullptr) { - throw std::runtime_error( - "Not enough memory: HNSW add failed to allocate a link list" - ); + throw std::bad_alloc(); } } link_lists.emplace_back(raw); diff --git a/tests/python/test_hnsw.py b/tests/python/test_hnsw.py index 1b4bd5e2..6af01f38 100644 --- a/tests/python/test_hnsw.py +++ b/tests/python/test_hnsw.py @@ -333,6 +333,13 @@ def test_resize_grows_capacity_and_allows_more(base_data, clusters): np.testing.assert_array_equal(ids[:, 0], np.arange(10)) +def test_zero_capacity_raises(): + with pytest.raises( + ValueError, match="^HNSW max_elements must be greater than zero$" + ): + HnswIndex(max_elements=0, dim=64) + + def test_resize_below_element_count_raises(partial_hnsw): with pytest.raises(ValueError, match="element count"): partial_hnsw.resize(10) diff --git a/tests/unit/rabitqlib/index/hnsw_test.cpp b/tests/unit/rabitqlib/index/hnsw_test.cpp index 8bd72a66..638096ad 100644 --- a/tests/unit/rabitqlib/index/hnsw_test.cpp +++ b/tests/unit/rabitqlib/index/hnsw_test.cpp @@ -576,6 +576,25 @@ TEST(HnswUninitializedTest, RejectsSearchAndSave) { EXPECT_FALSE(std::filesystem::exists(std::filesystem::u8path(path_file.path()))); } +TEST(HnswResizeTest, RejectsZeroCapacity) { + // A zero-byte level-0 allocation yields no storage; neither path may accept it. + try { + HierarchicalNSW index(0, AddFixture::kDim, 4, 8, 50); + FAIL() << "zero-capacity construction succeeded"; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ(error.what(), "HNSW max_elements must be greater than zero"); + } + + HierarchicalNSW index(16, AddFixture::kDim, 4, 8, 50); + try { + index.resize(0); + FAIL() << "zero-capacity resize succeeded"; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ(error.what(), "HNSW resize capacity must be greater than zero"); + } + EXPECT_EQ(index.max_elements(), 16U); +} + TEST(HnswResizeTest, AllowsResizeBeforeConstruction) { constexpr size_t kBuilt = 64; constexpr size_t kCapacity = 128; From 0f7175f2038773a7919b7cc93c63c477e2411289 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 14:08:45 +0800 Subject: [PATCH 4/5] fix(simd): give helpers shared across ISA builds internal linkage Headers compiled into generic, AVX2 and AVX-512 translation units defined external-linkage inline functions and templates (pack_4bit_block, load_excode_tail, bitops::load_word and countr_zero32, contribute_ip*, BoundedRescaleState's constructor). When the compiler does not inline them, each object emits a weak copy and the linker keeps one for every caller, so a lower-ISA path could run code compiled for a higher ISA. docs/development.md already requires internal linkage for these headers. Use anonymous namespaces or static inline. Release object code is unchanged; at -O0 the shared weak symbols drop from 10 (GCC) and 11 (Clang) to 0. --- include/rabitqlib/utils/bitops.hpp | 9 ++++++--- src/simd/hnsw_neon_kernels.hpp | 3 +++ src/simd/pack_excode_kernels.hpp | 3 +++ src/simd/packed_tail_avx2.hpp | 9 +++++---- src/simd/rescale_search.hpp | 4 ++++ src/simd/space_excode_avx2.cpp | 6 +++--- 6 files changed, 24 insertions(+), 10 deletions(-) diff --git a/include/rabitqlib/utils/bitops.hpp b/include/rabitqlib/utils/bitops.hpp index f75d4356..d299f3ff 100644 --- a/include/rabitqlib/utils/bitops.hpp +++ b/include/rabitqlib/utils/bitops.hpp @@ -10,8 +10,11 @@ namespace rabitqlib::bitops { +// ISA-specific translation units include these helpers; internal linkage keeps each +// backend's compiled body from being merged with another backend's. + // Load 1..8 bytes of a compact word, zero-extending a final partial word. -inline uint64_t load_word(const void* data, size_t bytes) { +static inline uint64_t load_word(const void* data, size_t bytes) { if (bytes == 8) { uint64_t word; std::memcpy(&word, data, sizeof(word)); @@ -27,7 +30,7 @@ inline uint64_t load_word(const void* data, size_t bytes) { return word; } -inline unsigned popcount32(uint32_t value) { +static inline unsigned popcount32(uint32_t value) { #if defined(_MSC_VER) value -= (value >> 1) & 0x55555555U; value = (value & 0x33333333U) + ((value >> 2) & 0x33333333U); @@ -53,7 +56,7 @@ static inline unsigned popcount64(uint64_t value) { } // Callers pass nonzero masks. -inline unsigned countr_zero32(uint32_t value) { +static inline unsigned countr_zero32(uint32_t value) { #if defined(_MSC_VER) unsigned long index = 0; _BitScanForward(&index, value); diff --git a/src/simd/hnsw_neon_kernels.hpp b/src/simd/hnsw_neon_kernels.hpp index e2600201..22a09cd6 100644 --- a/src/simd/hnsw_neon_kernels.hpp +++ b/src/simd/hnsw_neon_kernels.hpp @@ -11,6 +11,8 @@ #include "rabitqlib/core/detail/simd/space_dispatch.hpp" namespace rabitqlib::simd::detail { +// Internal linkage: backends compiled with different ISA flags keep their own bodies. +namespace { inline float mask_ip_neon(const float* query, const uint8_t* data, size_t dim) { float32x4_t sum0 = vdupq_n_f32(0), sum1 = sum0, sum2 = sum0, sum3 = sum0; @@ -96,4 +98,5 @@ inline float warmup_ip_neon( } return delta * static_cast(ip) + vl * static_cast(count); } +} // namespace } // namespace rabitqlib::simd::detail diff --git a/src/simd/pack_excode_kernels.hpp b/src/simd/pack_excode_kernels.hpp index a69f8542..88cdf712 100644 --- a/src/simd/pack_excode_kernels.hpp +++ b/src/simd/pack_excode_kernels.hpp @@ -9,6 +9,8 @@ #include namespace rabitqlib::simd::detail { +// Internal linkage: backends compiled with different ISA flags keep their own bodies. +namespace { template inline void pack_excode_tail_intrinsics(const uint8_t* raw, uint8_t* compact, size_t dim) { @@ -353,4 +355,5 @@ inline void packing_7bit_excode_intrinsics( pack_excode_tail_intrinsics<7>(o_raw, o_compact, dim % 64); } +} // namespace } // namespace rabitqlib::simd::detail diff --git a/src/simd/packed_tail_avx2.hpp b/src/simd/packed_tail_avx2.hpp index 267e9dfe..6e7e40f0 100644 --- a/src/simd/packed_tail_avx2.hpp +++ b/src/simd/packed_tail_avx2.hpp @@ -8,6 +8,8 @@ #include namespace rabitqlib::simd::detail { +// Internal linkage: backends compiled with different ISA flags keep their own bodies. +namespace { // Read exactly 2 * Bits bytes for 16 codes. Staging short stores through a // 16-byte stack buffer blocks store forwarding on the following vector load. template @@ -84,15 +86,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) { +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 -) { +inline float mask_ip_avx2(const float* query, const uint8_t* data, size_t padded_dim) { const size_t num_blk = padded_dim / 64; const auto* it_data = data; const float* it_query = query; @@ -150,4 +150,5 @@ static inline float mask_ip_avx2( _mm256_add_ps(_mm256_add_ps(sum0, sum1), _mm256_add_ps(sum2, sum3)) ); } +} // namespace } // namespace rabitqlib::simd::detail diff --git a/src/simd/rescale_search.hpp b/src/simd/rescale_search.hpp index f800e40e..e0a44991 100644 --- a/src/simd/rescale_search.hpp +++ b/src/simd/rescale_search.hpp @@ -104,6 +104,9 @@ static inline double next_double_down(double value) { return value; } +// Its default member initializers give it a constructor, so keep the type +// TU-local as well. +namespace { struct BoundedRescaleState { RescaleSearchState state{}; // Enclose the true dot product despite SIMD summation rounding. S is exact @@ -113,6 +116,7 @@ struct BoundedRescaleState { double score_lower = 0; double score_upper = 0; }; +} // namespace // Upper bound on squared cosine similarity for every code between two evaluated // states. Here a_i = input[i] is a rotated, normalized residual magnitude, and diff --git a/src/simd/space_excode_avx2.cpp b/src/simd/space_excode_avx2.cpp index a6a5d20c..270661cd 100644 --- a/src/simd/space_excode_avx2.cpp +++ b/src/simd/space_excode_avx2.cpp @@ -22,7 +22,6 @@ namespace { std::memcpy(&signed_low, &low, sizeof(signed_low)); return _mm_set_epi64x(signed_high, signed_low); } -} // namespace // Accumulate the two halves independently to shorten the FMA dependency chain. inline void contribute_ip( @@ -35,7 +34,7 @@ inline void contribute_ip( q = _mm256_loadu_ps(query + 8); cf = _mm256_cvtepi32_ps(_mm256_cvtepu8_epi32(_mm_srli_si128(vec, 8))); sum_hi = _mm256_fmadd_ps(q, cf, sum_hi); -}; +} inline void contribute_ip_signed( __m128i vec, const float* __restrict__ query, __m256& sum, __m256& sum_hi @@ -47,7 +46,8 @@ inline void contribute_ip_signed( q = _mm256_loadu_ps(query + 8); cf = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128(vec, 8))); sum_hi = _mm256_fmadd_ps(cf, q, sum_hi); -}; +} +} // namespace // 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. From d6b357571ecffa2fd9566b89c42121d1ecccdee6 Mon Sep 17 00:00:00 2001 From: gouyt13clear Date: Sat, 31 Oct 2026 14:08:45 +0800 Subject: [PATCH 5/5] fix: correct comments and messages that contradict the code - docs/development.md: HNSW has kDefaultEfSearch, not default_ef_search_, and query search inserts removed points into search_pool, not candidate_set. - HNSW: the quantized-construction, live-node and bounded-queue comments described old behavior or the wrong declaration. - IVF: the save/load format comments omitted the initializer payload written first, and ex_bits_ is 0 with raw reranking. - FastScan: kPos, the 1-bit code offsets and the HACC LUT layout comments were misleading or stale. - Core: unpack_sign_bits accepts any 8-coordinate tail, not only 32. - The k-means overflow message named size_t although the check uses ptrdiff_t, and an AVX-512 static_assert named a constant that does not exist. --- docs/development.md | 4 ++-- include/rabitqlib/clustering/detail/kmeans.hpp | 2 +- include/rabitqlib/core/detail/rabitq_impl.hpp | 3 ++- include/rabitqlib/fastscan/data_layout.hpp | 16 ++++++++-------- include/rabitqlib/fastscan/fastscan.hpp | 2 +- include/rabitqlib/index/hnsw/hnsw.hpp | 8 ++++---- include/rabitqlib/index/ivf/ivf.hpp | 2 +- src/fastscan/fastscan_generic.cpp | 2 +- src/index/hnsw_dispatch.cpp | 2 +- src/index/hnsw_search_impl.hpp | 2 +- src/index/ivf.cpp | 6 ++++-- src/simd/rotator_avx512.cpp | 2 +- 12 files changed, 27 insertions(+), 24 deletions(-) diff --git a/docs/development.md b/docs/development.md index b0d91b97..34714d56 100644 --- a/docs/development.md +++ b/docs/development.md @@ -602,7 +602,7 @@ preprocessing objects, `visited` for visited state. **Search windows** -`ef_search` for a per-call graph window, `default_ef_search_` for a stored default, and +`ef_search` for a per-call graph window, `kDefaultEfSearch` for HNSW's fixed fallback, and `ef_construction` for graph-building search breadth. IVF's `nprobe` counts clusters and remains distinct. @@ -809,7 +809,7 @@ This mark is unrelated to the `SearchBuffer` ID marker. Keep these rules when to - Read cluster IDs only through `get_clusterid_by_internalid`, which masks the bit. Estimators, routing, and the `add` reconstruction index centroids with it. -- A removed point stays in the graph: base-layer search still inserts it into `candidate_set`, and +- A removed point stays in the graph: base-layer search still inserts it into `search_pool`, and `add` may link to it. Only the result buffer skips it, including when it is the entry point. - `load` masks the bit before range-checking the cluster ID. - `HnswRemoveTest.SearchReachesLivePointsThroughRemovedOnes`, diff --git a/include/rabitqlib/clustering/detail/kmeans.hpp b/include/rabitqlib/clustering/detail/kmeans.hpp index 714b2035..40ec3912 100644 --- a/include/rabitqlib/clustering/detail/kmeans.hpp +++ b/include/rabitqlib/clustering/detail/kmeans.hpp @@ -244,7 +244,7 @@ class LloydKMeans : public Parameters { } if (n > static_cast(std::numeric_limits::max()) / d) { throw std::invalid_argument( - std::string(Assignment::kName) + " input dimensions overflow size_t" + std::string(Assignment::kName) + " input dimensions overflow ptrdiff_t" ); } const uint32_t num_workers = effective_num_threads(); diff --git a/include/rabitqlib/core/detail/rabitq_impl.hpp b/include/rabitqlib/core/detail/rabitq_impl.hpp index 0fa4f4ee..bb1bcd0b 100644 --- a/include/rabitqlib/core/detail/rabitq_impl.hpp +++ b/include/rabitqlib/core/detail/rabitq_impl.hpp @@ -102,7 +102,8 @@ inline T direct_rescale(T l2_sqr, T residual_code_ip, MetricType metric_type) { /** * Inverse of pack_binary_to_bytes: write each coordinate's sign bit (0 or - * 1) to signs. Full 64-coordinate words come first, followed by a final 32-bit word. + * 1) to signs. Full 64-coordinate words come first, followed by a final partial word + * of any multiple of 8 coordinates. */ inline void unpack_sign_bits(const uint8_t* packed, size_t dim, uint16_t* signs) { for (size_t start = 0; start < dim; start += 64) { diff --git a/include/rabitqlib/fastscan/data_layout.hpp b/include/rabitqlib/fastscan/data_layout.hpp index ac305ff2..faf2b078 100644 --- a/include/rabitqlib/fastscan/data_layout.hpp +++ b/include/rabitqlib/fastscan/data_layout.hpp @@ -28,8 +28,8 @@ template struct BatchDataMap { public: explicit BatchDataMap(char* data, size_t padded_dim) - : batch_bin_code_(reinterpret_cast(data)) - , f_add_(data + detail::BatchCodeOffsets::add(padded_dim)) // 1 bit code + : batch_bin_code_(reinterpret_cast(data)) // 1-bit codes + , f_add_(data + detail::BatchCodeOffsets::add(padded_dim)) , f_rescale_(data + detail::BatchCodeOffsets::rescale(padded_dim)) , f_error_(data + detail::BatchCodeOffsets::error(padded_dim)) {} @@ -54,8 +54,8 @@ template struct ConstBatchDataMap { public: explicit ConstBatchDataMap(const char* data, size_t padded_dim) - : batch_bin_code_(reinterpret_cast(data)) - , f_add_(data + detail::BatchCodeOffsets::add(padded_dim)) // 1 bit code + : batch_bin_code_(reinterpret_cast(data)) // 1-bit codes + , f_add_(data + detail::BatchCodeOffsets::add(padded_dim)) , f_rescale_(data + detail::BatchCodeOffsets::rescale(padded_dim)) , f_error_(data + detail::BatchCodeOffsets::error(padded_dim)) {} @@ -76,8 +76,8 @@ template struct QGBatchDataMap { public: explicit QGBatchDataMap(char* data, size_t padded_dim) - : batch_bin_code_(reinterpret_cast(data)) - , f_add_(data + detail::BatchCodeOffsets::add(padded_dim)) // 1 bit code + : batch_bin_code_(reinterpret_cast(data)) // 1-bit codes + , f_add_(data + detail::BatchCodeOffsets::add(padded_dim)) , f_rescale_(data + detail::BatchCodeOffsets::rescale(padded_dim)) {} [[nodiscard]] uint8_t* bin_code() { return batch_bin_code_; } @@ -107,8 +107,8 @@ template struct ConstQGBatchDataMap { public: explicit ConstQGBatchDataMap(const char* data, size_t padded_dim) - : batch_bin_code_(reinterpret_cast(data)) - , f_add_(data + detail::BatchCodeOffsets::add(padded_dim)) // 1 bit code + : batch_bin_code_(reinterpret_cast(data)) // 1-bit codes + , f_add_(data + detail::BatchCodeOffsets::add(padded_dim)) , f_rescale_(data + detail::BatchCodeOffsets::rescale(padded_dim)) {} [[nodiscard]] const uint8_t* bin_code() const { return batch_bin_code_; } diff --git a/include/rabitqlib/fastscan/fastscan.hpp b/include/rabitqlib/fastscan/fastscan.hpp index e63957dd..bc6546aa 100644 --- a/include/rabitqlib/fastscan/fastscan.hpp +++ b/include/rabitqlib/fastscan/fastscan.hpp @@ -26,7 +26,7 @@ constexpr static std::array kPos = { 3 /*1101*/, 2 /*1110*/, 3 /*1111*/, -}; // all possible combination for a 4 bit string +}; // kPos[j]: query coordinate of j's lowest set bit (bit 3 is q0, bit 0 is q3) // Accumulate one batch into int32_t results. Throws if a result exceeds int32_t. // dim must be a positive multiple of 16. diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index a520c3b2..28a3f9c0 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -169,12 +169,12 @@ class HierarchicalNSW { // Borrows construct() input only while its workers run; cleared on success or failure. const float* construction_data_{nullptr}; - // Bounded priority queue implemented as a sorted vector. struct Candidate { float est_dist; PID id; }; + // Bounded priority queue implemented as a sorted vector. class BoundedKNN { public: explicit BoundedKNN(size_t capacity) : capacity_(capacity) {} @@ -268,7 +268,8 @@ class HierarchicalNSW { // Layout: (# of edges + edges) + (cluster_id) + (External_id) + (BinData) + (ExData) AlignedStorage data_level0_memory_; - // Reserve capacity without initializing unused slots; size tracks live nodes. + // Reserve capacity without initializing unused slots; size tracks stored nodes, + // including removed ones. std::vector link_lists_; std::vector element_levels_; @@ -447,8 +448,7 @@ class HierarchicalNSW { BoundedKNN& nearest_candidates ) const; - // Construction - // Currently only support index construction with non-quantized vectors + // Raw-vector distance used by construct; add scores links with get_quant_dist. float get_data_dist(PID obj1, PID obj2) const { PID label1 = get_external_label(obj1); PID label2 = get_external_label(obj2); diff --git a/include/rabitqlib/index/ivf/ivf.hpp b/include/rabitqlib/index/ivf/ivf.hpp index 49127902..eea36ae7 100644 --- a/include/rabitqlib/index/ivf/ivf.hpp +++ b/include/rabitqlib/index/ivf/ivf.hpp @@ -83,7 +83,7 @@ class IVF { size_t num_clusters_ = 0; bool raw_reranking_ = false; // raw vectors replace extra-bit codes bool ready_ = false; // construction or loading completed - size_t ex_bits_ = 0; // total bits = ex_bits_ + 1 + size_t ex_bits_ = 0; // total bits = ex_bits_ + 1; 0 with raw reranking RotatorType rotator_type_ = RotatorType::FhtKacRotator; std::unique_ptr> rotator_; std::vector clusters_; diff --git a/src/fastscan/fastscan_generic.cpp b/src/fastscan/fastscan_generic.cpp index d39908e6..22ea98d5 100644 --- a/src/fastscan/fastscan_generic.cpp +++ b/src/fastscan/fastscan_generic.cpp @@ -60,7 +60,7 @@ void pack_lut_range_generic( namespace rabitqlib::fastscan::simd { void transfer_lut_hacc_generic(const uint16_t* lut, size_t dim, uint8_t* hc_lut) { - // Match the AVX2 layout: two low-byte tables, then two high-byte tables. + // HaccLutLayout::Pair32: two low-byte tables, then two high-byte tables. for (size_t group = 0; group < dim / 4; ++group) { const size_t offset = group / 2 * 64 + group % 2 * 16; for (size_t i = 0; i < 16; ++i) { diff --git a/src/index/hnsw_dispatch.cpp b/src/index/hnsw_dispatch.cpp index 9f85a6ae..5446725b 100644 --- a/src/index/hnsw_dispatch.cpp +++ b/src/index/hnsw_dispatch.cpp @@ -13,12 +13,12 @@ namespace rabitqlib::hnsw::detail { namespace { using SearchKnnFn = std::vector> (*)(HierarchicalNSW&, const float*, size_t, size_t); -// The core variant uses AVX2 warmup; the popcount variant has its own stricter tier. } // namespace std::vector> search_knn_sorted( HierarchicalNSW& index, const float* query, size_t k, size_t ef_search ) { #if defined(__x86_64__) || defined(_M_X64) + // The core variant uses AVX2 warmup; the popcount variant has its own stricter tier. static const SearchKnnFn kSearchKnnFn = cpu::has_avx512_popcnt() ? search_knn_sorted_avx512_popcnt : RABITQ_RESOLVE( diff --git a/src/index/hnsw_search_impl.hpp b/src/index/hnsw_search_impl.hpp index 0995e4a3..22331535 100644 --- a/src/index/hnsw_search_impl.hpp +++ b/src/index/hnsw_search_impl.hpp @@ -203,7 +203,7 @@ inline void HierarchicalNSW::search_base_layer_direct( ) const { auto visited = visited_list_pool_->get_free_vislist(); - // Use our bounded priority queue instead of the maxheap. + // Keep the ef_search best estimates in a bounded sorted buffer. buffer::SearchBuffer search_pool(ef_search); float dist_k = 1e10; diff --git a/src/index/ivf.cpp b/src/index/ivf.cpp index 1efbb382..67ccf355 100644 --- a/src/index/ivf.cpp +++ b/src/index/ivf.cpp @@ -754,7 +754,8 @@ void IVF::write_file(const char* filename) const { /* Save rotator */ this->rotator_->save(output); - /* Save data: all batch data, then all ex data, then all IDs, each in cluster order */ + /* Save data: initializer payload, then all batch data, all ex data and all IDs, + each in cluster order */ this->initializer_->save(output, filename); for (size_t i = 0; i < num_clusters_; ++i) { output.write( @@ -920,7 +921,8 @@ void IVF::load(const char* filename) { /* Load rotator */ loaded.rotator_->load(input); - /* Load data: all batch data, then all ex data, then all IDs, each in cluster order */ + /* Load data: initializer payload, then all batch data, all ex data and all IDs, + each in cluster order */ loaded.allocate_memory(cluster_sizes); loaded.initializer_->load(input, filename); for (size_t i = 0; i < loaded.num_clusters_; ++i) { diff --git a/src/simd/rotator_avx512.cpp b/src/simd/rotator_avx512.cpp index 82fbc9cd..6c0fc5ee 100644 --- a/src/simd/rotator_avx512.cpp +++ b/src/simd/rotator_avx512.cpp @@ -38,7 +38,7 @@ void flip_sign_avx512(const uint8_t* flip, float* data, size_t dim) { static_assert( kFloatsPerChunk % 16 == 0, - "floats_per_chunk must be divisible by AVX512 register width" + "kFloatsPerChunk must be divisible by AVX512 register width" ); size_t i = 0;