diff --git a/Makefile b/Makefile index 66701ac5c..23b78eb74 100644 --- a/Makefile +++ b/Makefile @@ -64,12 +64,13 @@ endif # which requires std::partial_ordering from (a C++20 feature). CONAN_SETTINGS := -s compiler.libcxx=$(LIBCXX) -s build_type=$(BUILD_TYPE) -s compiler.cppstd=20 -s:b compiler.cppstd=20 -# DiskANN and liburing require libaio (Linux-only). +# DiskANN is enabled for Linux builds. CONAN_INSTALL_FLAGS already contains +# --build=missing, which builds liburing when it is present in the graph and +# lacks a binary package. A separate --build=liburing pattern is unsafe with +# Conan because it is a hard error when a profile/option removes liburing from +# the resolved graph. ifneq ($(UNAME_S),Darwin) CONAN_SETTINGS += -o \&:with_diskann=True - ifndef WITH_GPU - CONAN_INSTALL_FLAGS += --build=liburing - endif endif # GPU builds use cuVS. diff --git a/include/knowhere/sparse_utils.h b/include/knowhere/sparse_utils.h index da73584a8..dc267892e 100644 --- a/include/knowhere/sparse_utils.h +++ b/include/knowhere/sparse_utils.h @@ -15,9 +15,11 @@ #pragma once #include +#include #include #include #include +#include #include #include #include @@ -29,6 +31,133 @@ namespace knowhere::sparse { +// DSP instrumentation (compile with -DKNOWHERE_DSP_INSTRUMENTATION to enable) +#ifdef KNOWHERE_DSP_INSTRUMENTATION +struct SeekStats { + std::atomic bucket_0{0}; // delta = 0 + std::atomic bucket_1_3{0}; // delta 1-3 + std::atomic bucket_4_15{0}; // delta 4-15 + std::atomic bucket_16_63{0}; // delta 16-63 + std::atomic bucket_64_255{0}; // delta 64-255 + std::atomic bucket_256_plus{0}; // delta 256+ + std::atomic seek_hits{0}; // seek found target doc_id + std::atomic seek_misses{0}; // seek did NOT find target doc_id + + void + record(size_t delta) { + if (delta == 0) + bucket_0++; + else if (delta <= 3) + bucket_1_3++; + else if (delta <= 15) + bucket_4_15++; + else if (delta <= 63) + bucket_16_63++; + else if (delta <= 255) + bucket_64_255++; + else + bucket_256_plus++; + } + + void + record_hit() { + seek_hits++; + } + void + record_miss() { + seek_misses++; + } + + void + print(const char* label = nullptr) const { + if (label) + printf("\n[Seek Stats: %s]\n", label); + else + printf("\n[Seek Distance Distribution]\n"); + uint64_t total = bucket_0 + bucket_1_3 + bucket_4_15 + bucket_16_63 + bucket_64_255 + bucket_256_plus; + printf(" delta=0: %lu (%.1f%%)\n", bucket_0.load(), total ? 100.0 * bucket_0 / total : 0); + printf(" delta 1-3: %lu (%.1f%%)\n", bucket_1_3.load(), total ? 100.0 * bucket_1_3 / total : 0); + printf(" delta 4-15: %lu (%.1f%%)\n", bucket_4_15.load(), total ? 100.0 * bucket_4_15 / total : 0); + printf(" delta 16-63: %lu (%.1f%%)\n", bucket_16_63.load(), total ? 100.0 * bucket_16_63 / total : 0); + printf(" delta 64-255: %lu (%.1f%%)\n", bucket_64_255.load(), total ? 100.0 * bucket_64_255 / total : 0); + printf(" delta 256+: %lu (%.1f%%)\n", bucket_256_plus.load(), total ? 100.0 * bucket_256_plus / total : 0); + printf(" total seeks: %lu\n", total); + uint64_t h = seek_hits.load(), m = seek_misses.load(); + uint64_t hm = h + m; + printf(" seek hits: %lu (%.1f%%)\n", h, hm ? 100.0 * h / hm : 0); + printf(" seek misses: %lu (%.1f%%)\n", m, hm ? 100.0 * m / hm : 0); + } + + void + reset() { + bucket_0.store(0, std::memory_order_relaxed); + bucket_1_3.store(0, std::memory_order_relaxed); + bucket_4_15.store(0, std::memory_order_relaxed); + bucket_16_63.store(0, std::memory_order_relaxed); + bucket_64_255.store(0, std::memory_order_relaxed); + bucket_256_plus.store(0, std::memory_order_relaxed); + seek_hits.store(0, std::memory_order_relaxed); + seek_misses.store(0, std::memory_order_relaxed); + } +}; + +inline SeekStats g_seek_stats; + +struct DspStats { + std::atomic total_superblocks{0}; // total superblocks considered + std::atomic surviving_superblocks{0}; // superblocks surviving coarse pruning + std::atomic candidate_blocks{0}; // subblocks passing the initial UB threshold + std::atomic blocks_processed{0}; // candidate subblocks actually scored + std::atomic saturated_ubs{0}; // surviving subblock UBs saturated at uint16 max + std::atomic entries_scored{0}; // posting list entries iterated + std::atomic docs_pushed{0}; // docs pushed to heap + std::atomic queries{0}; // number of queries + std::atomic workspace_pool_misses{0}; // searches that allocate because the per-index pool is empty + + void + print(const char* label = nullptr) const { + if (label) + printf("\n[DSP Block Stats: %s]\n", label); + else + printf("\n[DSP Block Stats]\n"); + uint64_t q = queries.load(); + uint64_t total_spb = total_superblocks.load(); + uint64_t surviving_spb = surviving_superblocks.load(); + uint64_t candidates = candidate_blocks.load(); + uint64_t processed = blocks_processed.load(); + printf(" queries: %lu\n", q); + printf(" superblocks total: %lu (avg %.1f/q)\n", total_spb, q ? (double)total_spb / q : 0); + printf(" superblocks surviving:%lu (avg %.1f/q, %.1f%%)\n", surviving_spb, q ? (double)surviving_spb / q : 0, + total_spb ? 100.0 * surviving_spb / total_spb : 0); + printf(" candidate blocks: %lu (avg %.1f/q)\n", candidates, q ? (double)candidates / q : 0); + printf(" blocks processed: %lu (avg %.1f/q, %.1f%% of candidates)\n", processed, + q ? (double)processed / q : 0, candidates ? 100.0 * processed / candidates : 0); + printf(" saturated UBs: %lu (avg %.1f/q)\n", saturated_ubs.load(), q ? (double)saturated_ubs / q : 0); + printf(" entries scored: %lu (avg %.1f/q)\n", entries_scored.load(), q ? (double)entries_scored / q : 0); + printf(" docs pushed: %lu (avg %.1f/q)\n", docs_pushed.load(), q ? (double)docs_pushed / q : 0); + printf(" workspace pool misses:%lu\n", workspace_pool_misses.load()); + if (processed > 0) { + printf(" entries/block: %.1f\n", (double)entries_scored / processed); + } + } + + void + reset() { + total_superblocks.store(0, std::memory_order_relaxed); + surviving_superblocks.store(0, std::memory_order_relaxed); + candidate_blocks.store(0, std::memory_order_relaxed); + blocks_processed.store(0, std::memory_order_relaxed); + saturated_ubs.store(0, std::memory_order_relaxed); + entries_scored.store(0, std::memory_order_relaxed); + docs_pushed.store(0, std::memory_order_relaxed); + queries.store(0, std::memory_order_relaxed); + workspace_pool_misses.store(0, std::memory_order_relaxed); + } +}; + +inline DspStats g_dsp_stats; +#endif + enum class SparseMetricType { METRIC_IP = 1, METRIC_BM25 = 2, diff --git a/src/index/sparse/sparse_dsp_config.h b/src/index/sparse/sparse_dsp_config.h index b4c8ddd9f..8c6c28dfa 100644 --- a/src/index/sparse/sparse_dsp_config.h +++ b/src/index/sparse/sparse_dsp_config.h @@ -36,7 +36,7 @@ class SparseDspConfig : public BaseConfig { CFG_INT dsp_gamma; CFG_BOOL dsp_kth_init; CFG_FLOAT dsp_kth_alpha; - KNOHWERE_DECLARE_CONFIG(SparseDspConfig) { + KNOWHERE_DECLARE_CONFIG(SparseDspConfig) { KNOWHERE_CONFIG_DECLARE_FIELD(drop_ratio_search) .description("drop ratio for search") .set_default(0.0f) diff --git a/src/index/sparse/sparse_dsp_index.h b/src/index/sparse/sparse_dsp_index.h index ac3117ed5..a51dd5091 100644 --- a/src/index/sparse/sparse_dsp_index.h +++ b/src/index/sparse/sparse_dsp_index.h @@ -22,8 +22,10 @@ #include #include #include +#include #include #include +#include #include #include #include @@ -32,6 +34,7 @@ #include "io/memory_io.h" #include "knowhere/bitsetview.h" #include "knowhere/comp/index_param.h" +#include "knowhere/comp/task.h" #include "knowhere/config.h" #include "knowhere/expected.h" #include "knowhere/heap.h" @@ -44,6 +47,8 @@ namespace knowhere::sparse { +// Keep score-only admission semantics here: equal-to-threshold scores do not enter the heap. Faiss CMin heaps also +// compare IDs on score ties, which would change DSP's established tie selection and bit-exact search results. using DspHeap = knowhere::ResultMinHeap; // Section types for DSP index serialization format @@ -129,6 +134,9 @@ class DspIndexBase { // - Counting sort (bucket sort) for block ordering by upper bound // - Forward index with two-pointer merge scoring // - Two-level hierarchy: superblocks for coarse pruning, subblocks for scoring +// +// The upper-bound construction assumes every corpus and query value is finite and non-negative. The index node +// validates this precondition at both build and search time; silently clamping invalid values would break rank safety. template class DspIndex : public DspIndexBase { public: @@ -231,25 +239,86 @@ class DspIndex : public DspIndexBase { if constexpr (mmapped) { throw std::invalid_argument("mmapped DspIndex does not support Add"); } else { - auto current_rows = n_rows_internal_; - if ((size_t)dim > max_dim_) { - max_dim_ = dim; + if (n_rows_internal_ != 0 || rows > std::numeric_limits::max()) { + return Status::invalid_args; } - - if (metric_type_ == SparseMetricType::METRIC_BM25) { - bm25_params_->row_sums.reserve(current_rows + rows); + n_rows_internal_ = rows; + max_dim_ = std::max(max_dim_, static_cast(dim)); + const bool is_bm25 = metric_type_ == SparseMetricType::METRIC_BM25; + std::vector dim_counts; + uint64_t total_entries = 0; + if (is_bm25) { + bm25_params_->row_sums.clear(); + bm25_params_->row_sums.reserve(rows); } - for (size_t i = 0; i < rows; ++i) { - add_row_to_index(data[i], current_rows + i); +#if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) + build_stats_.dataset_nnz_stats_.clear(); + build_stats_.dataset_nnz_stats_.reserve(rows); +#endif + for (uint32_t doc_id = 0; doc_id < rows; ++doc_id) { + float row_sum = 0.0f; + if (is_bm25) { + for (size_t j = 0; j < data[doc_id].size(); ++j) { + row_sum += data[doc_id][j].val; + } + bm25_params_->row_sums.push_back(row_sum); + } +#if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) + build_stats_.dataset_nnz_stats_.push_back(data[doc_id].size()); +#endif + for (size_t j = 0; j < data[doc_id].size(); ++j) { + const auto [raw_dim, val] = data[doc_id][j]; + if (val == 0) + continue; + auto [it, inserted] = dim_map_.try_emplace(raw_dim, next_dim_id_); + if (inserted) { + ++next_dim_id_; + dim_counts.push_back(0); + max_score_in_dim_.emplace_back(0.0f); + } + const uint32_t inner_dim = it->second; + ++dim_counts[inner_dim]; + ++total_entries; + const QType quantized = get_quant_val(val); + const float score = + is_bm25 ? bm25_params_->max_score_computer(quantized, row_sum) : static_cast(quantized); + max_score_in_dim_[inner_dim] = std::max(max_score_in_dim_[inner_dim], score); + } + } + if (total_entries > std::numeric_limits::max()) { + return Status::invalid_args; + } + nr_inner_dims_ = next_dim_id_; + inverted_index_ids_.resize(nr_inner_dims_); + inverted_index_vals_.resize(nr_inner_dims_); + for (uint32_t d = 0; d < nr_inner_dims_; ++d) { + inverted_index_ids_[d].resize(dim_counts[d]); + inverted_index_vals_[d].resize(dim_counts[d]); + } + std::vector write_pos(nr_inner_dims_, 0); + std::vector dim_spb_counts(nr_inner_dims_, 0); + std::vector last_spb(nr_inner_dims_, std::numeric_limits::max()); + for (uint32_t doc_id = 0; doc_id < rows; ++doc_id) { + for (size_t j = 0; j < data[doc_id].size(); ++j) { + const auto [raw_dim, val] = data[doc_id][j]; + if (val == 0) + continue; + const uint32_t inner_dim = dim_map_.at(raw_dim); + const uint64_t pos = write_pos[inner_dim]++; + inverted_index_ids_[inner_dim][pos] = doc_id; + inverted_index_vals_[inner_dim][pos] = get_quant_val(val); + const uint32_t spb = doc_id / kSuperblockSize; + if (last_spb[inner_dim] != spb) { + last_spb[inner_dim] = spb; + ++dim_spb_counts[inner_dim]; + } + } } - n_rows_internal_ += rows; - - nr_inner_dims_ = dim_map_.size(); #if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) build_stats_.posting_list_length_stats_.resize(nr_inner_dims_); for (size_t i = 0; i < nr_inner_dims_; ++i) { - build_stats_.posting_list_length_stats_[i] = inverted_index_ids_[i].size(); + build_stats_.posting_list_length_stats_[i] = dim_counts[i]; } #endif @@ -272,7 +341,7 @@ class DspIndex : public DspIndexBase { boost::span(bm25_params_->row_sums.data(), bm25_params_->row_sums.size()); } - build_dsp_metadata(); + build_dsp_metadata(data, rows, &dim_spb_counts); return Status::success; } } @@ -630,17 +699,20 @@ class DspIndex : public DspIndexBase { return; } - const size_t heap_capacity = k * approx_params.refine_factor; + if (approx_params.refine_factor > 1) { + static std::once_flag refine_warning_once; + std::call_once(refine_warning_once, []() { + LOG_KNOWHERE_WARNING_ << "DSP ignores refine_factor because its full-precision forward index already " + "performs exact scoring and build does not retain posting lists"; + }); + } + const size_t heap_capacity = k; DspHeap heap(heap_capacity); search_dsp(q_vec, heap, heap_capacity, bitset, computer, approx_params.dsp_mode, approx_params.dsp_mu, approx_params.dsp_eta, approx_params.dsp_gamma, approx_params.dsp_kth_init, approx_params.dsp_kth_alpha); - if (approx_params.refine_factor == 1) { - collect_result(heap, distances, labels); - } else { - refine_and_collect(query, heap, k, distances, labels, computer); - } + collect_result(heap, distances, labels); } std::vector @@ -1004,16 +1076,188 @@ class DspIndex : public DspIndexBase { uint32_t n_sb_padded_ = 0; bool dsp_loaded_ = false; + struct SearchWorkspace { + std::vector superblock_ub; + std::vector superblock_asc; + std::vector surviving_spb; + std::vector spb_alive; + std::vector block_ub; + std::vector spb_candidate_mask; + std::vector spb_in_batch; + }; + + // Roughly 5MB/workspace at 10M documents; the production cap of 32 bounds retained scratch near 160MB/index. + mutable std::once_flag search_workspace_pool_once_; + mutable size_t max_cached_search_workspaces_ = 2; + mutable std::mutex search_workspace_mutex_; + mutable std::vector> cached_search_workspaces_; + + struct SearchWorkspaceDeleter { + const DspIndex* index; + void + operator()(SearchWorkspace* workspace) const { + index->release_search_workspace(workspace); + } + }; + using SearchWorkspacePtr = std::unique_ptr; + + void + initialize_search_workspace_pool() const { + std::call_once(search_workspace_pool_once_, [this]() { + max_cached_search_workspaces_ = std::clamp(GetSearchThreadPoolSize(), 2, 32); + std::lock_guard lock(search_workspace_mutex_); + cached_search_workspaces_.reserve(max_cached_search_workspaces_); + }); + } + + SearchWorkspacePtr + acquire_search_workspace() const { + initialize_search_workspace_pool(); + { + std::lock_guard lock(search_workspace_mutex_); + if (!cached_search_workspaces_.empty()) { + auto workspace = std::move(cached_search_workspaces_.back()); + cached_search_workspaces_.pop_back(); + return SearchWorkspacePtr(workspace.release(), SearchWorkspaceDeleter{this}); + } + } +#ifdef KNOWHERE_DSP_INSTRUMENTATION + g_dsp_stats.workspace_pool_misses++; +#endif + return SearchWorkspacePtr(new SearchWorkspace(), SearchWorkspaceDeleter{this}); + } + + void + release_search_workspace(SearchWorkspace* workspace) const { + std::unique_ptr owned(workspace); + std::lock_guard lock(search_workspace_mutex_); + if (cached_search_workspaces_.size() < max_cached_search_workspaces_) + cached_search_workspaces_.push_back(std::move(owned)); + } + static constexpr float kDenseThreshold = 0.125f; static constexpr uint32_t kNumSegments = 8; static constexpr uint32_t kSegmentSize = kSuperblockSize / kNumSegments; + template + static void + run_build_ranges(uint32_t count, uint32_t min_grain, Function&& function) { + if (count == 0) + return; + const size_t pool_size = GetBuildThreadPoolSize(); + const uint32_t task_count = + static_cast(std::min(pool_size, (count + min_grain - 1) / min_grain)); + if (task_count <= 1) { + function(0, count); + return; + } + std::vector> tasks; + tasks.reserve(task_count); + for (uint32_t task = 0; task < task_count; ++task) { + const uint32_t begin = static_cast(static_cast(count) * task / task_count); + const uint32_t end = static_cast(static_cast(count) * (task + 1) / task_count); + tasks.emplace_back([begin, end, &function]() { function(begin, end); }); + } + ExecOverBuildThreadPool(tasks); + } + + Status + build_forward_index_from_rows(const SparseRow* data, size_t rows, uint32_t total_entries) { + struct BlockEntry { + uint32_t inner_dim; + uint8_t doc_offset; + float score; + }; + std::vector block_term_counts(n_subblocks_); + std::vector block_entry_counts(n_subblocks_); + auto count_blocks = [&](uint32_t begin, uint32_t end) { + std::vector block_dims; + block_dims.reserve(1024); + for (uint32_t sb = begin; sb < end; ++sb) { + block_dims.clear(); + const uint32_t doc_start = sb * kSubblockSize; + const uint32_t doc_end = std::min(doc_start + kSubblockSize, static_cast(rows)); + for (uint32_t doc_id = doc_start; doc_id < doc_end; ++doc_id) { + for (size_t j = 0; j < data[doc_id].size(); ++j) { + const auto [raw_dim, val] = data[doc_id][j]; + if (val != 0) + block_dims.push_back(dim_map_.at(raw_dim)); + } + } + block_entry_counts[sb] = block_dims.size(); + std::sort(block_dims.begin(), block_dims.end()); + block_term_counts[sb] = std::unique(block_dims.begin(), block_dims.end()) - block_dims.begin(); + } + }; + run_build_ranges(n_subblocks_, 1024, count_blocks); + + uint64_t total_terms = 0; + uint64_t counted_entries = 0; + fwd_block_term_offsets_.resize(n_subblocks_ + 1); + std::vector block_entry_offsets(n_subblocks_ + 1); + for (uint32_t sb = 0; sb < n_subblocks_; ++sb) { + fwd_block_term_offsets_[sb] = total_terms; + block_entry_offsets[sb] = counted_entries; + total_terms += block_term_counts[sb]; + counted_entries += block_entry_counts[sb]; + } + if (total_terms > std::numeric_limits::max() || counted_entries != total_entries) + return Status::invalid_args; + fwd_block_term_offsets_[n_subblocks_] = total_terms; + block_entry_offsets[n_subblocks_] = counted_entries; + fwd_term_ids_.resize(total_terms); + fwd_term_entry_offsets_.resize(total_terms + 1); + fwd_doc_offsets_.resize(total_entries); + fwd_scores_.resize(total_entries); + const bool is_bm25 = metric_type_ == SparseMetricType::METRIC_BM25; + auto fill_blocks = [&](uint32_t begin, uint32_t end) { + std::vector block_entries; + block_entries.reserve(1024); + for (uint32_t sb = begin; sb < end; ++sb) { + block_entries.clear(); + const uint32_t doc_start = sb * kSubblockSize; + const uint32_t doc_end = std::min(doc_start + kSubblockSize, static_cast(rows)); + for (uint32_t doc_id = doc_start; doc_id < doc_end; ++doc_id) { + const uint8_t doc_offset = static_cast(doc_id - doc_start); + const float row_sum = is_bm25 ? bm25_params_->row_sums[doc_id] : 0.0f; + for (size_t j = 0; j < data[doc_id].size(); ++j) { + const auto [raw_dim, val] = data[doc_id][j]; + if (val == 0) + continue; + const QType quantized = get_quant_val(val); + const float score = is_bm25 ? bm25_params_->max_score_computer(quantized, row_sum) + : static_cast(quantized); + block_entries.push_back({dim_map_.at(raw_dim), doc_offset, score}); + } + } + std::sort(block_entries.begin(), block_entries.end(), [](const auto& lhs, const auto& rhs) { + return lhs.inner_dim < rhs.inner_dim || + (lhs.inner_dim == rhs.inner_dim && lhs.doc_offset < rhs.doc_offset); + }); + uint32_t term_pos = fwd_block_term_offsets_[sb]; + uint32_t entry_pos = block_entry_offsets[sb]; + for (size_t i = 0; i < block_entries.size(); ++i) { + if (i == 0 || block_entries[i].inner_dim != block_entries[i - 1].inner_dim) { + fwd_term_ids_[term_pos] = block_entries[i].inner_dim; + fwd_term_entry_offsets_[term_pos++] = entry_pos; + } + fwd_doc_offsets_[entry_pos] = block_entries[i].doc_offset; + fwd_scores_[entry_pos++] = block_entries[i].score; + } + } + }; + run_build_ranges(n_subblocks_, 1024, fill_blocks); + fwd_term_entry_offsets_[total_terms] = total_entries; + return Status::success; + } + // ======================================================================== // Build DSP metadata from inverted index // ======================================================================== void - build_dsp_metadata() { + build_dsp_metadata(const SparseRow* source_rows = nullptr, size_t source_row_count = 0, + const std::vector* precomputed_spb_counts = nullptr) { if (n_rows_internal_ == 0 || nr_inner_dims_ == 0) { return; } @@ -1030,26 +1274,56 @@ class DspIndex : public DspIndexBase { uint32_t inner_dim = 0; float score = 0.0f; }; - std::vector> per_doc_fwd(n_rows_internal_); + std::vector> per_doc_fwd; + if (source_rows == nullptr) { + per_doc_fwd.resize(n_rows_internal_); + } std::vector tmp_sb_max(n_subblocks_, 0.0f); std::vector sb_touched(n_subblocks_, 0); std::vector touched_list; touched_list.reserve(n_subblocks_); - std::vector tmp_spb_max(n_superblocks_, 0.0f); - std::vector spb_touched(n_superblocks_, 0); - std::vector spb_touched_list; - spb_touched_list.reserve(n_superblocks_); - - std::vector tmp_seg_max(n_superblocks_ * kNumSegments, 0.0f); + // Build computes these counts while filling the doc-major -> dim-major CSC. Legacy deserialization has no + // source rows, so count distinct (dimension, superblock) pairs from its sorted posting lists here instead. + std::vector legacy_spb_counts; + const std::vector* spb_counts = precomputed_spb_counts; + if (spb_counts == nullptr) { + legacy_spb_counts.resize(nr_dims, 0); + for (uint32_t d = 0; d < nr_dims; ++d) { + uint32_t last_spb = std::numeric_limits::max(); + for (const uint32_t doc_id : inverted_index_ids_spans_[d]) { + const uint32_t spb = doc_id / kSuperblockSize; + if (spb != last_spb) { + if (last_spb != std::numeric_limits::max() && spb < last_spb) { + throw std::runtime_error("DSP posting list is not sorted by document ID"); + } + last_spb = spb; + ++legacy_spb_counts[d]; + } + } + } + spb_counts = &legacy_spb_counts; + } + if (spb_counts->size() != nr_dims) { + throw std::runtime_error("DSP superblock count size does not match dimension count"); + } - struct SpbEntry { - uint32_t block_id; - float max_score; - float asc; - }; - std::vector> per_dim_spb(nr_dims); + spb_dim_offsets_.resize(nr_dims + 1); + uint64_t total_spb = 0; + for (uint32_t d = 0; d < nr_dims; ++d) { + spb_dim_offsets_[d] = static_cast(total_spb); + if (max_score_in_dim_spans_[d] > 0.0f) { + total_spb += (*spb_counts)[d]; + } + if (total_spb > std::numeric_limits::max()) { + throw std::runtime_error("DSP superblock metadata exceeds uint32 capacity"); + } + } + spb_dim_offsets_[nr_dims] = static_cast(total_spb); + spb_block_ids_.resize(total_spb); + spb_max_vals_.resize(total_spb); + spb_asc_vals_.resize(total_spb); dim_block_max_.resize(nr_dims); @@ -1068,6 +1342,30 @@ class DspIndex : public DspIndexBase { const float inv_max_score_u8 = 255.0f / max_score_d; KthHeap kth_heaps[4]; + uint32_t current_spb = std::numeric_limits::max(); + uint32_t spb_write_pos = spb_dim_offsets_[d]; + float current_spb_max = 0.0f; + std::array current_seg_max{}; + auto emit_current_spb = [&]() { + if (current_spb == std::numeric_limits::max()) { + return; + } + if (spb_write_pos >= spb_dim_offsets_[d + 1]) { + throw std::runtime_error("DSP emitted more superblocks than counted"); + } + float seg_sum = 0.0f; + uint32_t seg_count = 0; + for (const float seg_max : current_seg_max) { + if (seg_max > 0.0f) { + seg_sum += seg_max; + ++seg_count; + } + } + spb_block_ids_[spb_write_pos] = current_spb; + spb_max_vals_[spb_write_pos] = current_spb_max; + spb_asc_vals_[spb_write_pos] = seg_count > 0 ? seg_sum / seg_count : 0.0f; + ++spb_write_pos; + }; for (size_t i = 0; i < plist_ids.size(); ++i) { const uint32_t doc_id = plist_ids[i]; @@ -1092,22 +1390,33 @@ class DspIndex : public DspIndexBase { const uint32_t sb = doc_id / kSubblockSize; const uint32_t spb = doc_id / kSuperblockSize; + if (spb != current_spb) { + if (current_spb != std::numeric_limits::max() && spb < current_spb) { + throw std::runtime_error("DSP posting list is not sorted by document ID"); + } + emit_current_spb(); + current_spb = spb; + current_spb_max = 0.0f; + current_seg_max.fill(0.0f); + } + if (!sb_touched[sb]) { touched_list.push_back(sb); sb_touched[sb] = 1; } tmp_sb_max[sb] = std::max(tmp_sb_max[sb], score); - if (!spb_touched[spb]) { - spb_touched_list.push_back(spb); - spb_touched[spb] = 1; - } - tmp_spb_max[spb] = std::max(tmp_spb_max[spb], score); - - const uint32_t seg = doc_id / kSegmentSize; - tmp_seg_max[seg] = std::max(tmp_seg_max[seg], score); + current_spb_max = std::max(current_spb_max, score); + const uint32_t segment_in_spb = (doc_id / kSegmentSize) % kNumSegments; + current_seg_max[segment_in_spb] = std::max(current_seg_max[segment_in_spb], score); - per_doc_fwd[doc_id].push_back({d, score}); + if (source_rows == nullptr) { + per_doc_fwd[doc_id].push_back({d, score}); + } + } + emit_current_spb(); + if (spb_write_pos != spb_dim_offsets_[d + 1]) { + throw std::runtime_error("DSP emitted fewer superblocks than counted"); } auto& bm = dim_block_max_[d]; @@ -1140,58 +1449,35 @@ class DspIndex : public DspIndexBase { } } - std::sort(spb_touched_list.begin(), spb_touched_list.end()); - per_dim_spb[d].reserve(spb_touched_list.size()); - for (uint32_t spb : spb_touched_list) { - float seg_sum = 0.0f; - uint32_t seg_count = 0; - for (uint32_t s = 0; s < kNumSegments; ++s) { - float seg_max = tmp_seg_max[spb * kNumSegments + s]; - if (seg_max > 0.0f) { - seg_sum += seg_max; - seg_count++; - } - } - float asc = (seg_count > 0) ? (seg_sum / seg_count) : 0.0f; - per_dim_spb[d].push_back({spb, tmp_spb_max[spb], asc}); - } - for (uint32_t sb : touched_list) { tmp_sb_max[sb] = 0.0f; sb_touched[sb] = 0; } touched_list.clear(); - for (uint32_t spb : spb_touched_list) { - tmp_spb_max[spb] = 0.0f; - spb_touched[spb] = 0; - for (uint32_t s = 0; s < kNumSegments; ++s) { - tmp_seg_max[spb * kNumSegments + s] = 0.0f; - } - } - spb_touched_list.clear(); } - // ---- Phase 2: Build superblock CSR ---- - { - uint32_t total_spb = 0; - spb_dim_offsets_.resize(nr_dims + 1); - for (uint32_t d = 0; d < nr_dims; ++d) { - spb_dim_offsets_[d] = total_spb; - total_spb += per_dim_spb[d].size(); - } - spb_dim_offsets_[nr_dims] = total_spb; - - spb_block_ids_.resize(total_spb); - spb_max_vals_.resize(total_spb); - spb_asc_vals_.resize(total_spb); - for (uint32_t d = 0; d < nr_dims; ++d) { - uint32_t off = spb_dim_offsets_[d]; - for (const auto& e : per_dim_spb[d]) { - spb_block_ids_[off] = e.block_id; - spb_max_vals_[off] = e.max_score; - spb_asc_vals_[off] = e.asc; - ++off; + if constexpr (!mmapped) { + if (source_rows != nullptr) { + uint64_t total_entries = 0; + for (uint32_t d = 0; d < nr_dims; ++d) { + total_entries += inverted_index_ids_spans_[d].size(); + } + // The transient exact CSC has served its only purpose. Release it before allocating the flat forward + // arrays so the two corpus-sized representations do not overlap at the build peak. + inverted_index_ids_spans_.clear(); + inverted_index_vals_spans_.clear(); + inverted_index_ids_.clear(); + inverted_index_vals_.clear(); + inverted_index_ids_.shrink_to_fit(); + inverted_index_vals_.shrink_to_fit(); + inverted_index_ids_spans_.resize(nr_dims); + inverted_index_vals_spans_.resize(nr_dims); + const auto status = + build_forward_index_from_rows(source_rows, source_row_count, static_cast(total_entries)); + if (status != Status::success) { + throw std::runtime_error("failed to build DSP forward index directly from sparse rows"); } + return; } } @@ -1309,6 +1595,12 @@ class DspIndex : public DspIndexBase { } std::sort(query.begin(), query.end(), [](const auto& a, const auto& b) { return a.inner_dim < b.inner_dim; }); const size_t n_query_terms = query.size(); + uint32_t hybrid_query_dims[kHybridMergeMaxQueryTerms] = {}; + if (n_query_terms <= kHybridMergeMaxQueryTerms) { + for (size_t i = 0; i < n_query_terms; ++i) { + hybrid_query_dims[i] = query[i].inner_dim; + } + } // ---- Step 1: Compute u8 query weights and scale factor ---- float S = 0.0f; @@ -1328,10 +1620,17 @@ class DspIndex : public DspIndexBase { // ---- Step 2: Initialize thresholds from kth scores ---- const bool has_filter = !filter.empty(); + const bool use_asc_survivor_guard = mu < 1.0f; + // With mu >= 1, an ASC-only survivor has max_ub <= theta/mu <= theta. Since max_ub bounds every document + // score and theta only rises, such a superblock cannot satisfy the strict score > theta admission condition. + const bool need_asc = + use_asc_survivor_guard && (mode == DspSearchMode::DSP || mode == DspSearchMode::LSP2 || gamma <= 0); bool bootstrap_mode = has_filter; float float_threshold = 0.0f; - if (kth_init && !has_filter) { +#ifndef DSP_DISABLE_KTH_INIT + if (kth_init && !has_filter && heap_capacity <= 10000) { + // Select kth bucket based on k int kth_bucket = (heap_capacity > 10) + (heap_capacity > 100) + (heap_capacity > 1000); for (const auto& qt : query) { const auto& bm = dim_block_max_[qt.inner_dim]; @@ -1342,31 +1641,58 @@ class DspIndex : public DspIndexBase { float term_thresh = qt.weight * kth_float; float_threshold = std::max(float_threshold, term_thresh); } - float_threshold *= kth_alpha; + float_threshold *= kth_alpha * (1.0f - 1e-6f); } +#endif float float_block_threshold = (eta > 0.0f) ? float_threshold / eta : float_threshold; uint16_t u16_block_threshold = static_cast(std::min(65535.0f, float_block_threshold * score_scale)); + auto workspace_owner = acquire_search_workspace(); + auto& workspace = *workspace_owner; + workspace.superblock_ub.resize(n_superblocks_); + if (need_asc) { + workspace.superblock_asc.resize(n_superblocks_); + } + workspace.surviving_spb.clear(); + workspace.spb_alive.resize(n_superblocks_); + workspace.block_ub.resize(n_sb_padded_); + workspace.spb_candidate_mask.resize(n_superblocks_); + workspace.spb_in_batch.resize(n_superblocks_); + std::fill(workspace.superblock_ub.begin(), workspace.superblock_ub.end(), 0.0f); + if (need_asc) { + std::fill(workspace.superblock_asc.begin(), workspace.superblock_asc.end(), 0.0f); + } + std::fill(workspace.spb_alive.begin(), workspace.spb_alive.end(), uint8_t{0}); + std::fill(workspace.spb_in_batch.begin(), workspace.spb_in_batch.end(), uint8_t{0}); + auto& superblock_ub = workspace.superblock_ub; + auto& superblock_asc = workspace.superblock_asc; + auto& surviving_spb = workspace.surviving_spb; + auto& spb_alive = workspace.spb_alive; + auto& block_ub = workspace.block_ub; + auto& spb_candidate_mask = workspace.spb_candidate_mask; + auto& spb_in_batch = workspace.spb_in_batch; + // ---- Step 3: Superblock pruning ---- - std::vector superblock_ub(n_superblocks_, 0.0f); - std::vector superblock_asc(n_superblocks_, 0.0f); for (const auto& qt : query) { const float qw = qt.weight; const uint32_t start = spb_dim_offsets_[qt.inner_dim]; const uint32_t end = spb_dim_offsets_[qt.inner_dim + 1]; - for (uint32_t i = start; i < end; ++i) { - superblock_ub[spb_block_ids_[i]] += qw * spb_max_vals_[i]; - superblock_asc[spb_block_ids_[i]] += qw * spb_asc_vals_[i]; + if (need_asc) { + for (uint32_t i = start; i < end; ++i) { + superblock_ub[spb_block_ids_[i]] += qw * spb_max_vals_[i]; + superblock_asc[spb_block_ids_[i]] += qw * spb_asc_vals_[i]; + } + } else { + for (uint32_t i = start; i < end; ++i) { + superblock_ub[spb_block_ids_[i]] += qw * spb_max_vals_[i]; + } } } const float theta = float_threshold; float mu_threshold = (mu > 0.0f) ? theta / mu : theta; float eta_threshold = (eta > 0.0f) ? theta / eta : theta; - - std::vector surviving_spb; surviving_spb.reserve(n_superblocks_); - std::vector spb_alive(n_superblocks_, 0); auto mark_alive = [&](uint32_t spb) { if (!spb_alive[spb]) { @@ -1405,7 +1731,8 @@ class DspIndex : public DspIndexBase { case DspSearchMode::DSP: { // dual-threshold (mu, eta) + optional top-gamma backstop for (uint32_t spb = 0; spb < n_superblocks_; ++spb) { - if (superblock_ub[spb] > mu_threshold || superblock_asc[spb] > eta_threshold) { + if (superblock_ub[spb] > mu_threshold || + (use_asc_survivor_guard && superblock_asc[spb] > eta_threshold)) { mark_alive(spb); } } @@ -1418,7 +1745,8 @@ class DspIndex : public DspIndexBase { if (gamma <= 0) { // fallback to DSP behavior for (uint32_t spb = 0; spb < n_superblocks_; ++spb) { - if (superblock_ub[spb] > mu_threshold || superblock_asc[spb] > eta_threshold) { + if (superblock_ub[spb] > mu_threshold || + (use_asc_survivor_guard && superblock_asc[spb] > eta_threshold)) { mark_alive(spb); } } @@ -1431,7 +1759,8 @@ class DspIndex : public DspIndexBase { if (gamma <= 0) { // fallback to DSP behavior for (uint32_t spb = 0; spb < n_superblocks_; ++spb) { - if (superblock_ub[spb] > mu_threshold || superblock_asc[spb] > eta_threshold) { + if (superblock_ub[spb] > mu_threshold || + (use_asc_survivor_guard && superblock_asc[spb] > eta_threshold)) { mark_alive(spb); } } @@ -1449,7 +1778,8 @@ class DspIndex : public DspIndexBase { if (gamma <= 0) { // fallback to DSP behavior for (uint32_t spb = 0; spb < n_superblocks_; ++spb) { - if (superblock_ub[spb] > mu_threshold || superblock_asc[spb] > eta_threshold) { + if (superblock_ub[spb] > mu_threshold || + (use_asc_survivor_guard && superblock_asc[spb] > eta_threshold)) { mark_alive(spb); } } @@ -1457,7 +1787,8 @@ class DspIndex : public DspIndexBase { } add_top_gamma(gamma, float_threshold, true); for (uint32_t spb = 0; spb < n_superblocks_; ++spb) { - if (superblock_ub[spb] > mu_threshold || superblock_asc[spb] > eta_threshold) { + if (superblock_ub[spb] > mu_threshold || + (use_asc_survivor_guard && superblock_asc[spb] > eta_threshold)) { mark_alive(spb); } } @@ -1466,7 +1797,8 @@ class DspIndex : public DspIndexBase { default: { // Default: same as DSP mode for (uint32_t spb = 0; spb < n_superblocks_; ++spb) { - if (superblock_ub[spb] > mu_threshold || superblock_asc[spb] > eta_threshold) { + if (superblock_ub[spb] > mu_threshold || + (use_asc_survivor_guard && superblock_asc[spb] > eta_threshold)) { mark_alive(spb); } } @@ -1478,23 +1810,33 @@ class DspIndex : public DspIndexBase { } }; - // ---- Compute block UBs for a set of superblocks ---- - std::vector block_ub(n_sb_padded_, 0); - std::vector spb_in_batch(n_superblocks_, 0); +#ifdef KNOWHERE_DSP_INSTRUMENTATION + g_dsp_stats.total_superblocks += n_superblocks_; + g_dsp_stats.queries++; +#endif + // ---- Compute block UBs for a set of superblocks ---- auto compute_block_ubs = [&](const std::vector& spbs) { - for (uint32_t spb : spbs) spb_in_batch[spb] = 1; +#ifdef KNOWHERE_DSP_INSTRUMENTATION + g_dsp_stats.surviving_superblocks += spbs.size(); +#endif + for (uint32_t spb : spbs) { + spb_in_batch[spb] = 1; + spb_candidate_mask[spb] = 0; + std::fill_n(block_ub.begin() + spb * kStride, kStride, uint16_t{0}); + } + std::vector dense_block_max_rows; + std::vector dense_query_weights; + dense_block_max_rows.reserve(n_query_terms); + dense_query_weights.reserve(n_query_terms); for (const auto& qt : query) { const auto& bm = dim_block_max_[qt.inner_dim]; if (bm.n_logical == 0) continue; if (bm.is_dense()) { - for (uint32_t spb : spbs) { - const uint32_t sb_start = spb * kStride; - accumulate_block_ub_dispatch(block_ub.data() + sb_start, bm.max_scores.data() + sb_start, - static_cast(qt.u8_weight), kStride); - } + dense_block_max_rows.push_back(bm.max_scores.data()); + dense_query_weights.push_back(static_cast(qt.u8_weight)); } else { const uint16_t u16w = static_cast(qt.u8_weight); for (size_t i = 0; i < bm.block_ids.size(); ++i) { @@ -1503,11 +1845,21 @@ class DspIndex : public DspIndexBase { continue; uint32_t prod = u16w * bm.max_scores[i]; uint32_t sum = static_cast(block_ub[sb]) + prod; - block_ub[sb] = static_cast(sum < 65535u ? sum : 65535u); + const uint16_t saturated_sum = static_cast(sum < 65535u ? sum : 65535u); + block_ub[sb] = saturated_sum; + if (saturated_sum > u16_block_threshold) { + spb_candidate_mask[sb / kStride] |= uint64_t{1} << (sb % kStride); + } } } } + if (!dense_block_max_rows.empty()) { + accumulate_dense_block_ubs_dispatch(block_ub.data(), spb_candidate_mask.data(), u16_block_threshold, + dense_block_max_rows.data(), dense_query_weights.data(), + dense_block_max_rows.size(), spbs.data(), spbs.size(), kStride); + } + for (uint32_t spb : spbs) spb_in_batch[spb] = 0; }; @@ -1516,18 +1868,29 @@ class DspIndex : public DspIndexBase { std::vector cands; cands.reserve(spbs.size() * kStride / 4); uint16_t local_max_ub = 0; +#ifdef KNOWHERE_DSP_INSTRUMENTATION + uint64_t local_saturated_ubs = 0; +#endif for (uint32_t spb : spbs) { const uint32_t sb_start = spb * kStride; - if (!scan_block_ub_any_above_dispatch(block_ub.data() + sb_start, u16_block_threshold, kStride)) - continue; - const uint32_t sb_end = std::min(sb_start + kStride, n_subblocks_); - for (uint32_t sb = sb_start; sb < sb_end; ++sb) { - if (block_ub[sb] > u16_block_threshold) { - cands.push_back(sb); - local_max_ub = std::max(local_max_ub, block_ub[sb]); - } + uint64_t candidate_mask = spb_candidate_mask[spb]; + while (candidate_mask != 0) { + const uint32_t lane = static_cast(__builtin_ctzll(candidate_mask)); + candidate_mask &= candidate_mask - 1; + const uint32_t sb = sb_start + lane; + if (sb >= n_subblocks_) + continue; +#ifdef KNOWHERE_DSP_INSTRUMENTATION + local_saturated_ubs += block_ub[sb] == std::numeric_limits::max(); +#endif + cands.push_back(sb); + local_max_ub = std::max(local_max_ub, block_ub[sb]); } } +#ifdef KNOWHERE_DSP_INSTRUMENTATION + g_dsp_stats.candidate_blocks += cands.size(); + g_dsp_stats.saturated_ubs += local_saturated_ubs; +#endif if (cands.empty()) return {}; const uint32_t rng = local_max_ub - u16_block_threshold; @@ -1549,16 +1912,39 @@ class DspIndex : public DspIndexBase { auto score_blocks = [&](const std::vector& sorted_blocks) -> bool { bool bootstrap_completed = false; +#ifdef KNOWHERE_DSP_INSTRUMENTATION + uint64_t local_entries = 0; + uint64_t local_blocks = 0; + uint64_t local_docs = 0; +#endif for (size_t ci = 0; ci < sorted_blocks.size(); ++ci) { const uint32_t sb_id = sorted_blocks[ci]; if (block_ub[sb_id] <= u16_block_threshold) break; + const uint32_t doc_base = sb_id * kSubblockSize; + const uint32_t doc_end = std::min(doc_base + kSubblockSize, static_cast(n_rows_internal_)); + if (has_filter) { + bool all_docs_filtered = true; + for (uint32_t doc_id = doc_base; doc_id < doc_end; ++doc_id) { + if (!filter.test(doc_id)) { + all_docs_filtered = false; + break; + } + } + if (all_docs_filtered) + continue; + } + const uint32_t block_term_start = fwd_block_term_offsets_[sb_id]; const uint32_t block_term_end = fwd_block_term_offsets_[sb_id + 1]; if (block_term_start == block_term_end) continue; +#ifdef KNOWHERE_DSP_INSTRUMENTATION + local_blocks++; +#endif + if (ci + 1 < sorted_blocks.size()) { const uint32_t next_sb = sorted_blocks[ci + 1]; const uint32_t next_start = fwd_block_term_offsets_[next_sb]; @@ -1571,32 +1957,68 @@ class DspIndex : public DspIndexBase { std::memset(scores, 0, sizeof(scores)); size_t qi = 0; uint32_t bi = block_term_start; - while (qi < n_query_terms && bi < block_term_end) { - const uint32_t q_dim = query[qi].inner_dim; - const uint32_t b_dim = fwd_term_ids_[bi]; - if (q_dim < b_dim) { - ++qi; - } else if (q_dim > b_dim) { - ++bi; - } else { - const float q_weight = query[qi].weight; - const uint32_t e_start = fwd_term_entry_offsets_[bi]; - const uint32_t e_end = fwd_term_entry_offsets_[bi + 1]; - for (uint32_t j = e_start; j < e_end; ++j) { - scores[fwd_doc_offsets_[j]] += q_weight * fwd_scores_[j]; + auto accumulate_match = [&](size_t query_idx, uint32_t block_term_idx) { + const float q_weight = query[query_idx].weight; + const uint32_t e_start = fwd_term_entry_offsets_[block_term_idx]; + const uint32_t e_end = fwd_term_entry_offsets_[block_term_idx + 1]; +#ifdef KNOWHERE_DSP_INSTRUMENTATION + local_entries += e_end - e_start; +#endif + for (uint32_t j = e_start; j < e_end; ++j) { + scores[fwd_doc_offsets_[j]] += q_weight * fwd_scores_[j]; + } + }; + uint32_t match_positions[kHybridMergeMaxQueryTerms]; + const bool used_hybrid = find_terms_hybrid_dispatch( + fwd_term_ids_.data() + block_term_start, block_term_end - block_term_start, hybrid_query_dims, + static_cast(n_query_terms), match_positions); + if (used_hybrid) { + for (qi = 0; qi < n_query_terms; ++qi) { + if (match_positions[qi] != std::numeric_limits::max()) { + accumulate_match(qi, block_term_start + match_positions[qi]); + } + } + } else { + while (qi < n_query_terms && bi < block_term_end) { + const uint32_t q_dim = query[qi].inner_dim; + const uint32_t b_dim = fwd_term_ids_[bi]; + if (q_dim < b_dim) { + ++qi; + } else if (q_dim > b_dim) { + if (metric_type_ == SparseMetricType::METRIC_BM25) { + // BM25 queries are short while a block's term list is comparatively long. + // Gallop to the first term that can match q_dim instead of advancing one + // term at a time. Long-query IP/SPLADE searches retain the linear merge. + const uint32_t base = bi; + uint32_t step = 1; + while (base + step < block_term_end) { + if (fwd_term_ids_[base + step] >= q_dim) + break; + step <<= 1; + } + uint32_t lo = base + (step >> 1) + 1; + uint32_t hi = std::min(block_term_end, uint64_t(base) + step + 1); + bi = static_cast( + std::lower_bound(fwd_term_ids_.begin() + lo, fwd_term_ids_.begin() + hi, q_dim) - + fwd_term_ids_.begin()); + } else { + ++bi; + } + } else { + accumulate_match(qi, bi); + ++qi; + ++bi; } - ++qi; - ++bi; } } - - const uint32_t doc_base = sb_id * kSubblockSize; - const uint32_t doc_end = std::min(doc_base + kSubblockSize, static_cast(n_rows_internal_)); for (uint32_t i = 0; i < doc_end - doc_base; ++i) { if (scores[i] > float_threshold) { const uint32_t doc_id = doc_base + i; if (has_filter && filter.test(doc_id)) continue; +#ifdef KNOWHERE_DSP_INSTRUMENTATION + local_docs++; +#endif heap.Push(scores[i], doc_id); if (heap.Full()) { float new_thresh = heap.Results().front().first; @@ -1616,6 +2038,11 @@ class DspIndex : public DspIndexBase { } } } +#ifdef KNOWHERE_DSP_INSTRUMENTATION + g_dsp_stats.blocks_processed += local_blocks; + g_dsp_stats.entries_scored += local_entries; + g_dsp_stats.docs_pushed += local_docs; +#endif return bootstrap_completed; }; diff --git a/src/index/sparse/sparse_index_node.cc b/src/index/sparse/sparse_index_node.cc index dd01f73d3..54a29e301 100644 --- a/src/index/sparse/sparse_index_node.cc +++ b/src/index/sparse/sparse_index_node.cc @@ -14,6 +14,7 @@ #include #include #include +#include #include #include @@ -44,6 +45,20 @@ namespace knowhere { +template +bool +HasInvalidDspValues(const sparse::SparseRow* rows, size_t row_count) { + for (size_t row = 0; row < row_count; ++row) { + for (size_t entry = 0; entry < rows[row].size(); ++entry) { + const T value = rows[row][entry].val; + if (!std::isfinite(value) || value < T{0}) { + return true; + } + } + } + return false; +} + // Peek at the encoding_type stored in serialized index data without fully parsing. // Returns nullopt if the data is too small or the first section is not POSTING_LISTS. // @@ -1083,14 +1098,19 @@ class SparseDspIndexNode : public IndexNode { LOG_KNOWHERE_ERROR_ << "Could not add data to empty " << Type(); return Status::empty_index; } + const auto* rows = static_cast*>(dataset->GetTensor()); + if (HasInvalidDspValues(rows, dataset->GetRows())) { + LOG_KNOWHERE_ERROR_ << Type() << " requires finite, non-negative sparse values"; + return Status::invalid_args; + } + if (use_knowhere_build_pool) { + // DSP partitions its own forward-index build on the global pool. Do not occupy a worker while waiting for + // nested work from that same pool. + return index_->Add(rows, dataset->GetRows(), dataset->GetDim()); + } auto build_pool_wrapper = std::make_shared(build_pool_, use_knowhere_build_pool); auto tryObj = - build_pool_wrapper - ->push([&] { - return index_->Add(static_cast*>(dataset->GetTensor()), - dataset->GetRows(), dataset->GetDim()); - }) - .getTry(); + build_pool_wrapper->push([&] { return index_->Add(rows, dataset->GetRows(), dataset->GetDim()); }).getTry(); if (!tryObj.hasValue()) { LOG_KNOWHERE_WARNING_ << "failed to add data to index " << Type() << ": " << tryObj.exception().what(); return Status::sparse_inner_error; @@ -1115,6 +1135,10 @@ class SparseDspIndexNode : public IndexNode { auto queries = static_cast*>(dataset->GetTensor()); auto nq = dataset->GetRows(); + if (HasInvalidDspValues(queries, nq)) { + return expected::Err(Status::invalid_args, + "DSP requires finite, non-negative sparse query values"); + } auto k = cfg.k.value(); auto p_id = std::make_unique(nq * k); auto p_dist = std::make_unique(nq * k); @@ -1370,7 +1394,7 @@ class SparseDspIndexNodeCC : public SparseDspIndexNode { auto res = SparseDspIndexNode::Add(dataset, config, use_knowhere_build_pool); auto cfg = static_cast(*config); - if (IsMetricType(cfg.metric_type.value(), metric::IP)) { + if (res == Status::success && IsMetricType(cfg.metric_type.value(), metric::IP)) { auto data = static_cast*>(dataset->GetTensor()); auto rows = dataset->GetRows(); raw_data_.insert(raw_data_.end(), data, data + rows); diff --git a/src/simd/sparse_simd.h b/src/simd/sparse_simd.h index ea2d80b53..e25f62586 100644 --- a/src/simd/sparse_simd.h +++ b/src/simd/sparse_simd.h @@ -1,8 +1,11 @@ #ifndef KNOWHERE_SIMD_SPARSE_SIMD_H #define KNOWHERE_SIMD_SPARSE_SIMD_H +#include +#include #include #include +#include #include #include "knowhere/sparse_utils.h" @@ -40,12 +43,47 @@ accumulate_block_ub_avx512_generic(uint16_t* __restrict ub, const uint8_t* __res void accumulate_block_ub_avx512(uint16_t* ub, const uint8_t* block_max, uint16_t query_weight, uint32_t n); +// Accumulate all dense query-term rows one superblock at a time, keeping the 64 u16 accumulators resident across +// terms. block_max_rows point to full per-term arrays indexed by subblock ID. +void +accumulate_dense_block_ubs_avx512(uint16_t* block_ub, uint64_t* spb_candidate_mask, uint16_t threshold, + const uint8_t* const* block_max_rows, const uint16_t* query_weights, uint32_t n_terms, + const uint32_t* superblock_ids, uint32_t n_superblocks, uint32_t stride); + +// Intersect a short sorted query with one sorted block term list. The kernel gallops over +// 16-term chunk maxima and uses one vector equality comparison in the selected chunk. +uint32_t +find_terms_hybrid_avx512(const uint32_t* terms, uint32_t count, const uint32_t* query_dims, uint32_t query_count, + uint32_t* positions); + // ---- AVX512: Posting list IP accumulation ---- void accumulate_posting_list_ip_avx512(const uint32_t* doc_ids, const float* doc_vals, size_t list_size, float q_weight, float* scores); #endif +inline constexpr uint32_t kHybridMergeMaxQueryTerms = 16; +inline constexpr uint32_t kHybridMergeMinBlockTerms = 64; + +inline bool +find_terms_hybrid_dispatch(const uint32_t* terms, uint32_t count, const uint32_t* query_dims, uint32_t query_count, + uint32_t* positions, uint32_t* probes = nullptr) { +#if defined(__x86_64__) || defined(_M_X64) + if (query_count <= kHybridMergeMaxQueryTerms && count >= kHybridMergeMinBlockTerms && + faiss::cppcontrib::knowhere::InstructionSet::GetInstance().AVX512BW()) { + const uint32_t local_probes = find_terms_hybrid_avx512(terms, count, query_dims, query_count, positions); + if (probes != nullptr) { + *probes = local_probes; + } + return true; + } +#endif + if (probes != nullptr) { + *probes = 0; + } + return false; +} + // Scalar fallback for SIMD block UB scan: check if any of n u16 values > threshold inline bool scan_block_ub_any_above_scalar(const uint16_t* block_ub, uint16_t threshold, uint32_t n) { @@ -107,6 +145,48 @@ accumulate_block_ub_dispatch(uint16_t* __restrict ub, const uint8_t* __restrict accumulate_block_ub_scalar(ub, block_max, query_weight, n); } +inline void +accumulate_dense_block_ubs_scalar(uint16_t* block_ub, uint64_t* spb_candidate_mask, uint16_t threshold, + const uint8_t* const* block_max_rows, const uint16_t* query_weights, uint32_t n_terms, + const uint32_t* superblock_ids, uint32_t n_superblocks, uint32_t stride) { + assert(stride == 64 && "accumulate_dense_block_ubs_scalar expects 64 subblocks per superblock"); + for (uint32_t spb_index = 0; spb_index < n_superblocks; ++spb_index) { + const uint32_t offset = superblock_ids[spb_index] * stride; + uint16_t accumulators[64]; + std::copy_n(block_ub + offset, stride, accumulators); + for (uint32_t term = 0; term < n_terms; ++term) { + const uint8_t* block_max = block_max_rows[term] + offset; + const uint32_t query_weight = query_weights[term]; + for (uint32_t lane = 0; lane < stride; ++lane) { + const uint32_t sum = static_cast(accumulators[lane]) + query_weight * block_max[lane]; + accumulators[lane] = static_cast(sum < 65535u ? sum : 65535u); + } + } + uint64_t candidate_mask = 0; + for (uint32_t lane = 0; lane < stride; ++lane) { + candidate_mask |= static_cast(accumulators[lane] > threshold) << lane; + } + spb_candidate_mask[superblock_ids[spb_index]] = candidate_mask; + std::copy_n(accumulators, stride, block_ub + offset); + } +} + +inline void +accumulate_dense_block_ubs_dispatch(uint16_t* block_ub, uint64_t* spb_candidate_mask, uint16_t threshold, + const uint8_t* const* block_max_rows, const uint16_t* query_weights, + uint32_t n_terms, const uint32_t* superblock_ids, uint32_t n_superblocks, + uint32_t stride) { +#if defined(__x86_64__) || defined(_M_X64) + if (faiss::cppcontrib::knowhere::InstructionSet::GetInstance().AVX512BW()) { + accumulate_dense_block_ubs_avx512(block_ub, spb_candidate_mask, threshold, block_max_rows, query_weights, + n_terms, superblock_ids, n_superblocks, stride); + return; + } +#endif + accumulate_dense_block_ubs_scalar(block_ub, spb_candidate_mask, threshold, block_max_rows, query_weights, n_terms, + superblock_ids, n_superblocks, stride); +} + template inline void accumulate_posting_list_contribution_ip_dispatch(const uint32_t* doc_ids, const QType* doc_vals, size_t list_size, diff --git a/src/simd/sparse_simd_avx512.cc b/src/simd/sparse_simd_avx512.cc index c80ea3a92..afdc3d4a7 100644 --- a/src/simd/sparse_simd_avx512.cc +++ b/src/simd/sparse_simd_avx512.cc @@ -14,12 +14,74 @@ #include +#include #include +#include #include "sparse_simd.h" namespace knowhere::sparse { +uint32_t +find_terms_hybrid_avx512(const uint32_t* terms, uint32_t count, const uint32_t* query_dims, uint32_t query_count, + uint32_t* positions) { + const uint32_t full_chunks = count / 16; + uint32_t chunk_cursor = 0; + uint32_t tail_cursor = full_chunks * 16; + uint32_t probes = 0; + for (uint32_t qi = 0; qi < query_count; ++qi) { + const uint32_t target = query_dims[qi]; + positions[qi] = std::numeric_limits::max(); + + if (chunk_cursor < full_chunks) { + ++probes; + if (terms[chunk_cursor * 16 + 15] < target) { + const uint32_t base = chunk_cursor; + uint32_t step = 1; + while (base + step < full_chunks) { + ++probes; + if (terms[(base + step) * 16 + 15] >= target) { + break; + } + step <<= 1; + } + uint32_t lo = base + (step >> 1) + 1; + uint32_t hi = std::min(full_chunks, uint64_t(base) + step + 1); + const uint32_t* first = std::lower_bound(terms + lo * 16, terms + hi * 16, target, + [&probes](uint32_t term, uint32_t needle) { + ++probes; + return term < needle; + }); + chunk_cursor = static_cast(first - terms) / 16; + } + } + + if (chunk_cursor < full_chunks) { + ++probes; + const uint32_t base = chunk_cursor * 16; + const __m512i values = _mm512_loadu_si512(reinterpret_cast(terms + base)); + const __m512i needle = _mm512_set1_epi32(static_cast(target)); + const __mmask16 matches = _mm512_cmpeq_epi32_mask(values, needle); + if (matches != 0) { + positions[qi] = base + static_cast(__builtin_ctz(matches)); + } + continue; + } + + while (tail_cursor < count && terms[tail_cursor] < target) { + ++tail_cursor; + ++probes; + } + if (tail_cursor < count) { + ++probes; + if (terms[tail_cursor] == target) { + positions[qi] = tail_cursor; + } + } + } + return probes; +} + // ============================================================================ // AVX512 BW: Block UB Threshold Scan — Stride-Specific Specializations // ============================================================================ @@ -198,4 +260,31 @@ accumulate_block_ub_avx512(uint16_t* ub, const uint8_t* block_max, uint16_t quer accumulate_block_ub_avx512_generic(ub, block_max, query_weight, n); } +void +accumulate_dense_block_ubs_avx512(uint16_t* block_ub, uint64_t* spb_candidate_mask, uint16_t threshold, + const uint8_t* const* block_max_rows, const uint16_t* query_weights, uint32_t n_terms, + const uint32_t* superblock_ids, uint32_t n_superblocks, uint32_t stride) { + assert(stride == 64 && "accumulate_dense_block_ubs_avx512 expects 64 subblocks per superblock"); + const __m512i threshold_vec = _mm512_set1_epi16(static_cast(threshold)); + for (uint32_t spb_index = 0; spb_index < n_superblocks; ++spb_index) { + const uint32_t offset = superblock_ids[spb_index] * stride; + __m512i accum0 = _mm512_loadu_si512(reinterpret_cast(block_ub + offset)); + __m512i accum1 = _mm512_loadu_si512(reinterpret_cast(block_ub + offset + 32)); + for (uint32_t term = 0; term < n_terms; ++term) { + const uint8_t* block_max = block_max_rows[term] + offset; + const __m512i query_weight = _mm512_set1_epi16(static_cast(query_weights[term])); + const __m256i max0 = _mm256_loadu_si256(reinterpret_cast(block_max)); + const __m256i max1 = _mm256_loadu_si256(reinterpret_cast(block_max + 32)); + accum0 = _mm512_adds_epu16(accum0, _mm512_mullo_epi16(_mm512_cvtepu8_epi16(max0), query_weight)); + accum1 = _mm512_adds_epu16(accum1, _mm512_mullo_epi16(_mm512_cvtepu8_epi16(max1), query_weight)); + } + _mm512_storeu_si512(reinterpret_cast<__m512i*>(block_ub + offset), accum0); + _mm512_storeu_si512(reinterpret_cast<__m512i*>(block_ub + offset + 32), accum1); + const __mmask32 above0 = _mm512_cmp_epu16_mask(accum0, threshold_vec, _MM_CMPINT_GT); + const __mmask32 above1 = _mm512_cmp_epu16_mask(accum1, threshold_vec, _MM_CMPINT_GT); + spb_candidate_mask[superblock_ids[spb_index]] = + static_cast(above0) | (static_cast(above1) << 32); + } +} + } // namespace knowhere::sparse diff --git a/tests/ut/test_sparse.cc b/tests/ut/test_sparse.cc index 22e36e826..de3e8b189 100644 --- a/tests/ut/test_sparse.cc +++ b/tests/ut/test_sparse.cc @@ -9,12 +9,19 @@ // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express // or implied. See the License for the specific language governing permissions and limitations under the License. +#include +#include +#include #include +#include +#include #include #include #include +#include #include +#include "catch2/catch_approx.hpp" #include "catch2/catch_test_macros.hpp" #include "catch2/generators/catch_generators.hpp" #include "index/sparse/inverted_index_format.h" @@ -439,8 +446,6 @@ TEST_CASE("Test Mem Sparse Index With Float Vector", "[float metrics]") { auto [name, gen] = GENERATE_REF(table>({ make_tuple(knowhere::IndexEnum::INDEX_SPARSE_INVERTED_INDEX, sparse_inverted_index_gen), make_tuple(knowhere::IndexEnum::INDEX_SPARSE_WAND, sparse_inverted_index_gen), - make_tuple(knowhere::IndexEnum::INDEX_SPARSE_DSP, sparse_dsp_gen), - make_tuple(knowhere::IndexEnum::INDEX_SPARSE_DSP_CC, sparse_dsp_gen), })); auto idx = knowhere::IndexFactory::Instance().Create(name, version).value(); auto cfg_json = gen().dump(); @@ -485,8 +490,6 @@ TEST_CASE("Test Mem Sparse Index With Float Vector", "[float metrics]") { auto [name, gen] = GENERATE_REF(table>({ make_tuple(knowhere::IndexEnum::INDEX_SPARSE_INVERTED_INDEX, sparse_inverted_index_gen), make_tuple(knowhere::IndexEnum::INDEX_SPARSE_WAND, sparse_inverted_index_gen), - make_tuple(knowhere::IndexEnum::INDEX_SPARSE_DSP, sparse_dsp_gen), - make_tuple(knowhere::IndexEnum::INDEX_SPARSE_DSP_CC, sparse_dsp_gen), })); auto idx = knowhere::IndexFactory::Instance().Create(name, version).value(); @@ -678,6 +681,552 @@ TEST_CASE("Test Mem Sparse Index Handle Empty Vector", "[float metrics]") { } } +TEST_CASE("Test DSP Sparse Index Large K Is Rank Safe", "[float metrics][sparse][dsp]") { + constexpr int64_t nb = 10001; + constexpr int64_t topk = 10001; + constexpr int32_t dim = 1; + + std::vector> base_data(nb); + for (int64_t i = 0; i < nb - 2; ++i) { + base_data[i][0] = 255.0f; + } + base_data[nb - 2][0] = 200.0f; + base_data[nb - 1][0] = 1.0f; + const auto train_ds = GenSparseDataSet(base_data, dim); + const auto query_ds = GenSparseDataSet(std::vector>{{{0, 1.0f}}}, dim); + + knowhere::Json json = { + {knowhere::meta::DIM, dim}, + {knowhere::meta::METRIC_TYPE, knowhere::metric::IP}, + {knowhere::meta::TOPK, topk}, + {knowhere::indexparam::DROP_RATIO_SEARCH, 0.0f}, + {"dsp_mu", 1.0f}, + {"dsp_eta", 1.0f}, + }; + + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_SPARSE_DSP_CC, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + REQUIRE(index.Build(train_ds, json) == knowhere::Status::success); + + auto expected = knowhere::BruteForce::SearchSparse(train_ds, query_ds, json, nullptr); + REQUIRE(expected.has_value()); + auto actual = index.Search(query_ds, json, nullptr); + REQUIRE(actual.has_value()); + REQUIRE(GetKNNRecall(*expected.value(), *actual.value()) == 1.0f); +} + +TEST_CASE("Test DSP BM25 Kth Init Includes Tied Maximum Scores", "[float metrics][sparse][dsp]") { + constexpr int64_t nb = 2000; + constexpr int64_t topk = 10; + constexpr int32_t dim = 1; + + std::vector> base_data(nb, {{{0, 1.0f}}}); + const auto train_ds = GenSparseDataSet(base_data, dim); + const auto query_ds = GenSparseDataSet(std::vector>{{{0, 1.0f}}}, dim); + + knowhere::Json json = { + {knowhere::meta::DIM, dim}, + {knowhere::meta::METRIC_TYPE, knowhere::metric::BM25}, + {knowhere::meta::TOPK, topk}, + {knowhere::meta::BM25_K1, 1.2f}, + {knowhere::meta::BM25_B, 0.75f}, + {knowhere::meta::BM25_AVGDL, 1.0f}, + {knowhere::indexparam::DROP_RATIO_SEARCH, 0.0f}, + {"dsp_mu", 1.0f}, + {"dsp_eta", 1.0f}, + }; + + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_SPARSE_DSP_CC, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + REQUIRE(index.Build(train_ds, json) == knowhere::Status::success); + + auto actual = index.Search(query_ds, json, nullptr); + REQUIRE(actual.has_value()); + REQUIRE(actual.value()->GetDim() == topk); + const auto* ids = actual.value()->GetIds(); + for (int64_t i = 0; i < topk; ++i) { + REQUIRE(ids[i] != -1); + } +} + +TEST_CASE("Test DSP Kth Init Ignores Partially Filled Heaps", "[float metrics][sparse][dsp]") { + auto [topk, head_count, nb] = GENERATE(table({ + {100, 50, 101}, + {1000, 100, 1001}, + })); + constexpr int32_t dim = 2; + constexpr int64_t tail_id = 0; + + std::vector> base_data(nb); + base_data[tail_id][1] = 0.01f; + for (int64_t i = 1; i <= head_count; ++i) { + base_data[i][0] = 1.0f; + } + const auto train_ds = GenSparseDataSet(base_data, dim); + const auto query_ds = GenSparseDataSet(std::vector>{{{0, 1.0f}, {1, 1.0f}}}, dim); + + knowhere::Json json = { + {knowhere::meta::DIM, dim}, + {knowhere::meta::METRIC_TYPE, knowhere::metric::IP}, + {knowhere::meta::TOPK, topk}, + {knowhere::indexparam::DROP_RATIO_SEARCH, 0.0f}, + {"dsp_mu", 1.0f}, + {"dsp_eta", 1.0f}, + }; + + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_SPARSE_DSP_CC, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + REQUIRE(index.Build(train_ds, json) == knowhere::Status::success); + + auto actual = index.Search(query_ds, json, nullptr); + REQUIRE(actual.has_value()); + REQUIRE(actual.value()->GetDim() == topk); + const auto* ids = actual.value()->GetIds(); + REQUIRE(std::find(ids, ids + topk, tail_id) != ids + topk); +} + +TEST_CASE("Test DSP Kth Init Is Disabled With Bitset Filter", "[float metrics][sparse][dsp]") { + constexpr int64_t nb = 20; + constexpr int64_t topk = 10; + constexpr int32_t dim = 1; + + // The corpus-wide 10th score is 100, but all ten documents supporting that threshold are filtered. The filtered + // top-10 consists entirely of the remaining score-1 documents, so using the unfiltered kth initializer would + // incorrectly prune every valid result. IDs 0..7 also form a fully filtered DSP block, covering its fast skip. + std::vector> base_data(nb); + for (int64_t i = 0; i < topk; ++i) { + base_data[i][0] = 100.0f; + } + for (int64_t i = topk; i < nb; ++i) { + base_data[i][0] = 1.0f; + } + const auto train_ds = GenSparseDataSet(base_data, dim); + const auto query_ds = GenSparseDataSet(std::vector>{{{0, 1.0f}}}, dim); + + knowhere::Json json = { + {knowhere::meta::DIM, dim}, + {knowhere::meta::METRIC_TYPE, knowhere::metric::IP}, + {knowhere::meta::TOPK, topk}, + {knowhere::indexparam::DROP_RATIO_SEARCH, 0.0f}, + {"dsp_mu", 1.0f}, + {"dsp_eta", 1.0f}, + }; + + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_SPARSE_DSP_CC, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + REQUIRE(index.Build(train_ds, json) == knowhere::Status::success); + + const auto bitset_data = GenerateBitsetWithFirstTbitsSet(nb, topk); + const knowhere::BitsetView bitset(bitset_data.data(), nb); + auto expected = knowhere::BruteForce::SearchSparse(train_ds, query_ds, json, bitset); + REQUIRE(expected.has_value()); + auto actual = index.Search(query_ds, json, bitset); + REQUIRE(actual.has_value()); + REQUIRE(GetKNNRecall(*expected.value(), *actual.value()) == 1.0f); + for (int64_t rank = 0; rank < topk; ++rank) { + REQUIRE(actual.value()->GetIds()[rank] >= topk); + REQUIRE(actual.value()->GetDistance()[rank] == 1.0f); + } +} + +TEST_CASE("Test DSP Native Serialization Round Trip", "[float metrics][sparse][dsp]") { + constexpr int64_t nb = 2000; + constexpr int64_t nq = 10; + constexpr int64_t topk = 100; + constexpr int32_t dim = 300; + const auto metric = GENERATE(knowhere::metric::IP, knowhere::metric::BM25); + const bool use_mmap = GENERATE(false, true); + const auto train_ds = GenSparseDataSet(nb, dim, 0.95f); + const auto query_ds = GenSparseDataSet(nq, dim, 0.97f); + + knowhere::Json json = { + {knowhere::meta::DIM, dim}, + {knowhere::meta::METRIC_TYPE, metric}, + {knowhere::meta::TOPK, topk}, + {knowhere::meta::BM25_K1, 1.2f}, + {knowhere::meta::BM25_B, 0.75f}, + {knowhere::meta::BM25_AVGDL, 100.0f}, + {knowhere::indexparam::DROP_RATIO_SEARCH, 0.0f}, + {"dsp_mu", 1.0f}, + {"dsp_eta", 1.0f}, + }; + + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_SPARSE_DSP_CC, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + REQUIRE(index.Build(train_ds, json) == knowhere::Status::success); + auto before = index.Search(query_ds, json, nullptr); + REQUIRE(before.has_value()); + + knowhere::BinarySet binary_set; + REQUIRE(index.Serialize(binary_set) == knowhere::Status::success); + if (use_mmap) { + const std::string filename = "/tmp/knowhere_dsp_native_serialization_test"; + WriteBinaryToFile(filename, binary_set.GetByName(index.Type())); + REQUIRE(index.DeserializeFromFile(filename, json) == knowhere::Status::success); + REQUIRE(std::remove(filename.c_str()) == 0); + } else { + REQUIRE(index.Deserialize(binary_set, json) == knowhere::Status::success); + } + + auto after = index.Search(query_ds, json, nullptr); + REQUIRE(after.has_value()); + REQUIRE(std::memcmp(after.value()->GetIds(), before.value()->GetIds(), nq * topk * sizeof(int64_t)) == 0); + REQUIRE(std::memcmp(after.value()->GetDistance(), before.value()->GetDistance(), nq * topk * sizeof(float)) == 0); +} + +TEST_CASE("Test DSP Loads Legacy Sparse Serialization", "[float metrics][sparse][dsp]") { + constexpr int64_t nb = 2048; + constexpr int64_t nq = 10; + constexpr int64_t topk = 100; + constexpr int32_t dim = 32; + const bool use_mmap = GENERATE(false, true); + + // A v1 DSP file without a DSP_METADATA section is the legacy format supported by the rebuild fallback. Keep the + // corpus deterministic and make the first row contain every dimension so that raw and inner dimension IDs match. + std::vector> base_data(nb); + for (int32_t d = 0; d < dim; ++d) { + base_data[0][d] = 1.0f + static_cast(d % 7) * 0.1f; + } + for (int64_t doc = 1; doc < nb; ++doc) { + const int32_t d0 = static_cast(doc % dim); + const int32_t d1 = static_cast((doc * 7 + 3) % dim); + base_data[doc][d0] = 0.5f + static_cast(doc % 11) * 0.03f; + base_data[doc][d1] = 0.7f + static_cast(doc % 13) * 0.02f; + } + std::vector> query_data(nq); + for (int64_t query = 0; query < nq; ++query) { + query_data[query][static_cast(query % dim)] = 1.0f; + query_data[query][static_cast((query * 5 + 1) % dim)] = 0.8f; + } + const auto train_ds = GenSparseDataSet(base_data, dim); + const auto query_ds = GenSparseDataSet(query_data, dim); + + knowhere::Json json = { + {knowhere::meta::DIM, dim}, + {knowhere::meta::METRIC_TYPE, knowhere::metric::IP}, + {knowhere::meta::TOPK, topk}, + {knowhere::indexparam::DROP_RATIO_SEARCH, 0.0f}, + {"dsp_mu", 1.0f}, + {"dsp_eta", 1.0f}, + }; + + auto fresh_dsp = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_SPARSE_DSP_CC, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + REQUIRE(fresh_dsp.Build(train_ds, json) == knowhere::Status::success); + auto expected = fresh_dsp.Search(query_ds, json, nullptr); + REQUIRE(expected.has_value()); + + std::vector> posting_ids(dim); + std::vector> posting_vals(dim); + std::vector max_scores(dim, 0.0f); + for (uint32_t doc = 0; doc < nb; ++doc) { + for (const auto& [raw_dim, value] : base_data[doc]) { + posting_ids[raw_dim].push_back(doc); + posting_vals[raw_dim].push_back(value); + max_scores[raw_dim] = std::max(max_scores[raw_dim], value); + } + } + + auto append_bytes = [](std::vector& output, const void* data, size_t size) { + const auto* first = static_cast(data); + output.insert(output.end(), first, first + size); + }; + auto append_value = [&](std::vector& output, const auto& value) { + append_bytes(output, &value, sizeof(value)); + }; + + std::vector posting_section; + const uint32_t encoding_type = 0; + append_value(posting_section, encoding_type); + std::vector posting_offsets(dim + 1, 0); + for (int32_t d = 0; d < dim; ++d) { + posting_offsets[d + 1] = posting_offsets[d] + posting_ids[d].size(); + } + append_bytes(posting_section, posting_offsets.data(), posting_offsets.size() * sizeof(uint64_t)); + for (int32_t d = 0; d < dim; ++d) { + append_bytes(posting_section, posting_ids[d].data(), posting_ids[d].size() * sizeof(uint32_t)); + } + for (int32_t d = 0; d < dim; ++d) { + append_bytes(posting_section, posting_vals[d].data(), posting_vals[d].size() * sizeof(float)); + } + + std::vector dim_map(dim); + std::iota(dim_map.begin(), dim_map.end(), 0); + struct LegacySectionHeader { + uint32_t type; + uint32_t padding = 0; + uint64_t offset; + uint64_t size; + }; + static_assert(sizeof(LegacySectionHeader) == 24); + constexpr uint32_t kPostingListsSection = 0; + constexpr uint32_t kDimMapSection = 2; + constexpr uint32_t kMaxScoresSection = 4; + constexpr uint32_t kHeaderSize = 32; + constexpr uint32_t kSectionCount = 3; + uint64_t next_offset = kHeaderSize + sizeof(uint32_t) + kSectionCount * sizeof(LegacySectionHeader); + std::array section_headers = { + LegacySectionHeader{kPostingListsSection, 0, next_offset, posting_section.size()}, + LegacySectionHeader{kDimMapSection, 0, next_offset + posting_section.size(), dim_map.size() * sizeof(uint32_t)}, + LegacySectionHeader{kMaxScoresSection, 0, + next_offset + posting_section.size() + dim_map.size() * sizeof(uint32_t), + max_scores.size() * sizeof(float)}, + }; + + std::vector legacy_blob; + const uint32_t format_version = 1; + const uint32_t row_count = nb; + const uint32_t max_dim = dim; + const uint32_t inner_dim_count = dim; + append_value(legacy_blob, format_version); + append_value(legacy_blob, row_count); + append_value(legacy_blob, max_dim); + append_value(legacy_blob, inner_dim_count); + const std::array reserved{}; + append_bytes(legacy_blob, reserved.data(), reserved.size()); + append_value(legacy_blob, kSectionCount); + append_bytes(legacy_blob, section_headers.data(), section_headers.size() * sizeof(LegacySectionHeader)); + append_bytes(legacy_blob, posting_section.data(), posting_section.size()); + append_bytes(legacy_blob, dim_map.data(), dim_map.size() * sizeof(uint32_t)); + append_bytes(legacy_blob, max_scores.data(), max_scores.size() * sizeof(float)); + + auto legacy_data = std::shared_ptr(new uint8_t[legacy_blob.size()]); + std::memcpy(legacy_data.get(), legacy_blob.data(), legacy_blob.size()); + knowhere::BinarySet legacy_binary; + legacy_binary.Append(knowhere::IndexEnum::INDEX_SPARSE_DSP_CC, legacy_data, legacy_blob.size()); + + auto loaded_dsp = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_SPARSE_DSP_CC, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + if (use_mmap) { + const std::string filename = "/tmp/knowhere_dsp_legacy_serialization_test"; + WriteBinaryToFile(filename, legacy_binary.GetByName(knowhere::IndexEnum::INDEX_SPARSE_DSP_CC)); + REQUIRE(loaded_dsp.DeserializeFromFile(filename, json) == knowhere::Status::success); + REQUIRE(std::remove(filename.c_str()) == 0); + } else { + REQUIRE(loaded_dsp.Deserialize(legacy_binary, json) == knowhere::Status::success); + } + auto actual = loaded_dsp.Search(query_ds, json, nullptr); + REQUIRE(actual.has_value()); + REQUIRE(std::memcmp(actual.value()->GetIds(), expected.value()->GetIds(), nq * topk * sizeof(int64_t)) == 0); + REQUIRE(std::memcmp(actual.value()->GetDistance(), expected.value()->GetDistance(), nq * topk * sizeof(float)) == + 0); +} + +TEST_CASE("Test DSP Parallel Build Is Byte Identical", "[float metrics][sparse][dsp]") { + constexpr int64_t nb = 65536; + constexpr int32_t dim = 300; + const auto metric = GENERATE(knowhere::metric::IP, knowhere::metric::BM25); + const auto train_ds = GenSparseDataSet(nb, dim, 0.99f); + knowhere::Json json = { + {knowhere::meta::DIM, dim}, {knowhere::meta::METRIC_TYPE, metric}, {knowhere::meta::TOPK, 100}, + {knowhere::meta::BM25_K1, 1.2f}, {knowhere::meta::BM25_B, 0.75f}, {knowhere::meta::BM25_AVGDL, 100.0f}, + }; + + struct BuildPoolSizeGuard { + size_t original = knowhere::KnowhereConfig::GetBuildThreadPoolSize(); + ~BuildPoolSizeGuard() { + // Zero means the global pool had not been initialized yet; zero is not a valid size to restore. + if (original != 0) { + knowhere::KnowhereConfig::SetBuildThreadPoolSize(original); + } + } + } pool_size_guard; + + knowhere::KnowhereConfig::SetBuildThreadPoolSize(1); + auto serial_index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_SPARSE_DSP_CC, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + REQUIRE(serial_index.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet serial_binary; + REQUIRE(serial_index.Serialize(serial_binary) == knowhere::Status::success); + + knowhere::KnowhereConfig::SetBuildThreadPoolSize(8); + auto parallel_index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_SPARSE_DSP_CC, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + REQUIRE(parallel_index.Build(train_ds, json) == knowhere::Status::success); + knowhere::BinarySet parallel_binary; + REQUIRE(parallel_index.Serialize(parallel_binary) == knowhere::Status::success); + + const auto serial_blob = serial_binary.GetByName(serial_index.Type()); + const auto parallel_blob = parallel_binary.GetByName(parallel_index.Type()); + REQUIRE(serial_blob->size == parallel_blob->size); + REQUIRE(std::memcmp(serial_blob->data.get(), parallel_blob->data.get(), serial_blob->size) == 0); +} + +TEST_CASE("Test DSP Concurrent Search Reuses Workspaces", "[float metrics][sparse][dsp][concurrent]") { + constexpr int64_t nb = 4096; + constexpr int64_t nq = 4; + constexpr int64_t topk = 100; + constexpr int32_t dim = 300; + constexpr int32_t num_threads = 8; + constexpr int32_t repetitions = 50; + const auto train_ds = GenSparseDataSet(nb, dim, 0.95f); + const auto query_ds = GenSparseDataSet(nq, dim, 0.97f); + knowhere::Json json = { + {knowhere::meta::DIM, dim}, + {knowhere::meta::METRIC_TYPE, knowhere::metric::IP}, + {knowhere::meta::TOPK, topk}, + {knowhere::indexparam::DROP_RATIO_SEARCH, 0.0f}, + {"dsp_mu", 1.0f}, + {"dsp_eta", 1.0f}, + }; + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_SPARSE_DSP_CC, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + REQUIRE(index.Build(train_ds, json) == knowhere::Status::success); + auto expected = index.Search(query_ds, json, nullptr); + REQUIRE(expected.has_value()); + + std::atomic all_equal{true}; + std::vector> futures; + futures.reserve(num_threads); + for (int32_t thread = 0; thread < num_threads; ++thread) { + futures.emplace_back(std::async(std::launch::async, [&]() { + for (int32_t repetition = 0; repetition < repetitions; ++repetition) { + auto actual = index.Search(query_ds, json, nullptr); + if (!actual.has_value() || + std::memcmp(actual.value()->GetIds(), expected.value()->GetIds(), nq * topk * sizeof(int64_t)) != + 0 || + std::memcmp(actual.value()->GetDistance(), expected.value()->GetDistance(), + nq * topk * sizeof(float)) != 0) { + all_equal = false; + return; + } + } + })); + } + for (auto& future : futures) { + future.get(); + } + REQUIRE(all_equal.load()); +} + +TEST_CASE("Test DSP Rejects Invalid Sparse Values", "[float metrics][sparse][dsp]") { + const auto index_type = GENERATE(knowhere::IndexEnum::INDEX_SPARSE_DSP, knowhere::IndexEnum::INDEX_SPARSE_DSP_CC); + const float invalid_value = + GENERATE(-1.0f, std::numeric_limits::quiet_NaN(), std::numeric_limits::infinity()); + constexpr int32_t dim = 4; + knowhere::Json json = { + {knowhere::meta::DIM, dim}, + {knowhere::meta::METRIC_TYPE, knowhere::metric::IP}, + {knowhere::meta::TOPK, 1}, + {knowhere::indexparam::DROP_RATIO_SEARCH, 0.0f}, + }; + + std::vector> invalid_base = {{{0, 1.0f}}, {{1, invalid_value}}}; + const auto invalid_train_ds = GenSparseDataSet(invalid_base, dim); + auto invalid_index = + knowhere::IndexFactory::Instance() + .Create(index_type, knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + REQUIRE(invalid_index.Build(invalid_train_ds, json) == knowhere::Status::invalid_args); + + std::vector> valid_base = {{{0, 1.0f}}, {{1, 2.0f}}}; + const auto valid_train_ds = GenSparseDataSet(valid_base, dim); + auto valid_index = + knowhere::IndexFactory::Instance() + .Create(index_type, knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + REQUIRE(valid_index.Build(valid_train_ds, json) == knowhere::Status::success); + + std::vector> invalid_query = {{{0, 1.0f}, {1, invalid_value}}}; + const auto invalid_query_ds = GenSparseDataSet(invalid_query, dim); + auto result = valid_index.Search(invalid_query_ds, json, nullptr); + REQUIRE_FALSE(result.has_value()); + REQUIRE(result.error() == knowhere::Status::invalid_args); +} + +TEST_CASE("Test DSP Safe Mode Matches Brute Force", "[float metrics][sparse][dsp]") { + using Catch::Approx; + constexpr int64_t nb = 2000; + constexpr int64_t nq = 10; + constexpr int32_t dim = 300; + const int64_t topk = GENERATE(10, 100, 1000); + INFO("topk=" << topk); + const auto train_ds = GenSparseDataSet(nb, dim, 0.95f); + const auto query_ds = GenSparseDataSet(nq, dim, 0.97f); + + knowhere::Json json = { + {knowhere::meta::DIM, dim}, + {knowhere::meta::METRIC_TYPE, knowhere::metric::IP}, + {knowhere::meta::TOPK, topk}, + {knowhere::indexparam::DROP_RATIO_SEARCH, 0.0f}, + {"dsp_mu", 1.0f}, + {"dsp_eta", 1.0f}, + }; + + auto index = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_SPARSE_DSP_CC, + knowhere::Version::GetCurrentVersion().VersionNumber()) + .value(); + REQUIRE(index.Build(train_ds, json) == knowhere::Status::success); + + auto expected = knowhere::BruteForce::SearchSparse(train_ds, query_ds, json, nullptr); + REQUIRE(expected.has_value()); + auto actual = index.Search(query_ds, json, nullptr); + REQUIRE(actual.has_value()); + REQUIRE(expected.value()->GetDim() == topk); + REQUIRE(actual.value()->GetDim() == topk); + + const auto* expected_ids = expected.value()->GetIds(); + const auto* expected_scores = expected.value()->GetDistance(); + const auto* actual_ids = actual.value()->GetIds(); + const auto* actual_scores = actual.value()->GetDistance(); + for (int64_t query = 0; query < nq; ++query) { + const int64_t offset = query * topk; + int64_t positive_count = 0; + while (positive_count < topk && expected_scores[offset + positive_count] > 0.0f) { + ++positive_count; + } + CAPTURE(query, positive_count); + + // Brute force may fill the tail with arbitrary zero-score IDs, while DSP (like the other DAAT paths) only + // emits positive-score matches. Compare the sorted scores only where a positive match exists. The two paths + // accumulate floats in a different order, hence the same relative tolerance used by the brute-force tests; + // tied IDs may legitimately appear in a different order. + for (int64_t rank = 0; rank < positive_count; ++rank) { + REQUIRE(actual_scores[offset + rank] == Approx(expected_scores[offset + rank]).epsilon(0.00001)); + } + + // IDs above the kth-score tie boundary are unique members of the exact top-k result. IDs at the boundary may + // be exchanged with other equal-score documents, so comparing them would make this assertion tie-sensitive. + // Use the same tolerance as the score comparison above when identifying that boundary: DSP and brute force + // accumulate in different orders, and a tied score can otherwise land a few ulps above kth_score. + const float kth_score = expected_scores[offset + topk - 1]; + const auto is_kth_tie = [kth_score](float score) { return score == Approx(kth_score).epsilon(0.00001); }; + std::unordered_set expected_strict_ids; + std::unordered_set actual_strict_ids; + for (int64_t rank = 0; rank < topk; ++rank) { + if (expected_scores[offset + rank] > kth_score && !is_kth_tie(expected_scores[offset + rank])) { + expected_strict_ids.insert(expected_ids[offset + rank]); + } + if (actual_ids[offset + rank] >= 0 && actual_scores[offset + rank] > kth_score && + !is_kth_tie(actual_scores[offset + rank])) { + actual_strict_ids.insert(actual_ids[offset + rank]); + } + } + REQUIRE(actual_strict_ids == expected_strict_ids); + } +} + TEST_CASE("Test Mem Sparse Index CC", "[float metrics]") { std::atomic value_base(0); // each time a new batch of vectors are generated, the base value is increased by 1. diff --git a/tests/ut/test_sparse_simd.cc b/tests/ut/test_sparse_simd.cc index 102edc536..a596cd64e 100644 --- a/tests/ut/test_sparse_simd.cc +++ b/tests/ut/test_sparse_simd.cc @@ -9,6 +9,8 @@ // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express // or implied. See the License for the specific language governing permissions and limitations under the License. +#include +#include #include #include #include @@ -54,6 +56,64 @@ accumulate_posting_list_ip_scalar_ref(const uint32_t* doc_ids, const float* doc_ } } +static std::vector +membership_reference(const std::vector& terms, const std::vector& query) { + std::vector positions(query.size(), std::numeric_limits::max()); + for (size_t i = 0; i < query.size(); ++i) { + const auto it = std::lower_bound(terms.begin(), terms.end(), query[i]); + if (it != terms.end() && *it == query[i]) { + positions[i] = static_cast(it - terms.begin()); + } + } + return positions; +} + +TEST_CASE("DSP hybrid SIMD membership matches scalar", "[sparse simd avx512][dsp]") { +#if defined(__x86_64__) || defined(_M_X64) + if (!faiss::cppcontrib::knowhere::InstructionSet::GetInstance().AVX512BW()) { + SKIP("AVX512BW not available on this CPU"); + } + + auto check = [](const std::vector& terms, const std::vector& query) { + const auto expected = membership_reference(terms, query); + std::vector actual(query.size()); + find_terms_hybrid_avx512(terms.data(), static_cast(terms.size()), query.data(), + static_cast(query.size()), actual.data()); + REQUIRE(actual == expected); + }; + + SECTION("edge cases") { + check({}, {}); + check({}, {1, 2, 3}); + + std::vector terms(80); + std::iota(terms.begin(), terms.end(), 100); + check(terms, terms); // all match + check(terms, {164, 165, 178, 179}); // all hits live in the final partial chunk + + terms[15] = 115; + terms[16] = 115; // duplicate spanning a chunk boundary + check(terms, {114, 115, 115, 116}); + } + + SECTION("randomized sorted intersections") { + std::mt19937 rng(20260720); + for (int iteration = 0; iteration < 500; ++iteration) { + const uint32_t term_count = 64 + rng() % 512; + const uint32_t query_count = rng() % 17; + std::set term_set; + while (term_set.size() < term_count) term_set.insert(rng() % 100000); + std::set query_set; + while (query_set.size() < query_count) query_set.insert(rng() % 100000); + check(std::vector(term_set.begin(), term_set.end()), + std::vector(query_set.begin(), query_set.end())); + } + } +#else + SKIP("Test only runs on x86_64 platforms"); +#endif +} + TEST_CASE("Test Sparse SIMD AVX512 - Basic Correctness", "[sparse simd avx512]") { #if defined(__x86_64__) || defined(_M_X64) if (!faiss::cppcontrib::knowhere::InstructionSet::GetInstance().AVX512F()) { @@ -279,6 +339,64 @@ TEST_CASE("Test Sparse SIMD AVX512 - Special Values", "[sparse simd avx512]") { #endif } +TEST_CASE("Test DSP Superblock-Major UB Accumulation", "[sparse simd avx512][dsp]") { +#if defined(__x86_64__) || defined(_M_X64) + if (!faiss::cppcontrib::knowhere::InstructionSet::GetInstance().AVX512BW()) { + SKIP("AVX512BW not available on this CPU"); + } + + constexpr uint32_t stride = 64; + constexpr uint32_t n_superblocks = 3; + constexpr uint32_t n_terms = 7; + const std::vector surviving_superblocks = {0, 2}; + + std::mt19937 generator(24680); + std::uniform_int_distribution max_distribution(0, 255); + std::uniform_int_distribution weight_distribution(1, 255); + std::vector> rows(n_terms, std::vector(stride * n_superblocks)); + std::vector row_pointers; + std::vector weights(n_terms); + for (uint32_t term = 0; term < n_terms; ++term) { + for (auto& value : rows[term]) { + value = static_cast(max_distribution(generator)); + } + row_pointers.push_back(rows[term].data()); + weights[term] = static_cast(weight_distribution(generator)); + } + + for (const uint16_t threshold : {uint16_t{0}, uint16_t{1234}, uint16_t{30000}, uint16_t{65534}, uint16_t{65535}}) { + CAPTURE(threshold); + std::vector expected(stride * n_superblocks, 1234); + std::vector actual = expected; + std::vector expected_masks(n_superblocks, 0xdeadbeef); + std::vector actual_masks = expected_masks; + accumulate_dense_block_ubs_scalar(expected.data(), expected_masks.data(), threshold, row_pointers.data(), + weights.data(), n_terms, surviving_superblocks.data(), + surviving_superblocks.size(), stride); + accumulate_dense_block_ubs_avx512(actual.data(), actual_masks.data(), threshold, row_pointers.data(), + weights.data(), n_terms, surviving_superblocks.data(), + surviving_superblocks.size(), stride); + + REQUIRE(actual == expected); + REQUIRE(actual_masks == expected_masks); + for (uint32_t spb : surviving_superblocks) { + const auto stripe_begin = expected.begin() + spb * stride; + uint64_t expected_mask = 0; + for (uint32_t lane = 0; lane < stride; ++lane) { + expected_mask |= static_cast(stripe_begin[lane] > threshold) << lane; + } + REQUIRE(expected_masks[spb] == expected_mask); + } + REQUIRE(actual_masks[1] == 0xdeadbeef); + for (uint32_t lane = stride; lane < 2 * stride; ++lane) { + REQUIRE(actual[lane] == 1234); + } + } +#else + SKIP("Test only runs on x86_64 platforms"); +#endif +} + TEST_CASE("Test Sparse SIMD AVX512 - Multiple Accumulations", "[sparse simd avx512]") { #if defined(__x86_64__) || defined(_M_X64) if (!faiss::cppcontrib::knowhere::InstructionSet::GetInstance().AVX512F()) {