diff --git a/include/rabitqlib/fastscan/detail/qg_estimator.hpp b/include/rabitqlib/fastscan/detail/qg_estimator.hpp new file mode 100644 index 00000000..b761eda5 --- /dev/null +++ b/include/rabitqlib/fastscan/detail/qg_estimator.hpp @@ -0,0 +1,42 @@ +#pragma once + +#include +#include +#include + +#include "rabitqlib/core/detail/eigen.hpp" +#include "rabitqlib/fastscan/data_layout.hpp" +#include "rabitqlib/fastscan/pack_codes.hpp" +#include "rabitqlib/fastscan/query.hpp" + +namespace rabitqlib::fastscan::detail { +// Share QG accumulation and correction while callers select the accumulator. +// TU-local linkage keeps each ISA's Eigen expression compiled with its own flags. +template +static inline void qg_batch_estdist_impl( + const char* batch_data, + const BatchQuery& query_wrapper, + size_t padded_dim, + T* estimated_distances +) { + ConstQGBatchDataMap cur_batch(batch_data, padded_dim); + std::array accu_values{}; + Accumulate(cur_batch.bin_code(), query_wrapper.lut(), accu_values.data(), padded_dim); + + ConstRowMajorArrayMap ip_arr(accu_values.data(), 1, fastscan::kBatchSize); + std::array f_add_values; + std::array f_rescale_values; + cur_batch.f_add().copy_to(f_add_values.data(), f_add_values.size()); + cur_batch.f_rescale().copy_to(f_rescale_values.data(), f_rescale_values.size()); + ConstRowMajorArrayMap f_add_arr(f_add_values.data(), 1, fastscan::kBatchSize); + ConstRowMajorArrayMap f_rescale_arr( + f_rescale_values.data(), 1, fastscan::kBatchSize + ); + + RowMajorArrayMap est_dist_arr(estimated_distances, 1, fastscan::kBatchSize); + + est_dist_arr = f_add_arr + query_wrapper.g_add() + + (f_rescale_arr * (query_wrapper.delta() * (ip_arr.template cast()) + + query_wrapper.sum_vl_lut() + query_wrapper.k1xsumq())); +} +} // namespace rabitqlib::fastscan::detail diff --git a/include/rabitqlib/fastscan/estimator.hpp b/include/rabitqlib/fastscan/estimator.hpp index 9d96d3ce..e6c7ff10 100644 --- a/include/rabitqlib/fastscan/estimator.hpp +++ b/include/rabitqlib/fastscan/estimator.hpp @@ -1,16 +1,13 @@ #pragma once -#include #include #include #include #include -#include "rabitqlib/core/detail/eigen.hpp" -#include "rabitqlib/fastscan/data_layout.hpp" #include "rabitqlib/fastscan/detail/estimator_dispatch.hpp" +#include "rabitqlib/fastscan/detail/qg_estimator.hpp" #include "rabitqlib/fastscan/fastscan.hpp" -#include "rabitqlib/fastscan/pack_codes.hpp" #include "rabitqlib/fastscan/query.hpp" namespace rabitqlib { @@ -81,26 +78,9 @@ inline void qg_batch_estdist( return; } - ConstQGBatchDataMap cur_batch(batch_data, padded_dim); - std::array accumulated_values{}; - fastscan::accumulate( - cur_batch.bin_code(), query_wrapper.lut(), accumulated_values.data(), padded_dim + fastscan::detail::qg_batch_estdist_impl( + batch_data, query_wrapper, padded_dim, estimated_distances ); - ConstRowMajorArrayMap ip_arr(accumulated_values.data(), 1, fastscan::kBatchSize); - std::array f_add_values; - std::array f_rescale_values; - cur_batch.f_add().copy_to(f_add_values.data(), f_add_values.size()); - cur_batch.f_rescale().copy_to(f_rescale_values.data(), f_rescale_values.size()); - ConstRowMajorArrayMap f_add_arr(f_add_values.data(), 1, fastscan::kBatchSize); - ConstRowMajorArrayMap f_rescale_arr( - f_rescale_values.data(), 1, fastscan::kBatchSize - ); - - RowMajorArrayMap est_dist_arr(estimated_distances, 1, fastscan::kBatchSize); - - est_dist_arr = f_add_arr + query_wrapper.g_add() + - (f_rescale_arr * (query_wrapper.delta() * (ip_arr.template cast()) + - query_wrapper.sum_vl_lut() + query_wrapper.k1xsumq())); } } // namespace rabitqlib diff --git a/include/rabitqlib/index/detail/visited_set_hash.hpp b/include/rabitqlib/index/detail/visited_set_hash.hpp index a79478f1..7fdb39f2 100644 --- a/include/rabitqlib/index/detail/visited_set_hash.hpp +++ b/include/rabitqlib/index/detail/visited_set_hash.hpp @@ -34,7 +34,6 @@ class HashBasedVisitedSet { static constexpr size_t kNoSizeHint = std::numeric_limits::max(); private: - size_t table_size_ = 0; PID mask_ = 0; std::vector> table_; std::unordered_set stl_hash_; @@ -44,8 +43,7 @@ class HashBasedVisitedSet { void initialize_table(const size_t table_size) { // Allocate before changing the mask so failure preserves the previous table. table_ = std::vector>(table_size); - table_size_ = table_size; - mask_ = static_cast(table_size_ - 1); + mask_ = static_cast(table_size - 1); assert(hash1(static_cast(table_size)) == 0); std::fill(table_.begin(), table_.end(), kPidMax); diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index 3b7a3de4..5c0ed5c5 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -432,7 +432,6 @@ class HierarchicalNSW { size_t k, const SplitSingleQuery& query_wrapper, const std::vector& q_to_centroids, - const float* query, BoundedKNN& nearest_candidates ) const; diff --git a/include/rabitqlib/utils/space.hpp b/include/rabitqlib/utils/space.hpp index af064813..8c6ec87b 100644 --- a/include/rabitqlib/utils/space.hpp +++ b/include/rabitqlib/utils/space.hpp @@ -15,6 +15,7 @@ #include "rabitqlib/utils/bitops.hpp" #include "rabitqlib/utils/openmp.hpp" // IWYU pragma: keep (configuration macros) #include "rabitqlib/utils/tools.hpp" +#include "rabitqlib/utils/warmup_space.hpp" #if RABITQ_USE_OPENMP #include @@ -319,35 +320,7 @@ inline float ip_x0_q( size_t padded_dim, size_t b_query ) { - auto num_blk = padded_dim / 64; - const auto* it_data = reinterpret_cast(data); - const auto* it_query = query; - - size_t ip = 0; - size_t ppc = 0; - - for (size_t i = 0; i < num_blk; ++i) { - const uint64_t x = rabitqlib::bitops::load_word(it_data, sizeof(uint64_t)); - ppc += bitops::popcount64(x); - - for (size_t j = 0; j < b_query; ++j) { - uint64_t y = *static_cast(it_query); - ip += (bitops::popcount64(x & y) << j); - it_query++; - } - it_data += sizeof(uint64_t); - } - - if (padded_dim % 64 != 0) { - const uint64_t x = rabitqlib::bitops::load_word(it_data, (padded_dim % 64) / 8); - ppc += bitops::popcount64(x); - for (size_t j = 0; j < b_query; ++j) { - ip += (bitops::popcount64(x & *it_query) << j); - ++it_query; - } - } - - return (delta * static_cast(ip)) + (vl * static_cast(ppc)); + return detail::ip_x0_q_impl(data, query, delta, vl, padded_dim, b_query); } static inline uint32_t ip_bin_bin(const uint64_t* q, const uint64_t* d, size_t padded_dim) { diff --git a/include/rabitqlib/utils/warmup_space.hpp b/include/rabitqlib/utils/warmup_space.hpp index 30544ec3..050581b3 100644 --- a/include/rabitqlib/utils/warmup_space.hpp +++ b/include/rabitqlib/utils/warmup_space.hpp @@ -25,14 +25,16 @@ float warmup_ip_x0_q_512( size_t b_query ); -template -inline float warmup_ip_x0_q( +namespace detail { +// Scalar word-interleaved query planes; distinct from the 512-coordinate layout. +// Keep the shared loop local to each translation unit's ISA configuration. +static inline float ip_x0_q_impl( const uint64_t* data, const uint64_t* query, float delta, float vl, size_t padded_dim, - [[maybe_unused]] size_t unused_query_bits = 0 // not used + size_t b_query ) { auto num_blk = padded_dim / 64; const auto* it_data = reinterpret_cast(data); @@ -64,5 +66,18 @@ inline float warmup_ip_x0_q( return (delta * static_cast(ip)) + (vl * static_cast(ppc)); } +} // namespace detail + +template +inline float warmup_ip_x0_q( + const uint64_t* data, + const uint64_t* query, + float delta, + float vl, + size_t padded_dim, + [[maybe_unused]] size_t unused_query_bits = 0 // not used +) { + return detail::ip_x0_q_impl(data, query, delta, vl, padded_dim, b_query); +} } // namespace rabitqlib diff --git a/src/fastscan/estimator_avx2.cpp b/src/fastscan/estimator_avx2.cpp index f41a28ed..4050fd4c 100644 --- a/src/fastscan/estimator_avx2.cpp +++ b/src/fastscan/estimator_avx2.cpp @@ -5,6 +5,7 @@ #include "estimator_kernels.hpp" #include "rabitqlib/fastscan/detail/estimator_dispatch.hpp" +#include "rabitqlib/fastscan/detail/qg_estimator.hpp" #include "rabitqlib/fastscan/fastscan.hpp" #include "rabitqlib/fastscan/highacc_fastscan.hpp" @@ -50,7 +51,7 @@ void qg_batch_estdist_avx2( size_t padded_dim, float* estimated_distances ) { - qg_batch_estdist_impl( + fastscan::detail::qg_batch_estdist_impl( batch_data, query_wrapper, padded_dim, estimated_distances ); } @@ -66,7 +67,7 @@ uint32_t qg_batch_estdist_mask_avx2( float* estimated_distances, float threshold ) { - qg_batch_estdist_impl( + fastscan::detail::qg_batch_estdist_impl( batch_data, query_wrapper, padded_dim, estimated_distances ); return candidate_mask_impl(estimated_distances, threshold); diff --git a/src/fastscan/estimator_avx512.cpp b/src/fastscan/estimator_avx512.cpp index a7e8a092..1536b4cc 100644 --- a/src/fastscan/estimator_avx512.cpp +++ b/src/fastscan/estimator_avx512.cpp @@ -3,6 +3,7 @@ #include "estimator_kernels.hpp" #include "rabitqlib/fastscan/detail/estimator_dispatch.hpp" +#include "rabitqlib/fastscan/detail/qg_estimator.hpp" #include "rabitqlib/fastscan/detail/simd.hpp" namespace rabitqlib::simd { @@ -45,7 +46,7 @@ void qg_batch_estdist_avx512( if (padded_dim == 0 || padded_dim % 16 != 0) { throw std::invalid_argument("FastScan dimension must be a positive multiple of 16"); } - qg_batch_estdist_impl( + fastscan::detail::qg_batch_estdist_impl( batch_data, query_wrapper, padded_dim, estimated_distances ); } diff --git a/src/fastscan/estimator_generic.cpp b/src/fastscan/estimator_generic.cpp index 5601ecde..d0d3e325 100644 --- a/src/fastscan/estimator_generic.cpp +++ b/src/fastscan/estimator_generic.cpp @@ -3,6 +3,7 @@ #include "estimator_kernels.hpp" #include "rabitqlib/fastscan/detail/estimator_dispatch.hpp" +#include "rabitqlib/fastscan/detail/qg_estimator.hpp" #include "rabitqlib/fastscan/fastscan.hpp" #include "rabitqlib/fastscan/highacc_fastscan.hpp" @@ -45,7 +46,7 @@ void qg_batch_estdist_generic( size_t padded_dim, float* estimated_distances ) { - qg_batch_estdist_impl( + fastscan::detail::qg_batch_estdist_impl( batch_data, query_wrapper, padded_dim, estimated_distances ); } @@ -61,7 +62,7 @@ uint32_t qg_batch_estdist_mask_generic( float* estimated_distances, float threshold ) { - qg_batch_estdist_impl( + fastscan::detail::qg_batch_estdist_impl( batch_data, query_wrapper, padded_dim, estimated_distances ); return candidate_mask_impl(estimated_distances, threshold); diff --git a/src/fastscan/estimator_kernels.hpp b/src/fastscan/estimator_kernels.hpp index 8d3cedae..2454ec26 100644 --- a/src/fastscan/estimator_kernels.hpp +++ b/src/fastscan/estimator_kernels.hpp @@ -63,34 +63,5 @@ inline void split_batch_estdist_impl( low_dist_arr = est_dist_arr - f_error_arr * query_wrapper.g_error(); } -template -inline void qg_batch_estdist_impl( - const char* batch_data, - const BatchQuery& query_wrapper, - size_t padded_dim, - float* estimated_distances -) { - using T = float; - ConstQGBatchDataMap cur_batch(batch_data, padded_dim); - std::array accu_values{}; - Accumulate(cur_batch.bin_code(), query_wrapper.lut(), accu_values.data(), padded_dim); - - ConstRowMajorArrayMap ip_arr(accu_values.data(), 1, fastscan::kBatchSize); - std::array f_add_values; - std::array f_rescale_values; - cur_batch.f_add().copy_to(f_add_values.data(), f_add_values.size()); - cur_batch.f_rescale().copy_to(f_rescale_values.data(), f_rescale_values.size()); - ConstRowMajorArrayMap f_add_arr(f_add_values.data(), 1, fastscan::kBatchSize); - ConstRowMajorArrayMap f_rescale_arr( - f_rescale_values.data(), 1, fastscan::kBatchSize - ); - - RowMajorArrayMap est_dist_arr(estimated_distances, 1, fastscan::kBatchSize); - - est_dist_arr = f_add_arr + query_wrapper.g_add() + - (f_rescale_arr * (query_wrapper.delta() * (ip_arr.template cast()) + - query_wrapper.sum_vl_lut() + query_wrapper.k1xsumq())); -} - } // namespace } // namespace rabitqlib::simd diff --git a/src/index/hnsw_search_impl.hpp b/src/index/hnsw_search_impl.hpp index 48f85f34..fc6759dd 100644 --- a/src/index/hnsw_search_impl.hpp +++ b/src/index/hnsw_search_impl.hpp @@ -177,7 +177,6 @@ inline std::vector> HierarchicalNSW::search_knn_sorted_dir k, query_wrapper, q_to_centroids, - rotated_query, nearest_candidates ); result.reserve(nearest_candidates.size()); @@ -204,7 +203,6 @@ inline void HierarchicalNSW::search_base_layer_direct( size_t k, const SplitSingleQuery& query_wrapper, const std::vector& q_to_centroids, - [[maybe_unused]] const float* query, BoundedKNN& nearest_candidates ) const { auto visited = visited_list_pool_->get_free_vislist();