diff --git a/Makefile b/Makefile index 66701ac5c..b187b7c79 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_BASE_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 1 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 + CONAN_FLAGS += -o with_diskann=True endif # GPU builds use cuVS. diff --git a/benchmark/CMakeLists.txt b/benchmark/CMakeLists.txt index 2c70642d6..474e52004 100644 --- a/benchmark/CMakeLists.txt +++ b/benchmark/CMakeLists.txt @@ -70,6 +70,14 @@ endif() benchmark_test(gen_hdf5_file hdf5/gen_hdf5_file.cpp) benchmark_test(gen_fbin_file hdf5/gen_fbin_file.cpp) +# Sparse DSP benchmark (standalone, no HDF5/GTest required) +add_executable(benchmark_sparse_dsp benchmark_sparse_dsp.cpp) +target_link_libraries(benchmark_sparse_dsp knowhere) +if(NOT APPLE) + target_link_libraries(benchmark_sparse_dsp atomic) +endif() +install(TARGETS benchmark_sparse_dsp DESTINATION unittest) + # Sparse SIMD benchmark (x86_64 only, standalone, no HDF5 required) # Only build on x86_64/AMD64, skip on ARM/aarch64/arm64 if(CMAKE_SYSTEM_PROCESSOR MATCHES "^(x86_64|AMD64|amd64|X86_64)$") diff --git a/benchmark/benchmark_sparse_dsp.cpp b/benchmark/benchmark_sparse_dsp.cpp new file mode 100644 index 000000000..46e13fa96 --- /dev/null +++ b/benchmark/benchmark_sparse_dsp.cpp @@ -0,0 +1,714 @@ +// Copyright (C) 2019-2023 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// 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. + +// +// Sparse DSP benchmark: measures QPS, latency percentiles, recall, and result +// coverage for SPARSE_DSP_CC across a parameter sweep of (mode, mu, eta, gamma). +// +// Supports SPLADE/IP and MSMARCO/BM25 datasets in CSR binary format. +// +// Usage: +// ./benchmark_sparse_dsp --data-dir ~/data/splade_full --metric IP +// --gt ~/data/splade_full/base_small.dev.ip.gt +// ./benchmark_sparse_dsp --data-dir ~/data/msmarco_full_bm25_v2 --metric BM25 +// --gt ~/data/msmarco_full_bm25_v2/base_small.dev.bm25.gt +// + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include "knowhere/bitsetview.h" +#include "knowhere/comp/index_param.h" +#include "knowhere/comp/knowhere_config.h" +#include "knowhere/index/index_factory.h" +#include "knowhere/operands.h" +#include "knowhere/sparse_utils.h" + +// ============================================================================ +// CSR binary file loader +// ============================================================================ +struct CSRDataset { + std::vector indptr; + std::vector indices; + std::vector data; + int64_t n_rows = 0; + int64_t n_cols = 0; + int64_t nnz = 0; + + bool + load(const std::string& path) { + std::ifstream file(path, std::ios::binary); + if (!file) { + fprintf(stderr, "Error: Cannot open %s\n", path.c_str()); + return false; + } + file.read(reinterpret_cast(&n_rows), sizeof(int64_t)); + file.read(reinterpret_cast(&n_cols), sizeof(int64_t)); + file.read(reinterpret_cast(&nnz), sizeof(int64_t)); + printf(" Loading CSR: %ld rows, %ld cols, %ld nnz\n", n_rows, n_cols, nnz); + + indptr.resize(n_rows + 1); + file.read(reinterpret_cast(indptr.data()), (n_rows + 1) * sizeof(int64_t)); + indices.resize(nnz); + file.read(reinterpret_cast(indices.data()), nnz * sizeof(int32_t)); + data.resize(nnz); + file.read(reinterpret_cast(data.data()), nnz * sizeof(float)); + return file.good(); + } + + std::unique_ptr[]> + to_sparse_rows() const { + auto rows = std::make_unique[]>(n_rows); + for (int64_t i = 0; i < n_rows; ++i) { + int64_t start = indptr[i]; + int64_t end = indptr[i + 1]; + int64_t len = end - start; + rows[i] = knowhere::sparse::SparseRow(len); + for (int64_t j = 0; j < len; ++j) { + rows[i].set_at(j, indices[start + j], data[start + j]); + } + } + return rows; + } + + // Compute avgdl as sum of all values / n_rows. + // For raw term-frequency CSR, this is the average document length. + // For impact-weighted CSR, this is the average sum of impact scores. + double compute_avgdl() const { + double total = 0.0; + for (int64_t i = 0; i < n_rows; ++i) { + for (int64_t j = indptr[i]; j < indptr[i + 1]; ++j) { + total += data[j]; + } + } + return total / n_rows; + } + + void + free_raw() { + indptr.clear(); + indptr.shrink_to_fit(); + indices.clear(); + indices.shrink_to_fit(); + data.clear(); + data.shrink_to_fit(); + } +}; + +// ============================================================================ +// Ground truth loader (binary: int32 nq, int32 k, then nq*k int32 IDs) +// ============================================================================ +struct GroundTruth { + std::vector> ids; + int64_t nq = 0; + int64_t k = 0; + + bool + load(const std::string& path, int64_t max_nq) { + std::ifstream file(path, std::ios::binary); + if (!file) { + fprintf(stderr, "Error: Cannot open GT file %s\n", path.c_str()); + return false; + } + int32_t nq32, k32; + file.read(reinterpret_cast(&nq32), sizeof(int32_t)); + file.read(reinterpret_cast(&k32), sizeof(int32_t)); + nq = std::min(static_cast(nq32), max_nq); + k = k32; + printf(" Loading GT: %ld queries, k=%ld (file has %d queries)\n", nq, k, nq32); + ids.resize(nq); + for (int64_t i = 0; i < nq; ++i) { + ids[i].resize(k); + file.read(reinterpret_cast(ids[i].data()), k * sizeof(int32_t)); + } + return true; + } + + // Compute recall for a single query. Returns 0 if query has no results. + float + recall(const int64_t* result, int64_t qi, int64_t result_k) const { + if (qi >= nq) + return 0.0f; + int64_t check_k = std::min(result_k, k); + int matches = 0; + for (int64_t i = 0; i < check_k; ++i) { + if (result[i] == -1) + continue; + for (int64_t j = 0; j < check_k; ++j) { + if (result[i] == ids[qi][j]) { + ++matches; + break; + } + } + } + return static_cast(matches) / check_k; + } +}; + +// ============================================================================ +// Latency percentile helper +// ============================================================================ +struct LatencyStats { + double mean_ms; + double p50_ms; + double p95_ms; + double p99_ms; + double max_ms; + + static LatencyStats + compute(std::vector& latencies_ms) { + std::sort(latencies_ms.begin(), latencies_ms.end()); + size_t n = latencies_ms.size(); + double sum = 0; + for (double v : latencies_ms) sum += v; + return { + .mean_ms = sum / n, + .p50_ms = latencies_ms[n / 2], + .p95_ms = latencies_ms[static_cast(n * 0.95)], + .p99_ms = latencies_ms[static_cast(n * 0.99)], + .max_ms = latencies_ms[n - 1], + }; + } +}; + +// ============================================================================ +// Bitset generation for filtered benchmarks +// ============================================================================ + +// Generate a bitset where each bit is set with probability filter_rate. +// In knowhere, a set bit means the doc is FILTERED OUT (excluded from results). +std::vector +generate_random_bitset(int64_t n_docs, float filter_rate, uint64_t seed) { + size_t n_bytes = (n_docs + 7) / 8; + std::vector bitset(n_bytes, 0); + std::mt19937_64 rng(seed); + std::uniform_real_distribution dist(0.0f, 1.0f); + int64_t n_set = 0; + for (int64_t i = 0; i < n_docs; ++i) { + if (dist(rng) < filter_rate) { + bitset[i / 8] |= (1u << (i % 8)); + ++n_set; + } + } + printf(" Generated random bitset: %ld/%ld docs filtered (%.1f%%)\n", n_set, n_docs, 100.0 * n_set / n_docs); + return bitset; +} + +// Generate a bitset that masks the docs with highest nnz (most non-zero dimensions). +// These dense docs appear in many posting lists and are likely to score well across +// diverse queries, making this an adversarial filter for pruning-based indexes. +std::vector +generate_global_topk_bitset(const knowhere::sparse::SparseRow* base_rows, int64_t n_docs, float filter_rate) { + std::vector doc_nnz(n_docs, 0); + for (int64_t i = 0; i < n_docs; ++i) { + doc_nnz[i] = static_cast(base_rows[i].size()); + } + // Sort indices by nnz descending — mask the densest docs first + std::vector order(n_docs); + std::iota(order.begin(), order.end(), 0); + std::sort(order.begin(), order.end(), [&](int64_t a, int64_t b) { return doc_nnz[a] > doc_nnz[b]; }); + + int64_t n_mask = static_cast(filter_rate * n_docs); + size_t n_bytes = (n_docs + 7) / 8; + std::vector bitset(n_bytes, 0); + for (int64_t i = 0; i < n_mask; ++i) { + int64_t doc = order[i]; + bitset[doc / 8] |= (1u << (doc % 8)); + } + printf(" Generated global-dense-nnz bitset: %ld/%ld docs filtered (%.1f%%)\n", n_mask, n_docs, + 100.0 * n_mask / n_docs); + return bitset; +} + +// Compute filtered ground truth using safe (exact) search. +GroundTruth +compute_filtered_gt(const knowhere::Index& brute_force_index, + const knowhere::sparse::SparseRow* query_rows, int64_t nq, int64_t n_cols, int64_t topk, + const knowhere::Json& search_conf, const knowhere::BitsetView& bitset) { + GroundTruth gt; + gt.nq = nq; + gt.k = topk; + gt.ids.resize(nq); + + auto query_ds = knowhere::GenDataSet(1, n_cols, nullptr); + query_ds->SetIsSparse(true); + + for (int64_t q = 0; q < nq; ++q) { + gt.ids[q].resize(topk, -1); + query_ds->SetTensor(&query_rows[q]); + auto result = brute_force_index.Search(query_ds, search_conf, bitset); + if (result.has_value()) { + auto ids = result.value()->GetIds(); + for (int64_t i = 0; i < topk; ++i) { + gt.ids[q][i] = static_cast(ids[i]); + } + } + } + return gt; +} + +// ============================================================================ +// Benchmark result with coverage metrics +// ============================================================================ +struct BenchResult { + std::vector latencies_ms; + std::vector result_ids; + double total_ms; + // Recall: averaged over ALL nq queries (failed queries contribute 0) + float avg_recall; + // Coverage: how many queries returned at least 1 result + int64_t n_queries; + int64_t n_failed; // queries where first result is -1 + float avg_filled; // average number of non-(-1) results per query out of topk + std::vector failed_indices; // query indices that returned zero results +}; + +BenchResult +run_search(const knowhere::Index& index, const knowhere::sparse::SparseRow* query_rows, + int64_t nq, int64_t n_cols, int64_t topk, const knowhere::Json& search_conf, const GroundTruth& gt, + const knowhere::BitsetView& bitset = knowhere::BitsetView()) { + BenchResult res; + res.latencies_ms.resize(nq); + res.result_ids.resize(nq * topk, -1); + res.n_queries = nq; + + auto query_ds = knowhere::GenDataSet(1, n_cols, nullptr); + query_ds->SetIsSparse(true); + + auto t_total_start = std::chrono::high_resolution_clock::now(); + for (int64_t q = 0; q < nq; ++q) { + query_ds->SetTensor(&query_rows[q]); + auto t0 = std::chrono::high_resolution_clock::now(); + auto result = index.Search(query_ds, search_conf, bitset); + auto t1 = std::chrono::high_resolution_clock::now(); + res.latencies_ms[q] = std::chrono::duration(t1 - t0).count(); + if (result.has_value()) { + memcpy(&res.result_ids[q * topk], result.value()->GetIds(), topk * sizeof(int64_t)); + } + } + auto t_total_end = std::chrono::high_resolution_clock::now(); + res.total_ms = std::chrono::duration(t_total_end - t_total_start).count(); + + // Compute recall over ALL queries (failed queries get recall=0). + // Also compute coverage metrics. + int64_t eval_nq = std::min(nq, gt.nq); + float recall_sum = 0; + int64_t n_failed = 0; + int64_t total_filled = 0; + + for (int64_t q = 0; q < nq; ++q) { + // Count filled (non -1) slots + int64_t filled = 0; + for (int64_t i = 0; i < topk; ++i) { + if (res.result_ids[q * topk + i] != -1) + ++filled; + } + total_filled += filled; + + if (filled == 0) { + ++n_failed; + res.failed_indices.push_back(q); + } + + // Recall: every query in [0, eval_nq) contributes, even if failed (=0 recall) + if (q < eval_nq) { + recall_sum += gt.recall(&res.result_ids[q * topk], q, topk); + } + } + + res.avg_recall = (eval_nq > 0) ? recall_sum / eval_nq : 0.0f; + res.n_failed = n_failed; + res.avg_filled = static_cast(total_filled) / nq; + return res; +} + +// ============================================================================ +// Main +// ============================================================================ +void +print_usage(const char* prog) { + printf( + "Usage: %s --data-dir --metric --gt [options]\n" + "\n" + "Required:\n" + " --data-dir Directory containing base and query CSR files\n" + " --metric Metric type\n" + " --gt Ground truth file (binary: int32 nq, int32 k, nq*k int32 IDs)\n" + "\n" + "Options:\n" + " --topk Top-k results (default: 10)\n" + " --nq Number of queries, 0=all (default: 0)\n" + " --warmup Warmup runs before timed run (default: 1)\n" + " --base Base vectors file (default: base_small.csr)\n" + " --query Query vectors file (default: queries.dev.csr)\n" + " --bm25-k1 BM25 k1 (default: 1.2)\n" + " --bm25-b BM25 b (default: 0.75)\n" + " --avgdl Override avgdl (default: computed from base data)\n" + " --default-only Only run default DSP config (no sweep)\n" + "\n" + "Filter options:\n" + " --filter-rate Corpus-level filter rate 0.0-1.0 (default: 0.0 = no filter)\n" + " --filter-mode Filter mode (default: random):\n" + " random uniform random docs\n" + " global-topk docs with highest nnz (most dimensions)\n" + " --filter-seed Random seed for filter generation (default: 42)\n" + "\n", + prog); +} + +int +main(int argc, char** argv) { + std::string data_dir; + std::string metric; + std::string gt_path; + std::string base_file = "base_small.csr"; + std::string query_file = "queries.dev.csr"; + int64_t topk = 10; + int64_t nq = 0; + int warmup = 1; + float bm25_k1 = 1.2f; + float bm25_b = 0.75f; + float avgdl_override = -1.0f; + bool default_only = false; + float filter_rate = 0.0f; + std::string filter_mode = "random"; + uint64_t filter_seed = 42; + + for (int i = 1; i < argc; ++i) { + if (strcmp(argv[i], "--data-dir") == 0 && i + 1 < argc) + data_dir = argv[++i]; + else if (strcmp(argv[i], "--metric") == 0 && i + 1 < argc) + metric = argv[++i]; + else if (strcmp(argv[i], "--gt") == 0 && i + 1 < argc) + gt_path = argv[++i]; + else if (strcmp(argv[i], "--topk") == 0 && i + 1 < argc) + topk = atoi(argv[++i]); + else if (strcmp(argv[i], "--nq") == 0 && i + 1 < argc) + nq = atoi(argv[++i]); + else if (strcmp(argv[i], "--warmup") == 0 && i + 1 < argc) + warmup = atoi(argv[++i]); + else if (strcmp(argv[i], "--base") == 0 && i + 1 < argc) + base_file = argv[++i]; + else if (strcmp(argv[i], "--query") == 0 && i + 1 < argc) + query_file = argv[++i]; + else if (strcmp(argv[i], "--bm25-k1") == 0 && i + 1 < argc) + bm25_k1 = atof(argv[++i]); + else if (strcmp(argv[i], "--bm25-b") == 0 && i + 1 < argc) + bm25_b = atof(argv[++i]); + else if (strcmp(argv[i], "--avgdl") == 0 && i + 1 < argc) + avgdl_override = atof(argv[++i]); + else if (strcmp(argv[i], "--default-only") == 0) + default_only = true; + else if (strcmp(argv[i], "--filter-rate") == 0 && i + 1 < argc) + filter_rate = atof(argv[++i]); + else if (strcmp(argv[i], "--filter-mode") == 0 && i + 1 < argc) + filter_mode = argv[++i]; + else if (strcmp(argv[i], "--filter-seed") == 0 && i + 1 < argc) + filter_seed = strtoull(argv[++i], nullptr, 10); + else if (strcmp(argv[i], "-h") == 0 || strcmp(argv[i], "--help") == 0) { + print_usage(argv[0]); + return 0; + } + } + + if (data_dir.empty() || metric.empty() || gt_path.empty()) { + print_usage(argv[0]); + return 1; + } + + bool is_bm25 = (metric == "BM25" || metric == "bm25"); + + printf("==========================================================\n"); + printf(" Sparse DSP Benchmark\n"); + printf("==========================================================\n\n"); + + knowhere::KnowhereConfig::SetSimdType(knowhere::KnowhereConfig::SimdType::AUTO); + + // ---- Load data ---- + printf("[Loading Data]\n"); + CSRDataset base, queries; + if (!base.load(data_dir + "/" + base_file)) + return 1; + if (!queries.load(data_dir + "/" + query_file)) + return 1; + + if (nq == 0 || nq > queries.n_rows) + nq = queries.n_rows; + + auto base_rows = base.to_sparse_rows(); + auto query_rows = queries.to_sparse_rows(); + + double avgdl = 0.0; + if (is_bm25) { + if (avgdl_override > 0) { + avgdl = avgdl_override; + printf(" avgdl: %.2f (user-provided override)\n", avgdl); + } else { + avgdl = base.compute_avgdl(); + printf(" avgdl: %.2f (computed from base data)\n", avgdl); + } + } + base.free_raw(); + + // ---- Load ground truth ---- + printf("[Loading Ground Truth]\n"); + printf(" GT path: %s\n", gt_path.c_str()); + GroundTruth gt; + if (!gt.load(gt_path, queries.n_rows)) { + fprintf(stderr, "Error: failed to load ground truth from %s\n", gt_path.c_str()); + return 1; + } + + printf("\n[Config]\n"); + printf(" base=%ld nq=%ld k=%ld metric=%s warmup=%d\n", base.n_rows, nq, topk, metric.c_str(), warmup); + printf(" base_file=%s query_file=%s\n", base_file.c_str(), query_file.c_str()); + printf(" gt=%s (nq=%ld, k=%ld)\n", gt_path.c_str(), gt.nq, gt.k); + if (is_bm25) + printf(" bm25: k1=%.2f b=%.2f avgdl=%.2f\n", bm25_k1, bm25_b, avgdl); + if (filter_rate > 0) + printf(" filter: rate=%.2f mode=%s seed=%lu\n", filter_rate, filter_mode.c_str(), filter_seed); + printf("\n"); + + // ---- Helper: populate BM25 params into JSON ---- + auto make_bm25_json = [&](knowhere::Json& json) { + if (is_bm25) { + json["bm25_k1"] = bm25_k1; + json["bm25_b"] = bm25_b; + json["bm25_avgdl"] = static_cast(avgdl); + } + }; + + // ---- Helper: print result row ---- + // Columns: config, params, recall, QPS, failed, avg_filled/k, latency percentiles + auto print_header = [&]() { + printf(" %-22s %-26s %-7s %-8s %-6s %-8s %-50s\n", "Config", "Params", "Recall", "QPS", "Fail", "Fill/k", + "Latency(ms): mean / p50 / p95 / p99 / max"); + printf(" %s\n", std::string(170, '-').c_str()); + }; + + auto print_row = [&](const char* label, float mu, float eta, int gamma, const BenchResult& res, + const LatencyStats& lat) { + char params_buf[64]; + snprintf(params_buf, sizeof(params_buf), "mu=%.2f eta=%.2f g=%-5d", mu, eta, gamma); + printf( + " %-22s %-26s %.4f %-8.1f %-6ld %.1f/%-3ld " + "%.2f / %.2f / %.2f / %.2f / %.2f\n", + label, params_buf, res.avg_recall, res.n_queries * 1000.0 / res.total_ms, res.n_failed, res.avg_filled, + topk, lat.mean_ms, lat.p50_ms, lat.p95_ms, lat.p99_ms, lat.max_ms); + }; + + // ---- Helper: diagnose failed queries ---- + auto print_failed_diag = [&](const char* label, const BenchResult& res) { + if (res.failed_indices.empty()) + return; + printf("\n [%s] %ld failed queries (zero results):\n", label, res.n_failed); + size_t show = std::min(res.failed_indices.size(), static_cast(20)); + for (size_t i = 0; i < show; ++i) { + int64_t qi = res.failed_indices[i]; + int64_t nnz = query_rows[qi].size(); + printf(" q[%ld]: nnz=%ld", qi, nnz); + // Show first few dims for context + if (nnz > 0) { + printf(" dims=["); + for (int64_t j = 0; j < std::min(nnz, static_cast(5)); ++j) { + if (j > 0) + printf(","); + printf("%d:%.2f", query_rows[qi][j].id, query_rows[qi][j].val); + } + if (nnz > 5) + printf(",..."); + printf("]"); + } + printf("\n"); + } + if (res.failed_indices.size() > show) { + printf(" ... and %ld more\n", res.failed_indices.size() - show); + } + printf("\n"); + }; + + // ---- Build base dataset ---- + auto ds = knowhere::GenDataSet(base.n_rows, base.n_cols, nullptr); + ds->SetIsSparse(true); + ds->SetTensor(base_rows.get()); + + // DSP safe-search config for filtered ground truth + // (mu=1, eta=1, mode=0, gamma=0, no kth-init gives exact results) + knowhere::Json safe_search_json; + if (filter_rate > 0) { + safe_search_json["metric_type"] = metric; + safe_search_json["topk"] = topk; + safe_search_json["drop_ratio_search"] = 0.0f; + safe_search_json["dsp_mode"] = 0; + safe_search_json["dsp_mu"] = 1.0f; + safe_search_json["dsp_eta"] = 1.0f; + safe_search_json["dsp_gamma"] = 0; + safe_search_json["dsp_kth_init"] = false; + make_bm25_json(safe_search_json); + } + + // ============================================================ + // Filter setup: generate bitset (filtered GT computed after DSP build) + // ============================================================ + std::vector filter_bitset_data; + knowhere::BitsetView filter_bitset; + + if (filter_rate > 0) { + printf("\n[Generating Filter]\n"); + if (filter_mode == "random") { + filter_bitset_data = generate_random_bitset(base.n_rows, filter_rate, filter_seed); + } else if (filter_mode == "global-topk") { + filter_bitset_data = generate_global_topk_bitset(base_rows.get(), base.n_rows, filter_rate); + } else { + fprintf(stderr, "Error: unknown filter mode '%s' (supported: random, global-topk)\n", filter_mode.c_str()); + return 1; + } + filter_bitset = knowhere::BitsetView(filter_bitset_data.data(), base.n_rows); + } + + // ============================================================ + // DSP: Build once, sweep params + // ============================================================ + printf("[DSP Index]\n"); + auto dsp_or = knowhere::IndexFactory::Instance().Create( + knowhere::IndexEnum::INDEX_SPARSE_DSP_CC, knowhere::Version::GetCurrentVersion().VersionNumber()); + if (!dsp_or.has_value()) { + fprintf(stderr, "Error: failed to create DSP index\n"); + return 1; + } + auto dsp = dsp_or.value(); + + { + knowhere::Json build_json; + build_json["metric_type"] = metric; + make_bm25_json(build_json); + + auto t0 = std::chrono::high_resolution_clock::now(); + auto st = dsp.Build(ds, build_json); + auto t1 = std::chrono::high_resolution_clock::now(); + if (st != knowhere::Status::success) { + fprintf(stderr, "Error: DSP build failed\n"); + return 1; + } + printf(" Build: %.1f ms\n", std::chrono::duration(t1 - t0).count()); + } + + // Compute filtered GT using DSP safe search (exact results with mu=1, eta=1) + if (filter_rate > 0) { + printf("[Computing Filtered Ground Truth via DSP safe search]\n"); + gt = compute_filtered_gt(dsp, query_rows.get(), nq, queries.n_cols, topk, safe_search_json, filter_bitset); + printf(" Filtered GT: %ld queries, k=%ld\n", gt.nq, gt.k); + } + + // ---- Parameter sweep ---- + struct ParamSet { + const char* label; + int mode; // 0=dsp, 1=lsp0, 2=lsp1, 3=lsp2 + float mu; + float eta; + int gamma; + bool kth_init = true; + float kth_alpha = 1.0f; + }; + + // clang-format off + std::vector params; + if (default_only) { + params = { + {"dsp default", 0, 1.0f, 1.0f, 0}, + {"dsp a=0.50", 0, 1.0f, 1.0f, 0, true, 0.50f}, + {"dsp a=0.25", 0, 1.0f, 1.0f, 0, true, 0.25f}, + {"dsp no-kth", 0, 1.0f, 1.0f, 0, false}, + }; + } else { + params = { + // DSP mode (mode=0): dual-threshold (mu, eta) + optional top-gamma backstop + {"dsp default", 0, 1.0f, 1.0f, 0}, + {"dsp mu=0.3", 0, 0.3f, 1.0f, 0}, + {"dsp mu=0.5 eta=1.0", 0, 0.5f, 1.0f, 0}, + {"dsp mu=0.5 eta=0.8", 0, 0.5f, 0.8f, 0}, + {"dsp mu=0.5 eta=0.5", 0, 0.5f, 0.5f, 0}, + {"dsp mu=0.3 g=100", 0, 0.3f, 1.0f, 100}, + + // LSP/0 (mode=1): top-gamma from ub>=theta, no mu/asc gate + {"lsp0 g=50", 1, 1.0f, 1.0f, 50}, + {"lsp0 g=100", 1, 1.0f, 1.0f, 100}, + {"lsp0 g=500", 1, 1.0f, 1.0f, 500}, + {"lsp0 g=1000", 1, 1.0f, 1.0f, 1000}, + + // LSP/1 (mode=2): lsp0 safe set + mu gate + {"lsp1 g=100", 2, 1.0f, 1.0f, 100}, + {"lsp1 g=100 mu=0.3", 2, 0.3f, 1.0f, 100}, + {"lsp1 g=100 mu=0.5", 2, 0.5f, 1.0f, 100}, + + // LSP/2 (mode=3): lsp1 + asc gate + {"lsp2 g=100", 3, 1.0f, 1.0f, 100}, + {"lsp2 g=100 mu=0.3", 3, 0.3f, 1.0f, 100}, + {"lsp2 g=100 mu=0.5", 3, 0.5f, 1.0f, 100}, + {"lsp2 g=100 mu=0.5 eta=0.8", 3, 0.5f, 0.8f, 100}, + + // kth-init OFF: isolate hierarchy-only contribution + {"dsp no-kth", 0, 1.0f, 1.0f, 0, false}, + {"lsp0 g=100 no-kth", 1, 1.0f, 1.0f, 100, false}, + {"lsp1 g=100 no-kth", 2, 1.0f, 1.0f, 100, false}, + + // Alpha-clamped kth threshold (DSP mode) + {"dsp a=0.25", 0, 1.0f, 1.0f, 0, true, 0.25f}, + {"dsp a=0.50", 0, 1.0f, 1.0f, 0, true, 0.50f}, + {"dsp a=0.75", 0, 1.0f, 1.0f, 0, true, 0.75f}, + }; + } + // clang-format on + + // ============================================================ + // Parameter sweep + // ============================================================ + { + printf("\n"); + print_header(); + + for (const auto& p : params) { + knowhere::Json search_json; + search_json["metric_type"] = metric; + search_json["topk"] = topk; + search_json["drop_ratio_search"] = 0.0f; + search_json["dsp_mode"] = p.mode; + search_json["dsp_mu"] = p.mu; + search_json["dsp_eta"] = p.eta; + search_json["dsp_gamma"] = p.gamma; + search_json["dsp_kth_init"] = p.kth_init; + search_json["dsp_kth_alpha"] = p.kth_alpha; + make_bm25_json(search_json); + + for (int w = 0; w < warmup; ++w) { + run_search(dsp, query_rows.get(), nq, queries.n_cols, topk, search_json, gt, filter_bitset); + } + + auto res = run_search(dsp, query_rows.get(), nq, queries.n_cols, topk, search_json, gt, filter_bitset); + auto lat = LatencyStats::compute(res.latencies_ms); + print_row(p.label, p.mu, p.eta, p.gamma, res, lat); + if (strstr(p.label, "default") != nullptr || strstr(p.label, "no-kth") != nullptr || res.n_failed > 0) { + print_failed_diag(p.label, res); + } + } + } + + printf("\n=== Done ===\n"); + return 0; +} diff --git a/cmake/libs/libfaiss.cmake b/cmake/libs/libfaiss.cmake index 59fc1f26d..b37d3967c 100644 --- a/cmake/libs/libfaiss.cmake +++ b/cmake/libs/libfaiss.cmake @@ -243,7 +243,7 @@ if(__X86_64) -mavx512bw -mpopcnt -mavx512vl) target_compile_options(utils_avx512icx PRIVATE -mfma -mf16c -mavx512f -mavx512dq -mavx512bw -mpopcnt -mavx512vl -mavx512vpopcntdq) - target_compile_options(sparse_simd_avx512 PRIVATE -mavx512f -mavx512dq) + target_compile_options(sparse_simd_avx512 PRIVATE -mavx512f -mavx512dq -mavx512bw -mavx512vl -mavx512cd) target_include_directories(sparse_simd_avx512 PRIVATE ${Boost_INCLUDE_DIRS}) target_link_libraries(sparse_simd_avx512 PRIVATE milvus-common::milvus-common) diff --git a/include/knowhere/comp/index_param.h b/include/knowhere/comp/index_param.h index 9d45d6cfb..2d30a9cdb 100644 --- a/include/knowhere/comp/index_param.h +++ b/include/knowhere/comp/index_param.h @@ -72,6 +72,8 @@ constexpr const char* INDEX_SVS_IVF_LEANVEC = "SVS_IVF_LEANVEC"; constexpr const char* INDEX_SPARSE_INVERTED_INDEX = "SPARSE_INVERTED_INDEX"; constexpr const char* INDEX_SPARSE_WAND = "SPARSE_WAND"; +constexpr const char* INDEX_SPARSE_DSP = "SPARSE_DSP"; +constexpr const char* INDEX_SPARSE_DSP_CC = "SPARSE_DSP_CC"; constexpr const char* INDEX_SPARSE_INVERTED_INDEX_CC = "SPARSE_INVERTED_INDEX_CC"; constexpr const char* INDEX_SPARSE_WAND_CC = "SPARSE_WAND_CC"; diff --git a/include/knowhere/index/index_table.h b/include/knowhere/index/index_table.h index b5b4863bb..fac5736f3 100644 --- a/include/knowhere/index/index_table.h +++ b/include/knowhere/index/index_table.h @@ -140,6 +140,8 @@ static std::set> legal_knowhere_index = { // sparse index {IndexEnum::INDEX_SPARSE_INVERTED_INDEX, VecType::VECTOR_SPARSE_FLOAT}, {IndexEnum::INDEX_SPARSE_WAND, VecType::VECTOR_SPARSE_FLOAT}, + {IndexEnum::INDEX_SPARSE_DSP, VecType::VECTOR_SPARSE_FLOAT}, + {IndexEnum::INDEX_SPARSE_DSP_CC, VecType::VECTOR_SPARSE_FLOAT}, // minhash index {IndexEnum::INDEX_MINHASH_LSH, VecType::VECTOR_BINARY}, }; @@ -168,6 +170,8 @@ static std::set legal_support_mmap_knowhere_index = { // sparse index IndexEnum::INDEX_SPARSE_INVERTED_INDEX, IndexEnum::INDEX_SPARSE_WAND, + IndexEnum::INDEX_SPARSE_DSP, + IndexEnum::INDEX_SPARSE_DSP_CC, }; static std::set legal_support_emb_list_knowhere_index = { diff --git a/include/knowhere/sparse_utils.h b/include/knowhere/sparse_utils.h index da73584a8..80000169a 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,121 @@ namespace knowhere::sparse { +// Seek distance instrumentation (compile with -DSEEK_INSTRUMENTATION to enable) +#ifdef SEEK_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 = bucket_1_3 = bucket_4_15 = 0; + bucket_16_63 = bucket_64_255 = bucket_256_plus = 0; + seek_hits = seek_misses = 0; + } +}; + +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 = surviving_superblocks = candidate_blocks = 0; + blocks_processed = saturated_ubs = entries_scored = docs_pushed = queries = workspace_pool_misses = 0; + } +}; + +inline DspStats g_dsp_stats; +#endif + enum class SparseMetricType { METRIC_IP = 1, METRIC_BM25 = 2, @@ -253,6 +370,72 @@ class SparseRow { bool own_data_; }; +// When pushing new elements into a MaxMinHeap, only `capacity` elements with the +// largest val are kept. pop()/top() returns the smallest element out of them. +template +class MaxMinHeap { + public: + explicit MaxMinHeap(int capacity) : capacity_(capacity), pool_(capacity) { + } + void + push(table_t id, T val) { + if (size_ < capacity_) { + pool_[size_] = {id, val}; + size_ += 1; + std::push_heap(pool_.begin(), pool_.begin() + size_, std::greater>()); + } else if (val > pool_[0].val) { + sift_down(id, val); + } + } + table_t + pop() { + std::pop_heap(pool_.begin(), pool_.begin() + size_, std::greater>()); + size_ -= 1; + return pool_[size_].id; + } + [[nodiscard]] size_t + size() const { + return size_; + } + [[nodiscard]] bool + empty() const { + return size() == 0; + } + SparseIdVal + top() const { + return pool_[0]; + } + [[nodiscard]] bool + full() const { + return size_ == capacity_; + } + + private: + void + sift_down(table_t id, T val) { + size_t i = 0; + for (; 2 * i + 1 < size_;) { + size_t j = i; + size_t l = 2 * i + 1, r = 2 * i + 2; + if (pool_[l].val < val) { + j = l; + } + if (r < size_ && pool_[r].val < std::min(pool_[l].val, val)) { + j = r; + } + if (i == j) { + break; + } + pool_[i] = pool_[j]; + i = j; + } + pool_[i] = {id, val}; + } + + size_t size_ = 0, capacity_; + std::vector> pool_; +}; // class MaxMinHeap + // A std::vector like container but uses fixed size free memory(typically from // mmap) as backing store and can only be appended at the end. // diff --git a/src/index/sparse/sparse_dsp_config.h b/src/index/sparse/sparse_dsp_config.h new file mode 100644 index 000000000..8c6c28dfa --- /dev/null +++ b/src/index/sparse/sparse_dsp_config.h @@ -0,0 +1,117 @@ +// Copyright (C) 2019-2023 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// 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. + +#ifndef SPARSE_DSP_CONFIG_H +#define SPARSE_DSP_CONFIG_H + +#include "knowhere/config.h" + +namespace knowhere { + +// Search parameters for the DSP (Dynamic Superblock Pruning) index. +// +// Mode selection (dsp_mode): values map to DspSearchMode enum in sparse_dsp_index.h. +// 0 = DSP: dual-threshold (mu, eta) superblock pruning. Safe with default mu=1, eta=1. +// 1 = LSP/0: top-gamma superblock inclusion only, no mu/eta gates. +// 2 = LSP/1: LSP/0 + mu-overestimation gate (ub > theta/mu). +// 3 = LSP/2: LSP/1 + ASC gate (ub > theta/mu || asc > theta/eta). +// +// For all modes, dsp_eta also controls subblock BoundSum pruning. +// dsp_kth_init seeds the pruning threshold from per-dimension kth-largest scores (orthogonal to mode). +class SparseDspConfig : public BaseConfig { + public: + CFG_FLOAT drop_ratio_search; + CFG_INT refine_factor; + CFG_INT dsp_mode; + CFG_FLOAT dsp_mu; + CFG_FLOAT dsp_eta; + CFG_INT dsp_gamma; + CFG_BOOL dsp_kth_init; + CFG_FLOAT dsp_kth_alpha; + KNOWHERE_DECLARE_CONFIG(SparseDspConfig) { + KNOWHERE_CONFIG_DECLARE_FIELD(drop_ratio_search) + .description("drop ratio for search") + .set_default(0.0f) + .set_range(0.0f, 1.0f, true, false) + .for_search() + .for_range_search() + .for_iterator(); + KNOWHERE_CONFIG_DECLARE_FIELD(refine_factor) + .description("refine factor for approximate search") + .set_default(1) + .for_search(); + KNOWHERE_CONFIG_DECLARE_FIELD(dsp_mode) + .set_range(0, 3) + .set_default(0) + .description( + "superblock selection mode: " + "0=dsp (dual-threshold mu/eta + optional top-gamma backstop), " + "1=lsp0 (top-gamma from ub>=theta, no mu/asc gate), " + "2=lsp1 (lsp0 + mu-overestimation gate: ub>theta/mu), " + "3=lsp2 (lsp1 + asc gate: ub>theta/mu || asc>theta/eta)") + .for_search(); + KNOWHERE_CONFIG_DECLARE_FIELD(dsp_mu) + .set_range(0.0, 1.0, false, true) + .set_default(1.0) + .description( + "superblock max-based threshold relaxation factor (used by dsp/lsp1/lsp2). " + "Paper-aligned range is 0 < mu <= 1, with 1.0 as the safe default.") + .for_search(); + KNOWHERE_CONFIG_DECLARE_FIELD(dsp_eta) + .set_range(0.0, 1.0, false, true) + .set_default(1.0) + .description( + "threshold relaxation for superblock ASC pruning (dsp/lsp2) " + "and subblock BoundSum pruning (all modes). " + "Paper-aligned range is 0 < eta <= 1, with 1.0 as the safe default.") + .for_search(); + KNOWHERE_CONFIG_DECLARE_FIELD(dsp_gamma) + .set_range(0, 100000) + .set_default(0) + .description( + "always visit top-gamma superblocks by UB score " + "(0 = disabled, higher = safer but slower)") + .for_search(); + KNOWHERE_CONFIG_DECLARE_FIELD(dsp_kth_init) + .set_default(true) + .description( + "enable kth-score threshold initialization before pruning " + "(false = start threshold at 0, orthogonal to mode)") + .for_search(); + KNOWHERE_CONFIG_DECLARE_FIELD(dsp_kth_alpha) + .set_range(0.0, 1.0) + .set_default(1.0) + .description( + "scale factor for kth-score threshold seed: threshold *= alpha " + "(1.0 = full seed, 0.0 = no seed, intermediate = calibrated)") + .for_search(); + } + + Status + CheckAndAdjust(PARAM_TYPE param_type, std::string* err_msg) override { + if (param_type == PARAM_TYPE::SEARCH) { + const int mode = dsp_mode.value_or(0); + const float mu = dsp_mu.value_or(1.0f); + const float eta = dsp_eta.value_or(1.0f); + // Paper constraint 0 < mu <= eta <= 1 applies to DSP (mode=0) and LSP/2 (mode=3) + // which both use the dual-threshold (mu, eta) pruning. + // LSP/0 (mode=1) has no mu gate; LSP/1 (mode=2) uses mu but not the eta inequality. + if ((mode == 0 || mode == 3) && mu > eta) { + return HandleError(err_msg, "dsp_mu must be <= dsp_eta for DSP/LSP2 modes", Status::invalid_args); + } + } + return Status::success; + } +}; // class SparseDspConfig + +} // namespace knowhere + +#endif // SPARSE_DSP_CONFIG_H diff --git a/src/index/sparse/sparse_dsp_index.h b/src/index/sparse/sparse_dsp_index.h new file mode 100644 index 000000000..c960d9875 --- /dev/null +++ b/src/index/sparse/sparse_dsp_index.h @@ -0,0 +1,2123 @@ +// Copyright (C) 2019-2023 Zilliz. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software distributed under the License +// 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. + +#ifndef SPARSE_DSP_INDEX_H +#define SPARSE_DSP_INDEX_H + +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#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" +#include "knowhere/log.h" +#include "knowhere/prometheus_client.h" +#include "knowhere/sparse_utils.h" +#include "knowhere/utils.h" +#include "simd/instruction_set.h" +#include "simd/sparse_simd.h" + +namespace knowhere::sparse { + +using DspHeap = knowhere::ResultMinHeap; + +// Section types for DSP index serialization format +enum class DspSectionType : uint32_t { + POSTING_LISTS = 0, + METRIC_PARAMS = 1, + DIM_MAP = 2, + ROW_SUMS = 3, + MAX_SCORES_PER_DIM = 4, + PROMETHEUS_BUILD_STATS = 5, + DSP_METADATA = 6 +}; + +struct DspSectionHeader { + DspSectionType type = DspSectionType::POSTING_LISTS; + uint64_t offset = 0; + uint64_t size = 0; +}; + +struct DspBuildStats { + std::vector dataset_nnz_stats_; + std::vector posting_list_length_stats_; +}; + +// Superblock selection modes for DSP search. +enum class DspSearchMode : int { + DSP = 0, // dual-threshold (mu, eta) + optional top-gamma backstop + LSP0 = 1, // top-gamma from ub>=theta, no mu/asc gate (recommended for SPLADE) + LSP1 = 2, // LSP/0 + mu-overestimation gate (ub>theta/mu) + LSP2 = 3, // LSP/1 + asc gate (ub>theta/mu || asc>theta/eta) +}; + +struct DspSearchParams { + int refine_factor = 1; + float drop_ratio_search = 0.0f; + float dim_max_score_ratio = 1.0f; + DspSearchMode dsp_mode = DspSearchMode::DSP; + float dsp_mu = 1.0f; + float dsp_eta = 1.0f; + int dsp_gamma = 0; + bool dsp_kth_init = true; + float dsp_kth_alpha = 1.0f; +}; + +// Type-erased base for DspIndex so that the index node can hold either mmapped or non-mmapped variant. +template +class DspIndexBase { + public: + virtual ~DspIndexBase() = default; + virtual Status + Serialize(MemoryIOWriter& writer) const = 0; + virtual Status + Deserialize(MemoryIOReader& reader) = 0; + virtual Status + Train(const SparseRow* data, size_t rows) = 0; + virtual Status + Add(const SparseRow* data, size_t rows, int64_t dim) = 0; + virtual void + Search(const SparseRow& query, size_t k, float* distances, label_t* labels, const BitsetView& bitset, + const DocValueComputer& computer, DspSearchParams& params) const = 0; + virtual std::vector + GetAllDistances(const SparseRow& query, float drop_ratio_search, const BitsetView& bitset, + const DocValueComputer& computer) const = 0; + virtual float + GetRawDistance(const label_t vec_id, const SparseRow& query, const DocValueComputer& computer) const = 0; + virtual expected> + GetDocValueComputer(const BaseConfig& cfg) const = 0; + [[nodiscard]] virtual size_t + size() const = 0; + [[nodiscard]] virtual size_t + n_rows() const = 0; + [[nodiscard]] virtual size_t + n_cols() const = 0; + virtual void + SetBM25Params(float k1, float b, float avgdl) = 0; +}; + +// DSP (Dynamic Superblock Pruning) index for fast sparse vector search. +// +// DSP index structure: +// - u8 quantized block max scores +// - u16 upper bound accumulators with AVX-512 SIMD +// - 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 +template +class DspIndex : public DspIndexBase { + public: + template + using Vector = std::conditional_t, std::vector>; + + static constexpr uint32_t kSubblockSize = 8; + static constexpr uint32_t kSuperblockSize = 512; + static constexpr uint32_t kStride = kSuperblockSize / kSubblockSize; // 64 + static constexpr uint32_t kSimdWidth = 32; // AVX-512 processes 32 u16 values + + explicit DspIndex(SparseMetricType metric_type) : metric_type_(metric_type) { +#if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) + // for now, use timestamp as index_id + index_id_ = std::to_string( + std::chrono::duration_cast(std::chrono::system_clock::now().time_since_epoch()) + .count()); + index_size_gauge_ = &sparse_inverted_index_size_family.Add({{"index_id", index_id_}, {"index_type", "dsp"}}); + index_dataset_nnz_len_histogram_ = + &sparse_dataset_nnz_len_family.Add({{"index_id", index_id_}, {"index_type", "dsp"}}, defaultBuckets); + index_posting_list_len_histogram_ = &sparse_inverted_index_posting_list_len_family.Add( + {{"index_id", index_id_}, {"index_type", "dsp"}}, defaultBuckets); +#endif + } + + ~DspIndex() { + if constexpr (mmapped) { + if (map_ != nullptr) { + auto res = munmap(map_, map_byte_size_); + if (res != 0) { + LOG_KNOWHERE_ERROR_ << "Failed to munmap when deleting sparse DspIndex: " << strerror(errno); + } + map_ = nullptr; + map_byte_size_ = 0; + } + if (map_fd_ != -1) { + close(map_fd_); + map_fd_ = -1; + } + } +#if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) + if (index_size_gauge_ != nullptr) { + sparse_inverted_index_size_family.Remove(index_size_gauge_); + } + if (index_dataset_nnz_len_histogram_ != nullptr) { + sparse_dataset_nnz_len_family.Remove(index_dataset_nnz_len_histogram_); + } + if (index_posting_list_len_histogram_ != nullptr) { + sparse_inverted_index_posting_list_len_family.Remove(index_posting_list_len_histogram_); + } +#endif + } + + void + SetBM25Params(float k1, float b, float avgdl) override { + bm25_params_ = std::make_unique(k1, b, avgdl); + } + + expected> + GetDocValueComputer(const BaseConfig& cfg) const override { + auto metric_type = cfg.metric_type; + if (metric_type_ != SparseMetricType::METRIC_BM25) { + if (metric_type.has_value() && !IsMetricType(metric_type.value(), metric::IP)) { + auto msg = + "metric type not match, expected: " + std::string(metric::IP) + ", got: " + metric_type.value(); + return expected>::Err(Status::invalid_metric_type, msg); + } + return GetDocValueOriginalComputer(); + } + if (metric_type.has_value() && !IsMetricType(metric_type.value(), metric::BM25)) { + auto msg = + "metric type not match, expected: " + std::string(metric::BM25) + ", got: " + metric_type.value(); + return expected>::Err(Status::invalid_metric_type, msg); + } + if (!cfg.bm25_avgdl.has_value()) { + return expected>::Err(Status::invalid_args, + "avgdl must be supplied during searching"); + } + auto avgdl = cfg.bm25_avgdl.value(); + avgdl = std::max(avgdl, 1.0f); + if ((cfg.bm25_k1.has_value() && cfg.bm25_k1.value() != bm25_params_->k1) || + ((cfg.bm25_b.has_value() && cfg.bm25_b.value() != bm25_params_->b))) { + return expected>::Err(Status::invalid_args, + "search time k1/b must equal load time config."); + } + return GetDocValueBM25Computer(bm25_params_->k1, bm25_params_->b, avgdl); + } + + Status + Train(const SparseRow* data, size_t rows) override { + if constexpr (mmapped) { + throw std::invalid_argument("mmapped DspIndex does not support Train"); + } else { + return Status::success; + } + } + + Status + Add(const SparseRow* data, size_t rows, int64_t dim) override { + if constexpr (mmapped) { + throw std::invalid_argument("mmapped DspIndex does not support Add"); + } else { + if (n_rows_internal_ != 0 || rows > std::numeric_limits::max()) { + return Status::invalid_args; + } + 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); + } +#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]; + } + } + } + +#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] = dim_counts[i]; + } +#endif + + inverted_index_ids_spans_.clear(); + inverted_index_vals_spans_.clear(); + inverted_index_ids_spans_.reserve(nr_inner_dims_); + inverted_index_vals_spans_.reserve(nr_inner_dims_); + + for (size_t i = 0; i < nr_inner_dims_; ++i) { + inverted_index_ids_spans_.emplace_back(inverted_index_ids_[i].data(), inverted_index_ids_[i].size()); + inverted_index_vals_spans_.emplace_back(inverted_index_vals_[i].data(), inverted_index_vals_[i].size()); + } + + if (max_score_in_dim_.size() > 0) { + max_score_in_dim_spans_ = boost::span(max_score_in_dim_.data(), max_score_in_dim_.size()); + } + + if (metric_type_ == SparseMetricType::METRIC_BM25) { + bm25_params_->row_sums_spans_ = + boost::span(bm25_params_->row_sums.data(), bm25_params_->row_sums.size()); + } + + build_dsp_metadata(data, rows, &dim_spb_counts); + return Status::success; + } + } + + Status + Serialize(MemoryIOWriter& writer) const override { + const uint32_t index_format_version = 1; + + // Index File Header (v1) + writer.write(&index_format_version, sizeof(uint32_t)); + writer.write(&n_rows_internal_, sizeof(uint32_t)); + writer.write(&max_dim_, sizeof(uint32_t)); + writer.write(&nr_inner_dims_, sizeof(uint32_t)); + auto reserved = std::array(); + writer.write(reserved.data(), reserved.size()); + + // Phase 1: Collect section metadata + std::vector> section_meta; + + uint64_t posting_lists_size = sizeof(uint32_t); // encoding type + posting_lists_size += sizeof(uint64_t) * (nr_inner_dims_ + 1); // dim offsets + for (size_t i = 0; i < nr_inner_dims_; ++i) { + posting_lists_size += inverted_index_ids_spans_[i].size() * sizeof(uint32_t) + + inverted_index_vals_spans_[i].size() * sizeof(QType); + } + section_meta.emplace_back(DspSectionType::POSTING_LISTS, posting_lists_size); + section_meta.emplace_back(DspSectionType::DIM_MAP, sizeof(uint32_t) * nr_inner_dims_); + + if (metric_type_ == SparseMetricType::METRIC_BM25) { + section_meta.emplace_back(DspSectionType::ROW_SUMS, sizeof(float) * n_rows_internal_); + } + if (max_score_in_dim_spans_.size() > 0) { + section_meta.emplace_back(DspSectionType::MAX_SCORES_PER_DIM, sizeof(float) * nr_inner_dims_); + } +#if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) + section_meta.emplace_back(DspSectionType::PROMETHEUS_BUILD_STATS, + sizeof(uint32_t) * n_rows_internal_ + sizeof(uint32_t) * nr_inner_dims_); +#endif + + // Append DSP metadata section + if (n_subblocks_ > 0) { + uint64_t dsp_size = 0; + dsp_size += 4 * sizeof(uint32_t); // header: version, n_subblocks, n_superblocks, n_sb_padded + + for (uint32_t d = 0; d < nr_inner_dims_; ++d) { + const auto& bm = dim_block_max_[d]; + dsp_size += sizeof(uint32_t); // n_block_ids + dsp_size += sizeof(uint32_t); // n_logical + dsp_size += 4; // kth[4] + dsp_size += bm.block_ids.size() * sizeof(uint32_t); // block_ids + dsp_size += sizeof(uint32_t); // packed_size + dsp_size += bm.max_scores.size() * sizeof(uint8_t); // u8 max_scores + } + + uint32_t spb_total = spb_block_ids_.size(); + dsp_size += sizeof(uint32_t); // spb_total + dsp_size += (nr_inner_dims_ + 1) * sizeof(uint32_t); // spb_dim_offsets + dsp_size += spb_total * sizeof(uint32_t); // spb_block_ids + dsp_size += spb_total * sizeof(float); // spb_max_vals + dsp_size += spb_total * sizeof(float); // spb_asc_vals + + uint32_t total_terms = fwd_term_ids_.size(); + uint32_t total_entries = fwd_doc_offsets_.size(); + dsp_size += sizeof(uint32_t); // total_terms + dsp_size += sizeof(uint32_t); // total_entries + dsp_size += (n_subblocks_ + 1) * sizeof(uint32_t); // fwd_block_term_offsets + dsp_size += total_terms * sizeof(uint32_t); // fwd_term_ids + dsp_size += (total_terms + 1) * sizeof(uint32_t); // fwd_term_entry_offsets + dsp_size += total_entries * sizeof(uint8_t); // fwd_doc_offsets + dsp_size += total_entries * sizeof(float); // fwd_scores + + section_meta.emplace_back(DspSectionType::DSP_METADATA, dsp_size); + } + + // Phase 2: Build headers with offsets and write section table + uint32_t nr_sections = static_cast(section_meta.size()); + writer.write(&nr_sections, sizeof(uint32_t)); + + std::vector section_headers(nr_sections); + uint64_t used_offset = index_file_v1_header_size + sizeof(uint32_t) + sizeof(DspSectionHeader) * nr_sections; + for (uint32_t i = 0; i < nr_sections; ++i) { + section_headers[i].type = section_meta[i].first; + section_headers[i].offset = used_offset; + section_headers[i].size = section_meta[i].second; + used_offset += section_meta[i].second; + } + writer.write(section_headers.data(), sizeof(DspSectionHeader), nr_sections); + + // Write posting lists + uint32_t index_encoding_type = 0; + writer.write(&index_encoding_type, sizeof(uint32_t)); + std::vector inverted_index_offsets(nr_inner_dims_ + 1); + inverted_index_offsets[0] = 0; + for (size_t i = 1; i <= nr_inner_dims_; ++i) { + inverted_index_offsets[i] = inverted_index_offsets[i - 1] + inverted_index_ids_spans_[i - 1].size(); + } + writer.write(inverted_index_offsets.data(), sizeof(uint64_t), inverted_index_offsets.size()); + for (size_t i = 0; i < nr_inner_dims_; ++i) { + writer.write(inverted_index_ids_spans_[i].data(), sizeof(uint32_t), inverted_index_ids_spans_[i].size()); + } + for (size_t i = 0; i < nr_inner_dims_; ++i) { + writer.write(inverted_index_vals_spans_[i].data(), sizeof(QType), inverted_index_vals_spans_[i].size()); + } + + // Write dim map + auto dim_map_reverse = std::vector(nr_inner_dims_); + for (const auto& [dim, dim_id] : dim_map_) { + dim_map_reverse[dim_id] = dim; + } + writer.write(dim_map_reverse.data(), sizeof(uint32_t), nr_inner_dims_); + + // Write row sums (BM25) + if (metric_type_ == SparseMetricType::METRIC_BM25) { + writer.write(bm25_params_->row_sums_spans_.data(), sizeof(float), n_rows_internal_); + } + + // Write max scores per dim + if (max_score_in_dim_spans_.size() > 0) { + writer.write(max_score_in_dim_spans_.data(), sizeof(float), nr_inner_dims_); + } + + // Write prometheus build stats +#if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) + writer.write(build_stats_.dataset_nnz_stats_.data(), sizeof(uint32_t), n_rows_internal_); + writer.write(build_stats_.posting_list_length_stats_.data(), sizeof(uint32_t), nr_inner_dims_); +#endif + + // Write DSP metadata section + if (n_subblocks_ > 0) { + uint32_t dsp_version = 1; + writer.write(&dsp_version, sizeof(uint32_t)); + writer.write(&n_subblocks_, sizeof(uint32_t)); + writer.write(&n_superblocks_, sizeof(uint32_t)); + writer.write(&n_sb_padded_, sizeof(uint32_t)); + + for (uint32_t d = 0; d < nr_inner_dims_; ++d) { + const auto& bm = dim_block_max_[d]; + uint32_t n_block_ids = bm.block_ids.size(); + uint32_t n_logical = bm.n_logical; + writer.write(&n_block_ids, sizeof(uint32_t)); + writer.write(&n_logical, sizeof(uint32_t)); + writer.write(bm.kth, 4); + if (n_block_ids > 0) { + writer.write(bm.block_ids.data(), sizeof(uint32_t), n_block_ids); + } + uint32_t packed_size = bm.max_scores.size(); + writer.write(&packed_size, sizeof(uint32_t)); + if (packed_size > 0) { + writer.write(bm.max_scores.data(), sizeof(uint8_t), packed_size); + } + } + + uint32_t spb_total = spb_block_ids_.size(); + writer.write(&spb_total, sizeof(uint32_t)); + writer.write(spb_dim_offsets_.data(), sizeof(uint32_t), nr_inner_dims_ + 1); + writer.write(spb_block_ids_.data(), sizeof(uint32_t), spb_total); + writer.write(spb_max_vals_.data(), sizeof(float), spb_total); + writer.write(spb_asc_vals_.data(), sizeof(float), spb_total); + + uint32_t total_terms = fwd_term_ids_.size(); + uint32_t total_entries = fwd_doc_offsets_.size(); + writer.write(&total_terms, sizeof(uint32_t)); + writer.write(&total_entries, sizeof(uint32_t)); + writer.write(fwd_block_term_offsets_.data(), sizeof(uint32_t), n_subblocks_ + 1); + writer.write(fwd_term_ids_.data(), sizeof(uint32_t), total_terms); + writer.write(fwd_term_entry_offsets_.data(), sizeof(uint32_t), total_terms + 1); + writer.write(fwd_doc_offsets_.data(), sizeof(uint8_t), total_entries); + writer.write(fwd_scores_.data(), sizeof(float), total_entries); + } + + return Status::success; + } + + Status + Deserialize(MemoryIOReader& reader) override { + dsp_loaded_ = false; + + // Read file header + uint32_t index_format_version = 0; + reader.read(&index_format_version, sizeof(uint32_t)); + if (index_format_version != 1) { + return Status::invalid_serialized_index_type; + } + + reader.read(&n_rows_internal_, sizeof(uint32_t)); + reader.read(&max_dim_, sizeof(uint32_t)); + reader.read(&nr_inner_dims_, sizeof(uint32_t)); + reader.advance(index_file_v1_header_reserved_size); + + // Read sections + uint32_t nr_sections = 0; + reader.read(&nr_sections, sizeof(uint32_t)); + size_t sec_table_offset = reader.tellg(); + + for (uint32_t i = 0; i < nr_sections; ++i) { + DspSectionHeader section_header; + reader.seekg(sec_table_offset); + reader.read(§ion_header, sizeof(DspSectionHeader)); + sec_table_offset += sizeof(DspSectionHeader); + + switch (section_header.type) { + case DspSectionType::POSTING_LISTS: { + reader.seekg(section_header.offset); + uint32_t index_encoding_type = 0; + reader.read(&index_encoding_type, sizeof(uint32_t)); + if (index_encoding_type != 0) { + return Status::invalid_serialized_index_type; + } + auto inverted_index_offsets_span = boost::span( + reinterpret_cast(reader.data() + reader.tellg()), nr_inner_dims_ + 1); + reader.advance(sizeof(uint64_t) * (nr_inner_dims_ + 1)); + inverted_index_ids_spans_.resize(nr_inner_dims_); + inverted_index_vals_spans_.resize(nr_inner_dims_); + for (size_t j = 0; j < nr_inner_dims_; ++j) { + inverted_index_ids_spans_[j] = boost::span( + reinterpret_cast(reader.data() + reader.tellg()), + inverted_index_offsets_span[j + 1] - inverted_index_offsets_span[j]); + reader.advance(inverted_index_ids_spans_[j].size() * sizeof(uint32_t)); + } + for (size_t j = 0; j < nr_inner_dims_; ++j) { + inverted_index_vals_spans_[j] = boost::span( + reinterpret_cast(reader.data() + reader.tellg()), + inverted_index_offsets_span[j + 1] - inverted_index_offsets_span[j]); + reader.advance(inverted_index_vals_spans_[j].size() * sizeof(QType)); + } + break; + } + case DspSectionType::DIM_MAP: { + reader.seekg(section_header.offset); + for (uint32_t j = 0; j < nr_inner_dims_; ++j) { + uint32_t dim = 0; + reader.read(&dim, sizeof(uint32_t)); + dim_map_[dim] = j; + } + break; + } + case DspSectionType::ROW_SUMS: { + reader.seekg(section_header.offset); + bm25_params_->row_sums_spans_ = boost::span( + reinterpret_cast(reader.data() + section_header.offset), n_rows_internal_); + reader.advance(sizeof(float) * n_rows_internal_); + break; + } + case DspSectionType::MAX_SCORES_PER_DIM: { + reader.seekg(section_header.offset); + max_score_in_dim_spans_ = boost::span( + reinterpret_cast(reader.data() + section_header.offset), nr_inner_dims_); + reader.advance(sizeof(float) * nr_inner_dims_); + break; + } + case DspSectionType::PROMETHEUS_BUILD_STATS: { +#if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) + reader.seekg(section_header.offset); + auto dataset_nnz_stats = std::vector(n_rows_internal_); + reader.read(dataset_nnz_stats.data(), sizeof(uint32_t), n_rows_internal_); + auto posting_list_length_stats = std::vector(nr_inner_dims_); + reader.read(posting_list_length_stats.data(), sizeof(uint32_t), nr_inner_dims_); + for (size_t j = 0; j < n_rows_internal_; ++j) { + index_dataset_nnz_len_histogram_->Observe(dataset_nnz_stats[j]); + } + for (size_t j = 0; j < nr_inner_dims_; ++j) { + index_posting_list_len_histogram_->Observe(posting_list_length_stats[j]); + } +#endif + break; + } + case DspSectionType::DSP_METADATA: { + reader.seekg(section_header.offset); + + uint32_t dsp_version = 0; + reader.read(&dsp_version, sizeof(uint32_t)); + if (dsp_version != 1) { + return Status::invalid_serialized_index_type; + } + reader.read(&n_subblocks_, sizeof(uint32_t)); + reader.read(&n_superblocks_, sizeof(uint32_t)); + reader.read(&n_sb_padded_, sizeof(uint32_t)); + + const uint32_t nr_dims = nr_inner_dims_; + dim_block_max_.resize(nr_dims); + for (uint32_t d = 0; d < nr_dims; ++d) { + auto& bm = dim_block_max_[d]; + uint32_t n_block_ids = 0, n_logical = 0; + reader.read(&n_block_ids, sizeof(uint32_t)); + reader.read(&n_logical, sizeof(uint32_t)); + reader.read(bm.kth, 4); + bm.n_logical = n_logical; + if (n_block_ids > 0) { + bm.block_ids.resize(n_block_ids); + reader.read(bm.block_ids.data(), sizeof(uint32_t), n_block_ids); + } + uint32_t packed_size = 0; + reader.read(&packed_size, sizeof(uint32_t)); + if (packed_size > 0) { + bm.max_scores.resize(packed_size); + reader.read(bm.max_scores.data(), sizeof(uint8_t), packed_size); + } + } + + uint32_t spb_total = 0; + reader.read(&spb_total, sizeof(uint32_t)); + spb_dim_offsets_.resize(nr_dims + 1); + reader.read(spb_dim_offsets_.data(), sizeof(uint32_t), nr_dims + 1); + spb_block_ids_.resize(spb_total); + reader.read(spb_block_ids_.data(), sizeof(uint32_t), spb_total); + spb_max_vals_.resize(spb_total); + reader.read(spb_max_vals_.data(), sizeof(float), spb_total); + spb_asc_vals_.resize(spb_total); + reader.read(spb_asc_vals_.data(), sizeof(float), spb_total); + + uint32_t total_terms = 0, total_entries = 0; + reader.read(&total_terms, sizeof(uint32_t)); + reader.read(&total_entries, sizeof(uint32_t)); + fwd_block_term_offsets_.resize(n_subblocks_ + 1); + reader.read(fwd_block_term_offsets_.data(), sizeof(uint32_t), n_subblocks_ + 1); + fwd_term_ids_.resize(total_terms); + reader.read(fwd_term_ids_.data(), sizeof(uint32_t), total_terms); + fwd_term_entry_offsets_.resize(total_terms + 1); + reader.read(fwd_term_entry_offsets_.data(), sizeof(uint32_t), total_terms + 1); + fwd_doc_offsets_.resize(total_entries); + reader.read(fwd_doc_offsets_.data(), sizeof(uint8_t), total_entries); + fwd_scores_.resize(total_entries); + reader.read(fwd_scores_.data(), sizeof(float), total_entries); + + dsp_loaded_ = true; + break; + } + default: + break; + } + } + +#if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) + index_size_gauge_->Set((double)size() / 1024.0 / 1024.0); +#endif + + if (!dsp_loaded_) { + build_dsp_metadata(); + } + + return Status::success; + } + + void + Search(const SparseRow& query, size_t k, float* distances, label_t* labels, const BitsetView& bitset, + const DocValueComputer& computer, DspSearchParams& approx_params) const override { + std::fill(distances, distances + k, std::numeric_limits::quiet_NaN()); + std::fill(labels, labels + k, -1); + if (query.size() == 0) { + return; + } + + auto q_vec = parse_query(query, approx_params.drop_ratio_search); + if (q_vec.empty()) { + return; + } + + 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); + + collect_result(heap, distances, labels); + } + + std::vector + GetAllDistances(const SparseRow& query, float drop_ratio_search, const BitsetView& bitset, + const DocValueComputer& computer) const override { + if (query.size() == 0) { + return {}; + } + std::vector values(query.size()); + for (size_t i = 0; i < query.size(); ++i) { + values[i] = std::abs(query[i].val); + } + auto q_vec = parse_query(query, drop_ratio_search); + + auto distances = compute_all_distances(q_vec, computer); + if (!bitset.empty()) { + for (size_t i = 0; i < distances.size(); ++i) { + if (bitset.test(i)) { + distances[i] = 0.0f; + } + } + } + return distances; + } + + float + GetRawDistance(const label_t vec_id, const SparseRow& query, + const DocValueComputer& computer) const override { + float distance = 0.0f; + + for (size_t i = 0; i < query.size(); ++i) { + auto [dim, val] = query[i]; + auto dim_it = dim_map_.find(dim); + if (dim_it == dim_map_.cend()) { + continue; + } + auto& plist_ids = inverted_index_ids_spans_[dim_it->second]; + auto it = std::lower_bound(plist_ids.begin(), plist_ids.end(), vec_id, + [](const auto& x, table_t y) { return x < y; }); + if (it != plist_ids.end() && *it == vec_id) { + distance += + val * + computer(inverted_index_vals_spans_[dim_it->second][it - plist_ids.begin()], + metric_type_ == SparseMetricType::METRIC_BM25 ? bm25_params_->row_sums_spans_[vec_id] : 0); + } + } + + return distance; + } + + [[nodiscard]] size_t + size() const override { + size_t res = sizeof(*this); + res += dim_map_.size() * + (sizeof(typename decltype(dim_map_)::key_type) + sizeof(typename decltype(dim_map_)::mapped_type)); + + if constexpr (mmapped) { + res += map_byte_size_; + } else { + // Posting list data: use owning vectors if populated (Train path), + // otherwise fall back to spans (Deserialize path where spans point into binary_). + if (inverted_index_ids_.size() == nr_inner_dims_) { + for (uint32_t i = 0; i < nr_inner_dims_; ++i) { + res += inverted_index_ids_[i].capacity() * sizeof(table_t); + res += inverted_index_vals_[i].capacity() * sizeof(QType); + } + } else { + for (uint32_t i = 0; i < inverted_index_ids_spans_.size(); ++i) { + res += inverted_index_ids_spans_[i].size() * sizeof(table_t); + res += inverted_index_vals_spans_[i].size() * sizeof(QType); + } + } + // Span metadata (lightweight views into the owning vectors) + res += inverted_index_ids_spans_.capacity() * sizeof(boost::span); + res += inverted_index_vals_spans_.capacity() * sizeof(boost::span); + // Max score per dimension + res += max_score_in_dim_.capacity() * sizeof(float); + } + + // BM25 row sums + if (bm25_params_) { + if constexpr (mmapped) { + // row_sums data is in the mmap region, already counted above + } else { + res += bm25_params_->row_sums.capacity() * sizeof(float); + } + } + + // DSP metadata: dim_block_max_, superblock CSR, forward index + for (const auto& bm : dim_block_max_) { + res += bm.block_ids.capacity() * sizeof(uint32_t); + res += bm.max_scores.capacity() * sizeof(uint8_t); + } + res += dim_block_max_.capacity() * sizeof(DimBlockMax); + res += spb_dim_offsets_.capacity() * sizeof(uint32_t); + res += spb_block_ids_.capacity() * sizeof(uint32_t); + res += spb_max_vals_.capacity() * sizeof(float); + res += spb_asc_vals_.capacity() * sizeof(float); + res += fwd_block_term_offsets_.capacity() * sizeof(uint32_t); + res += fwd_term_ids_.capacity() * sizeof(uint32_t); + res += fwd_term_entry_offsets_.capacity() * sizeof(uint32_t); + res += fwd_doc_offsets_.capacity() * sizeof(uint8_t); + res += fwd_scores_.capacity() * sizeof(float); + + return res; + } + + [[nodiscard]] size_t + n_rows() const override { + return n_rows_internal_; + } + + [[nodiscard]] size_t + n_cols() const override { + return max_dim_; + } + + private: + // ======================================================================== + // Storage members (self-contained, no longer inherited from SparseInvertedStorage) + // ======================================================================== + std::unordered_map dim_map_; + uint32_t nr_inner_dims_ = 0; + + Vector> inverted_index_ids_; + Vector> inverted_index_vals_; + std::vector> inverted_index_ids_spans_; + std::vector> inverted_index_vals_spans_; + Vector max_score_in_dim_; + boost::span max_score_in_dim_spans_; + + SparseMetricType metric_type_; + + size_t n_rows_internal_ = 0; + size_t max_dim_ = 0; + uint32_t next_dim_id_ = 0; + + char* map_ = nullptr; + size_t map_byte_size_ = 0; + int map_fd_ = -1; + + struct BM25Params { + float k1; + float b; + Vector row_sums; + boost::span row_sums_spans_; + + DocValueComputer max_score_computer; + + BM25Params(float k1, float b, float avgdl) + : k1(k1), b(b), max_score_computer(GetDocValueBM25Computer(k1, b, avgdl)) { + } + }; + + std::unique_ptr bm25_params_; + + static constexpr uint32_t index_file_v1_header_size = 32; + static constexpr uint32_t index_file_v1_header_reserved_size = 16; + +#if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) + DspBuildStats build_stats_; + + std::string index_id_{}; + prometheus::Gauge* index_size_gauge_{nullptr}; + prometheus::Histogram* index_dataset_nnz_len_histogram_{nullptr}; + prometheus::Histogram* index_posting_list_len_histogram_{nullptr}; +#endif + + // ======================================================================== + // Helper methods from SparseInvertedStorage + // ======================================================================== + inline DType + get_threshold(std::vector& values, float drop_ratio) const { + auto drop_count = static_cast(drop_ratio * values.size()); + if (drop_count == 0) { + return 0; + } + auto pos = values.begin() + drop_count; + std::nth_element(values.begin(), pos, values.end()); + return *pos; + } + + std::vector + compute_all_distances(const std::vector>& q_vec, + const DocValueComputer& computer) const { + std::vector scores(n_rows_internal_, 0.0f); + + if (metric_type_ == SparseMetricType::METRIC_IP) { + for (const auto& [dim_idx, q_weight] : q_vec) { + const auto& plist_ids = inverted_index_ids_spans_[dim_idx]; + const auto& plist_vals = inverted_index_vals_spans_[dim_idx]; + + accumulate_posting_list_contribution_ip_dispatch( + plist_ids.data(), plist_vals.data(), plist_ids.size(), static_cast(q_weight), scores.data()); + } + } else { + const auto& doc_len_ratios = bm25_params_->row_sums_spans_; + for (const auto& [dim_idx, q_weight] : q_vec) { + const auto& plist_ids = inverted_index_ids_spans_[dim_idx]; + const auto& plist_vals = inverted_index_vals_spans_[dim_idx]; + const float q_weight_float = static_cast(q_weight); + for (size_t j = 0; j < plist_ids.size(); ++j) { + const auto doc_id = plist_ids[j]; + const float doc_val = computer(plist_vals[j], doc_len_ratios[doc_id]); + scores[doc_id] += q_weight_float * doc_val; + } + } + } + + return scores; + } + + std::vector> + parse_query(const SparseRow& query, float drop_ratio_search) const { + DType q_threshold = 0; + if (drop_ratio_search != 0) { + std::vector values(query.size()); + for (size_t i = 0; i < query.size(); ++i) { + values[i] = std::abs(query[i].val); + } + q_threshold = get_threshold(values, drop_ratio_search); + } + + std::vector> filtered_query; + for (size_t i = 0; i < query.size(); ++i) { + auto [dim, val] = query[i]; + auto dim_it = dim_map_.find(dim); + if (dim_it == dim_map_.cend() || std::abs(val) < q_threshold) { + continue; + } + filtered_query.emplace_back(dim_it->second, val); + } + + return filtered_query; + } + + template + void + collect_result(HeapType& heap, float* distances, label_t* labels) const { + int cnt = heap.Size(); + for (auto i = cnt - 1; i >= 0; --i) { + const auto& entry = heap.Results().front(); + labels[i] = entry.second; + distances[i] = entry.first; + heap.Pop(); + } + } + + inline void + add_row_to_index(const SparseRow& row, table_t vec_id) { + [[maybe_unused]] float row_sum = 0; + for (size_t j = 0; j < row.size(); ++j) { + auto [dim, val] = row[j]; + if (metric_type_ == SparseMetricType::METRIC_BM25) { + row_sum += val; + } + if (val == 0) { + continue; + } + auto dim_it = dim_map_.find(dim); + if (dim_it == dim_map_.cend()) { + if constexpr (mmapped) { + throw std::runtime_error("unexpected vector dimension in mmapped DspIndex"); + } + dim_it = dim_map_.insert({dim, next_dim_id_++}).first; + inverted_index_ids_.emplace_back(); + inverted_index_vals_.emplace_back(); + max_score_in_dim_.emplace_back(0.0f); + } + inverted_index_ids_[dim_it->second].emplace_back(vec_id); + inverted_index_vals_[dim_it->second].emplace_back(get_quant_val(val)); + } +#if defined(NOT_COMPILE_FOR_SWIG) && !defined(KNOWHERE_WITH_LIGHT) + build_stats_.dataset_nnz_stats_.push_back(row.size()); +#endif + // update max_score_in_dim_ + for (size_t j = 0; j < row.size(); ++j) { + auto [dim, val] = row[j]; + if (val == 0) { + continue; + } + auto dim_it = dim_map_.find(dim); + if (dim_it == dim_map_.cend()) { + throw std::runtime_error("unexpected vector dimension in DspIndex"); + } + auto score = static_cast(val); + if (metric_type_ == SparseMetricType::METRIC_BM25) { + score = bm25_params_->max_score_computer(val, row_sum); + } + max_score_in_dim_[dim_it->second] = std::max(max_score_in_dim_[dim_it->second], score); + } + if (metric_type_ == SparseMetricType::METRIC_BM25) { + bm25_params_->row_sums.emplace_back(row_sum); + } + } + + inline QType + get_quant_val(DType val) const { + if constexpr (!std::is_same_v) { + const DType max_val = static_cast(std::numeric_limits::max()); + if (val >= max_val) { + return std::numeric_limits::max(); + } else if (val <= std::numeric_limits::min()) { + return std::numeric_limits::min(); + } else { + return static_cast(val); + } + } else { + return val; + } + } + + // ======================================================================== + // DSP-specific members + // ======================================================================== + void + refine_and_collect(const SparseRow& query, DspHeap& inacc_heap, size_t k, float* distances, label_t* labels, + const DocValueComputer& computer) const { + DspHeap heap(k); + while (inacc_heap.Size() > 0) { + table_t doc_id = inacc_heap.Pop()->second; + float score = GetRawDistance(doc_id, query, computer); + heap.Push(score, doc_id); + } + collect_result(heap, distances, labels); + } + + struct DimBlockMax { + std::vector block_ids; + std::vector max_scores; + uint32_t n_logical = 0; + uint8_t kth[4] = {0, 0, 0, 0}; + bool + is_dense() const { + return block_ids.empty() && n_logical > 0; + } + }; + std::vector dim_block_max_; + + // ======================================================================== + // Superblock max + ASC (sparse CSR format, float -- used for coarse pruning) + // ======================================================================== + // TODO: Consider span/view-backed storage for spb_* and fwd_* in deserialize/mmapped + // paths if these arrays become a meaningful memory overhead. + std::vector spb_dim_offsets_; + std::vector spb_block_ids_; + std::vector spb_max_vals_; + std::vector spb_asc_vals_; + + // ======================================================================== + // Forward index (flat layout for cache-friendly scoring) + // ======================================================================== + std::vector fwd_block_term_offsets_; + std::vector fwd_term_ids_; + std::vector fwd_term_entry_offsets_; + std::vector fwd_doc_offsets_; + std::vector fwd_scores_; + + uint32_t n_subblocks_ = 0; + uint32_t n_superblocks_ = 0; + 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 SEEK_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(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; + } + + n_subblocks_ = (n_rows_internal_ + kSubblockSize - 1) / kSubblockSize; + n_superblocks_ = (n_rows_internal_ + kSuperblockSize - 1) / kSuperblockSize; + n_sb_padded_ = (n_subblocks_ + kStride - 1) / kStride * kStride; + + const uint32_t nr_dims = nr_inner_dims_; + const bool is_bm25 = metric_type_ == SparseMetricType::METRIC_BM25; + + // Per-doc forward index: (inner_dim, score) pairs appended per doc. + struct DocFwdEntry { + uint32_t inner_dim = 0; + float score = 0.0f; + }; + 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_); + + // 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"); + } + + 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); + + static constexpr uint32_t kKthSizes[4] = {10, 100, 1000, 10000}; + using KthHeap = std::priority_queue, std::greater>; + + for (uint32_t d = 0; d < nr_dims; ++d) { + const auto& plist_ids = inverted_index_ids_spans_[d]; + const auto& plist_vals = inverted_index_vals_spans_[d]; + const float max_score_d = max_score_in_dim_spans_[d]; + + if (plist_ids.size() == 0 || max_score_d <= 0.0f) { + continue; + } + + 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]; + const QType val = plist_vals[i]; + + float score; + if (is_bm25) { + score = bm25_params_->max_score_computer(val, bm25_params_->row_sums_spans_[doc_id]); + } else { + score = static_cast(val); + } + + for (int h = 0; h < 4; ++h) { + if (kth_heaps[h].size() < kKthSizes[h]) { + kth_heaps[h].push(score); + } else if (score > kth_heaps[h].top()) { + kth_heaps[h].pop(); + kth_heaps[h].push(score); + } + } + + 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); + + 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); + + 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]; + const size_t posting_len = plist_ids.size(); + for (int h = 0; h < 4; ++h) { + if (posting_len >= kKthSizes[h] && !kth_heaps[h].empty()) { + float kth_f = kth_heaps[h].top(); + bm.kth[h] = static_cast(std::min(255.0f, std::floor(kth_f * inv_max_score_u8))); + } + } + + const uint32_t nnz_blocks = touched_list.size(); + if (nnz_blocks > static_cast(n_subblocks_ * kDenseThreshold)) { + bm.n_logical = n_sb_padded_; + bm.max_scores.resize(n_sb_padded_, 0); + for (uint32_t sb : touched_list) { + bm.max_scores[sb] = + static_cast(std::min(255.0f, std::ceil(tmp_sb_max[sb] * inv_max_score_u8))); + } + } else { + std::sort(touched_list.begin(), touched_list.end()); + bm.block_ids.resize(nnz_blocks); + bm.n_logical = nnz_blocks; + bm.max_scores.resize(nnz_blocks); + for (uint32_t i = 0; i < nnz_blocks; ++i) { + uint32_t sb = touched_list[i]; + bm.block_ids[i] = sb; + bm.max_scores[i] = + static_cast(std::min(255.0f, std::ceil(tmp_sb_max[sb] * inv_max_score_u8))); + } + } + + for (uint32_t sb : touched_list) { + tmp_sb_max[sb] = 0.0f; + sb_touched[sb] = 0; + } + touched_list.clear(); + } + + 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; + } + } + + // ---- Phase 3: Build flat forward index from per-doc data ---- + { + for (uint32_t doc = 0; doc < n_rows_internal_; ++doc) { + auto& entries = per_doc_fwd[doc]; + if (entries.size() > 1) { + std::sort(entries.begin(), entries.end(), + [](const DocFwdEntry& a, const DocFwdEntry& b) { return a.inner_dim < b.inner_dim; }); + } + } + + uint32_t total_terms = 0; + uint32_t total_entries = 0; + + struct BlockEntry { + uint32_t inner_dim = 0; + uint8_t doc_offset = 0; + float score = 0.0f; + }; + std::vector block_buf; + block_buf.reserve(1024); + + for (uint32_t sb = 0; sb < n_subblocks_; ++sb) { + block_buf.clear(); + const uint32_t doc_start = sb * kSubblockSize; + const uint32_t doc_end = std::min(doc_start + kSubblockSize, static_cast(n_rows_internal_)); + for (uint32_t doc = doc_start; doc < doc_end; ++doc) { + const uint8_t doc_off = static_cast(doc - doc_start); + for (const auto& e : per_doc_fwd[doc]) { + block_buf.push_back({e.inner_dim, doc_off, e.score}); + } + } + if (block_buf.empty()) + continue; + std::sort(block_buf.begin(), block_buf.end(), [](const BlockEntry& a, const BlockEntry& b) { + return a.inner_dim < b.inner_dim || (a.inner_dim == b.inner_dim && a.doc_offset < b.doc_offset); + }); + total_entries += block_buf.size(); + total_terms++; + for (size_t i = 1; i < block_buf.size(); ++i) { + if (block_buf[i].inner_dim != block_buf[i - 1].inner_dim) { + total_terms++; + } + } + } + + fwd_block_term_offsets_.resize(n_subblocks_ + 1); + 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); + + uint32_t term_pos = 0; + uint32_t entry_pos = 0; + + for (uint32_t sb = 0; sb < n_subblocks_; ++sb) { + fwd_block_term_offsets_[sb] = term_pos; + block_buf.clear(); + const uint32_t doc_start = sb * kSubblockSize; + const uint32_t doc_end = std::min(doc_start + kSubblockSize, static_cast(n_rows_internal_)); + for (uint32_t doc = doc_start; doc < doc_end; ++doc) { + const uint8_t doc_off = static_cast(doc - doc_start); + for (const auto& e : per_doc_fwd[doc]) { + block_buf.push_back({e.inner_dim, doc_off, e.score}); + } + per_doc_fwd[doc].clear(); + per_doc_fwd[doc].shrink_to_fit(); + } + if (block_buf.empty()) + continue; + std::sort(block_buf.begin(), block_buf.end(), [](const BlockEntry& a, const BlockEntry& b) { + return a.inner_dim < b.inner_dim || (a.inner_dim == b.inner_dim && a.doc_offset < b.doc_offset); + }); + + fwd_term_ids_[term_pos] = block_buf[0].inner_dim; + fwd_term_entry_offsets_[term_pos] = entry_pos; + + for (size_t i = 0; i < block_buf.size(); ++i) { + if (i > 0 && block_buf[i].inner_dim != block_buf[i - 1].inner_dim) { + term_pos++; + fwd_term_ids_[term_pos] = block_buf[i].inner_dim; + fwd_term_entry_offsets_[term_pos] = entry_pos; + } + fwd_doc_offsets_[entry_pos] = block_buf[i].doc_offset; + fwd_scores_[entry_pos] = block_buf[i].score; + entry_pos++; + } + term_pos++; + } + fwd_block_term_offsets_[n_subblocks_] = term_pos; + fwd_term_entry_offsets_[total_terms] = entry_pos; + } + } + + // ======================================================================== + // DSP Search + // ======================================================================== + template + void + search_dsp(const std::vector>& q_vec, DspHeap& heap, size_t heap_capacity, + DocIdFilter& filter, const DocValueComputer& computer, DspSearchMode mode, float mu, float eta, + int gamma, bool kth_init = true, float kth_alpha = 1.0f) const { + // ---- Step 0: Prepare sorted query ---- + struct QueryTerm { + uint32_t inner_dim = 0; + float weight = 0.0f; + uint8_t u8_weight = 0; + }; + std::vector query(q_vec.size()); + for (size_t i = 0; i < q_vec.size(); ++i) { + query[i].inner_dim = static_cast(q_vec[i].first); + query[i].weight = static_cast(q_vec[i].second); + } + 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; + for (const auto& qt : query) { + S += qt.weight * max_score_in_dim_spans_[qt.inner_dim]; + } + if (S <= 0.0f) + return; + + const float inv_S = 255.0f / S; + for (auto& qt : query) { + float w = qt.weight * max_score_in_dim_spans_[qt.inner_dim] * inv_S; + uint8_t u8w = static_cast(std::min(255.0f, std::max(1.0f, std::ceil(w)))); + qt.u8_weight = u8w; + } + const float score_scale = 65025.0f / S; + + // ---- Step 2: Initialize thresholds from kth scores ---- + const bool has_filter = !filter.empty(); + bool bootstrap_mode = has_filter; + + float float_threshold = 0.0f; +#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]; + uint8_t kth_u8 = bm.kth[kth_bucket]; + if (kth_u8 == 0) + continue; + float kth_float = kth_u8 / 255.0f * max_score_in_dim_spans_[qt.inner_dim]; + float term_thresh = qt.weight * kth_float; + float_threshold = std::max(float_threshold, term_thresh); + } + 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_); + 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); + 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 ---- + 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]; + } + } + + const float theta = float_threshold; + float mu_threshold = (mu > 0.0f) ? theta / mu : theta; + float eta_threshold = (eta > 0.0f) ? theta / eta : theta; + // 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 use_asc_survivor_guard = mu < 1.0f; + + surviving_spb.reserve(n_superblocks_); + + auto mark_alive = [&](uint32_t spb) { + if (!spb_alive[spb]) { + surviving_spb.push_back(spb); + spb_alive[spb] = 1; + } + }; + + auto add_top_gamma = [&](int g_val, float min_ub, bool inclusive = false) { + uint32_t g = static_cast(std::min(g_val, static_cast(n_superblocks_))); + std::vector eligible; + eligible.reserve(n_superblocks_); + for (uint32_t spb = 0; spb < n_superblocks_; ++spb) { + if (inclusive ? (superblock_ub[spb] >= min_ub) : (superblock_ub[spb] > min_ub)) { + eligible.push_back(spb); + } + } + if (eligible.size() <= g) { + for (uint32_t spb : eligible) { + mark_alive(spb); + } + } else { + std::nth_element(eligible.begin(), eligible.begin() + g, eligible.end(), + [&](uint32_t a, uint32_t b) { return superblock_ub[a] > superblock_ub[b]; }); + for (uint32_t i = 0; i < g; ++i) { + mark_alive(eligible[i]); + } + } + }; + + // ---- Mode-driven superblock selection ---- + // DspSearchMode enum values: + // DSP = 0, LSP0 = 1, LSP1 = 2, LSP2 = 3 + auto select_superblocks_by_mode = [&]() { + switch (mode) { + 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 || + (use_asc_survivor_guard && superblock_asc[spb] > eta_threshold)) { + mark_alive(spb); + } + } + if (gamma > 0) { + add_top_gamma(gamma, 0.0f); + } + break; + } + case DspSearchMode::LSP0: { + if (gamma <= 0) { + // fallback to DSP behavior + for (uint32_t spb = 0; spb < n_superblocks_; ++spb) { + if (superblock_ub[spb] > mu_threshold || + (use_asc_survivor_guard && superblock_asc[spb] > eta_threshold)) { + mark_alive(spb); + } + } + break; + } + add_top_gamma(gamma, float_threshold, true); + break; + } + case DspSearchMode::LSP1: { + if (gamma <= 0) { + // fallback to DSP behavior + for (uint32_t spb = 0; spb < n_superblocks_; ++spb) { + if (superblock_ub[spb] > mu_threshold || + (use_asc_survivor_guard && superblock_asc[spb] > eta_threshold)) { + mark_alive(spb); + } + } + break; + } + add_top_gamma(gamma, float_threshold, true); + for (uint32_t spb = 0; spb < n_superblocks_; ++spb) { + if (superblock_ub[spb] > mu_threshold) { + mark_alive(spb); + } + } + break; + } + case DspSearchMode::LSP2: { + if (gamma <= 0) { + // fallback to DSP behavior + for (uint32_t spb = 0; spb < n_superblocks_; ++spb) { + if (superblock_ub[spb] > mu_threshold || + (use_asc_survivor_guard && superblock_asc[spb] > eta_threshold)) { + mark_alive(spb); + } + } + break; + } + add_top_gamma(gamma, float_threshold, true); + for (uint32_t spb = 0; spb < n_superblocks_; ++spb) { + if (superblock_ub[spb] > mu_threshold || + (use_asc_survivor_guard && superblock_asc[spb] > eta_threshold)) { + mark_alive(spb); + } + } + break; + } + default: { + // Default: same as DSP mode + for (uint32_t spb = 0; spb < n_superblocks_; ++spb) { + if (superblock_ub[spb] > mu_threshold || + (use_asc_survivor_guard && superblock_asc[spb] > eta_threshold)) { + mark_alive(spb); + } + } + if (gamma > 0) { + add_top_gamma(gamma, 0.0f); + } + break; + } + } + }; + +#ifdef SEEK_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) { +#ifdef SEEK_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()) { + 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) { + const uint32_t sb = bm.block_ids[i]; + if (!spb_in_batch[sb / kStride]) + continue; + uint32_t prod = u16w * bm.max_scores[i]; + uint32_t sum = static_cast(block_ub[sb]) + prod; + 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; + }; + + // ---- Collect candidates from superblocks and sort by UB descending ---- + auto collect_and_sort = [&](const std::vector& spbs) -> std::vector { + std::vector cands; + cands.reserve(spbs.size() * kStride / 4); + uint16_t local_max_ub = 0; +#ifdef SEEK_INSTRUMENTATION + uint64_t local_saturated_ubs = 0; +#endif + for (uint32_t spb : spbs) { + const uint32_t sb_start = spb * kStride; + 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 SEEK_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 SEEK_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; + std::vector cnt(rng + 1, 0); + for (uint32_t sb : cands) cnt[block_ub[sb] - u16_block_threshold - 1]++; + uint32_t p = 0; + for (int b = static_cast(rng); b >= 0; --b) { + uint32_t c = cnt[b]; + cnt[b] = p; + p += c; + } + std::vector sorted(cands.size()); + for (uint32_t sb : cands) sorted[cnt[block_ub[sb] - u16_block_threshold - 1]++] = sb; + return sorted; + }; + + // ---- Score a list of sorted blocks, updating heap and thresholds ---- + float scores[kSubblockSize]; + + auto score_blocks = [&](const std::vector& sorted_blocks) -> bool { + bool bootstrap_completed = false; +#ifdef SEEK_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 SEEK_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]; + __builtin_prefetch(&fwd_term_ids_[next_start], 0, 1); + const uint32_t next_entry_start = fwd_term_entry_offsets_[next_start]; + __builtin_prefetch(&fwd_doc_offsets_[next_entry_start], 0, 0); + __builtin_prefetch(&fwd_scores_[next_entry_start], 0, 0); + } + + std::memset(scores, 0, sizeof(scores)); + size_t qi = 0; + uint32_t bi = block_term_start; + 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 SEEK_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); + while (lo < hi) { + const uint32_t mid = lo + (hi - lo) / 2; + if (fwd_term_ids_[mid] < q_dim) + lo = mid + 1; + else + hi = mid; + } + bi = lo; + } else { + ++bi; + } + } else { + accumulate_match(qi, bi); + ++qi; + ++bi; + } + } + } + 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 SEEK_INSTRUMENTATION + local_docs++; +#endif + heap.Push(scores[i], doc_id); + if (heap.Full()) { + float new_thresh = heap.Results().front().first; + if (new_thresh > float_threshold) { + float_threshold = new_thresh; + float_block_threshold = (eta > 0.0f) ? float_threshold / eta : float_threshold; + u16_block_threshold = + static_cast(std::min(65535.0f, float_block_threshold * score_scale)); + } + if (bootstrap_mode) { + bootstrap_mode = false; + bootstrap_completed = true; + mu_threshold = (mu > 0.0f) ? float_threshold / mu : float_threshold; + eta_threshold = (eta > 0.0f) ? float_threshold / eta : float_threshold; + } + } + } + } + } +#ifdef SEEK_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; + }; + + // ==================================================================== + // Two-phase filtered bootstrap + // ==================================================================== + if (has_filter) { + std::vector spb_by_ub; + spb_by_ub.reserve(n_superblocks_); + for (uint32_t spb = 0; spb < n_superblocks_; ++spb) { + if (superblock_ub[spb] > 0.0f) + spb_by_ub.push_back(spb); + } + std::sort(spb_by_ub.begin(), spb_by_ub.end(), + [&](uint32_t a, uint32_t b) { return superblock_ub[a] > superblock_ub[b]; }); + + std::vector spb_processed(n_superblocks_, 0); + + const uint32_t batch_sizes[] = {64, 256}; + const int n_batches = sizeof(batch_sizes) / sizeof(batch_sizes[0]); + uint32_t cursor = 0; + + for (int batch_idx = 0; batch_idx <= n_batches && cursor < spb_by_ub.size(); ++batch_idx) { + uint32_t batch_end; + if (batch_idx < n_batches) { + batch_end = std::min(static_cast(spb_by_ub.size()), batch_sizes[batch_idx]); + } else { + batch_end = static_cast(spb_by_ub.size()); + } + if (batch_end <= cursor) + continue; + + std::vector batch_spbs; + batch_spbs.reserve(batch_end - cursor); + for (uint32_t i = cursor; i < batch_end; ++i) { + mark_alive(spb_by_ub[i]); + spb_processed[spb_by_ub[i]] = 1; + batch_spbs.push_back(spb_by_ub[i]); + } + cursor = batch_end; + + compute_block_ubs(batch_spbs); + auto sorted = collect_and_sort(batch_spbs); + if (sorted.empty()) + continue; + bool done = score_blocks(sorted); + if (done) + break; + } + + if (!bootstrap_mode) { + surviving_spb.clear(); + std::fill(spb_alive.begin(), spb_alive.end(), 0); + select_superblocks_by_mode(); + + std::vector new_spbs; + new_spbs.reserve(surviving_spb.size()); + for (uint32_t spb : surviving_spb) { + if (!spb_processed[spb]) + new_spbs.push_back(spb); + } + + if (!new_spbs.empty()) { + compute_block_ubs(new_spbs); + auto sorted = collect_and_sort(new_spbs); + if (!sorted.empty()) + score_blocks(sorted); + } + } + } else { + // ==================================================================== + // Unfiltered path: original single-pass logic + // ==================================================================== + select_superblocks_by_mode(); + if (surviving_spb.empty()) + return; + + compute_block_ubs(surviving_spb); + auto sorted = collect_and_sort(surviving_spb); + if (sorted.empty()) + return; + score_blocks(sorted); + } + } +}; + +} // namespace knowhere::sparse + +#endif // SPARSE_DSP_INDEX_H diff --git a/src/index/sparse/sparse_index_node.cc b/src/index/sparse/sparse_index_node.cc index 62e3d6e02..d2224c8fb 100644 --- a/src/index/sparse/sparse_index_node.cc +++ b/src/index/sparse/sparse_index_node.cc @@ -25,11 +25,14 @@ #include "index/sparse/growable_inverted_index.h" #include "index/sparse/inverted_index.h" #include "index/sparse/sindi_inverted_index.h" +#include "index/sparse/sparse_dsp_config.h" +#include "index/sparse/sparse_dsp_index.h" #include "index/sparse/sparse_index_config.h" #include "io/file_io.h" #include "io/memory_io.h" #include "knowhere/comp/index_param.h" #include "knowhere/config.h" +#include "knowhere/context.h" #include "knowhere/dataset.h" #include "knowhere/expected.h" #include "knowhere/index/index_factory.h" @@ -1084,6 +1087,469 @@ class SparseInvertedIndexNodeCC : public SparseInvertedIndexNode { mutable std::vector> raw_data_ = {}; }; // class SparseInvertedIndexNodeCC +// ============================================================================ +// SparseDspIndexNode: standalone DSP (Dynamic Superblock Pruning) sparse index. +// Does not share code with SparseInvertedIndexNode — fully self-contained. +// ============================================================================ +template +class SparseDspIndexNode : public IndexNode { + static_assert(std::is_same_v, "SparseDspIndexNode only support sparse_u32_f32"); + using value_type = typename T::ValueType; + + public: + explicit SparseDspIndexNode(const int32_t& version, const Object& /*object*/) + : search_pool_(ThreadPool::GetGlobalSearchThreadPool()), + build_pool_(ThreadPool::GetGlobalBuildThreadPool()), + index_version_(version) { + } + + ~SparseDspIndexNode() override { + // index_ must be destroyed before mmap_guard_ since index_ may hold spans into mmap'd memory + delete index_; + index_ = nullptr; + mmap_guard_.reset(); + } + + Status + Train(const DataSetPtr dataset, std::shared_ptr config, bool use_knowhere_build_pool) override { + auto& cfg = static_cast(*config); + if (!IsMetricType(cfg.metric_type.value(), metric::IP) && + !IsMetricType(cfg.metric_type.value(), metric::BM25)) { + LOG_KNOWHERE_ERROR_ << Type() << " only support metric_type IP or BM25"; + return Status::invalid_metric_type; + } + auto index = CreateDspIndex(cfg); + if (!index.has_value()) { + return index.error(); + } + auto* idx = index.value(); + idx->Train(static_cast*>(dataset->GetTensor()), dataset->GetRows()); + if (index_) { + LOG_KNOWHERE_WARNING_ << Type() << " already created, deleting old"; + delete index_; + } + index_ = idx; + return Status::success; + } + + Status + Add(const DataSetPtr dataset, std::shared_ptr config, bool use_knowhere_build_pool) override { + if (!index_) { + LOG_KNOWHERE_ERROR_ << "Could not add data to empty " << Type(); + return Status::empty_index; + } + 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(static_cast*>(dataset->GetTensor()), + 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(); + if (!tryObj.hasValue()) { + LOG_KNOWHERE_WARNING_ << "failed to add data to index " << Type() << ": " << tryObj.exception().what(); + return Status::sparse_inner_error; + } + return tryObj.value(); + } + + [[nodiscard]] expected + Search(const DataSetPtr dataset, std::unique_ptr config, const BitsetView& bitset, + milvus::OpContext* op_context) const override { + if (!index_) { + LOG_KNOWHERE_ERROR_ << "Could not search empty " << Type(); + return expected::Err(Status::empty_index, "index not loaded"); + } + auto& cfg = static_cast(*config); + auto computer_or = index_->GetDocValueComputer(cfg); + if (!computer_or.has_value()) { + return expected::Err(computer_or.error(), computer_or.what()); + } + auto computer = computer_or.value(); + auto approx_params = ExtractDspParams(cfg); + + auto queries = static_cast*>(dataset->GetTensor()); + auto nq = dataset->GetRows(); + auto k = cfg.k.value(); + auto p_id = std::make_unique(nq * k); + auto p_dist = std::make_unique(nq * k); + + std::vector> futs; + futs.reserve(nq); + for (int64_t idx = 0; idx < nq; ++idx) { + futs.emplace_back(search_pool_->push([&, idx = idx, p_id = p_id.get(), p_dist = p_dist.get()]() { + knowhere::checkCancellation(op_context); + index_->Search(queries[idx], k, p_dist + idx * k, p_id + idx * k, bitset, computer, approx_params); + })); + } + WaitAllSuccess(futs); + return GenResultDataSet(nq, k, p_id.release(), p_dist.release()); + } + + [[nodiscard]] expected> + AnnIterator(const DataSetPtr dataset, std::unique_ptr config, const BitsetView& bitset, + bool use_knowhere_search_pool, milvus::OpContext* op_context) const override { + if (!index_) { + return expected>::Err(Status::empty_index, "index not loaded"); + } + auto nq = dataset->GetRows(); + auto& cfg = static_cast(*config); + auto computer_or = index_->GetDocValueComputer(cfg); + if (!computer_or.has_value()) { + return expected>::Err(computer_or.error(), computer_or.what()); + } + auto computer = computer_or.value(); + auto drop_ratio_search = cfg.drop_ratio_search.value_or(0.0f); + + auto vec = std::vector>(nq, nullptr); + try { + for (int i = 0; i < nq; ++i) { + auto compute_dist_func = [=]() -> std::vector { + auto queries = static_cast*>(dataset->GetTensor()); + std::vector distances = + index_->GetAllDistances(queries[i], drop_ratio_search, bitset, computer); + std::vector distances_ids; + distances_ids.reserve(distances.size() * 0.3); + for (size_t i = 0; i < distances.size(); i++) { + if (distances[i] != 0) { + distances_ids.emplace_back((int64_t)i, distances[i]); + } + } + return distances_ids; + }; + vec[i] = + std::make_shared(compute_dist_func, true, use_knowhere_search_pool); + } + } catch (const std::exception& e) { + return expected>::Err(Status::sparse_inner_error, e.what()); + } + return vec; + } + + [[nodiscard]] expected + GetVectorByIds(const DataSetPtr dataset, milvus::OpContext* op_context) const override { + return expected::Err(Status::not_implemented, "GetVectorByIds not implemented"); + } + + static bool + StaticHasRawData(const knowhere::BaseConfig& /*config*/, const IndexVersion& /*version*/) { + return false; + } + + [[nodiscard]] bool + HasRawData(const std::string& metric_type) const override { + return false; + } + + [[nodiscard]] expected + GetIndexMeta(std::unique_ptr cfg) const override { + return expected::Err(Status::not_implemented, "GetIndexMeta not supported"); + } + + Status + Serialize(BinarySet& binset) const override { + if (!index_) { + LOG_KNOWHERE_ERROR_ << "Could not serialize empty " << Type(); + return Status::empty_index; + } + MemoryIOWriter writer; + RETURN_IF_ERROR(index_->Serialize(writer)); + std::shared_ptr data(writer.data()); + binset.Append(Type(), data, writer.tellg()); + return Status::success; + } + + Status + Deserialize(const BinarySet& binset, std::shared_ptr config) override { + if (index_) { + LOG_KNOWHERE_WARNING_ << Type() << " already created, deleting old"; + delete index_; // destroy index before unmapping memory it may reference + index_ = nullptr; + mmap_guard_.reset(); + } + auto binary = binset.GetByName(Type()); + if (binary == nullptr) { + LOG_KNOWHERE_ERROR_ << "Invalid BinarySet."; + return Status::invalid_binary_set; + } + MemoryIOReader reader(binary->data.get(), binary->size); + auto index = CreateDspIndex(static_cast(*config)); + if (!index.has_value()) { + return index.error(); + } + index_ = index.value(); + binary_ = binary; + return index_->Deserialize(reader); + } + + Status + DeserializeFromFile(const std::string& filename, std::shared_ptr config) override { + if (index_) { + LOG_KNOWHERE_WARNING_ << Type() << " already created, deleting old"; + delete index_; // destroy index before unmapping memory it may reference + index_ = nullptr; + mmap_guard_.reset(); + } + auto& base_cfg = static_cast(*config); + auto index = CreateDspIndex(static_cast(*config)); + if (!index.has_value()) { + return index.error(); + } + index_ = index.value(); + + auto reader = knowhere::FileReader(filename); + size_t map_size = reader.size(); + int map_flags = MAP_SHARED; +#ifdef MAP_POPULATE + if (base_cfg.enable_mmap_pop.has_value() && base_cfg.enable_mmap_pop.value()) { + map_flags |= MAP_POPULATE; + } +#endif + void* mapped_memory = mmap(nullptr, map_size, PROT_READ, map_flags, reader.descriptor(), 0); + if (mapped_memory == MAP_FAILED) { + LOG_KNOWHERE_ERROR_ << "Failed to mmap file " << filename << ": " << strerror(errno); + return Status::disk_file_error; + } + mmap_guard_ = std::make_unique(map_size, filename, mapped_memory); + MemoryIOReader map_reader(reinterpret_cast(mapped_memory), map_size); + return index_->Deserialize(map_reader); + } + + static std::unique_ptr + StaticCreateConfig() { + return std::make_unique(); + } + + [[nodiscard]] std::unique_ptr + CreateConfig() const override { + return StaticCreateConfig(); + } + + [[nodiscard]] std::string + Type() const override { + return knowhere::IndexEnum::INDEX_SPARSE_DSP; + } + + [[nodiscard]] int64_t + Dim() const override { + return index_ ? index_->n_cols() : 0; + } + + [[nodiscard]] int64_t + Size() const override { + return index_ ? index_->size() : 0; + } + + [[nodiscard]] int64_t + Count() const override { + return index_ ? index_->n_rows() : 0; + } + + protected: + static sparse::DspSearchParams + ExtractDspParams(const SparseDspConfig& c) { + auto drop_ratio_search = c.drop_ratio_search.value_or(0.0f); + auto refine_factor = c.refine_factor.value_or(1); + if (drop_ratio_search == 0) { + refine_factor = 1; + } + return { + .refine_factor = refine_factor, + .drop_ratio_search = drop_ratio_search, + .dim_max_score_ratio = 1.0f, + .dsp_mode = static_cast(c.dsp_mode.value_or(0)), + .dsp_mu = c.dsp_mu.value_or(1.0f), + .dsp_eta = c.dsp_eta.value_or(1.0f), + .dsp_gamma = static_cast(c.dsp_gamma.value_or(0)), + .dsp_kth_init = c.dsp_kth_init.value_or(true), + .dsp_kth_alpha = c.dsp_kth_alpha.value_or(1.0f), + }; + } + + template + expected*> + CreateDspIndex(const SparseDspConfig& cfg) const { + if (IsMetricType(cfg.metric_type.value(), metric::BM25)) { + if (!cfg.bm25_k1.has_value() || !cfg.bm25_b.has_value() || !cfg.bm25_avgdl.has_value()) { + return expected*>::Err(Status::invalid_args, + "BM25 parameters k1, b, and avgdl must be set"); + } + auto idx = new sparse::DspIndex(sparse::SparseMetricType::METRIC_BM25); + idx->SetBM25Params(cfg.bm25_k1.value(), cfg.bm25_b.value(), std::max(cfg.bm25_avgdl.value(), 1.0f)); + return idx; + } + return new sparse::DspIndex(sparse::SparseMetricType::METRIC_IP); + } + + struct MmapGuard { + size_t map_size; + std::string filename; + void* map_addr; + + MmapGuard(size_t size, const std::string& fname, void* addr) : map_size(size), filename(fname), map_addr(addr) { + } + + ~MmapGuard() { + if (munmap(map_addr, map_size) != 0) { + LOG_KNOWHERE_ERROR_ << "Failed to munmap file " << filename << ": " << strerror(errno); + } + } + }; + + sparse::DspIndexBase* index_{nullptr}; + std::shared_ptr search_pool_; + std::shared_ptr build_pool_; + const int32_t index_version_; + BinaryPtr binary_{nullptr}; + std::unique_ptr mmap_guard_{nullptr}; +}; // class SparseDspIndexNode + +// Concurrent version of SparseDspIndexNode +template +class SparseDspIndexNodeCC : public SparseDspIndexNode { + static_assert(std::is_same_v, "SparseDspIndexNodeCC only support sparse_u32_f32"); + using value_type = typename T::ValueType; + + public: + explicit SparseDspIndexNodeCC(const int32_t& version, const Object& object) + : SparseDspIndexNode(version, object) { + } + + Status + Add(const DataSetPtr dataset, std::shared_ptr config, bool use_knowhere_build_pool) override { + std::unique_lock lock(mutex_); + uint64_t task_id = next_task_id_++; + add_tasks_.push(task_id); + cv_.wait(lock, [this, task_id]() { return current_task_id_ == task_id && active_readers_ == 0; }); + + auto res = SparseDspIndexNode::Add(dataset, config, use_knowhere_build_pool); + + auto cfg = static_cast(*config); + if (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); + } + + add_tasks_.pop(); + current_task_id_++; + lock.unlock(); + cv_.notify_all(); + return res; + } + + expected + Search(const DataSetPtr dataset, std::unique_ptr cfg, const BitsetView& bitset, + milvus::OpContext* op_context) const override { + ReadPermission permission(*this); + return SparseDspIndexNode::Search(dataset, std::move(cfg), bitset, op_context); + } + + expected> + AnnIterator(const DataSetPtr dataset, std::unique_ptr cfg, const BitsetView& bitset, + bool use_knowhere_search_pool, milvus::OpContext* op_context) const override { + ReadPermission permission(*this); + static_cast(*cfg).drop_ratio_search = 0.0f; + return SparseDspIndexNode::AnnIterator(dataset, std::move(cfg), bitset, use_knowhere_search_pool, + op_context); + } + + expected + RangeSearch(const DataSetPtr dataset, std::unique_ptr cfg, const BitsetView& bitset, + milvus::OpContext* op_context) const override { + ReadPermission permission(*this); + return SparseDspIndexNode::RangeSearch(dataset, std::move(cfg), bitset, op_context); + } + + int64_t + Dim() const override { + ReadPermission p(*this); + return SparseDspIndexNode::Dim(); + } + int64_t + Size() const override { + ReadPermission p(*this); + return SparseDspIndexNode::Size(); + } + int64_t + Count() const override { + ReadPermission p(*this); + return SparseDspIndexNode::Count(); + } + + std::string + Type() const override { + return knowhere::IndexEnum::INDEX_SPARSE_DSP_CC; + } + + expected + GetVectorByIds(const DataSetPtr dataset, milvus::OpContext* op_context) const override { + ReadPermission permission(*this); + if (raw_data_.empty()) { + return expected::Err(Status::invalid_args, "GetVectorByIds failed: raw data is empty"); + } + auto rows = dataset->GetRows(); + auto ids = dataset->GetIds(); + auto data = std::make_unique[]>(rows); + int64_t dim = 0; + try { + for (int64_t i = 0; i < rows; ++i) { + data[i] = raw_data_[ids[i]]; + dim = std::max(dim, data[i].dim()); + } + } catch (std::exception& e) { + return expected::Err(Status::invalid_args, "GetVectorByIds failed: " + std::string(e.what())); + } + auto res = GenResultDataSet(rows, dim, data.release()); + res->SetIsSparse(true); + return res; + } + + static bool + StaticHasRawData(const knowhere::BaseConfig& config, const IndexVersion& version) { + return config.metric_type.has_value() && IsMetricType(config.metric_type.value(), metric::IP); + } + + [[nodiscard]] bool + HasRawData(const std::string& metric_type) const override { + return IsMetricType(metric_type, metric::IP); + } + + private: + struct ReadPermission { + ReadPermission(const SparseDspIndexNodeCC& node) : node_(node) { + std::unique_lock lock(node_.mutex_); + uint64_t task_id = node_.next_task_id_++; + if (!node_.add_tasks_.empty() && task_id > node_.add_tasks_.front()) { + node_.cv_.wait( + lock, [this, task_id]() { return node_.add_tasks_.empty() || task_id < node_.add_tasks_.front(); }); + } + node_.active_readers_++; + } + ~ReadPermission() { + std::unique_lock lock(node_.mutex_); + node_.active_readers_--; + node_.current_task_id_++; + node_.cv_.notify_all(); + } + const SparseDspIndexNodeCC& node_; + }; + + mutable std::mutex mutex_; + mutable std::condition_variable cv_; + mutable int64_t active_readers_ = 0; + mutable std::queue add_tasks_; + mutable uint64_t next_task_id_ = 0; + mutable uint64_t current_task_id_ = 0; + mutable std::vector> raw_data_ = {}; +}; // class SparseDspIndexNodeCC + KNOWHERE_SIMPLE_REGISTER_SPARSE_FLOAT_GLOBAL(SPARSE_INVERTED_INDEX, SparseInvertedIndexNode, knowhere::feature::MMAP, /*use_wand=*/false) KNOWHERE_SIMPLE_REGISTER_SPARSE_FLOAT_GLOBAL(SPARSE_WAND, SparseInvertedIndexNode, knowhere::feature::MMAP, @@ -1093,4 +1559,6 @@ KNOWHERE_SIMPLE_REGISTER_SPARSE_FLOAT_GLOBAL(SPARSE_INVERTED_INDEX_CC, SparseInv /*use_wand=*/false) KNOWHERE_SIMPLE_REGISTER_SPARSE_FLOAT_GLOBAL(SPARSE_WAND_CC, SparseInvertedIndexNodeCC, knowhere::feature::MMAP, /*use_wand=*/true) +KNOWHERE_SIMPLE_REGISTER_SPARSE_FLOAT_GLOBAL(SPARSE_DSP, SparseDspIndexNode, knowhere::feature::MMAP) +KNOWHERE_SIMPLE_REGISTER_SPARSE_FLOAT_GLOBAL(SPARSE_DSP_CC, SparseDspIndexNodeCC, knowhere::feature::MMAP) } // namespace knowhere diff --git a/src/simd/sparse_simd.h b/src/simd/sparse_simd.h index 8b010e346..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" @@ -11,11 +14,179 @@ namespace knowhere::sparse { #if defined(__x86_64__) || defined(_M_X64) +// ---- AVX512 BW: Block UB threshold scan ---- +// Stride-specific specializations (no loop counter overhead) +bool +scan_block_ub_any_above_avx512_32(const uint16_t* block_ub, uint16_t threshold); +bool +scan_block_ub_any_above_avx512_64(const uint16_t* block_ub, uint16_t threshold); +// Generic loop fallback for non-standard sizes (n must be a multiple of 32). +// In the current DSP code path, callers use kStride = 64, so this precondition holds. +bool +scan_block_ub_any_above_avx512_generic(const uint16_t* block_ub, uint16_t threshold, uint32_t n); +// Legacy entry point — dispatches internally to stride-specific or generic +bool +scan_block_ub_any_above_avx512(const uint16_t* block_ub, uint16_t threshold, uint32_t n); + +// ---- AVX512 BW: Block max UB accumulation ---- +// Stride-specific specializations (no loop counter overhead) +void +accumulate_block_ub_avx512_32(uint16_t* __restrict ub, const uint8_t* __restrict block_max, uint16_t query_weight); +void +accumulate_block_ub_avx512_64(uint16_t* __restrict ub, const uint8_t* __restrict block_max, uint16_t query_weight); +// Generic loop fallback for non-standard sizes (n must be a multiple of 32). +// In the current DSP code path, callers use kStride = 64, so this precondition holds. +void +accumulate_block_ub_avx512_generic(uint16_t* __restrict ub, const uint8_t* __restrict block_max, uint16_t query_weight, + uint32_t n); +// Legacy entry point — dispatches internally to stride-specific or generic +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) { + for (uint32_t i = 0; i < n; ++i) { + if (block_ub[i] > threshold) { + return true; + } + } + return false; +} + +// Dispatch for block UB scan with runtime CPU detection. +// Routes to stride-specific AVX512 kernels for n=32/64 (the DSP hot path). +inline bool +scan_block_ub_any_above_dispatch(const uint16_t* block_ub, uint16_t threshold, uint32_t n) { +#if defined(__x86_64__) || defined(_M_X64) + if (faiss::cppcontrib::knowhere::InstructionSet::GetInstance().AVX512BW()) { + if (n == 64) { + return scan_block_ub_any_above_avx512_64(block_ub, threshold); + } + if (n == 32) { + return scan_block_ub_any_above_avx512_32(block_ub, threshold); + } + return scan_block_ub_any_above_avx512_generic(block_ub, threshold, n); + } +#endif + return scan_block_ub_any_above_scalar(block_ub, threshold, n); +} + +// Scalar fallback for u8 block max to u16 UB accumulation +inline void +accumulate_block_ub_scalar(uint16_t* ub, const uint8_t* block_max, uint16_t query_weight, uint32_t n) { + for (uint32_t i = 0; i < n; ++i) { + uint32_t prod = static_cast(query_weight) * block_max[i]; + uint32_t sum = static_cast(ub[i]) + prod; + ub[i] = static_cast(sum < 65535u ? sum : 65535u); + } +} + +// Dispatch for u8 block max to u16 UB accumulation. +// Routes to stride-specific AVX512 kernels for n=32/64 (the DSP hot path). +inline void +accumulate_block_ub_dispatch(uint16_t* __restrict ub, const uint8_t* __restrict block_max, uint16_t query_weight, + uint32_t n) { +#if defined(__x86_64__) || defined(_M_X64) + if (faiss::cppcontrib::knowhere::InstructionSet::GetInstance().AVX512BW()) { + if (n == 64) { + accumulate_block_ub_avx512_64(ub, block_max, query_weight); + return; + } + if (n == 32) { + accumulate_block_ub_avx512_32(ub, block_max, query_weight); + return; + } + accumulate_block_ub_avx512_generic(ub, block_max, query_weight, n); + return; + } +#endif + 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 612d071ed..ee2a7058b 100644 --- a/src/simd/sparse_simd_avx512.cc +++ b/src/simd/sparse_simd_avx512.cc @@ -9,15 +9,133 @@ // 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. -// This file is compiled with -mavx512f flag to enable AVX512 intrinsics +// This file is compiled with -mavx512f -mavx512cd -mavx512bw flags to enable AVX512 intrinsics // Runtime CPU detection ensures it's only called on CPUs with AVX512 support #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); + while (lo < hi) { + const uint32_t mid = lo + (hi - lo) / 2; + ++probes; + if (terms[mid * 16 + 15] < target) { + lo = mid + 1; + } else { + hi = mid; + } + } + chunk_cursor = lo; + } + } + + 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 +// ============================================================================ +// Check if any u16 value in block_ub exceeds threshold. +// DSP currently uses kStride = 64, so the 64-element specialization is the main hot path. + +// n == 32: single AVX-512 register (32 u16 lanes) +bool +scan_block_ub_any_above_avx512_32(const uint16_t* block_ub, uint16_t threshold) { + const __m512i thresh_v = _mm512_set1_epi16(static_cast(threshold)); + __m512i v0 = _mm512_loadu_si512(reinterpret_cast(block_ub)); + return _mm512_cmp_epu16_mask(v0, thresh_v, _MM_CMPINT_NLE) != 0; +} + +// n == 64: two AVX-512 registers (DSP kStride = 64, the primary hot path) +bool +scan_block_ub_any_above_avx512_64(const uint16_t* block_ub, uint16_t threshold) { + const __m512i thresh_v = _mm512_set1_epi16(static_cast(threshold)); + __m512i v0 = _mm512_loadu_si512(reinterpret_cast(block_ub)); + __m512i v1 = _mm512_loadu_si512(reinterpret_cast(block_ub + 32)); + __mmask32 mask = + _mm512_cmp_epu16_mask(v0, thresh_v, _MM_CMPINT_NLE) | _mm512_cmp_epu16_mask(v1, thresh_v, _MM_CMPINT_NLE); + return mask != 0; +} + +// Generic: loop over n elements in 32-element chunks. n must be a multiple of 32. +bool +scan_block_ub_any_above_avx512_generic(const uint16_t* block_ub, uint16_t threshold, uint32_t n) { + assert(n % 32 == 0 && "n must be a multiple of 32 for AVX-512 u16 processing"); + const __m512i thresh_v = _mm512_set1_epi16(static_cast(threshold)); + for (uint32_t i = 0; i < n; i += 32) { + __m512i v = _mm512_loadu_si512(reinterpret_cast(block_ub + i)); + if (_mm512_cmp_epu16_mask(v, thresh_v, _MM_CMPINT_NLE) != 0) { + return true; + } + } + return false; +} + +// Kept for backward compatibility — routes to stride-64 specialization +bool +scan_block_ub_any_above_avx512(const uint16_t* block_ub, uint16_t threshold, uint32_t n) { + if (n == 64) { + return scan_block_ub_any_above_avx512_64(block_ub, threshold); + } + if (n == 32) { + return scan_block_ub_any_above_avx512_32(block_ub, threshold); + } + return scan_block_ub_any_above_avx512_generic(block_ub, threshold, n); +} + // ============================================================================ // AVX512 SIMD Implementation (16-wide vectorization with hardware scatter) // ============================================================================ @@ -80,4 +198,96 @@ accumulate_posting_list_ip_avx512(const uint32_t* doc_ids, const float* doc_vals } } +// ============================================================================ +// AVX512 BW: Block Max UB Accumulation — Stride-Specific Specializations +// ============================================================================ +// Accumulates u8 block max scores into u16 upper bound array with saturating add: +// ub[i] = sat_add_u16(ub[i], query_weight * block_max[i]) +// DSP currently uses kStride = 64, so the 64-element specialization is the main hot path. + +// Helper macro for one 32-element iteration (load u8, zero-extend, multiply, saturating add) +#define ACCUMULATE_BLOCK_UB_32(ub_ptr, bm_ptr, qw_vec) \ + do { \ + __m256i bm8 = _mm256_loadu_si256(reinterpret_cast(bm_ptr)); \ + __m512i bm16 = _mm512_cvtepu8_epi16(bm8); \ + __m512i prod = _mm512_mullo_epi16(bm16, qw_vec); \ + __m512i cur = _mm512_loadu_si512(reinterpret_cast(ub_ptr)); \ + cur = _mm512_adds_epu16(cur, prod); \ + _mm512_storeu_si512(reinterpret_cast<__m512i*>(ub_ptr), cur); \ + } while (0) + +// n == 32: single iteration, no loop +void +accumulate_block_ub_avx512_32(uint16_t* __restrict ub, const uint8_t* __restrict block_max, uint16_t query_weight) { + const __m512i qw = _mm512_set1_epi16(static_cast(query_weight)); + ACCUMULATE_BLOCK_UB_32(ub, block_max, qw); +} + +// n == 64: two iterations fully unrolled (DSP kStride = 64, the primary hot path) +void +accumulate_block_ub_avx512_64(uint16_t* __restrict ub, const uint8_t* __restrict block_max, uint16_t query_weight) { + const __m512i qw = _mm512_set1_epi16(static_cast(query_weight)); + ACCUMULATE_BLOCK_UB_32(ub, block_max, qw); + ACCUMULATE_BLOCK_UB_32(ub + 32, block_max + 32, qw); +} + +#undef ACCUMULATE_BLOCK_UB_32 + +// Generic: loop over n elements in 32-element chunks. n must be a multiple of 32. +void +accumulate_block_ub_avx512_generic(uint16_t* __restrict ub, const uint8_t* __restrict block_max, uint16_t query_weight, + uint32_t n) { + assert(n % 32 == 0 && "n must be a multiple of 32 for AVX-512 u16 processing"); + const __m512i qw = _mm512_set1_epi16(static_cast(query_weight)); + for (uint32_t i = 0; i < n; i += 32) { + __m256i bm8 = _mm256_loadu_si256(reinterpret_cast(block_max + i)); + __m512i bm16 = _mm512_cvtepu8_epi16(bm8); + __m512i prod = _mm512_mullo_epi16(bm16, qw); + __m512i cur = _mm512_loadu_si512(reinterpret_cast(ub + i)); + cur = _mm512_adds_epu16(cur, prod); + _mm512_storeu_si512(reinterpret_cast<__m512i*>(ub + i), cur); + } +} + +// Kept for backward compatibility — routes to stride-specific specializations +void +accumulate_block_ub_avx512(uint16_t* ub, const uint8_t* block_max, uint16_t query_weight, uint32_t n) { + if (n == 64) { + accumulate_block_ub_avx512_64(ub, block_max, query_weight); + return; + } + if (n == 32) { + accumulate_block_ub_avx512_32(ub, block_max, query_weight); + return; + } + 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 1e515f9f6..85fd47e00 100644 --- a/tests/ut/test_sparse.cc +++ b/tests/ut/test_sparse.cc @@ -9,13 +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" @@ -111,6 +117,12 @@ TEST_CASE("Test Mem Sparse Index With Float Vector", "[float metrics]") { return json; }; + auto sparse_dsp_gen = [base_gen, drop_ratio_search = drop_ratio_search]() { + knowhere::Json json = base_gen(); + json[knowhere::indexparam::DROP_RATIO_SEARCH] = drop_ratio_search; + return json; + }; + auto sparse_dataset_gen = [&](int nr, int dim, float sparsity) -> knowhere::DataSetPtr { if (metric == knowhere::metric::BM25) { return GenSparseDataSetWithMaxVal(nr, dim, sparsity, 256, true); @@ -161,6 +173,8 @@ 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 gt = knowhere::BruteForce::SearchSparse(train_ds, query_ds, conf, nullptr); check_distance_decreasing(*gt.value()); @@ -206,11 +220,194 @@ TEST_CASE("Test Mem Sparse Index With Float Vector", "[float metrics]") { } } + SECTION("Test DSP Params") { + // Build one DSP index, then search with different param combos to prove + // that eta and gamma actually affect pruning behavior. + auto gt = knowhere::BruteForce::SearchSparse(train_ds, query_ds, conf, nullptr); + REQUIRE(gt.has_value()); + + auto use_mmap = GENERATE(true, false); + auto tmp_file = "/tmp/knowhere_sparse_dsp_param_test"; + + auto idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_SPARSE_DSP_CC, version) + .value(); + knowhere::Json build_json = base_gen(); + build_json[knowhere::indexparam::DROP_RATIO_SEARCH] = 0.0; + REQUIRE(idx.Build(train_ds, build_json) == knowhere::Status::success); + + knowhere::BinarySet bs; + REQUIRE(idx.Serialize(bs) == knowhere::Status::success); + if (use_mmap) { + WriteBinaryToFile(tmp_file, bs.GetByName(idx.Type())); + REQUIRE(idx.DeserializeFromFile(tmp_file, build_json) == knowhere::Status::success); + } else { + REQUIRE(idx.Deserialize(bs, build_json) == knowhere::Status::success); + } + + // Helper: search with given params and return recall vs gt + auto search_recall = [&](int mode, float mu, float eta, int gamma, bool kth_init = true) -> float { + knowhere::Json json = base_gen(); + json[knowhere::indexparam::DROP_RATIO_SEARCH] = 0.0; + json["dsp_mode"] = mode; + json["dsp_mu"] = mu; + json["dsp_eta"] = eta; + json["dsp_gamma"] = gamma; + json["dsp_kth_init"] = kth_init; + auto results = idx.Search(query_ds, json, nullptr); + REQUIRE(results.has_value()); + check_distance_decreasing(*results.value()); + return GetKNNRecall(*gt.value(), *results.value()); + }; + + // 1. DSP mode (mode=0): default params → perfect recall + float recall_dsp = search_recall(0, 1.0f, 1.0f, 0); + REQUIRE(recall_dsp == 1.0f); + + // 2. DSP with mu=0.7, eta=0.7 → exercises both mu and eta pruning paths + float recall_mu_eta_07 = search_recall(0, 0.7f, 0.7f, 0); + REQUIRE(recall_mu_eta_07 >= 0.0f); + + // 2b. DSP with mu < eta → more aggressive mu pruning, eta still active + float recall_mu03_eta1 = search_recall(0, 0.3f, 1.0f, 0); + REQUIRE(recall_mu03_eta1 >= 0.0f); + + // 3. DSP gamma=100000 → perfect recall + float recall_gamma_all = search_recall(0, 1.0f, 1.0f, 100000); + REQUIRE(recall_gamma_all == 1.0f); + + // 4. DSP aggressive mu with gamma backstop + float recall_aggressive_no_gamma = search_recall(0, 0.3f, 1.0f, 0); + float recall_aggressive_with_gamma = search_recall(0, 0.3f, 1.0f, 50); + REQUIRE(recall_aggressive_with_gamma >= recall_aggressive_no_gamma); + REQUIRE(recall_aggressive_no_gamma >= 0.0f); + REQUIRE(recall_aggressive_with_gamma >= 0.5f); + + // 5. LSP/0 (mode=1): top-gamma only, no mu/asc gate + float recall_lsp0 = search_recall(1, 1.0f, 1.0f, 100); + REQUIRE(recall_lsp0 >= 0.5f); + // lsp0 ignores mu: changing mu should not affect recall + float recall_lsp0_mu03 = search_recall(1, 0.3f, 1.0f, 100); + REQUIRE(recall_lsp0_mu03 >= recall_lsp0 - 1e-6f); + REQUIRE(recall_lsp0_mu03 <= recall_lsp0 + 1e-6f); + + // 6. LSP/1 (mode=2): lsp0 safe set + mu gate + float recall_lsp1 = search_recall(2, 1.0f, 1.0f, 100); + REQUIRE(recall_lsp1 >= recall_lsp0); // lsp1 includes lsp0 + more + + // 7. LSP/2 (mode=3): lsp1 + asc gate → recall >= lsp1 + float recall_lsp2 = search_recall(3, 1.0f, 1.0f, 100); + REQUIRE(recall_lsp2 >= recall_lsp1); + + // 8. LSP modes with gamma=0 fall back to DSP (not silent empty results) + { + float recall_lsp0_g0 = search_recall(1, 1.0f, 1.0f, 0); + REQUIRE(recall_lsp0_g0 == recall_dsp); + float recall_lsp1_g0 = search_recall(2, 1.0f, 1.0f, 0); + REQUIRE(recall_lsp1_g0 == recall_dsp); + float recall_lsp2_g0 = search_recall(3, 1.0f, 1.0f, 0); + REQUIRE(recall_lsp2_g0 == recall_dsp); + } + + // 9. dsp_mode defaults to DSP (mode=0) with gamma=0 + { + knowhere::Json json = base_gen(); + json[knowhere::indexparam::DROP_RATIO_SEARCH] = 0.0; + auto results = idx.Search(query_ds, json, nullptr); + REQUIRE(results.has_value()); + } + + // 10. kth_init=false is orthogonal to mode + float recall_dsp_nokth = search_recall(0, 1.0f, 1.0f, 0, false); + REQUIRE(recall_dsp_nokth == 1.0f); + + if (use_mmap) { + REQUIRE(std::remove(tmp_file) == 0); + } + } + + SECTION("Test DSP Filtered Search with kth-init Safety") { + // Adversarial test: mask the exact top-k unfiltered results so the kth-init + // seeded threshold is maximally wrong. Verifies that the filtered bootstrap + // bypasses kth-init and still achieves perfect recall. + auto idx = knowhere::IndexFactory::Instance() + .Create(knowhere::IndexEnum::INDEX_SPARSE_DSP_CC, version) + .value(); + knowhere::Json build_json = base_gen(); + build_json[knowhere::indexparam::DROP_RATIO_SEARCH] = 0.0; + REQUIRE(idx.Build(train_ds, build_json) == knowhere::Status::success); + + // Find top-k results per query (unfiltered) + auto gt_unfiltered = knowhere::BruteForce::SearchSparse(train_ds, query_ds, conf, nullptr); + REQUIRE(gt_unfiltered.has_value()); + + // Create adversarial bitset: mask exactly the unfiltered top-k results. + // This maximizes the kth-init failure: the seeded threshold reflects scores + // of docs that are all filtered out. + auto bitset_data = std::vector((nb + 7) / 8, 0); + auto* gt_ids = gt_unfiltered.value()->GetIds(); + int64_t gt_k = gt_unfiltered.value()->GetDim(); + for (int64_t q = 0; q < nq; ++q) { + for (int64_t j = 0; j < gt_k; ++j) { + int64_t id = gt_ids[q * gt_k + j]; + if (id >= 0 && id < nb) { + bitset_data[id / 8] |= (1u << (id % 8)); + } + } + } + knowhere::BitsetView bitset(bitset_data.data(), nb); + + // Compute filtered ground truth + auto gt_filtered = knowhere::BruteForce::SearchSparse(train_ds, query_ds, conf, bitset); + REQUIRE(gt_filtered.has_value()); + check_result_match_filter(*gt_filtered.value(), bitset); + + // Search with kth_init=true (the potentially unsafe case without bootstrap fix) + knowhere::Json search_json = base_gen(); + search_json[knowhere::indexparam::DROP_RATIO_SEARCH] = 0.0; + search_json["dsp_kth_init"] = true; + auto results_kth = idx.Search(query_ds, search_json, bitset); + REQUIRE(results_kth.has_value()); + check_result_match_filter(*results_kth.value(), bitset); + float recall_kth = GetKNNRecall(*gt_filtered.value(), *results_kth.value()); + REQUIRE(recall_kth == 1.0f); + + // Search with kth_init=false (baseline: no seeded threshold) + search_json["dsp_kth_init"] = false; + auto results_nokth = idx.Search(query_ds, search_json, bitset); + REQUIRE(results_nokth.has_value()); + check_result_match_filter(*results_nokth.value(), bitset); + float recall_nokth = GetKNNRecall(*gt_filtered.value(), *results_nokth.value()); + REQUIRE(recall_nokth == 1.0f); + + // Both should produce identical results (kth_init is bypassed under filter) + auto* ids_kth = results_kth.value()->GetIds(); + auto* ids_nokth = results_nokth.value()->GetIds(); + for (int64_t q = 0; q < nq; ++q) { + for (int64_t j = 0; j < topk; ++j) { + REQUIRE(ids_kth[q * topk + j] == ids_nokth[q * topk + j]); + } + } + + // Test with aggressive pruning params under adversarial filter. + // Bootstrap recovery should re-prune superblocks after heap fills. + search_json["dsp_mu"] = 0.3; + search_json["dsp_eta"] = 0.85; + search_json["dsp_kth_init"] = true; + auto results_aggressive = idx.Search(query_ds, search_json, bitset); + REQUIRE(results_aggressive.has_value()); + check_result_match_filter(*results_aggressive.value(), bitset); + float recall_aggressive = GetKNNRecall(*gt_filtered.value(), *results_aggressive.value()); + REQUIRE(recall_aggressive >= 0.5f); + } + SECTION("Test Search with Bitset") { using std::make_tuple; 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(); @@ -587,6 +784,513 @@ 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 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. + const float kth_score = expected_scores[offset + topk - 1]; + 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) { + expected_strict_ids.insert(expected_ids[offset + rank]); + } + if (actual_ids[offset + rank] >= 0 && actual_scores[offset + rank] > kth_score) { + 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 c36bc3282..ec74bf7c4 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 @@ -138,6 +140,64 @@ TEST_CASE("Test simdcomp bit-packing kernels", "[sparse][simd][simdcomp]") { } } +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()) { @@ -363,6 +423,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()) {