Skip to content

Commit 756fb1a

Browse files
lyang24lyang24
committed
feat: optimize sparse DSP index
Improve DSP search correctness and performance with rank-safe threshold handling, superblock-major bound accumulation, candidate masks, conditional ASC computation, and hybrid SIMD intersections. Flatten and parallelize DSP metadata construction, persist native metadata, bound search workspace reuse, extend instrumentation, and add correctness, serialization, concurrency, and SIMD coverage. Signed-off-by: lyang24 <lanqingy93@gmail.com>
1 parent 810efce commit 756fb1a

9 files changed

Lines changed: 1633 additions & 154 deletions

File tree

‎Makefile‎

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -64,12 +64,13 @@ endif
6464
# which requires std::partial_ordering from <compare> (a C++20 feature).
6565
CONAN_SETTINGS := -s compiler.libcxx=$(LIBCXX) -s build_type=$(BUILD_TYPE) -s compiler.cppstd=20 -s:b compiler.cppstd=20
6666

67-
# DiskANN and liburing require libaio (Linux-only).
67+
# DiskANN is enabled for Linux builds. CONAN_INSTALL_FLAGS already contains
68+
# --build=missing, which builds liburing when it is present in the graph and
69+
# lacks a binary package. A separate --build=liburing pattern is unsafe with
70+
# Conan because it is a hard error when a profile/option removes liburing from
71+
# the resolved graph.
6872
ifneq ($(UNAME_S),Darwin)
6973
CONAN_SETTINGS += -o \&:with_diskann=True
70-
ifndef WITH_GPU
71-
CONAN_INSTALL_FLAGS += --build=liburing
72-
endif
7374
endif
7475

7576
# GPU builds use cuVS.

‎include/knowhere/sparse_utils.h‎

Lines changed: 195 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,11 @@
1515
#pragma once
1616

1717
#include <algorithm>
18+
#include <atomic>
1819
#include <boost/iterator/iterator_facade.hpp>
1920
#include <cstddef>
2021
#include <cstdint>
22+
#include <cstdio>
2123
#include <cstring>
2224
#include <functional>
2325
#include <type_traits>
@@ -29,6 +31,133 @@
2931

