diff --git a/docs/docs/index/ivf.md b/docs/docs/index/ivf.md index 1643b9ab..0e60a0cc 100644 --- a/docs/docs/index/ivf.md +++ b/docs/docs/index/ivf.md @@ -83,6 +83,10 @@ vectors, and compute one-bit codes and factors. Quantized mode also computes float32 vectors instead. In raw mode, the index owns a copy of the input, so the original data can be released after construction. +Construction reuses a rotation buffer of at most 32 vectors per worker. This +temporary storage uses at most `workers * 32 * padded_dim * sizeof(float)` bytes, +independently of cluster sizes. + After construction, you can directly save the index file to disk: ```c++ ivf.save(outoput_index_file); @@ -157,6 +161,25 @@ squared L2 or `1 - dot(query, vector)` in the original coordinates; the search API is unchanged. Cluster selection and filtering remain approximate in both modes. Search returns the top `k` results after scanning the selected clusters. +For multiple queries, `search_batch` distributes queries across workers. +Input queries are contiguous row-major float32 +vectors with shape `(num_queries, dimension())`; output buffers have shape +`(num_queries, k)`. + +```cpp +ivf.search_batch(queries, num_queries, k, nprobe, ids, distances, + std::nullopt, 4); // Automatic precision, up to four workers +``` + +The last two arguments are `std::optional use_hacc` (default +`std::nullopt`) and `size_t num_threads` (default `1`, using OpenMP for parallel +batches). Pass `true` or `false` +for an explicit precision override. Distances may be `nullptr`. An empty batch +does no work. Input and output buffers must not overlap. Simultaneous searches +need separate output buffers, and must not overlap index updates. +Python's existing `index.search(queries, k, nprobe, num_threads=4)` uses this +batch path with the same output and precision rules. + ## Updating an Index Starting with 0.5.1, C++ `IVF` and Python `IvfIndex` support `add()` and `remove()` diff --git a/docs/docs/index/qg.md b/docs/docs/index/qg.md index 21d9b47b..ef344618 100644 --- a/docs/docs/index/qg.md +++ b/docs/docs/index/qg.md @@ -29,7 +29,7 @@ index.build(data, ef_construction=200, num_threads=32, init="pipnn") index.save("qg_example.index") loaded = SymqgIndex.load("qg_example.index") -ids, distances = loaded.search(queries, k=10, ef=100, num_threads=1) +ids, distances = loaded.search_batch(queries, k=10, ef=100, num_threads=1) ``` `init` defaults to `"pipnn"`; use `"random"` for random initialization. @@ -118,5 +118,31 @@ std::vector distances(10); qg.search(query.data(), 10, ids.data(), distances.data()); ``` +For contiguous row-major batches, `search_batch()` reuses scratch storage within +each worker and preserves the results of independent `search()` calls: + +```cpp +// queries contains num_queries * qg.dimension() floats in the original dimension. +std::vector batch_ids(num_queries * 10); +std::vector batch_distances(num_queries * 10); +qg.search_batch( + queries.data(), num_queries, 10, batch_ids.data(), batch_distances.data(), 4 +); +``` + +The C++ signature is `search_batch(queries, num_queries, k, ids, distances, +num_threads=1)`. Set `ef >= k` with `set_ef()` before searching. Both output +buffers are required and contain `num_queries * k` elements. Empty C++ batches +do no work and may use null buffers. Inputs and outputs must not overlap. +Concurrent searches need separate outputs; do not modify the index or call +`set_ef()` while searches are running. + +Python provides `search_batch(queries, k, ef, num_threads=1)`; the existing +`search()` accepts the same two-dimensional array and remains an equivalent +entry point. Both return `(ids, distances)` arrays of shape `(num_queries, k)`, +including empty batches. Both APIs default to one worker; `num_threads=0` +selects the available hardware thread count, capped by the number of queries. +Queries use the original dimension, including when the internal rotation pads it. + See `sample/cpp/symqg_indexing.cpp`, `sample/cpp/symqg_querying.cpp`, and their Python counterparts for complete examples. diff --git a/include/rabitqlib/clustering/detail/kmeans.hpp b/include/rabitqlib/clustering/detail/kmeans.hpp index ef783719..f8efac06 100644 --- a/include/rabitqlib/clustering/detail/kmeans.hpp +++ b/include/rabitqlib/clustering/detail/kmeans.hpp @@ -21,6 +21,7 @@ #include "rabitqlib/defines.hpp" #include "rabitqlib/simd/matrix_dispatch.hpp" +#include "rabitqlib/utils/cpu_features.hpp" #include "rabitqlib/utils/rotator.hpp" #include "rabitqlib/utils/space.hpp" #include "rabitqlib/utils/tools.hpp" @@ -55,6 +56,8 @@ namespace detail { // Inputs have already passed clustering validation. GEMM only screens candidates; // final distances and tie breaking retain the scalar double-precision arithmetic. +// Optional previous assignments contain n valid centroid IDs and may alias the +// output. Recompute their distances against the current, updated centroids. inline void exact_assign( const float* x, const float* centroids, @@ -64,27 +67,41 @@ inline void exact_assign( bool spherical, uint32_t num_threads, PID* assignments, - float* distances + float* distances, + const PID* previous_assignments = nullptr ) { - constexpr size_t kPointBlock = 64; + constexpr size_t kMaxPointBlock = 128; + // Larger row blocks amortize packing on the AVX-512 Eigen backend. Retain + // smaller blocks on other backends and require at least two blocks per + // worker to avoid doubling the work of a worker handling a small tail. + const size_t point_block = + cpu::has_avx512_core() && k >= 256 && n / num_threads >= 2 * kMaxPointBlock + ? kMaxPointBlock + : 64; constexpr size_t kCentroidBlock = 256; - const size_t blocks = 1 + (n - 1) / kPointBlock; + const size_t blocks = 1 + (n - 1) / point_block; const auto threads = static_cast(std::min(num_threads, blocks)); std::vector centroid_norms(k); std::vector> products( - threads, std::vector(kPointBlock * std::min(k, kCentroidBlock)) + threads, std::vector(point_block * std::min(k, kCentroidBlock)) ); std::vector errors(threads); -#pragma omp parallel for num_threads(threads) schedule(static) - for (std::ptrdiff_t cluster = 0; cluster < static_cast(k); ++cluster) { - double norm = 0.0; - for (size_t dim = 0; dim < d; ++dim) { - const double value = centroids[static_cast(cluster) * d + dim]; - norm += value * value; + const auto exact_distance = [d, spherical](const float* vector, const float* centroid) { + double distance = 0.0; + if (spherical) { + for (size_t dim = 0; dim < d; ++dim) { + distance -= static_cast(vector[dim]) * centroid[dim]; + } + distance += 1.0; + } else { + for (size_t dim = 0; dim < d; ++dim) { + const double difference = static_cast(vector[dim]) - centroid[dim]; + distance += difference * difference; + } } - centroid_norms[static_cast(cluster)] = norm; - } + return distance; + }; // For d <= 65536, this exceeds the float dot-product error bound // gamma_d * sum(abs(x_i*c_i)), using 2*abs(x_i*c_i) <= x_i^2+c_i^2. @@ -101,8 +118,20 @@ inline void exact_assign( omp_set_num_threads(1); const auto thread = static_cast(omp_get_thread_num()); float* dots = products[thread].data(); - std::array point_norms{}; - std::array best_distances{}; + std::array point_norms{}; + std::array best_distances{}; + // Reuse this worker team for norms and assignment. The implicit barrier + // publishes all centroid norms before any matrix tile is screened. +#pragma omp for schedule(static) + for (std::ptrdiff_t cluster = 0; cluster < static_cast(k); + ++cluster) { + double norm = 0.0; + for (size_t dim = 0; dim < d; ++dim) { + const double value = centroids[static_cast(cluster) * d + dim]; + norm += value * value; + } + centroid_norms[static_cast(cluster)] = norm; + } #pragma omp for schedule(static) for (std::ptrdiff_t block = 0; block < static_cast(blocks); ++block) { @@ -110,8 +139,8 @@ inline void exact_assign( continue; } try { - const size_t first_point = static_cast(block) * kPointBlock; - const size_t rows = std::min(kPointBlock, n - first_point); + const size_t first_point = static_cast(block) * point_block; + const size_t rows = std::min(point_block, n - first_point); for (size_t row = 0; row < rows; ++row) { const float* vector = x + (first_point + row) * d; double norm = 0.0; @@ -121,6 +150,14 @@ inline void exact_assign( } point_norms[row] = norm; best_distances[row] = std::numeric_limits::infinity(); + if (previous_assignments != nullptr) { + const size_t point = first_point + row; + const PID previous = previous_assignments[point]; + assert(previous < k); + best_distances[row] = + exact_distance(vector, centroids + previous * d); + assignments[point] = previous; + } } for (size_t first_cluster = 0; first_cluster < k; first_cluster += kCentroidBlock) { @@ -148,22 +185,12 @@ inline void exact_assign( continue; } const float* centroid = centroids + cluster * d; - double distance = 0.0; - if (spherical) { - for (size_t dim = 0; dim < d; ++dim) { - distance -= - static_cast(vector[dim]) * centroid[dim]; - } - distance += 1.0; - } else { - for (size_t dim = 0; dim < d; ++dim) { - const double difference = - static_cast(vector[dim]) - centroid[dim]; - distance += difference * difference; - } - } - // Visit in ascending ID order, retaining the first exact tie. - if (distance < best_distances[row]) { + const double distance = exact_distance(vector, centroid); + // The hint may have a higher ID than an equally good + // centroid. Preserve the original lowest-ID tie rule. + if (distance < best_distances[row] || + (distance == best_distances[row] && + cluster < assignments[point])) { assignments[point] = static_cast(cluster); best_distances[row] = distance; } @@ -289,10 +316,10 @@ class LloydKMeans : public Parameters { double previous_obj = std::numeric_limits::infinity(); for (size_t iteration = 0; iteration < niter; ++iteration) { - previous_centroids = centroids; - assigner.assign( - previous_centroids.data(), assignments.data(), distances.data() - ); + assigner.assign(centroids.data(), assignments.data(), distances.data()); + // Assignment has finished borrowing the current centroids. Keep them + // for the shift calculation and overwrite the other buffer with means. + previous_centroids.swap(centroids); const double obj = std::accumulate(distances.begin(), distances.end(), 0.0); std::fill(sums.begin(), sums.end(), 0.0); @@ -332,7 +359,8 @@ class LloydKMeans : public Parameters { spherical, threads, assignments.data(), - distances.data() + distances.data(), + assignments.data() ); } else { assigner.assign(centroids.data(), assignments.data(), distances.data(), true); diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index c1b30433..e98b5093 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -919,8 +919,7 @@ inline void HierarchicalNSW::construct( num_cluster_ = cluster_num; const size_t centroids_bytes = num_cluster_ * padded_dim_ * sizeof(float); - centroids_memory_ = - static_cast(memory::aligned_allocate_bytes(64, centroids_bytes)); + centroids_memory_ = memory::huge_page_allocate(centroids_bytes); if (centroids_memory_ == nullptr) { throw std::runtime_error("Not enough memory: HNSW failed to allocate centroids"); } diff --git a/include/rabitqlib/index/ivf/ivf.hpp b/include/rabitqlib/index/ivf/ivf.hpp index 157b86d1..5168c20c 100644 --- a/include/rabitqlib/index/ivf/ivf.hpp +++ b/include/rabitqlib/index/ivf/ivf.hpp @@ -1,7 +1,10 @@ #pragma once +#include + #include #include +#include #include #include #include @@ -12,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -65,7 +69,7 @@ class IVF { static constexpr uint32_t kRawFormatVersion = 1; void - quantize_cluster(Cluster&, const std::vector&, const float*, const float*, float*, const quant::RabitqConfig&); + quantize_cluster(Cluster&, const std::vector&, const float*, const float*, float*, float*, const quant::RabitqConfig&); void extend_cluster(const Cluster&, const std::vector&, const float*, const float*, PID, char*, char*, PID*, const quant::RabitqConfig&) @@ -163,6 +167,8 @@ class IVF { scan_one_batch(const char* batch_data, const char* ex_data, const PID* ids, const SplitBatchQuery& q_obj, buffer::SearchBuffer& knns, size_t num_points, bool, const float*) const; + void validate_search(const float*, size_t, size_t, const PID*) const; + public: explicit IVF() = default; explicit IVF( @@ -238,6 +244,25 @@ class IVF { void search(const float*, size_t, size_t, PID*, float*, bool) const; + /** + * Search contiguous row-major queries using at most num_threads workers. + * queries contains num_queries * dimension() floats. results and optional + * dists each hold num_queries * k elements. Empty batches do no work. + * use_hacc defaults to the same automatic policy as search(). + * Input and output buffers must not overlap. Concurrent calls are safe if + * the index is not being modified and their output buffers are separate. + */ + void search_batch( + const float* queries, + size_t num_queries, + size_t k, + size_t nprobe, + PID* results, + float* dists = nullptr, + std::optional use_hacc = std::nullopt, + size_t num_threads = 1 + ) const; + [[nodiscard]] size_t padded_dim() const { return this->padded_dim_; } [[nodiscard]] size_t num_clusters() const { return this->num_cluster_; } @@ -333,7 +358,15 @@ inline void IVF::construct( config = quant::faster_config(padded_dim_, ex_bits_ + 1); } - num_threads = resolve_num_threads(num_threads); + num_threads = std::min(resolve_num_threads(num_threads), num_cluster_); + const size_t scratch_points = + std::min(fastscan::kBatchSize, *std::max_element(counts.begin(), counts.end())); + // Each worker reuses one batch, independently of the largest cluster size. + // Allocate before entering OpenMP so allocation failures reach the caller. + std::vector> rotated_blocks(num_threads); + for (auto& block : rotated_blocks) { + block.resize(scratch_points * padded_dim_); + } /* Quantize each cluster */ #pragma omp parallel for schedule(dynamic) num_threads(num_threads) for (std::ptrdiff_t index = 0; index < static_cast(num_cluster_); @@ -342,7 +375,15 @@ inline void IVF::construct( const float* cur_centroid = centroids + (i * dim_); float* cur_rotated_c = &rotated_centroids[i * padded_dim_]; Cluster& cp = cluster_lst_[i]; - quantize_cluster(cp, id_lists[i], data, cur_centroid, cur_rotated_c, config); + quantize_cluster( + cp, + id_lists[i], + data, + cur_centroid, + cur_rotated_c, + rotated_blocks[static_cast(omp_get_thread_num())].data(), + config + ); } this->initer_->add_vectors(rotated_centroids.data(), num_threads); @@ -407,6 +448,7 @@ inline void IVF::quantize_cluster( const float* data, const float* cur_centroid, float* rotated_centroid, + float* rotated_data, const quant::RabitqConfig& config ) { size_t num_points = IDs.size(); @@ -420,25 +462,24 @@ inline void IVF::quantize_cluster( // rotate centroid this->rotator_->rotate(cur_centroid, rotated_centroid); - // rotate vectors for this cluster - std::vector rotated_data(padded_dim_ * num_points); - for (size_t i = 0; i < num_points; ++i) { - const float* vector = data + (IDs[i] * dim_); - rotator_->rotate(vector, rotated_data.data() + (i * padded_dim_)); - if (raw_reranking_) { - std::memcpy( - cp.ex_data() + (i * rerank_vector_bytes()), vector, rerank_vector_bytes() - ); - } - } - char* batch_data = cp.batch_data(); char* ex_data = cp.ex_data(); for (size_t i = 0; i < num_points; i += fastscan::kBatchSize) { size_t n = std::min(fastscan::kBatchSize, num_points - i); + for (size_t j = 0; j < n; ++j) { + const float* vector = data + (IDs[i + j] * dim_); + rotator_->rotate(vector, rotated_data + (j * padded_dim_)); + if (raw_reranking_) { + std::memcpy( + cp.ex_data() + ((i + j) * rerank_vector_bytes()), + vector, + rerank_vector_bytes() + ); + } + } quant::quantize_split_batch( - rotated_data.data() + (i * padded_dim_), + rotated_data, rotated_centroid, n, padded_dim_, @@ -931,23 +972,7 @@ inline void IVF::search( float* __restrict__ dists, bool use_hacc ) const { - if (!ready_ || initer_ == nullptr || rotator_ == nullptr || - cluster_lst_.size() != num_cluster_ || batch_storage_.empty() || - id_storage_.size() != num_) { - throw std::logic_error("IVF index must be constructed or loaded before search"); - } - if (query == nullptr) { - throw std::invalid_argument("IVF search query must not be null"); - } - if (results == nullptr) { - throw std::invalid_argument("IVF search results must not be null"); - } - if (k == 0 || k > num_ || k > buffer::kSearchBufferMaxPointCount) { - throw std::invalid_argument("IVF search k must be between 1 and the point count"); - } - if (nprobe == 0) { - throw std::invalid_argument("IVF search nprobe must be positive"); - } + validate_search(query, k, nprobe, results); nprobe = std::min(nprobe, num_cluster_); // corner case std::vector rotated_query(padded_dim_); @@ -996,6 +1021,99 @@ inline void IVF::search( std::fill(results + found, results + k, kPidMax); } +inline void IVF::search_batch( + const float* queries, + size_t num_queries, + size_t k, + size_t nprobe, + PID* results, + float* dists, + std::optional use_hacc, + size_t num_threads +) const { + if (num_queries == 0) { + return; + } + validate_search(queries, k, nprobe, results); + const size_t max_values = std::numeric_limits::max() / sizeof(float); + if (num_queries > max_values / dim_ || num_queries > max_values / k) { + throw std::length_error("IVF query batch is too large"); + } + nprobe = std::min(nprobe, num_cluster_); + const bool hacc = use_hacc.value_or(!raw_reranking_ && ex_bits_ > 2); + const size_t workers = num_queries == 1 || num_threads == 1 + ? 1 + : std::min(resolve_num_threads(num_threads), num_queries); + if (workers == 1) { + for (size_t i = 0; i < num_queries; ++i) { + search( + queries + i * dim_, + k, + nprobe, + results + i * k, + dists == nullptr ? nullptr : dists + i * k, + hacc + ); + } + return; + } + std::atomic failed{false}; + std::exception_ptr error; + const auto capture_error = [&] { +#pragma omp critical(rabitq_ivf_search_error) + { + if (!error) { + error = std::current_exception(); + } + } + failed.store(true, std::memory_order_relaxed); + }; +#pragma omp parallel for num_threads(workers) schedule(dynamic) + for (std::ptrdiff_t query_index = 0; + query_index < static_cast(num_queries); + ++query_index) { + if (failed.load(std::memory_order_relaxed)) { + continue; + } + const size_t i = static_cast(query_index); + try { + search( + queries + i * dim_, + k, + nprobe, + results + i * k, + dists == nullptr ? nullptr : dists + i * k, + hacc + ); + } catch (...) { capture_error(); } + } + if (error) { + std::rethrow_exception(error); + } +} + +inline void IVF::validate_search( + const float* query, size_t k, size_t nprobe, const PID* results +) const { + if (!ready_ || initer_ == nullptr || rotator_ == nullptr || + cluster_lst_.size() != num_cluster_ || batch_storage_.empty() || + id_storage_.size() != num_) { + throw std::logic_error("IVF index must be constructed or loaded before search"); + } + if (query == nullptr) { + throw std::invalid_argument("IVF search query must not be null"); + } + if (results == nullptr) { + throw std::invalid_argument("IVF search results must not be null"); + } + if (k == 0 || k > num_ || k > buffer::kSearchBufferMaxPointCount) { + throw std::invalid_argument("IVF search k must be between 1 and the point count"); + } + if (nprobe == 0) { + throw std::invalid_argument("IVF search nprobe must be positive"); + } +} + inline void IVF::search_cluster( const Cluster& cur_cluster, const SplitBatchQuery& q_obj, diff --git a/include/rabitqlib/index/symqg/qg.hpp b/include/rabitqlib/index/symqg/qg.hpp index 67c07be5..a68e8ffc 100644 --- a/include/rabitqlib/index/symqg/qg.hpp +++ b/include/rabitqlib/index/symqg/qg.hpp @@ -59,11 +59,14 @@ class QuantizedGraph { std::vector centroid_; // rotated global centroid for qg-quant ex_ipfunc quantized_ip_func_ = nullptr; - using RowStorage = std::vector>; + using RowStorage = + std::vector>; // Complete rows are contiguous in both raw and quantized modes: // vector/code, neighbor quantization data, then packed neighbor IDs. // Typed storage establishes float lifetimes for raw vectors; packed portions // are accessed only as bytes, never as float values. + // Builders and load() fill rows before use; leave scalars uninitialized so + // construction workers first-touch their own pages instead of zeroing serially. RowStorage data_; std::unique_ptr> rotator_; // data rotator @@ -234,6 +237,25 @@ class QuantizedGraph { uint32_t* __restrict__ results, float* __restrict__ dists ); + + /** + * Search contiguous row-major queries, using the window set by set_ef(). + * queries contains num_queries * dimension() floats; results and dists + * each hold num_queries * knn elements. Empty batches do no work. + * Scratch is reused within each worker. num_threads defaults to one; + * zero selects the available hardware thread count, capped by the batch. + * Input and output buffers must not overlap. Concurrent searches require + * separate outputs and no index mutation, including set_ef(). + */ + void search_batch( + const float* queries, + size_t num_queries, + uint32_t knn, + uint32_t* results, + float* dists, + size_t num_threads = 1 + ); + /** * Search using caller-owned scratch buffers. * diff --git a/include/rabitqlib/utils/memory.hpp b/include/rabitqlib/utils/memory.hpp index 1c3c6b84..66b20fee 100644 --- a/include/rabitqlib/utils/memory.hpp +++ b/include/rabitqlib/utils/memory.hpp @@ -42,6 +42,43 @@ inline void advise_huge_pages(void* ptr, size_t size) { #endif } +inline constexpr size_t kHugePageSize = 2UL << 20; + +// Huge-page advice is opt-in and based on requested bytes, before rounding. +// Small buffers retain their required alignment; Linux buffers of at least +// 2 MiB also receive huge-page alignment. Other platforms avoid that padding. +template +inline T* align_allocate(size_t nbytes) { + static_assert(Alignment >= alignof(T)); + static_assert((Alignment & (Alignment - 1)) == 0, "Alignment must be a power of two"); + + if (nbytes == 0) { + return nullptr; + } + + const bool huge_pages = HugePage && nbytes >= kHugePageSize; + size_t alignment = Alignment; +#if defined(__linux__) + if (huge_pages && alignment < kHugePageSize) { + alignment = kHugePageSize; + } +#endif + const size_t remainder = nbytes % alignment; + const size_t padding = remainder == 0 ? 0 : alignment - remainder; + if (nbytes > std::numeric_limits::max() - padding) { + throw std::bad_array_new_length(); + } + const size_t size = nbytes + padding; + void* ptr = aligned_allocate_bytes(alignment, size); + if (ptr == nullptr) { + throw std::bad_alloc(); + } + if (huge_pages) { + advise_huge_pages(ptr, size); + } + return static_cast(ptr); +} + template class AlignedAllocator { private: @@ -51,18 +88,6 @@ class AlignedAllocator { template using ReboundAllocator = AlignedAllocator; - [[nodiscard]] static constexpr size_t aligned_size(size_t nbytes) { - const size_t remainder = nbytes % Alignment; - if (remainder == 0) { - return nbytes; - } - const size_t padding = Alignment - remainder; - if (nbytes > std::numeric_limits::max() - padding) { - throw std::bad_array_new_length(); - } - return nbytes + padding; - } - public: using value_type = T; using is_always_equal = std::true_type; @@ -94,24 +119,36 @@ class AlignedAllocator { throw std::bad_array_new_length(); } - if (n == 0) { - return nullptr; - } - - const auto nbytes = aligned_size(n * sizeof(T)); - auto* ptr = aligned_allocate_bytes(Alignment, nbytes); - if (ptr == nullptr) { - throw std::bad_alloc(); - } - if (HugePage) { - advise_huge_pages(ptr, nbytes); - } - return reinterpret_cast(ptr); + return align_allocate(n * sizeof(T)); } void deallocate(T* ptr, [[maybe_unused]] std::size_t n) { aligned_deallocate(ptr); } }; +// Retain alignment and object lifetimes without zero-filling scalar storage. +// Callers must initialize every element before reading it. +template +class DefaultInitAlignedAllocator : public AlignedAllocator { + template + using ReboundAllocator = DefaultInitAlignedAllocator; + + public: + template + struct rebind { + using other = ReboundAllocator; + }; + + constexpr DefaultInitAlignedAllocator() noexcept = default; + + template + constexpr explicit DefaultInitAlignedAllocator(const ReboundAllocator&) noexcept {} + + template + void construct(U* ptr) noexcept(std::is_nothrow_default_constructible_v) { + ::new (static_cast(ptr)) U; + } +}; + template struct Allocator { public: @@ -135,37 +172,12 @@ struct Allocator { } }; -template -inline T* align_allocate(size_t nbytes) { - static_assert(Alignment >= alignof(T)); - static_assert((Alignment & (Alignment - 1)) == 0, "Alignment must be a power of two"); - - if (nbytes == 0) { - return nullptr; - } - - const size_t remainder = nbytes % Alignment; - const size_t padding = remainder == 0 ? 0 : Alignment - remainder; - if (nbytes > std::numeric_limits::max() - padding) { - throw std::bad_array_new_length(); - } - const size_t size = nbytes + padding; - void* ptr = aligned_allocate_bytes(Alignment, size); - if (ptr == nullptr) { - throw std::bad_alloc(); - } - if (HugePage) { - advise_huge_pages(ptr, size); - } - return static_cast(ptr); -} - -inline constexpr size_t kHugePageSize = 2UL << 20; - -// Allocate a large buffer backed by transparent huge pages. +// Use the same size-aware policy as huge-page-enabled STL storage. Advice is +// best-effort; actual page sizes remain under the operating system's control. template inline T* huge_page_allocate(size_t nbytes) { - return align_allocate(nbytes); + constexpr size_t kAlignment = alignof(T) > 64 ? alignof(T) : 64; + return align_allocate(nbytes); } static inline void prefetch_l1(const void* addr) { diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index 3d3a212d..1073c890 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/python_bindings/ivf_bindings.cpp @@ -196,22 +196,15 @@ class IvfIndex { dists_data, dists_data + dists.size(), std::numeric_limits::infinity() ); - rabitqlib::ivf::parallel_for( - 0, + index_->search_batch( + query_array.data(), nq, - num_threads, - [&](size_t idx, size_t /*threadId*/) { - const float* query = query_array.data() + (idx * dim_); - auto* result_ids = ids_data + (idx * k); - auto* result_dists = dists_data + (idx * k); - if (high_accuracy.has_value()) { - index_->search( - query, k, nprobe, result_ids, result_dists, *high_accuracy - ); - } else { - index_->search(query, k, nprobe, result_ids, result_dists); - } - } + k, + nprobe, + ids_data, + dists_data, + high_accuracy, + num_threads ); return py::make_tuple(ids, dists); diff --git a/python_bindings/symqg_bindings.cpp b/python_bindings/symqg_bindings.cpp index 97cceec2..1d463517 100644 --- a/python_bindings/symqg_bindings.cpp +++ b/python_bindings/symqg_bindings.cpp @@ -4,11 +4,8 @@ #include #include -#include -#include #include #include -#include #include #include #include @@ -22,7 +19,6 @@ #include "rabitqlib/index/symqg/qg.hpp" #include "rabitqlib/index/symqg/qg_builder.hpp" #include "rabitqlib/utils/rotator.hpp" -#include "rabitqlib/utils/tools.hpp" namespace py = pybind11; @@ -118,45 +114,14 @@ class SymqgIndex { auto dists = py::array_t(shape); auto* ids_data = ids.mutable_data(); auto* dists_data = dists.mutable_data(); - - const auto* queries_data = query_array.data(); - const auto workers = static_cast(std::max( - 1, - std::min( - {resolve_num_threads(num_threads), - nq, - static_cast(std::numeric_limits::max())} - ) - )); - std::atomic failed{false}; - std::exception_ptr error; -#pragma omp parallel for num_threads(workers) if (workers > 1) schedule(dynamic) - for (std::ptrdiff_t query_index = 0; query_index < static_cast(nq); - ++query_index) { - const size_t idx = static_cast(query_index); - if (failed.load(std::memory_order_relaxed)) { - continue; - } - try { - index_->search( - queries_data + (idx * dim_), - static_cast(k), - ids_data + (idx * k), - dists_data + (idx * k) - ); - } catch (...) { -#pragma omp critical(rabitq_symqg_search_error) - { - if (!error) { - error = std::current_exception(); - } - } - failed.store(true, std::memory_order_relaxed); - } - } - if (error) { - std::rethrow_exception(error); - } + index_->search_batch( + query_array.data(), + nq, + static_cast(k), + ids_data, + dists_data, + num_threads + ); return py::make_tuple(ids, dists); } @@ -230,6 +195,14 @@ void register_symqg(py::module_& m) { py::arg("ef"), py::arg("num_threads") = 1 ) + .def( + "search_batch", + &SymqgIndex::search, + py::arg("queries"), + py::arg("k"), + py::arg("ef"), + py::arg("num_threads") = 1 + ) .def("save", &SymqgIndex::save, py::arg("path")) .def_static("load", &SymqgIndex::load, py::arg("path")) .def_property_readonly("dim", &SymqgIndex::dim) diff --git a/sample/cpp/ivf_rabitq_querying.cpp b/sample/cpp/ivf_rabitq_querying.cpp index 7c10f463..39a8fabe 100644 --- a/sample/cpp/ivf_rabitq_querying.cpp +++ b/sample/cpp/ivf_rabitq_querying.cpp @@ -95,21 +95,17 @@ int run(int argc, char** argv) { << ").\n"; } size_t total_correct = 0; - float total_time = 0; - std::vector results(topk); + std::vector results(nq * topk); + stopw.reset(); + ivf.search_batch( + query.data(), nq, topk, nprobe, results.data(), nullptr, use_hacc + ); + const float total_time = stopw.get_elapsed_micro(); for (size_t i = 0; i < nq; i++) { - stopw.reset(); - const float* vector = &query(static_cast(i), 0); - if (use_hacc.has_value()) { - ivf.search(vector, topk, nprobe, results.data(), *use_hacc); - } else { - ivf.search(vector, topk, nprobe, results.data()); - } - total_time += stopw.get_elapsed_micro(); for (size_t j = 0; j < topk; j++) { for (size_t k = 0; k < topk; k++) { if (gt(static_cast(i), - static_cast(k)) == results[j]) { + static_cast(k)) == results[i * topk + j]) { total_correct++; break; } diff --git a/sample/cpp/symqg_querying.cpp b/sample/cpp/symqg_querying.cpp index 4453bba2..6cc7c91f 100644 --- a/sample/cpp/symqg_querying.cpp +++ b/sample/cpp/symqg_querying.cpp @@ -68,21 +68,16 @@ int run(int argc, char** argv) { size_t total_correct = 0; float total_time = 0; qg.set_ef(ef); - std::vector results(topk); - std::vector dists(topk); + std::vector results(total_count); + std::vector dists(total_count); + stopw.reset(); + qg.search_batch(query.data(), nq, topk, results.data(), dists.data()); + total_time = stopw.get_elapsed_micro(); for (size_t z = 0; z < nq; z++) { - stopw.reset(); - qg.search( - &query(static_cast(z), 0), - topk, - results.data(), - dists.data() - ); - total_time += stopw.get_elapsed_micro(); for (size_t y = 0; y < topk; y++) { for (size_t k = 0; k < topk; k++) { if (gt(static_cast(z), - static_cast(k)) == results[y]) { + static_cast(k)) == results[z * topk + y]) { total_correct++; break; } diff --git a/sample/python/symqg_querying.py b/sample/python/symqg_querying.py index aab3c9d7..be839903 100644 --- a/sample/python/symqg_querying.py +++ b/sample/python/symqg_querying.py @@ -37,7 +37,7 @@ def main(args=None) -> None: for i_probe, ef in enumerate(EFS): for r in range(args.test_rounds): t0 = perf_counter() - ids, _ = idx.search( + ids, _ = idx.search_batch( queries, k=args.topk, ef=ef, num_threads=args.num_threads ) elapsed = perf_counter() - t0 # seconds diff --git a/src/index/qg.cpp b/src/index/qg.cpp index 9a93b80c..fd15130c 100644 --- a/src/index/qg.cpp +++ b/src/index/qg.cpp @@ -1,10 +1,14 @@ #include "rabitqlib/index/symqg/qg.hpp" +#include + #include +#include #include #include #include #include +#include #include #include #include @@ -29,6 +33,7 @@ #include "rabitqlib/utils/path.hpp" #include "rabitqlib/utils/rotator.hpp" #include "rabitqlib/utils/space.hpp" +#include "rabitqlib/utils/tools.hpp" #include "rabitqlib/utils/visited_set.hpp" namespace rabitqlib::symqg { @@ -164,7 +169,7 @@ void QuantizedGraph::copy_vectors(const float* data, size_t num_threads) { std::vector rotated_data(padded_dim_); std::vector quantized_data(padded_dim_); -#pragma omp for schedule(dynamic) +#pragma omp for schedule(static) for (std::ptrdiff_t index = 0; index < static_cast(num_points_); ++index) { const size_t i = static_cast(index); @@ -195,7 +200,7 @@ void QuantizedGraph::copy_vectors(const float* data, size_t num_threads) } return; } -#pragma omp parallel for schedule(dynamic) num_threads(thread_count) +#pragma omp parallel for schedule(static) num_threads(thread_count) for (std::ptrdiff_t index = 0; index < static_cast(num_points_); ++index) { const size_t i = static_cast(index); @@ -433,6 +438,99 @@ void QuantizedGraph::search( ); } +void QuantizedGraph::search_batch( + const float* queries, + size_t num_queries, + uint32_t knn, + uint32_t* results, + float* dists, + size_t num_threads +) { + if (num_queries == 0) { + return; + } + validate_search(queries, knn, results, dists); + const size_t max_values = std::numeric_limits::max() / sizeof(float); + if (num_queries > max_values / dim_ || num_queries > max_values / knn) { + throw std::length_error("QuantizedGraph query batch is too large"); + } + if (num_queries == 1) { + search(queries, knn, results, dists); + return; + } + + struct SearchScratch { + std::vector rotated_query; + std::vector estimated_distances; + std::vector lookup_table; + BatchQuery batch_query; + buffer::SearchBuffer search_pool; + buffer::SearchBuffer result_pool; + VisitedSet visited; + + SearchScratch(size_t padded_dim, size_t degree, size_t ef, size_t k, size_t points) + : rotated_query(padded_dim) + , estimated_distances(degree) + , lookup_table(padded_dim * 4) + , search_pool(ef) + , result_pool(k) + , visited(points, points / 10) {} + }; + + const auto workers = static_cast(std::min( + {resolve_num_threads(num_threads), + num_queries, + static_cast(std::numeric_limits::max())} + )); + std::atomic failed{false}; + std::exception_ptr error; + const auto capture_error = [&] { +#pragma omp critical(rabitq_symqg_search_error) + { + if (!error) { + error = std::current_exception(); + } + } + failed.store(true, std::memory_order_relaxed); + }; +#pragma omp parallel num_threads(workers) if (workers > 1) + { + // Allocate and release scratch on its worker, retaining it across queries. + std::optional scratch; + try { + scratch.emplace(padded_dim_, degree_bound_, ef_, knn, num_points_); + } catch (...) { capture_error(); } +#pragma omp for schedule(dynamic) + for (std::ptrdiff_t query_index = 0; + query_index < static_cast(num_queries); + ++query_index) { + if (failed.load(std::memory_order_relaxed)) { + continue; + } + const size_t i = static_cast(query_index); + try { + auto& local = *scratch; + search_with_scratch( + queries + i * dim_, + knn, + results + i * knn, + dists + i * knn, + local.rotated_query.data(), + local.estimated_distances.data(), + local.lookup_table.data(), + local.batch_query, + local.search_pool, + local.result_pool, + local.visited + ); + } catch (...) { capture_error(); } + } + } + if (error) { + std::rethrow_exception(error); + } +} + void QuantizedGraph::validate_search( const float* query, uint32_t k, const uint32_t* results, const float* dists ) const { diff --git a/tests/consumer/main.cpp b/tests/consumer/main.cpp index c4a68e5b..6040e951 100644 --- a/tests/consumer/main.cpp +++ b/tests/consumer/main.cpp @@ -7,6 +7,9 @@ #include "rabitqlib/clustering/qgkmeans.hpp" #include "rabitqlib/clustering/rabitqkmeans.hpp" +#include "rabitqlib/index/ivf/ivf.hpp" +#include "rabitqlib/index/symqg/qg.hpp" +#include "rabitqlib/index/symqg/qg_builder.hpp" #include "rabitqlib/utils/cpu_features.hpp" namespace { @@ -89,5 +92,35 @@ int main() { std::cerr << "Installed RaBitQKMeans returned invalid clustering results\n"; return 1; } + + rabitqlib::ivf::IVF index(kPoints, kDimension, flat.k, 4); + index.construct(data.data(), flat.centroids.data(), flat.assignments.data(), false, 2); + std::vector ids(kPoints); + std::vector distances(kPoints); + index.search_batch(data.data(), kPoints, 1, flat.k, ids.data(), distances.data()); + for (size_t i = 0; i < kPoints; ++i) { + rabitqlib::PID id = 0; + float distance = 0; + index.search(data.data() + i * kDimension, 1, flat.k, &id, &distance); + if (ids[i] != id || distances[i] != distance) { + std::cerr << "Installed IVF batch and single-query results differ\n"; + return 1; + } + } + + rabitqlib::symqg::QuantizedGraph symqg(kPoints, kDimension, 32); + rabitqlib::symqg::QGBuilder builder(symqg, 64, data.data(), 2); + builder.build(); + symqg.set_ef(64); + symqg.search_batch(data.data(), kPoints, 1, ids.data(), distances.data(), 2); + for (size_t i = 0; i < kPoints; ++i) { + rabitqlib::PID id = 0; + float distance = 0; + symqg.search(data.data() + i * kDimension, 1, &id, &distance); + if (ids[i] != id || distances[i] != distance) { + std::cerr << "Installed SymphonyQG batch and single-query results differ\n"; + return 1; + } + } return 0; } diff --git a/tests/python/test_examples.py b/tests/python/test_examples.py index f8ed6851..cc20e8ce 100644 --- a/tests/python/test_examples.py +++ b/tests/python/test_examples.py @@ -88,6 +88,7 @@ def test_query_timing_handles_unchanged_clock(kind, elapsed, monkeypatch, capsys index = SimpleNamespace( dim=2, num_clusters=2, search=lambda *a, **kw: (ids, np.zeros((2, 1))) ) + index.search_batch = index.search index_name = {"ivf": "IvfIndex", "hnsw": "HnswIndex", "symqg": "SymqgIndex"}[kind] monkeypatch.setattr(module, index_name, SimpleNamespace(load=lambda path: index)) monkeypatch.setattr(module, "read_fvecs", lambda path: np.zeros((2, 2))) diff --git a/tests/python/test_hnsw.py b/tests/python/test_hnsw.py index fd8e7b2e..fc6858f1 100644 --- a/tests/python/test_hnsw.py +++ b/tests/python/test_hnsw.py @@ -296,6 +296,30 @@ def test_resize_below_element_count_raises(partial_hnsw): assert partial_hnsw.max_elements == N_VECTORS +@pytest.mark.parametrize("nbits", [1, 9]) +@pytest.mark.parametrize("metric", ["l2", "ip"]) +def test_resize_across_allocation_threshold_preserves_search(tmp_path, nbits, metric): + rng = np.random.default_rng(73) + data = rng.standard_normal((97, 65)).astype(np.float32) + queries = rng.standard_normal((7, 65)).astype(np.float32) + centroids = data.mean(axis=0, keepdims=True) + index = HnswIndex(65, len(data), M=8, nbits=nbits, metric=metric) + index.build(data, centroids, np.zeros(len(data), dtype=np.uint32)) + expected_ids, expected_distances = index.search(queries, k=7, ef=97) + + # The level-0 allocation moves from below 2 MiB to above it and back. + for capacity in [32768, len(data)]: + index.resize(capacity) + path = str(tmp_path / "resized.index") + index.save(path) + index = HnswIndex.load(path) + assert index.max_elements == capacity + assert index.num_points == len(data) + ids, distances = index.search(queries, k=7, ef=97) + np.testing.assert_array_equal(ids, expected_ids) + np.testing.assert_array_equal(distances, expected_distances) + + def test_added_points_survive_save_and_load( partial_hnsw, base_data, clusters, tmp_path ): diff --git a/tests/python/test_ivf.py b/tests/python/test_ivf.py index 0f59a2f5..99fe6550 100644 --- a/tests/python/test_ivf.py +++ b/tests/python/test_ivf.py @@ -75,6 +75,41 @@ def test_single_query(built_ivf, query_data): assert dists.shape == (1, 1) +@pytest.mark.parametrize("bits", [1, 2, 3, 4, 5, 6, 7, 8, 9, 32]) +@pytest.mark.parametrize("metric", ["l2", "ip"]) +def test_batch_search_matches_individual_queries(bits, metric): + rng = np.random.default_rng(73) + data = rng.standard_normal((97, 65)).astype(np.float32) + queries = rng.standard_normal((13, 65)).astype(np.float32) + centroids = data[:4].copy() + labels = np.arange(len(data), dtype=np.uint32) % 3 + index = IvfIndex(65, len(data), 4, bits, metric=metric) + index.build(data, centroids, labels) + for hacc in [None, False, True]: + for nprobe in [1, 99]: + kwargs = dict(k=7, nprobe=nprobe, high_accuracy=hacc) + individual = [index.search(q[None, :], **kwargs) for q in queries] + expected_ids = np.concatenate([result[0] for result in individual]) + expected_distances = np.concatenate([result[1] for result in individual]) + for threads in [1, 4]: + ids, distances = index.search(queries, num_threads=threads, **kwargs) + np.testing.assert_array_equal(ids, expected_ids) + np.testing.assert_array_equal(distances, expected_distances) + one_ids, one_distances = index.search( + queries[:1], num_threads=threads, **kwargs + ) + np.testing.assert_array_equal(one_ids, expected_ids[:1]) + np.testing.assert_array_equal(one_distances, expected_distances[:1]) + empty_ids, empty_distances = index.search( + queries[:0], num_threads=threads, **kwargs + ) + assert empty_ids.shape == empty_distances.shape == (0, 7) + index.remove(np.arange(len(data))) + ids, distances = index.search(queries, k=7, nprobe=4, num_threads=4) + assert np.all(ids == np.iinfo(np.uint32).max) + assert np.all(np.isinf(distances)) + + def test_k_equals_one(built_ivf, query_data): ids, dists = built_ivf.search(query_data, k=1, nprobe=1) assert ids.shape == (N_QUERIES, 1) @@ -660,3 +695,31 @@ def test_remove_and_add_accept_any_integer_dtype_and_empty_input(): assert idx.remove(np.array([5], dtype=np.int16)) == 1 assert len(idx.add(new[:3], cluster_ids=[0, 1, 2])) == 3 assert len(idx.add(new[:3], cluster_ids=np.array([0, 1, 2], dtype=np.uint16))) == 3 + + +@pytest.mark.parametrize("bits", [1, 4, 9, 32]) +@pytest.mark.parametrize("metric", ["l2", "ip"]) +def test_rebuild_matches_incremental_encoding_with_empty_clusters( + tmp_path, bits, metric +): + rng = np.random.default_rng(731) + data = rng.standard_normal((130, 65), dtype=np.float32) + centers = rng.standard_normal((4, 65), dtype=np.float32) + labels = np.full(130, 3, dtype=np.uint32) + labels[0] = 1 + labels[1:33] = 2 + incremental = IvfIndex(65, 1, 4, nbits=bits, metric=metric) + incremental.build(data[:1], centers, labels[:1], num_threads=1) + incremental.add(data[1:], cluster_ids=labels[1:], num_threads=1) + path = tmp_path / "incremental.index" + incremental.save(str(path)) + expected_bytes = path.read_bytes() + expected_ids, expected_distances = incremental.search(data[:3], k=10, nprobe=4) + rebuilt = IvfIndex.load(str(path)) + for threads in (1, 4): + rebuilt.build(data, centers, labels, num_threads=threads) + rebuilt.save(str(path)) + assert path.read_bytes() == expected_bytes + ids, distances = rebuilt.search(data[:3], k=10, nprobe=4) + np.testing.assert_array_equal(ids, expected_ids) + np.testing.assert_array_equal(distances, expected_distances) diff --git a/tests/python/test_qgkmeans.py b/tests/python/test_qgkmeans.py index d29719fa..cc8d62b0 100644 --- a/tests/python/test_qgkmeans.py +++ b/tests/python/test_qgkmeans.py @@ -151,10 +151,11 @@ def test_exact_final_assignment_preserves_shifted_l2_accuracy(): @pytest.mark.parametrize("spherical", [False, True]) -def test_exact_final_assignment_uses_first_duplicate_centroid(spherical): +@pytest.mark.parametrize("model_type", [QGKMeans, RaBitQKMeans]) +def test_exact_final_assignment_uses_first_duplicate_centroid(spherical, model_type): x = np.zeros((65, 65), dtype=np.float32) x[:, 0] = 1.0 - clustering = QGKMeans( + clustering = model_type( 65, K, niter=1, diff --git a/tests/python/test_symqg.py b/tests/python/test_symqg.py index 117f14bc..67fc9f63 100644 --- a/tests/python/test_symqg.py +++ b/tests/python/test_symqg.py @@ -69,6 +69,61 @@ def test_single_query(built_symqg, query_data): assert ids.shape == (1, 1) +@pytest.mark.parametrize("bits", [0, 4, 8]) +@pytest.mark.parametrize("metric", ["l2", "ip"]) +@pytest.mark.parametrize("method", ["search", "search_batch"]) +def test_batch_search_matches_individual_queries(tmp_path, bits, metric, method): + rng = np.random.default_rng(73) + data = rng.standard_normal((97, 65)).astype(np.float32) + queries = rng.standard_normal((19, 65)).astype(np.float32) + queries[4] = queries[0] + index = SymqgIndex(65, 32, metric=metric, quantization_bits=bits) + index.build(data, ef_construction=48, num_threads=1) + path = str(tmp_path / "batch.index") + index.save(path) + restored = SymqgIndex.load(path) + search = getattr(restored, method) + + for k, ef in [(7, 64), (1, 1), (7, 64)]: + individual = [index.search(query[None, :], k=k, ef=ef) for query in queries] + expected_ids = np.concatenate([result[0] for result in individual]) + expected_distances = np.concatenate([result[1] for result in individual]) + for threads in [1, 4]: + ids, distances = search(queries, k=k, ef=ef, num_threads=threads) + np.testing.assert_array_equal(ids, expected_ids) + np.testing.assert_array_equal(distances, expected_distances) + empty_ids, empty_distances = search( + queries[:0], k=k, ef=ef, num_threads=threads + ) + assert empty_ids.shape == empty_distances.shape == (0, k) + + +def test_search_batch_accepts_noncontiguous_queries(built_symqg, query_data): + queries = query_data[::2] + assert not queries.flags.c_contiguous + expected = built_symqg.search(queries.copy(), k=5, ef=_EF) + actual = built_symqg.search_batch(queries, k=5, ef=_EF, num_threads=0) + np.testing.assert_array_equal(actual[0], expected[0]) + np.testing.assert_array_equal(actual[1], expected[1]) + + +def test_search_batch_validates_inputs(built_symqg, query_data): + with pytest.raises(ValueError, match="queries must be a 2D NumPy array"): + built_symqg.search_batch(query_data[0], k=1, ef=_EF) + with pytest.raises(ValueError, match="query dimension does not match index dim"): + built_symqg.search_batch(query_data[:, :-1], k=1, ef=_EF) + for k in [0, N_VECTORS + 1]: + with pytest.raises(ValueError, match="k must be between 1 and num_points"): + built_symqg.search_batch(query_data, k=k, ef=_EF) + with pytest.raises(ValueError, match="ef must be positive"): + built_symqg.search_batch(query_data, k=1, ef=0) + with pytest.raises(ValueError, match="ef must be at least k"): + built_symqg.search_batch(query_data, k=5, ef=4) + unbuilt = SymqgIndex(DIM, max_degree=_MAX_DEGREE) + with pytest.raises(RuntimeError, match="must be built or loaded before search"): + unbuilt.search_batch(query_data, k=1, ef=_EF) + + # ── search correctness ──────────────────────────────────────────────────────── diff --git a/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp b/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp index 36e8902a..4dc86a45 100644 --- a/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp +++ b/tests/unit/rabitqlib/clustering/qgkmeans_test.cpp @@ -379,28 +379,39 @@ void expect_exact_assignment_matches_scalar( } for (const uint32_t threads : {1U, 4U}) { SCOPED_TRACE(threads); - std::vector labels(points, kPidMax); - std::vector distances(points, std::numeric_limits::quiet_NaN()); - detail::exact_assign( - data.data(), - centroids.data(), - points, - dimension, - clusters, - spherical, - threads, - labels.data(), - distances.data() - ); - EXPECT_EQ(labels, expected_labels); - EXPECT_EQ(distances, expected_distances); + // Exercise no hint, separate hints, and in-place hints as used by Lloyd. + // The last centroid deliberately loses ties to earlier centroid IDs. + const std::vector previous(points, static_cast(clusters - 1)); + for (const int hint_mode : {0, 1, 2}) { + SCOPED_TRACE(hint_mode); + std::vector labels = + hint_mode == 2 ? previous : std::vector(points, kPidMax); + std::vector distances(points, std::numeric_limits::quiet_NaN()); + const PID* hints = hint_mode == 0 + ? nullptr + : (hint_mode == 1 ? previous.data() : labels.data()); + detail::exact_assign( + data.data(), + centroids.data(), + points, + dimension, + clusters, + spherical, + threads, + labels.data(), + distances.data(), + hints + ); + EXPECT_EQ(labels, expected_labels); + EXPECT_EQ(distances, expected_distances); + } } } TEST(QGKMeansTest, ExactAssignmentMatchesScalarAcrossMatrixTileTails) { std::mt19937 rng(741); std::uniform_real_distribution distribution(-2.0F, 2.0F); - for (const size_t size : {65U, 129U}) { + for (const size_t size : {65U, 129U, 257U}) { SCOPED_TRACE(size); std::vector data(size * size); std::vector centroids(257 * size); @@ -429,6 +440,20 @@ TEST(QGKMeansTest, ExactAssignmentMatchesScalarAtMaximumDimension) { } } +TEST(QGKMeansTest, ExactAssignmentMatchesScalarWithManyPointBlocks) { + constexpr size_t kDimension = 17; + std::mt19937 rng(385); + std::uniform_real_distribution distribution(-2.0F, 2.0F); + std::vector data(4097 * kDimension); + std::vector centroids(257 * kDimension); + std::generate(data.begin(), data.end(), [&] { return distribution(rng); }); + std::generate(centroids.begin(), centroids.end(), [&] { return distribution(rng); }); + for (const bool spherical : {false, true}) { + SCOPED_TRACE(spherical); + expect_exact_assignment_matches_scalar(data, centroids, kDimension, spherical); + } +} + TEST(QGKMeansTest, ExactAssignmentPreservesCallerOpenMPSettings) { struct SavedOpenMPSettings { int threads = omp_get_max_threads(); @@ -455,8 +480,8 @@ TEST(QGKMeansTest, ExactAssignmentPreservesCallerOpenMPSettings) { #endif constexpr size_t kDimension = 65; const std::vector centroids(3 * kDimension, 1.0F); - // Three rows reduce the requested team to one; 65 rows use two blocks. - for (const size_t points : {3U, 65U}) { + // Cover a single worker, a partial row block, and multiple workers. + for (const size_t points : {3U, 65U, 513U}) { SCOPED_TRACE(points); const std::vector data(points * kDimension, 0.0F); std::vector labels(points); diff --git a/tests/unit/rabitqlib/index/ivf_test.cpp b/tests/unit/rabitqlib/index/ivf_test.cpp index 07e4d04c..c79899db 100644 --- a/tests/unit/rabitqlib/index/ivf_test.cpp +++ b/tests/unit/rabitqlib/index/ivf_test.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -25,6 +26,112 @@ namespace rabitqlib::ivf { namespace { +TEST(IvfSearchTest, BatchSearchMatchesIndividualQueries) { + constexpr size_t kNum = 97, kDim = 65, kClusters = 4, kQueries = 13, k = 7; + std::mt19937 rng(73); + std::uniform_real_distribution random(-1, 1); + std::vector data(kNum * kDim), centroids(kClusters * kDim); + std::vector queries(kQueries * kDim); + std::vector labels(kNum); + std::generate(data.begin(), data.end(), [&] { return random(rng); }); + std::generate(centroids.begin(), centroids.end(), [&] { return random(rng); }); + std::generate(queries.begin(), queries.end(), [&] { return random(rng); }); + for (size_t i = 0; i < kNum; ++i) { + labels[i] = static_cast(i % (kClusters - 1)); + } + for (auto metric : {METRIC_L2, METRIC_IP}) { + for (size_t bits : {1U, 2U, 3U, 4U, 5U, 6U, 7U, 8U, 9U, 32U}) { + IVF index(kNum, kDim, kClusters, bits, metric); + index.construct(data.data(), centroids.data(), labels.data(), false, 1); + std::array expected{}, actual{}, ids_only{}; + std::array expected_distances{}, actual_distances{}; + for (size_t nprobe : {1U, 99U}) { + for (std::optional hacc : + {std::optional{}, + std::optional(false), + std::optional(true)}) { + for (size_t i = 0; i < kQueries; ++i) { + index.search( + queries.data() + i * kDim, + k, + nprobe, + expected.data() + i * k, + expected_distances.data() + i * k, + hacc.value_or(bits >= 4 && bits <= 9) + ); + } + for (size_t threads : {1U, 4U}) { + index.search_batch( + queries.data(), + kQueries, + k, + nprobe, + actual.data(), + actual_distances.data(), + hacc, + threads + ); + index.search_batch( + queries.data(), + kQueries, + k, + nprobe, + ids_only.data(), + nullptr, + hacc, + threads + ); + EXPECT_EQ(actual, expected); + EXPECT_EQ(ids_only, expected); + EXPECT_EQ(actual_distances, expected_distances); + } + } + } + EXPECT_NO_THROW(index.search_batch(nullptr, 0, k, 1, nullptr)); + EXPECT_THROW( + index.search_batch( + queries.data(), std::numeric_limits::max(), k, 1, actual.data() + ), + std::length_error + ); + EXPECT_THROW( + index.search_batch(nullptr, 1, k, 1, actual.data()), std::invalid_argument + ); + EXPECT_THROW( + index.search_batch(queries.data(), 1, 0, 1, actual.data()), + std::invalid_argument + ); + EXPECT_THROW( + index.search_batch(queries.data(), 1, k, 0, actual.data()), + std::invalid_argument + ); + EXPECT_THROW( + index.search_batch(queries.data(), 1, k, 1, nullptr), std::invalid_argument + ); + std::vector removed(kNum); + std::iota(removed.begin(), removed.end(), PID{0}); + index.remove(removed.data(), kNum); + index.search_batch( + queries.data(), + kQueries, + k, + 4, + actual.data(), + actual_distances.data(), + std::nullopt, + 4 + ); + for (size_t i = 0; i < actual.size(); ++i) { + EXPECT_EQ(actual[i], kPidMax); + EXPECT_EQ(actual_distances[i], std::numeric_limits::infinity()); + } + } + } + IVF empty; + PID id = 0; + EXPECT_THROW(empty.search_batch(queries.data(), 1, 1, 1, &id), std::logic_error); +} + TEST(IvfSearchTest, BatchCandidatesPreserveTiesAndTailCount) { buffer::SearchBuffer knns(3); const std::array ids{0, 1, 2, 3, 4, 5}; @@ -530,6 +637,48 @@ TEST(IvfAddTest, EveryWayOfAddingTheSamePointsGivesByteIdenticalIndexes) { std::remove(path.c_str()); } +TEST(IvfConstructionTest, RebuildMatchesIncrementalEncodingWithSkewedClustersAndTails) { + // Cluster sizes 0, 1, 32, 97 cover empty rows and multiple physical tail batches. + constexpr size_t kCount = 130; + const auto points = random_rows(kCount, 731); + const auto centroids = random_rows(kDynClusters, 732); + std::vector assigned(kCount, 3); + assigned[0] = 1; + std::fill(assigned.begin() + 1, assigned.begin() + 33, PID{2}); + const std::string path = ::testing::TempDir() + "rabitq_ivf_rebuild.index"; + for (auto metric : {METRIC_L2, METRIC_IP}) { + for (auto rotation : {RotatorType::FhtKacRotator, RotatorType::MatrixRotator}) { + for (size_t bits : {1UL, 2UL, 3UL, 4UL, 5UL, 6UL, 7UL, 8UL, 9UL, 32UL}) { + SCOPED_TRACE( + ::testing::Message() << "bits=" << bits << " metric=" << metric + ); + IVF incremental(1, kDynDim, kDynClusters, bits, metric, rotation); + incremental.construct( + points.data(), centroids.data(), assigned.data(), false, 1 + ); + incremental.add( + points.data() + kDynDim, kCount - 1, assigned.data() + 1, false, 1 + ); + const std::string expected = saved_bytes(incremental, path); + IVF rebuilt; + rebuilt.load(path.c_str()); + for (size_t threads : {1UL, 4UL}) { + rebuilt.construct( + points.data(), centroids.data(), assigned.data(), false, threads + ); + EXPECT_EQ(saved_bytes(rebuilt, path), expected); + const auto actual_hits = search_everything(rebuilt, points.data()); + const auto expected_hits = + search_everything(incremental, points.data()); + EXPECT_EQ(actual_hits.ids, expected_hits.ids); + EXPECT_EQ(actual_hits.distances, expected_hits.distances); + } + } + } + } + std::remove(path.c_str()); +} + TEST(IvfAddTest, KeepsExistingDistancesAndQuantizesNewPointsAccurately) { constexpr size_t kOld = 150; constexpr size_t kNew = 150; diff --git a/tests/unit/rabitqlib/index/qg_test.cpp b/tests/unit/rabitqlib/index/qg_test.cpp index f4f5ba1f..56bf8d87 100644 --- a/tests/unit/rabitqlib/index/qg_test.cpp +++ b/tests/unit/rabitqlib/index/qg_test.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -143,6 +144,9 @@ struct QGConstructionTestAccess { } return result; } + static void fill_storage(QuantizedGraph& graph, float value) { + std::fill(graph.data_.begin(), graph.data_.end(), value); + } static void check_contiguous_rows(QuantizedGraph& graph) { const size_t vector_bytes = graph.is_quantized() @@ -1051,6 +1055,42 @@ TEST(QGConstructionTest, InitializesUnusedPartialBatchFactors) { } } +TEST(QGConstructionTest, FullyInitializesRowsBeforeSearchAndSave) { + constexpr size_t kCount = 65, kDim = 65; + std::vector data(kCount * kDim); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = std::sin(static_cast(i) * 0.13F); + } + for (auto metric : {METRIC_L2, METRIC_IP}) { + for (size_t bits : {0U, 4U, 8U}) { + for (auto init : {QGInitialization::PiPNN, QGInitialization::Random}) { + SCOPED_TRACE( + ::testing::Message() + << metric << "/" << bits << "/" << static_cast(init) + ); + QuantizedGraph zeroed( + kCount, kDim, 32, metric, RotatorType::FhtKacRotator, bits, 42 + ); + QuantizedGraph poisoned( + kCount, kDim, 32, metric, RotatorType::FhtKacRotator, bits, 42 + ); + QGConstructionTestAccess::fill_storage(zeroed, 0.0F); + QGConstructionTestAccess::fill_storage( + poisoned, std::numeric_limits::quiet_NaN() + ); + QGBuilder zeroed_builder(zeroed, 48, data.data(), 1, init, 42); + QGBuilder poisoned_builder(poisoned, 48, data.data(), 1, init, 42); + zeroed_builder.build(); + poisoned_builder.build(); + EXPECT_EQ( + QGConstructionTestAccess::rows(zeroed), + QGConstructionTestAccess::rows(poisoned) + ); + } + } + } +} + TEST(QGConstructionTest, RefinesPartialSeedOnceAfterReleasingInputs) { constexpr size_t kCount = 97, kDim = 65, kDegree = 64; constexpr std::array kDegrees{0, 1, 31, 32, 33, 64}; @@ -1074,6 +1114,9 @@ TEST(QGConstructionTest, RefinesPartialSeedOnceAfterReleasingInputs) { QuantizedGraph graph( kCount, kDim, kDegree, metric, RotatorType::FhtKacRotator, bits ); + QGConstructionTestAccess::fill_storage( + graph, std::numeric_limits::quiet_NaN() + ); auto builder = QGConstructionTestAccess::from_graph( graph, 96, data.data(), offsets, edges, 2 ); @@ -1521,7 +1564,7 @@ TEST(QGSearchTest, RejectsInvalidKAndEfInsteadOfReturningPartialResults) { })); } -TEST(QGSearchTest, ParallelQueriesMatchSerialAcrossIndexesAndSettings) { +TEST(QGSearchTest, BatchAndParallelQueriesMatchSerialAcrossIndexesAndSettings) { constexpr size_t kNumQueries = 9; constexpr size_t kTopK = 5; for (size_t bits : {0U, 4U, 8U}) { @@ -1574,12 +1617,167 @@ TEST(QGSearchTest, ParallelQueriesMatchSerialAcrossIndexesAndSettings) { EXPECT_EQ(ids, serial_ids); EXPECT_EQ(distances, serial_distances); } + for (size_t threads : {0U, 1U, 4U}) { + SCOPED_TRACE(threads); + std::array ids{}; + std::array distances{}; + graph.search_batch( + queries.data(), + kNumQueries, + kTopK, + ids.data(), + distances.data(), + threads + ); + EXPECT_EQ(ids, serial_ids); + EXPECT_EQ(distances, serial_distances); + std::fill_n(ids.data(), kTopK, kPidMax); + std::fill_n(distances.data(), kTopK, -123.0F); + graph.search_batch( + queries.data(), 1, kTopK, ids.data(), distances.data(), threads + ); + EXPECT_EQ(ids, serial_ids); + EXPECT_EQ(distances, serial_distances); + } } } } } } +TEST(QGSearchTest, BatchSearchValidatesBuffersSizesAndState) { + QuantizedGraph graph(97, 65, 32, METRIC_L2, RotatorType::FhtKacRotator, 0, 42); + std::vector data(97 * 65, 0.1F), distances(97); + std::vector ids(97); + EXPECT_NO_THROW(graph.search_batch(nullptr, 0, 0, nullptr, nullptr)); + EXPECT_THROW( + graph.search_batch(data.data(), 1, 1, ids.data(), distances.data()), + std::logic_error + ); + QGBuilder builder(graph, 64, data.data(), 1); + builder.build(); + const auto reject = [&](const float* queries, + size_t count, + uint32_t k, + PID* output, + float* output_distances, + const char* message) { + try { + graph.search_batch(queries, count, k, output, output_distances, 4); + FAIL() << "Invalid batch search accepted"; + } catch (const std::invalid_argument& error) { + EXPECT_STREQ(error.what(), message); + } + }; + reject( + data.data(), + 2, + 1, + ids.data(), + distances.data(), + "QuantizedGraph ef must be at least k" + ); + graph.set_ef(64); + reject( + nullptr, + 2, + 1, + ids.data(), + distances.data(), + "QuantizedGraph search buffers must not be null" + ); + reject( + data.data(), + 2, + 1, + nullptr, + distances.data(), + "QuantizedGraph search buffers must not be null" + ); + reject( + data.data(), + 2, + 1, + ids.data(), + nullptr, + "QuantizedGraph search buffers must not be null" + ); + for (uint32_t k : {0U, 98U}) { + reject( + data.data(), + 2, + k, + ids.data(), + distances.data(), + "QuantizedGraph k must be between 1 and num_points" + ); + } + graph.set_ef(1); + reject( + data.data(), + 2, + 2, + ids.data(), + distances.data(), + "QuantizedGraph ef must be at least k" + ); + try { + graph.search_batch( + data.data(), std::numeric_limits::max(), 1, ids.data(), distances.data() + ); + FAIL() << "Oversized batch search accepted"; + } catch (const std::length_error& error) { + EXPECT_STREQ(error.what(), "QuantizedGraph query batch is too large"); + } + graph.set_ef(97); + const size_t oversized_output = + std::numeric_limits::max() / sizeof(float) / 97 + 1; + EXPECT_THROW( + graph.search_batch(data.data(), oversized_output, 97, ids.data(), distances.data()), + std::length_error + ); +} + +TEST(QGSearchTest, ConcurrentBatchCallsMatchIndividualQueries) { + constexpr size_t kDimension = 65, kPoints = 65, kQueries = 9, k = 3; + std::vector data(kPoints * kDimension); + for (size_t i = 0; i < data.size(); ++i) { + data[i] = std::sin(static_cast(i) * 0.19F); + } + QuantizedGraph graph( + kPoints, kDimension, 32, METRIC_IP, RotatorType::FhtKacRotator, 8, 42 + ); + QGBuilder builder(graph, 64, data.data(), 1); + builder.build(); + graph.set_ef(48); + std::vector expected(2 * kQueries * k), actual(expected.size(), kPidMax); + std::vector expected_distances(expected.size()), distances(expected.size()); + for (size_t i = 0; i < 2 * kQueries; ++i) { + graph.search( + data.data() + i * kDimension, + k, + expected.data() + i * k, + expected_distances.data() + i * k + ); + } + const auto search = [&](size_t start) { + graph.search_batch( + data.data() + start * kDimension, + kQueries, + k, + actual.data() + start * k, + distances.data() + start * k, + 2 + ); + }; + auto first = std::async(std::launch::async, search, 0); + auto second = std::async(std::launch::async, search, kQueries); + first.get(); + second.get(); + EXPECT_EQ(actual, expected); + EXPECT_EQ(distances, expected_distances); +} + } // namespace } // namespace rabitqlib::symqg @@ -1615,6 +1813,18 @@ TEST(QuantizedGraphSearchTest, ScratchSearchRejectsIncompleteResults) { ), std::runtime_error ); + std::vector batch_ids(2 * 66, kPidMax); + std::vector batch_distances(2 * 66, -1.0F); + for (size_t threads : {1U, 4U}) { + try { + graph.search_batch( + data.data(), 2, 66, batch_ids.data(), batch_distances.data(), threads + ); + FAIL() << "Incomplete batch results accepted"; + } catch (const std::runtime_error& error) { + EXPECT_STREQ(error.what(), "QuantizedGraph search could not produce k results"); + } + } EXPECT_TRUE(std::all_of(ids.begin(), ids.end(), [](PID id) { return id == kPidMax; })); EXPECT_TRUE(std::all_of(distances.begin(), distances.end(), [](float d) { return d == -1.0F; diff --git a/tests/unit/rabitqlib/utils/memory_test.cpp b/tests/unit/rabitqlib/utils/memory_test.cpp index d3266580..567b1cb2 100644 --- a/tests/unit/rabitqlib/utils/memory_test.cpp +++ b/tests/unit/rabitqlib/utils/memory_test.cpp @@ -2,10 +2,17 @@ #include +#if defined(__GLIBC__) +#include +#endif + +#include #include #include #include +#include #include +#include #include #include "rabitqlib/utils/array.hpp" @@ -21,6 +28,28 @@ TEST(AlignedAllocatorTest, RejectsSizeThatOverflowsAlignmentRounding) { ); } +TEST(AlignedAllocatorTest, ValueInitializesScalarsByDefault) { + const std::vector> values(17); + for (float value : values) { + EXPECT_EQ(value, 0.0F); + } +} + +TEST(AlignedAllocatorTest, DefaultInitializationPreservesObjectConstructionAndCopies) { + using Storage = std::vector>; + Storage values(3); + for (const auto& value : values) { + EXPECT_TRUE(value.empty()); + } + values[0] = "graph"; + values.resize(17, "row"); + const Storage copy(values); + EXPECT_EQ(copy, values); + EXPECT_EQ(copy[0], "graph"); + EXPECT_EQ(copy.back(), "row"); + EXPECT_EQ(reinterpret_cast(copy.data()) % 64, 0U); +} + TEST(AlignedAllocationTest, RejectsSizeThatOverflowsAlignmentRounding) { EXPECT_THROW( (align_allocate<64, char>(std::numeric_limits::max())), @@ -41,6 +70,81 @@ TEST(AlignedAllocationTest, UsesMatchingDeallocator) { allocator.deallocate(values, 3); } +TEST(AlignedAllocationTest, HugePageStorageHandlesSizeBoundaries) { + EXPECT_EQ(huge_page_allocate(0), nullptr); + AlignedAllocator allocator; + EXPECT_EQ(allocator.allocate(0), nullptr); + + for (size_t bytes : + {size_t{1}, size_t{65}, kHugePageSize - 1, kHugePageSize, kHugePageSize + 1}) { + std::unique_ptr storage( + huge_page_allocate(bytes), aligned_deallocate + ); + ASSERT_NE(storage, nullptr); + EXPECT_EQ(reinterpret_cast(storage.get()) % 64, 0U); + std::fill_n(storage.get(), bytes, 'q'); + EXPECT_EQ(storage.get()[bytes - 1], 'q'); + + // Exercise the element-count API as well as the byte-count API. + const size_t count = (bytes + sizeof(uint64_t) - 1) / sizeof(uint64_t); + std::vector> values(count, 42); + EXPECT_EQ(values.front(), 42U); + EXPECT_EQ(values.back(), 42U); + EXPECT_EQ(reinterpret_cast(values.data()) % 64, 0U); +#if defined(__linux__) + if (bytes >= kHugePageSize) { + EXPECT_EQ(reinterpret_cast(storage.get()) % kHugePageSize, 0U); + } + if (count * sizeof(uint64_t) >= kHugePageSize) { + EXPECT_EQ(reinterpret_cast(values.data()) % kHugePageSize, 0U); + } +#endif + } +} + +TEST(AlignedAllocationTest, SmallHugePageStorageAvoidsMegabytePadding) { +#if defined(__GLIBC__) + // Inspect the allocated capacity, not pointer alignment, which can happen + // to exceed the request. This catches rounding small buffers to 2/4 MiB. + std::unique_ptr storage( + huge_page_allocate(65), aligned_deallocate + ); + EXPECT_LT(malloc_usable_size(storage.get()), kHugePageSize); + std::vector> rows(17); + EXPECT_LT(malloc_usable_size(rows.data()), kHugePageSize); + std::vector> ids(17); + EXPECT_LT(malloc_usable_size(ids.data()), kHugePageSize); +#else + GTEST_SKIP() << "Allocated-capacity inspection requires glibc"; +#endif +} + +TEST(AlignedAllocationTest, HugePageStoragePreservesExplicitAlignmentAndOverflowChecks) { + struct alignas(128) OverAligned { + char value; + }; + std::unique_ptr storage( + huge_page_allocate(sizeof(OverAligned)), aligned_deallocate + ); + EXPECT_EQ(reinterpret_cast(storage.get()) % alignof(OverAligned), 0U); + EXPECT_THROW( + (huge_page_allocate(std::numeric_limits::max())), + std::bad_array_new_length + ); + AlignedAllocator allocator; + EXPECT_THROW( + (void)allocator.allocate(std::numeric_limits::max() / sizeof(uint64_t) + 1), + std::bad_array_new_length + ); +#if defined(__linux__) + // Fits 64-byte rounding but overflows when huge-page alignment is selected. + EXPECT_THROW( + (huge_page_allocate(std::numeric_limits::max() - 63)), + std::bad_array_new_length + ); +#endif +} + TEST(ArrayTest, RejectsDimensionProductOverflow) { EXPECT_THROW( (Array(std::vector{std::numeric_limits::max(), 2})),