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/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/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/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/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/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.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/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/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/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/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; 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. 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; 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; 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);