diff --git a/include/rabitqlib/index/detail/parallel_for.hpp b/include/rabitqlib/index/detail/parallel_for.hpp index b29691a7..4d5d6af1 100644 --- a/include/rabitqlib/index/detail/parallel_for.hpp +++ b/include/rabitqlib/index/detail/parallel_for.hpp @@ -13,7 +13,13 @@ namespace rabitqlib::index_detail { template -inline void parallel_for(size_t start, size_t end, size_t num_threads, Function fn) { +inline void parallel_for( + size_t start, + size_t end, + size_t num_threads, + Function fn, + [[maybe_unused]] bool reuse_workers = false +) { if (start >= end) { return; } @@ -25,9 +31,12 @@ inline void parallel_for(size_t start, size_t end, size_t num_threads, Function } } else { #if defined(_OPENMP) && _OPENMP >= 201307 - if (omp_get_proc_bind() != omp_proc_bind_false) { - // Native threads inherit only the caller's place. Let OpenMP honor - // the configured placement of the whole team when binding is enabled. + if ((reuse_workers && omp_get_level() == 0) || + omp_get_proc_bind() != omp_proc_bind_false) { + // Independent query batches can reuse workers; graph insertion keeps + // its existing scheduling. Binding requires OpenMP placement because + // native threads inherit only the caller's place. Keep unbound nested + // calls on the existing native-thread path. std::vector errors(num_threads); #pragma omp parallel for num_threads(num_threads) schedule(dynamic) for (size_t id = start; id < end; ++id) { diff --git a/include/rabitqlib/index/ivf/initializer.hpp b/include/rabitqlib/index/ivf/initializer.hpp index 8260a6fb..3da82fb2 100644 --- a/include/rabitqlib/index/ivf/initializer.hpp +++ b/include/rabitqlib/index/ivf/initializer.hpp @@ -7,7 +7,6 @@ #include #include #include -#include #include #include #include @@ -192,7 +191,6 @@ class HNSWInitializer : public Initializer { hnswlib::SpaceInterface* space_; MetricType metric_type_; std::vector label_to_internal_; - mutable std::mutex search_mutex_; void rebuild_label_lookup() { if (alg_hnsw_->cur_element_count != num_cluster_) { @@ -253,10 +251,8 @@ class HNSWInitializer : public Initializer { void centroids_distances( const float* query, size_t nprobe, std::vector>& candidates ) const override { - std::lock_guard lock(search_mutex_); - alg_hnsw_->setEf(std::max(size_t{768}, 2 * nprobe)); std::priority_queue> result = - alg_hnsw_->searchKnn(query, nprobe); + alg_hnsw_->searchKnnWithEf(query, nprobe, std::max(size_t{768}, 2 * nprobe)); if (result.size() != nprobe) { throw std::runtime_error("HNSW centroid graph returned too few candidates"); } diff --git a/include/rabitqlib/third/hnswlib/hnswalg.h b/include/rabitqlib/third/hnswlib/hnswalg.h index 7ad27883..c8ff85d1 100644 --- a/include/rabitqlib/third/hnswlib/hnswalg.h +++ b/include/rabitqlib/third/hnswlib/hnswalg.h @@ -1158,6 +1158,7 @@ class HierarchicalNSW : public AlgorithmInterface { tableint addPoint(const void *data_point, labeltype label, int level) { tableint cur_c = 0; + int curlevel = 0; { // Checking if the element with the same label already exists // if so, updating it *instead* of creating a new element. @@ -1187,10 +1188,11 @@ class HierarchicalNSW : public AlgorithmInterface { cur_c = cur_element_count; cur_element_count++; label_lookup_[label] = cur_c; + // RaBitQ: the shared generator must advance under the registration lock. + curlevel = getRandomLevel(mult_); } std::unique_lock lock_el(link_list_locks_[cur_c]); - int curlevel = getRandomLevel(mult_); if (level > 0) curlevel = level; @@ -1198,10 +1200,10 @@ class HierarchicalNSW : public AlgorithmInterface { std::unique_lock templock(global); int maxlevelcopy = maxlevel_; - if (curlevel <= maxlevelcopy) - templock.unlock(); tableint currObj = enterpoint_node_; tableint enterpoint_copy = enterpoint_node_; + if (curlevel <= maxlevelcopy) + templock.unlock(); memset(data_level0_memory_ + cur_c * size_data_per_element_ + offsetLevel0_, 0, size_data_per_element_); @@ -1275,6 +1277,13 @@ class HierarchicalNSW : public AlgorithmInterface { std::priority_queue> searchKnn(const void *query_data, size_t k, BaseFilterFunctor* isIdAllowed = nullptr) const { + return searchKnnWithEf(query_data, k, ef_, isIdAllowed); + } + + + // RaBitQ extension: concurrent searches can choose ef without changing shared state. + std::priority_queue> + searchKnnWithEf(const void *query_data, size_t k, size_t ef, BaseFilterFunctor* isIdAllowed = nullptr) const { std::priority_queue> result; if (cur_element_count == 0) return result; @@ -1312,10 +1321,10 @@ class HierarchicalNSW : public AlgorithmInterface { bool bare_bone_search = !num_deleted_ && !isIdAllowed; if (bare_bone_search) { top_candidates = searchBaseLayerST( - currObj, query_data, std::max(ef_, k), isIdAllowed); + currObj, query_data, std::max(ef, k), isIdAllowed); } else { top_candidates = searchBaseLayerST( - currObj, query_data, std::max(ef_, k), isIdAllowed); + currObj, query_data, std::max(ef, k), isIdAllowed); } while (top_candidates.size() > k) { diff --git a/include/rabitqlib/utils/space.hpp b/include/rabitqlib/utils/space.hpp index 4430f782..76186447 100644 --- a/include/rabitqlib/utils/space.hpp +++ b/include/rabitqlib/utils/space.hpp @@ -37,6 +37,10 @@ void scalar_quantize_normal( float lo, float delta ) { + if (delta == 0.0F) { + std::fill_n(result, dim, T{0}); + return; + } float one_over_delta = 1.0F / delta; ConstRowMajorArrayMap v0(vec0, 1, static_cast(dim)); @@ -188,6 +192,7 @@ inline void data_range(const T* __restrict__ vec0, size_t dim, T& lo, T& hi) { hi = v0.maxCoeff(); } +// A constant range (delta == 0) is represented by zero codes. template void scalar_quantize( T* __restrict__ result, const TD* __restrict__ vec0, size_t dim, TD lo, TD delta diff --git a/src/index/hnsw.cpp b/src/index/hnsw.cpp index 2e705e80..da678eba 100644 --- a/src/index/hnsw.cpp +++ b/src/index/hnsw.cpp @@ -951,13 +951,17 @@ std::vector>> HierarchicalNSW::search( } std::vector>> results(num_queries); rabitqlib::index_detail::parallel_for( - 0, num_queries, num_threads, [&](size_t idx, size_t /*threadId*/) { + 0, + num_queries, + num_threads, + [&](size_t idx, size_t /*threadId*/) { std::vector rotated_query(padded_dim_); this->rotator_->rotate(queries + (idx * dim_), rotated_query.data()); results[idx] = detail::search_knn_sorted( *this, rotated_query.data(), k, std::max(ef_search, k) ); - } + }, + /*reuse_workers=*/true ); return results; } diff --git a/src/index/pipnn.cpp b/src/index/pipnn.cpp index c0f4fdc7..d9794371 100644 --- a/src/index/pipnn.cpp +++ b/src/index/pipnn.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -28,6 +29,7 @@ namespace rabitqlib::symqg::detail { namespace pipnn_impl { constexpr size_t kTileSize = 2048; +constexpr size_t kMinTileSize = 64; constexpr size_t kHashBits = 12; constexpr size_t kLocalDegree = 2; constexpr float kAlpha = 1.1F; @@ -180,6 +182,33 @@ void pairwise(Workspace& work, size_t size, size_t dim, MetricType metric) { ); } +void select_local_neighbors(const Bucket& ids, Workspace& work) { + static_assert(kLocalDegree == 2); + work.candidates.assign( + ids.size() * kLocalDegree, {std::numeric_limits::infinity(), kPidMax} + ); + const auto keep_nearest = [&](size_t row, float distance, PID target) { + auto* best = work.candidates.data() + row * kLocalDegree; + const std::pair candidate(distance, target); + if (candidate < best[1]) { + best[1] = candidate; + if (best[1] < best[0]) { + std::swap(best[0], best[1]); + } + } + }; + // Both metrics are symmetric. Read each stored pair once, contiguously, + // instead of gathering the upper half through a strided column walk. + for (size_t i = 1; i < ids.size(); ++i) { + const float* distances = work.distances + i * ids.size(); + for (size_t j = 0; j < i; ++j) { + const float distance = distances[j]; + keep_nearest(i, distance, ids[j]); + keep_nearest(j, distance, ids[i]); + } + } +} + std::vector partition( const float* data, size_t dim, @@ -189,6 +218,15 @@ std::vector partition( ScratchPool& scratch ) { const auto& ids = job.ids; + // Small root partitions need enough tiles to occupy the requested team. + // Deeper levels already distribute independent buckets across workers. + const size_t tile_size = + job.depth == 0 + ? std::min( + kTileSize, + round_up_to_multiple(div_round_up(ids.size(), threads), kMinTileSize) + ) + : kTileSize; const size_t leader_count = std::min( ids.size(), std::clamp(job.depth == 0 ? 512 : ids.size() / 200, 3, 512) ); @@ -210,7 +248,10 @@ std::vector partition( simd::row_norms(leader_data.data(), leader_norms.data(), leader_count, dim); std::vector assignments(ids.size() * fanout); const auto assign_tile = [&](size_t begin, Workspace& work) { - const size_t count = std::min(kTileSize, ids.size() - begin); + // Both callers run inside a PiPNN OpenMP task, including one-thread teams. + // Keep Eigen within that worker without changing the caller's thread budget. + omp_set_num_threads(1); + const size_t count = std::min(tile_size, ids.size() - begin); RowMajorMatrixMap points( work.points, static_cast(count), static_cast(dim) ); @@ -251,13 +292,13 @@ std::vector partition( if (job.depth == 0) { { index_detail::ParallelException failure; -#pragma omp parallel num_threads(threads) +#pragma omp parallel num_threads(std::min(threads, div_round_up(ids.size(), tile_size))) { auto work = scratch.worker(static_cast(omp_get_thread_num())); #pragma omp for schedule(static) for (std::ptrdiff_t offset = 0; offset < static_cast(ids.size()); - offset += static_cast(kTileSize)) { + offset += static_cast(tile_size)) { if (failure.failed()) { continue; } @@ -273,7 +314,7 @@ std::vector partition( // Already inside the bucket-level team. Eigen must see that team rather // than a serialized nested region when deciding its own parallelism. auto work = scratch.worker(static_cast(omp_get_thread_num())); - for (size_t begin = 0; begin < ids.size(); begin += kTileSize) { + for (size_t begin = 0; begin < ids.size(); begin += tile_size) { assign_tile(begin, work); } } @@ -367,7 +408,7 @@ std::vector cluster( } else { { index_detail::ParallelException failure; -#pragma omp parallel for num_threads(threads) schedule(dynamic) +#pragma omp parallel for num_threads(std::min(threads, active.size())) schedule(dynamic) for (std::ptrdiff_t i = 0; i < static_cast(active.size()); ++i) { if (failure.failed()) { @@ -424,17 +465,22 @@ InitialGraph build_initial_graph( projections.data()[i] = normal(random); } RowMajorMatrix sketches(count, kHashBits); + const size_t sketch_tile = std::min( + kTileSize, round_up_to_multiple(div_round_up(count, threads), kMinTileSize) + ); { index_detail::ParallelException failure; -#pragma omp parallel for num_threads(threads) schedule(static) +#pragma omp parallel for num_threads(std::min(threads, div_round_up(count, sketch_tile))) \ + schedule(static) for (std::ptrdiff_t offset = 0; offset < static_cast(count); - offset += static_cast(kTileSize)) { + offset += static_cast(sketch_tile)) { if (failure.failed()) { continue; } try { const size_t begin = static_cast(offset); - const size_t rows = std::min(kTileSize, count - begin); + const size_t rows = std::min(sketch_tile, count - begin); + omp_set_num_threads(1); simd::matrix_product( data + begin * dim, projections.data(), @@ -450,7 +496,7 @@ InitialGraph build_initial_graph( CandidateTable table(count, degree * 5 / 2); { index_detail::ParallelException failure; -#pragma omp parallel num_threads(threads) +#pragma omp parallel num_threads(std::min(threads, leaves.size())) { auto work = scratch.worker(static_cast(omp_get_thread_num())); #pragma omp for schedule(dynamic) @@ -468,25 +514,14 @@ InitialGraph build_initial_graph( } gather(data, dim, ids, work); pairwise(work, ids.size(), dim, metric); + select_local_neighbors(ids, work); + const size_t keep = std::min(kLocalDegree, ids.size() - 1); + // Retain the original point order and (distance, ID) tie order + // when publishing forward and reverse proposals. for (size_t i = 0; i < ids.size(); ++i) { - work.candidates.clear(); - for (size_t j = 0; j < ids.size(); ++j) { - if (i != j) { - work.candidates.emplace_back( - work.distances - [std::max(i, j) * ids.size() + std::min(i, j)], - ids[j] - ); - } - } - const size_t keep = std::min(kLocalDegree, work.candidates.size()); - std::partial_sort( - work.candidates.begin(), - work.candidates.begin() + static_cast(keep), - work.candidates.end() - ); for (size_t j = 0; j < keep; ++j) { - const auto [distance, target] = work.candidates[j]; + const auto [distance, target] = + work.candidates[i * kLocalDegree + j]; table.merge(ids[i], target, distance, sketches); table.merge(target, ids[i], distance, sketches); } diff --git a/src/index/pipnn_internal.hpp b/src/index/pipnn_internal.hpp index 3634b6d0..d4ceb3f2 100644 --- a/src/index/pipnn_internal.hpp +++ b/src/index/pipnn_internal.hpp @@ -34,6 +34,9 @@ struct ScratchPool { void gather(const float*, size_t, const Bucket&, Workspace&); void pairwise(Workspace&, size_t, size_t, MetricType); +// After pairwise(), write two sorted (distance, ID) candidates per local row. +// Excludes self; a missing second neighbor retains kPidMax. +void select_local_neighbors(const Bucket&, Workspace&); std::vector cluster(const float*, size_t, size_t, MetricType, size_t, ScratchPool&); } // namespace rabitqlib::symqg::detail::pipnn_impl diff --git a/src/simd/portable_generic.cpp b/src/simd/portable_generic.cpp index 55201e50..3392abe8 100644 --- a/src/simd/portable_generic.cpp +++ b/src/simd/portable_generic.cpp @@ -135,6 +135,10 @@ float excode_ip(const float* query, const uint8_t* compact, size_t dim) { template void quantize(T* result, const float* data, size_t dim, float lo, float delta) { + if (delta == 0.0F) { + std::fill_n(result, dim, T{0}); + return; + } const float reciprocal = 1.0F / delta; for (size_t i = 0; i < dim; ++i) { result[i] = static_cast(std::round((data[i] - lo) * reciprocal)); diff --git a/src/simd/query_neon.cpp b/src/simd/query_neon.cpp index c10b5a8b..67940232 100644 --- a/src/simd/query_neon.cpp +++ b/src/simd/query_neon.cpp @@ -30,6 +30,10 @@ inline uint32x4_t quantize4(const float* data, float32x4_t lo, float32x4_t recip void scalar_quantize_uint8_neon( uint8_t* result, const float* data, size_t dim, float lo, float delta ) { + if (delta == 0.0F) { + std::fill_n(result, dim, uint8_t{0}); + return; + } const float reciprocal = 1.0F / delta; const float32x4_t low = vdupq_n_f32(lo), scale = vdupq_n_f32(reciprocal); size_t i = 0; @@ -48,6 +52,10 @@ void scalar_quantize_uint8_neon( void scalar_quantize_uint16_neon( uint16_t* result, const float* data, size_t dim, float lo, float delta ) { + if (delta == 0.0F) { + std::fill_n(result, dim, uint16_t{0}); + return; + } const float reciprocal = 1.0F / delta; const float32x4_t low = vdupq_n_f32(lo), scale = vdupq_n_f32(reciprocal); size_t i = 0; diff --git a/src/simd/space_avx2.cpp b/src/simd/space_avx2.cpp index aaf2178d..e7f28639 100644 --- a/src/simd/space_avx2.cpp +++ b/src/simd/space_avx2.cpp @@ -1,5 +1,6 @@ #include +#include #include #include #include @@ -44,6 +45,10 @@ float l2norm_sqr_avx2(const float* a, size_t dim) { void scalar_quantize_uint8_avx2( uint8_t* result, const float* vec0, size_t dim, float lo, float delta ) { + if (delta == 0.0F) { + std::fill_n(result, dim, uint8_t{0}); + return; + } size_t mul8 = dim - (dim & 0b111); size_t i = 0; float one_over_delta = 1.0F / delta; @@ -69,6 +74,10 @@ void scalar_quantize_uint8_avx2( void scalar_quantize_uint16_avx2( uint16_t* result, const float* vec0, size_t dim, float lo, float delta ) { + if (delta == 0.0F) { + std::fill_n(result, dim, uint16_t{0}); + return; + } size_t mul8 = dim - (dim & 0b111); size_t i = 0; float one_over_delta = 1.0F / delta; diff --git a/src/simd/space_avx512.cpp b/src/simd/space_avx512.cpp index 71c70710..5945e8c6 100644 --- a/src/simd/space_avx512.cpp +++ b/src/simd/space_avx512.cpp @@ -1,5 +1,6 @@ #include +#include #include #include #include @@ -44,6 +45,10 @@ float l2norm_sqr_avx512(const float* a, size_t dim) { void scalar_quantize_uint8_avx512( uint8_t* result, const float* vec0, size_t dim, float lo, float delta ) { + if (delta == 0.0F) { + std::fill_n(result, dim, uint8_t{0}); + return; + } size_t mul16 = dim - (dim & 0b1111); size_t i = 0; float one_over_delta = 1.0F / delta; @@ -63,6 +68,10 @@ void scalar_quantize_uint8_avx512( void scalar_quantize_uint16_avx512( uint16_t* result, const float* vec0, size_t dim, float lo, float delta ) { + if (delta == 0.0F) { + std::fill_n(result, dim, uint16_t{0}); + return; + } size_t mul16 = dim - (dim & 0b1111); size_t i = 0; float one_over_delta = 1.0F / delta; diff --git a/tests/python/test_ivf_initializers.py b/tests/python/test_ivf_initializers.py index 9ebf2b1f..44916408 100644 --- a/tests/python/test_ivf_initializers.py +++ b/tests/python/test_ivf_initializers.py @@ -1,6 +1,7 @@ """Explicit centroid routing, persistence, and legacy initializer inference.""" import struct +from concurrent.futures import ThreadPoolExecutor from pathlib import Path import numpy as np @@ -8,6 +9,51 @@ from rabitqlib import IvfIndex +@pytest.mark.parametrize("metric", ["l2", "ip"]) +@pytest.mark.parametrize("threads", [1, 4, 8]) +def test_hnsw_parallel_build_routes_and_reloads(tmp_path, metric, threads): + count, dim = 1024, 65 + data = np.random.default_rng(42).standard_normal((count, dim)).astype(np.float32) + data /= np.linalg.norm(data, axis=1, keepdims=True) + index = IvfIndex(dim, count, count, 4, metric, initializer="hnsw") + index.build(data, data, np.arange(count, dtype=np.uint32), num_threads=threads) + labels = np.array([0, 17, 511, 1023], dtype=np.uint32) + queries = data[labels] + expected = index.search(queries, 1, 8, num_threads=threads) + np.testing.assert_array_equal(expected[0][:, 0], labels) + assert np.isfinite(expected[1]).all() + path = tmp_path / "parallel-build.index" + index.save(str(path)) + loaded = IvfIndex.load(str(path)) + for actual, reference in zip( + loaded.search(queries, 1, 8, num_threads=threads), expected + ): + np.testing.assert_array_equal(actual, reference) + + +@pytest.mark.parametrize("metric", ["l2", "ip"]) +def test_hnsw_concurrent_routing_windows_match_serial(tmp_path, metric): + count, dim = 1024, 64 + rng = np.random.default_rng(42) + data = rng.standard_normal((count, dim)).astype(np.float32) + queries = rng.standard_normal((4, dim)).astype(np.float32) + index = IvfIndex(dim, count, count, 1, metric, initializer="hnsw") + index.build(data, data, np.arange(count, dtype=np.uint32), num_threads=1) + probes = (1, 16, 400, 500) + expected = [index.search(queries, 5, probe, num_threads=1) for probe in probes] + path = tmp_path / "concurrent.index" + index.save(str(path)) + for current in (index, IvfIndex.load(str(path))): + with ThreadPoolExecutor(max_workers=4) as pool: + futures = [ + pool.submit(current.search, queries, 5, probes[i % 4], num_threads=2) + for i in range(12) + ] + for i, future in enumerate(futures): + for actual, reference in zip(future.result(), expected[i % 4]): + np.testing.assert_array_equal(actual, reference) + + @pytest.mark.parametrize("metric", ["l2", "ip"]) @pytest.mark.parametrize("damage", ["entry", "neighbor", "level"]) def test_hnsw_rejects_corrupt_sidecar(tmp_path, metric, damage): diff --git a/tests/python/test_quantization_factors.py b/tests/python/test_quantization_factors.py index 39af396d..aebbb247 100644 --- a/tests/python/test_quantization_factors.py +++ b/tests/python/test_quantization_factors.py @@ -7,7 +7,8 @@ @pytest.mark.parametrize("nbits", [4, 9, 32]) @pytest.mark.parametrize("metric", ["l2", "ip"]) -def test_collinear_ivf_query_keeps_finite_distances(nbits, metric): +@pytest.mark.parametrize("high_accuracy", [False, True]) +def test_collinear_ivf_query_keeps_finite_distances(nbits, metric, high_accuracy): data = np.full((1, 512), 0.1, dtype=np.float32) centroid = np.zeros_like(data) index = IvfIndex(data.shape[1], len(data), 1, nbits=nbits, metric=metric) @@ -16,7 +17,9 @@ def test_collinear_ivf_query_keeps_finite_distances(nbits, metric): # The C++ regression fixes the rotated residual to trigger negative roundoff. # This exercises the normal Python boundary with its randomly sampled rotation. # A zero query avoids query-LUT approximation in the distance oracle. - ids, distances = index.search(centroid, k=1, nprobe=1, num_threads=1) + ids, distances = index.search( + centroid, k=1, nprobe=1, num_threads=1, high_accuracy=high_accuracy + ) np.testing.assert_array_equal(ids, [[0]]) assert np.all(np.isfinite(distances)) expected = np.sum(data.astype(np.float64) ** 2) if metric == "l2" else 1.0 diff --git a/tests/unit/rabitqlib/fastscan/fastscan_test.cpp b/tests/unit/rabitqlib/fastscan/fastscan_test.cpp index e7feafd7..11240944 100644 --- a/tests/unit/rabitqlib/fastscan/fastscan_test.cpp +++ b/tests/unit/rabitqlib/fastscan/fastscan_test.cpp @@ -126,6 +126,32 @@ TEST(FastScanLutTest, ResetOwnsBytesAndReusesStorageWithoutRetainingScratch) { } } +TEST(FastScanLutTest, ZeroQueryClearsCodesInStandardAndHighAccuracyModes) { + for (size_t dim : {32U, 64U, 96U}) { + for (bool hacc : {false, true}) { + SCOPED_TRACE(::testing::Message() << dim << "/" << hacc); + std::vector query(dim, 1.0F), scratch(dim * 4); + std::vector hacc_scratch(dim * 4); + Lut lut(query.data(), dim, hacc); + const size_t bytes = dim * 4 * (hacc ? 2 : 1); + ASSERT_TRUE(std::any_of(lut.lut(), lut.lut() + bytes, [](uint8_t value) { + return value != 0; + })); + std::fill(query.begin(), query.end(), 0.0F); + lut.reset(query.data(), dim, scratch.data(), hacc_scratch.data(), hacc); + EXPECT_EQ(lut.delta(), 0.0F); + EXPECT_EQ(lut.sum_vl(), 0.0F); + EXPECT_TRUE(std::all_of(lut.lut(), lut.lut() + bytes, [](uint8_t value) { + return value == 0; + })); + const Lut fresh(query.data(), dim, hacc); + EXPECT_EQ(fresh.delta(), 0.0F); + EXPECT_EQ(fresh.sum_vl(), 0.0F); + EXPECT_TRUE(std::equal(lut.lut(), lut.lut() + bytes, fresh.lut())); + } + } +} + TEST(FastScanLutTest, RejectedResetPreservesOwnedBytesAndFactors) { std::array query{}; query.fill(1); diff --git a/tests/unit/rabitqlib/index/initializer_test.cpp b/tests/unit/rabitqlib/index/initializer_test.cpp index eba48dcc..a805c60a 100644 --- a/tests/unit/rabitqlib/index/initializer_test.cpp +++ b/tests/unit/rabitqlib/index/initializer_test.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -147,6 +148,35 @@ TEST(HNSWInitializerTest, RoutesByInnerProductWithoutTakingSquareRoot) { EXPECT_FLOAT_EQ(candidates[0].distance, -99.0F); } +TEST(HNSWInitializerTest, ParallelInsertionPreservesRandomLevelDraws) { + constexpr size_t kDim = 17; + constexpr size_t kCount = 2048; + constexpr size_t kSeed = 100; + std::mt19937 rng(42); + std::uniform_real_distribution distribution(-1.0F, 1.0F); + std::vector centroids(kCount * kDim); + for (auto& value : centroids) { + value = distribution(rng); + } + CentroidL2Space space(kDim); + hnswlib::HierarchicalNSW reference(&space, 1, 16, 100, kSeed); + std::vector expected_levels(kCount); + for (auto& level : expected_levels) { + level = reference.getRandomLevel(reference.mult_); + } + for (size_t threads : {1U, 4U, 8U}) { + SCOPED_TRACE(threads); + hnswlib::HierarchicalNSW graph(&space, kCount, 16, 100, kSeed); + index_detail::parallel_for(0, kCount, threads, [&](size_t row, size_t) { + graph.addPoint(centroids.data() + row * kDim, row); + }); + EXPECT_EQ(graph.cur_element_count.load(), kCount); + // Registration assigns internal IDs and draws their levels in the same order. + EXPECT_EQ(graph.element_levels_, expected_levels); + EXPECT_EQ(graph.level_generator_, reference.level_generator_); + } +} + TEST(HNSWInitializerTest, ParallelInsertionPreservesCentroidLabelsAfterReload) { constexpr size_t kDim = 17; constexpr size_t kCount = 1024; @@ -203,37 +233,88 @@ TEST(HNSWInitializerTest, ReloadsInnerProductGraphForRouting) { std::remove((path + ".hnsw").c_str()); } -TEST(HNSWInitializerTest, ConcurrentSearchesKeepRoutingStateSynchronized) { +TEST(HNSWInitializerTest, PerCallSearchWindowsPreserveDefaultAndLegacyResults) { constexpr size_t kDim = 17; - constexpr size_t kCount = 32; - std::vector centroids(kCount * kDim, 0.0F); - for (size_t label = 0; label < kCount; ++label) { - centroids[label * kDim] = static_cast(label); + constexpr size_t kCount = 1024; + std::mt19937 rng(42); + std::uniform_real_distribution distribution(-1.0F, 1.0F); + std::vector centroids(kCount * kDim); + for (auto& value : centroids) { + value = distribution(rng); } - HNSWInitializer initializer(kDim, kCount); - initializer.add_vectors(centroids.data(), 4); + CentroidL2Space space(kDim); + hnswlib::HierarchicalNSW graph(&space, kCount, 16, 100); + for (size_t i = 0; i < kCount; ++i) { + graph.addPoint(centroids.data() + i * kDim, i); + } + struct EvenLabels : hnswlib::BaseFilterFunctor { + bool operator()(hnswlib::labeltype label) override { return label % 2 == 0; } + } even_labels; std::array query{}; - query[0] = 15.0F; - std::atomic invalid_result{false}; - std::vector threads; - for (size_t worker = 0; worker < 4; ++worker) { - threads.emplace_back([&, worker] { - for (size_t iteration = 0; iteration < 25; ++iteration) { - const size_t nprobe = 1 + (worker + iteration) % 4; - std::vector> candidates(nprobe); - initializer.centroids_distances(query.data(), nprobe, candidates); - for (const auto& candidate : candidates) { - if (candidate.id >= kCount || !std::isfinite(candidate.distance)) { - invalid_result = true; - } - } + for (auto* filter : + {static_cast(nullptr), + static_cast(&even_labels)}) { + for (size_t ef : {1U, 64U, 768U, 1000U}) { + graph.setEf(ef); + auto expected = graph.searchKnn(query.data(), 10, filter); + graph.setEf(11); + auto actual = graph.searchKnnWithEf(query.data(), 10, ef, filter); + EXPECT_EQ(graph.ef_, 11U); + ASSERT_EQ(actual.size(), expected.size()); + while (!expected.empty()) { + EXPECT_EQ(actual.top(), expected.top()); + actual.pop(); + expected.pop(); } - }); + } } - for (auto& thread : threads) { - thread.join(); +} + +TEST(HNSWInitializerTest, ConcurrentSearchesMatchSerialWithDifferentWindows) { + constexpr size_t kDim = 17; + constexpr size_t kCount = 1024; + const std::array probes{1, 16, 400, 500}; + std::mt19937 rng(42); + std::uniform_real_distribution distribution(-1.0F, 1.0F); + std::vector centroids(kCount * kDim); + for (auto& value : centroids) { + value = distribution(rng); + } + for (const auto metric : {METRIC_L2, METRIC_IP}) { + HNSWInitializer initializer(kDim, kCount, metric); + initializer.add_vectors(centroids.data(), 1); + std::array>, 4> expected; + for (size_t i = 0; i < probes.size(); ++i) { + expected[i].resize(probes[i]); + initializer.centroids_distances( + centroids.data() + i * kDim, probes[i], expected[i] + ); + } + std::atomic invalid_result{false}; + std::vector threads; + threads.reserve(4); + for (size_t worker = 0; worker < 4; ++worker) { + threads.emplace_back([&, worker] { + for (size_t iteration = 0; iteration < 8; ++iteration) { + const size_t test = (worker + iteration) % probes.size(); + std::vector> candidates(probes[test]); + initializer.centroids_distances( + centroids.data() + test * kDim, probes[test], candidates + ); + for (size_t i = 0; i < candidates.size(); ++i) { + if (candidates[i].id != expected[test][i].id || + candidates[i].distance != expected[test][i].distance) { + invalid_result = true; + } + } + } + }); + } + for (auto& thread : threads) { + thread.join(); + } + EXPECT_FALSE(invalid_result); } - EXPECT_FALSE(invalid_result); } TEST(ParallelForTest, JoinsWorkersBeforeRethrowingFunctionException) { diff --git a/tests/unit/rabitqlib/index/pipnn_test.cpp b/tests/unit/rabitqlib/index/pipnn_test.cpp index 60001776..1a82ad0b 100644 --- a/tests/unit/rabitqlib/index/pipnn_test.cpp +++ b/tests/unit/rabitqlib/index/pipnn_test.cpp @@ -8,7 +8,9 @@ #include #include #include +#include #include +#include #include #include "pipnn_internal.hpp" @@ -96,6 +98,55 @@ TEST(PipnnTest, LocalSeedRetainsNearestNeighborsAndValidIds) { } } +TEST(PipnnTest, LocalNeighborsMatchExactReferenceAndGlobalIdTies) { + constexpr size_t kDim = 65; + std::mt19937 generator(719); + std::uniform_int_distribution coordinate(-8, 8); + for (size_t count : {2U, 3U, 17U, 65U, 1024U}) { + std::vector data(count * kDim); + pipnn_impl::Bucket ids(count); + std::iota(ids.begin(), ids.end(), PID{0}); + std::shuffle(ids.begin(), ids.end(), generator); + pipnn_impl::ScratchPool scratch(count, kDim, 1); + auto work = scratch.worker(0); + for (bool ties : {false, true}) { + std::generate(data.begin(), data.end(), [&] { + return ties ? 1.0F : static_cast(coordinate(generator)); + }); + for (auto metric : {METRIC_L2, METRIC_IP}) { + SCOPED_TRACE(::testing::Message() << count << "/" << ties << "/" << metric); + pipnn_impl::gather(data.data(), kDim, ids, work); + pipnn_impl::pairwise(work, count, kDim, metric); + pipnn_impl::select_local_neighbors(ids, work); + ASSERT_EQ(work.candidates.size(), 2 * count); + for (size_t i = 0; i < count; ++i) { + std::vector> reference; + reference.reserve(count - 1); + for (PID target = 0; target < count; ++target) { + if (target == ids[i]) { + continue; + } + double distance = 0; + for (size_t d = 0; d < kDim; ++d) { + const double a = data[ids[i] * kDim + d]; + const double b = data[target * kDim + d]; + distance += metric == METRIC_L2 ? (a - b) * (a - b) : -a * b; + } + reference.emplace_back(static_cast(distance), target); + } + std::sort(reference.begin(), reference.end()); + for (size_t j = 0; j < std::min(size_t{2}, reference.size()); ++j) { + EXPECT_EQ(work.candidates[2 * i + j], reference[j]); + } + if (count == 2) { + EXPECT_EQ(work.candidates[2 * i + 1].second, kPidMax); + } + } + } + } + } +} + TEST(PipnnTest, OverlappingPartitionsTerminateForDuplicateVectors) { constexpr size_t kCount = 1100, kDim = 7; std::vector data(kCount * kDim, 0.5F); @@ -156,5 +207,33 @@ TEST(PipnnTest, RootTileTailAndRecursiveBucketsCoverDuplicateInputs) { } } } + +TEST(PipnnTest, AdaptiveRootTilesPreservePartitionsAcrossThreadCounts) { + constexpr size_t kDim = 65; + std::mt19937 generator(817); + // Integer coordinates keep these dot products exact on every SIMD backend. + std::uniform_int_distribution uniform(-8, 8); + // Exercise partial tiles on either side of the original 2048-row boundary. + for (size_t count : {1025U, 2049U, 4097U}) { + std::vector data(count * kDim); + std::generate(data.begin(), data.end(), [&] { + return static_cast(uniform(generator)); + }); + for (auto metric : {METRIC_L2, METRIC_IP}) { + pipnn_impl::ScratchPool serial_scratch(count, kDim, 1); + const auto expected = + pipnn_impl::cluster(data.data(), count, kDim, metric, 1, serial_scratch); + for (size_t threads : {2U, 4U}) { + SCOPED_TRACE( + ::testing::Message() << count << "/" << metric << "/" << threads + ); + pipnn_impl::ScratchPool scratch(count, kDim, threads); + const auto actual = + pipnn_impl::cluster(data.data(), count, kDim, metric, threads, scratch); + EXPECT_EQ(actual, expected); + } + } + } +} } // namespace } // namespace rabitqlib::symqg::detail diff --git a/tests/unit/rabitqlib/utils/query_preparation_test.cpp b/tests/unit/rabitqlib/utils/query_preparation_test.cpp index 21500ce6..ad831d77 100644 --- a/tests/unit/rabitqlib/utils/query_preparation_test.cpp +++ b/tests/unit/rabitqlib/utils/query_preparation_test.cpp @@ -58,6 +58,24 @@ void check_quantization(void (*quantize)(T*, const float*, size_t, float, float) } } +template +void check_zero_range(void (*quantize)(T*, const float*, size_t, float, float)) { + for (float lo : {0.0F, -3.25F, 7.5F}) { + for (float delta : {0.0F, -0.0F}) { + for (size_t dim : + {0U, 1U, 3U, 4U, 7U, 8U, 9U, 15U, 16U, 17U, 31U, 32U, 33U, 65U, 256U}) { + SCOPED_TRACE(::testing::Message() << lo << "/" << delta << "/" << dim); + std::vector input(dim + 2, lo); + // Offset both buffers to cover unaligned SIMD accesses and tails. + std::vector actual(dim + 2, 99), expected(dim + 2, 0); + expected.front() = expected.back() = 99; + quantize(actual.data() + 1, input.data() + 1, dim, lo, delta); + EXPECT_EQ(actual, expected); + } + } + } +} + template void check_transpose( void (*transpose)(const T*, uint64_t*, size_t, size_t), bool blocked512 @@ -134,6 +152,32 @@ TEST(QueryPreparation, QuantizationMatchesHalfAwayRoundingAndTails) { #endif } +TEST(QueryPreparation, ZeroRangeQuantizationMatchesAcrossBackendsAndTails) { + check_zero_range(simd::scalar_quantize_uint8_generic); + check_zero_range(simd::scalar_quantize_uint16_generic); + check_zero_range(scalar_quantize); + check_zero_range(scalar_quantize); + check_zero_range(scalar_quantize); + check_zero_range(scalar_impl::scalar_quantize_normal); + check_zero_range(scalar_impl::scalar_quantize_normal); +#if defined(__aarch64__) || defined(_M_ARM64) + if (cpu::has_neon()) { + check_zero_range(simd::scalar_quantize_uint8_neon); + check_zero_range(simd::scalar_quantize_uint16_neon); + } +#endif +#if defined(__x86_64__) || defined(_M_X64) + if (cpu::has_avx2()) { + check_zero_range(simd::scalar_quantize_uint8_avx2); + check_zero_range(simd::scalar_quantize_uint16_avx2); + } + if (cpu::has_avx512_core()) { + check_zero_range(simd::scalar_quantize_uint8_avx512); + check_zero_range(simd::scalar_quantize_uint16_avx512); + } +#endif +} + TEST(QueryPreparation, TranspositionMatchesCoordinateReferenceAndBounds) { check_transpose(simd::new_transpose_bin_generic, false); check_transpose(simd::new_transpose_bin_512_generic, true); diff --git a/tests/unit/rabitqlib/utils/thread_count_test.cpp b/tests/unit/rabitqlib/utils/thread_count_test.cpp index 78f38ac0..e945269f 100644 --- a/tests/unit/rabitqlib/utils/thread_count_test.cpp +++ b/tests/unit/rabitqlib/utils/thread_count_test.cpp @@ -138,3 +138,91 @@ TEST(ThreadCount, ParallelForVisitsEachItemOnceWithBoundedWorkerIds) { } } } + +TEST(ThreadCount, TopLevelParallelForReusesWorkersOnlyWhenRequested) { +#if defined(_OPENMP) && _OPENMP >= 201307 + if (rabitqlib::resolve_num_threads(2) < 2) { + GTEST_SKIP() << "Requires two available workers"; + } + const int max_threads = omp_get_max_threads(); + const int max_levels = omp_get_max_active_levels(); + std::vector levels(17, -1); + rabitqlib::index_detail::parallel_for(0, levels.size(), 2, [&](size_t i, size_t) { + levels[i] = omp_get_level(); + }); + const int build_level = omp_get_proc_bind() == omp_proc_bind_false ? 0 : 1; + for (int level : levels) { + EXPECT_EQ(level, build_level); + } + for (size_t repeat = 0; repeat < 3; ++repeat) { + rabitqlib::index_detail::parallel_for( + 0, + levels.size(), + 2, + [&](size_t i, size_t) { levels[i] = omp_get_level(); }, + true + ); + for (int level : levels) { + EXPECT_EQ(level, 1); + } + } + rabitqlib::index_detail::parallel_for( + 0, + levels.size(), + 1, + [&](size_t i, size_t worker) { + EXPECT_EQ(worker, 0U); + levels[i] = omp_get_level(); + }, + true + ); + for (int level : levels) { + EXPECT_EQ(level, 0); + } + EXPECT_THROW( + rabitqlib::index_detail::parallel_for( + 0, + levels.size(), + 2, + [](size_t, size_t) { throw std::runtime_error("failed"); }, + true + ), + std::runtime_error + ); + EXPECT_EQ(omp_get_max_threads(), max_threads); + EXPECT_EQ(omp_get_max_active_levels(), max_levels); +#else + GTEST_SKIP() << "Runtime worker reuse requires OpenMP 4.0"; +#endif +} + +TEST(ThreadCount, NestedParallelForPreservesBindingAndWorkerIds) { +#if defined(_OPENMP) && _OPENMP >= 201307 + if (rabitqlib::resolve_num_threads(2) < 2) { + GTEST_SKIP() << "Requires two available workers"; + } + const bool bound = omp_get_proc_bind() != omp_proc_bind_false; + std::vector levels(17, -1); + std::vector workers(levels.size(), 2); +#pragma omp parallel num_threads(2) + { +#pragma omp single + rabitqlib::index_detail::parallel_for( + 0, + levels.size(), + 2, + [&](size_t i, size_t worker) { + levels[i] = omp_get_level(); + workers[i] = worker; + }, + true + ); + } + for (size_t i = 0; i < levels.size(); ++i) { + EXPECT_EQ(levels[i], bound ? 2 : 0); + EXPECT_LT(workers[i], 2U); + } +#else + GTEST_SKIP() << "OpenMP binding requires OpenMP 4.0"; +#endif +}