3032
namespace knowhere::sparse {
3133

34+
// DSP instrumentation (compile with -DKNOWHERE_DSP_INSTRUMENTATION to enable)
35+
#ifdef KNOWHERE_DSP_INSTRUMENTATION
36+
struct SeekStats {
37+
std::atomic<uint64_t> bucket_0{0}; // delta = 0
38+
std::atomic<uint64_t> bucket_1_3{0}; // delta 1-3
39+
std::atomic<uint64_t> bucket_4_15{0}; // delta 4-15
40+
std::atomic<uint64_t> bucket_16_63{0}; // delta 16-63
41+
std::atomic<uint64_t> bucket_64_255{0}; // delta 64-255
42+
std::atomic<uint64_t> bucket_256_plus{0}; // delta 256+
43+
std::atomic<uint64_t> seek_hits{0}; // seek found target doc_id
44+
std::atomic<uint64_t> seek_misses{0}; // seek did NOT find target doc_id
45+
46+
void
47+
record(size_t delta) {
48+
if (delta == 0)
49+
bucket_0++;
50+
else if (delta <= 3)
51+
bucket_1_3++;
52+
else if (delta <= 15)
53+
bucket_4_15++;
54+
else if (delta <= 63)
55+
bucket_16_63++;
56+
else if (delta <= 255)
57+
bucket_64_255++;
58+
else
59+
bucket_256_plus++;
60+
}
61+
62+
void
63+
record_hit() {
64+
seek_hits++;
65+
}
66+
void
67+
record_miss() {
68+
seek_misses++;
69+
}
70+
71+
void
72+
print(const char* label = nullptr) const {
73+
if (label)
74+
printf("\n[Seek Stats: %s]\n", label);
75+
else
76+
printf("\n[Seek Distance Distribution]\n");
77+
uint64_t total = bucket_0 + bucket_1_3 + bucket_4_15 + bucket_16_63 + bucket_64_255 + bucket_256_plus;
78+
printf(" delta=0: %lu (%.1f%%)\n", bucket_0.load(), total ? 100.0 * bucket_0 / total : 0);
79+
printf(" delta 1-3: %lu (%.1f%%)\n", bucket_1_3.load(), total ? 100.0 * bucket_1_3 / total : 0);
80+
printf(" delta 4-15: %lu (%.1f%%)\n", bucket_4_15.load(), total ? 100.0 * bucket_4_15 / total : 0);
81+
printf(" delta 16-63: %lu (%.1f%%)\n", bucket_16_63.load(), total ? 100.0 * bucket_16_63 / total : 0);
82+
printf(" delta 64-255: %lu (%.1f%%)\n", bucket_64_255.load(), total ? 100.0 * bucket_64_255 / total : 0);
83+
printf(" delta 256+: %lu (%.1f%%)\n", bucket_256_plus.load(), total ? 100.0 * bucket_256_plus / total : 0);
84+
printf(" total seeks: %lu\n", total);
85+
uint64_t h = seek_hits.load(), m = seek_misses.load();
86+
uint64_t hm = h + m;
87+
printf(" seek hits: %lu (%.1f%%)\n", h, hm ? 100.0 * h / hm : 0);
88+
printf(" seek misses: %lu (%.1f%%)\n", m, hm ? 100.0 * m / hm : 0);
89+
}
90+
91+
void
92+
reset() {
93+
bucket_0.store(0, std::memory_order_relaxed);
94+
bucket_1_3.store(0, std::memory_order_relaxed);
95+
bucket_4_15.store(0, std::memory_order_relaxed);
96+
bucket_16_63.store(0, std::memory_order_relaxed);
97+
bucket_64_255.store(0, std::memory_order_relaxed);
98+
bucket_256_plus.store(0, std::memory_order_relaxed);
99+
seek_hits.store(0, std::memory_order_relaxed);
100+
seek_misses.store(0, std::memory_order_relaxed);
101+
}
102+
};
103+
104+
inline SeekStats g_seek_stats;
105+
106+
struct DspStats {
107+
std::atomic<uint64_t> total_superblocks{0}; // total superblocks considered
108+
std::atomic<uint64_t> surviving_superblocks{0}; // superblocks surviving coarse pruning
109+
std::atomic<uint64_t> candidate_blocks{0}; // subblocks passing the initial UB threshold
110+
std::atomic<uint64_t> blocks_processed{0}; // candidate subblocks actually scored
111+
std::atomic<uint64_t> saturated_ubs{0}; // surviving subblock UBs saturated at uint16 max
112+
std::atomic<uint64_t> entries_scored{0}; // posting list entries iterated
113+
std::atomic<uint64_t> docs_pushed{0}; // docs pushed to heap
114+
std::atomic<uint64_t> queries{0}; // number of queries
115+
std::atomic<uint64_t> workspace_pool_misses{0}; // searches that allocate because the per-index pool is empty
116+
117+
void
118+
print(const char* label = nullptr) const {
119+
if (label)
120+
printf("\n[DSP Block Stats: %s]\n", label);
121+
else
122+
printf("\n[DSP Block Stats]\n");
123+
uint64_t q = queries.load();
124+
uint64_t total_spb = total_superblocks.load();
125+
uint64_t surviving_spb = surviving_superblocks.load();
126+
uint64_t candidates = candidate_blocks.load();
127+
uint64_t processed = blocks_processed.load();
128+
printf(" queries: %lu\n", q);
129+
printf(" superblocks total: %lu (avg %.1f/q)\n", total_spb, q ? (double)total_spb / q : 0);
130+
printf(" superblocks surviving:%lu (avg %.1f/q, %.1f%%)\n", surviving_spb, q ? (double)surviving_spb / q : 0,
131+
total_spb ? 100.0 * surviving_spb / total_spb : 0);
132+
printf(" candidate blocks: %lu (avg %.1f/q)\n", candidates, q ? (double)candidates / q : 0);
133+
printf(" blocks processed: %lu (avg %.1f/q, %.1f%% of candidates)\n", processed,
134+
q ? (double)processed / q : 0, candidates ? 100.0 * processed / candidates : 0);
135+
printf(" saturated UBs: %lu (avg %.1f/q)\n", saturated_ubs.load(), q ? (double)saturated_ubs / q : 0);
136+
printf(" entries scored: %lu (avg %.1f/q)\n", entries_scored.load(), q ? (double)entries_scored / q : 0);
137+
printf(" docs pushed: %lu (avg %.1f/q)\n", docs_pushed.load(), q ? (double)docs_pushed / q : 0);
138+
printf(" workspace pool misses:%lu\n", workspace_pool_misses.load());
139+
if (processed > 0) {
140+
printf(" entries/block: %.1f\n", (double)entries_scored / processed);
141+
}
142+
}
143+
144+
void
145+
reset() {
146+
total_superblocks.store(0, std::memory_order_relaxed);
147+
surviving_superblocks.store(0, std::memory_order_relaxed);
148+
candidate_blocks.store(0, std::memory_order_relaxed);
149+
blocks_processed.store(0, std::memory_order_relaxed);
150+
saturated_ubs.store(0, std::memory_order_relaxed);
151+
entries_scored.store(0, std::memory_order_relaxed);
152+
docs_pushed.store(0, std::memory_order_relaxed);
153+
queries.store(0, std::memory_order_relaxed);
154+
workspace_pool_misses.store(0, std::memory_order_relaxed);
155+
}
156+
};
157+
158+
inline DspStats g_dsp_stats;
159+
#endif
160+
32161
enum class SparseMetricType {
33162
METRIC_IP = 1,
34163
METRIC_BM25 = 2,
@@ -253,6 +382,72 @@ class SparseRow {
253382
bool own_data_;
254383
};
255384

385+
// When pushing new elements into a MaxMinHeap, only `capacity` elements with the
386+
// largest val are kept. pop()/top() returns the smallest element out of them.
387+
template <typename T>
388+
class MaxMinHeap {
389+
public:
390+
explicit MaxMinHeap(int capacity) : capacity_(capacity), pool_(capacity) {
391+
}
392+
void
393+
push(table_t id, T val) {
394+
if (size_ < capacity_) {
395+
pool_[size_] = {id, val};
396+
size_ += 1;
397+
std::push_heap(pool_.begin(), pool_.begin() + size_, std::greater<SparseIdVal<T>>());
398+
} else if (val > pool_[0].val) {
399+
sift_down(id, val);
400+
}
401+
}
402+
table_t
403+
pop() {
404+
std::pop_heap(pool_.begin(), pool_.begin() + size_, std::greater<SparseIdVal<T>>());
405+
size_ -= 1;
406+
return pool_[size_].id;
407+
}
408+
[[nodiscard]] size_t
409+
size() const {
410+
return size_;
411+
}
412+
[[nodiscard]] bool
413+
empty() const {
414+
return size() == 0;
415+
}
416+
SparseIdVal<T>
417+
top() const {
418+
return pool_[0];
419+
}
420+
[[nodiscard]] bool
421+
full() const {
422+
return size_ == capacity_;
423+
}
424+
425+
private:
426+
void
427+
sift_down(table_t id, T val) {
428+
size_t i = 0;
429+
for (; 2 * i + 1 < size_;) {
430+
size_t j = i;
431+
size_t l = 2 * i + 1, r = 2 * i + 2;
432+
if (pool_[l].val < val) {
433+
j = l;
434+
}
435+
if (r < size_ && pool_[r].val < std::min(pool_[l].val, val)) {
436+
j = r;
437+
}
438+
if (i == j) {
439+
break;
440+
}
441+
pool_[i] = pool_[j];
442+
i = j;
443+
}
444+
pool_[i] = {id, val};
445+
}
446+
447+
size_t size_ = 0, capacity_;
448+
std::vector<SparseIdVal<T>> pool_;
449+
}; // class MaxMinHeap
450+
256451
// A std::vector like container but uses fixed size free memory(typically from
257452
// mmap) as backing store and can only be appended at the end.
258453
//

‎src/index/sparse/sparse_dsp_config.h‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ class SparseDspConfig : public BaseConfig {
3636
CFG_INT dsp_gamma;
3737
CFG_BOOL dsp_kth_init;
3838
CFG_FLOAT dsp_kth_alpha;
39-
KNOHWERE_DECLARE_CONFIG(SparseDspConfig) {
39+
KNOWHERE_DECLARE_CONFIG(SparseDspConfig) {
4040
KNOWHERE_CONFIG_DECLARE_FIELD(drop_ratio_search)
4141
.description("drop ratio for search")
4242
.set_default(0.0f)

0 commit comments

Comments
 (0)