Skip to content

Commit 2a4cab5

Browse files
committed
add dataLevel0BlocksMemory test case
1 parent 51450b0 commit 2a4cab5

7 files changed

+259
-16
lines changed

CMakeLists.txt

+3
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,9 @@ if(HNSWLIB_EXAMPLES)
100100
add_executable(multiThread_replace_test tests/cpp/multiThread_replace_test.cpp)
101101
target_link_libraries(multiThread_replace_test hnswlib)
102102

103+
add_executable(dataLevel0BlocksMemory_test tests/cpp/dataLevel0BlocksMemory_test.cpp)
104+
target_link_libraries(dataLevel0BlocksMemory_test hnswlib)
105+
103106
add_executable(main tests/cpp/main.cpp tests/cpp/sift_1b.cpp)
104107
target_link_libraries(main hnswlib)
105108
endif()
+175
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
1+
#include "../../hnswlib/hnswlib.h"
2+
3+
#include <assert.h>
4+
5+
#include <vector>
6+
#include <iostream>
7+
#include <cstdio>
8+
#include <thread>
9+
#include <chrono>
10+
11+
namespace {
12+
13+
const size_t M = 32;
14+
const size_t ef_construction = 500;
15+
const size_t random_seed = 100;
16+
const bool allow_replace_deleted = false;
17+
18+
const size_t dimension = 1024;
19+
const size_t total_items = 100 * 10000;
20+
const size_t num_query = 500 * 10000;
21+
size_t topk = 10;
22+
const size_t max_thread_num = 48;
23+
const std::string index_path = "./hnsw.index";
24+
25+
std::vector<float> data(total_items * dimension);
26+
std::vector<float> query(num_query * dimension);
27+
28+
29+
void check_knn_closer(hnswlib::AlgorithmInterface<float>* alg_hnsw) {
30+
for (size_t j = 0; j < num_query; ++j) {
31+
const void* p = query.data() + j * dimension;
32+
auto gd = alg_hnsw->searchKnn(p, topk);
33+
auto res = alg_hnsw->searchKnnCloserFirst(p, topk);
34+
assert(gd.size() == res.size());
35+
size_t t = gd.size();
36+
while (!gd.empty()) {
37+
assert(gd.top() == res[--t]);
38+
gd.pop();
39+
}
40+
}
41+
std::cout << "test hnsw search knn closer first success..." << std::endl;
42+
}
43+
44+
void test_compatibility(bool hnsw_first_use_blocks_memory,
45+
bool hnsw_second_use_blocks_memory) {
46+
47+
std::cout << "================== test compatibility ==================" << std::endl;
48+
hnswlib::L2Space space(dimension);
49+
hnswlib::AlgorithmInterface<float>* alg_hnsw_first = new hnswlib::HierarchicalNSW<float>(&space, 2 * total_items,
50+
M, ef_construction, random_seed, allow_replace_deleted, hnsw_first_use_blocks_memory);
51+
52+
for (size_t i = 0; i < total_items; ++i) {
53+
alg_hnsw_first->addPoint(data.data() + dimension * i, i);
54+
}
55+
check_knn_closer(alg_hnsw_first);
56+
57+
// save hnsw index
58+
std::remove(index_path.data());
59+
alg_hnsw_first->saveIndex(index_path);
60+
std::cout << "save hnsw(use_small_blocks_memory = " << hnsw_first_use_blocks_memory << ") index success" << std::endl;
61+
delete alg_hnsw_first;
62+
63+
// load hnsw index
64+
hnswlib::AlgorithmInterface<float>* alg_hnsw_second = new hnswlib::HierarchicalNSW<float>(&space, false,
65+
0, allow_replace_deleted, hnsw_second_use_blocks_memory);
66+
std::cout << "load hnsw(use_small_blocks_memory = " << hnsw_second_use_blocks_memory << ") index success" << std::endl;
67+
std::remove(index_path.data());
68+
check_knn_closer(alg_hnsw_second);
69+
70+
delete alg_hnsw_second;
71+
}
72+
73+
void test_performace(bool use_small_blocks_memory) {
74+
if (total_items == 0) {
75+
return;
76+
}
77+
78+
std::cout << "================== test preformace("
79+
<< "dimension: " << dimension
80+
<< ", M: " << M
81+
<< ", ef_construction: " << ef_construction
82+
<< ", topk: " << topk
83+
<< ", use_small_blocks_memory: " << (use_small_blocks_memory ? "ture" : "false" )
84+
<< ") ==================" << std::endl;
85+
hnswlib::L2Space space(dimension);
86+
hnswlib::HierarchicalNSW<float>* alg_hnsw = new hnswlib::HierarchicalNSW<float>(&space, 2 * total_items,
87+
M, ef_construction, random_seed, allow_replace_deleted, use_small_blocks_memory);
88+
89+
std::vector<std::thread> threads;
90+
size_t num_threads = (total_items >= max_thread_num ? max_thread_num : total_items);
91+
size_t batch_num = (total_items / (num_threads <= 1 ? 1 : (num_threads - 1))) + 1;
92+
auto start_time = std::chrono::system_clock::now();
93+
for (size_t idx = 0; idx < total_items; idx += batch_num) {
94+
size_t start = idx;
95+
size_t end = std::min(idx + batch_num, total_items);
96+
threads.push_back(
97+
std::thread(
98+
[alg_hnsw, start, end] {
99+
for (size_t i = start; i < end; i++) {
100+
alg_hnsw->addPoint(data.data() + i * dimension, i);
101+
}
102+
}
103+
)
104+
);
105+
}
106+
for (auto &thread : threads) {
107+
thread.join();
108+
}
109+
threads.clear();
110+
auto end_time = std::chrono::system_clock::now();
111+
double duration_in_ms = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time).count();
112+
double duration_in_seconds = static_cast<double>((std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time)).count()) / 1000.0;
113+
size_t qps = (duration_in_seconds == 0 ? total_items : total_items / duration_in_seconds);
114+
double latency = (total_items == 0 ? 0 : duration_in_ms / total_items);
115+
std::cout << "Start " << num_threads << " thread to add " << total_items << " items to hnsw index, cost "
116+
<< duration_in_seconds << " seconds, qps: " << qps << ", latency: " << latency << "ms" << std::endl;
117+
118+
119+
num_threads = (num_query >= max_thread_num ? max_thread_num : num_query);
120+
batch_num = (num_query / (num_threads <= 1 ? 1 : (num_threads - 1))) + 1;
121+
start_time = std::chrono::system_clock::now();
122+
for (size_t idx = 0; idx < num_query; idx += batch_num) {
123+
size_t start = idx;
124+
size_t end = std::min(idx + batch_num, num_query);
125+
threads.push_back(
126+
std::thread(
127+
[alg_hnsw, start, end] {
128+
for (size_t i = start; i < end; i++) {
129+
const void* p = query.data() + i * dimension;
130+
auto gd = alg_hnsw->searchKnn(p, topk);
131+
}
132+
}
133+
)
134+
);
135+
}
136+
for (auto &thread : threads) {
137+
thread.join();
138+
}
139+
threads.clear();
140+
end_time = std::chrono::system_clock::now();
141+
duration_in_ms = std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time).count();
142+
duration_in_seconds = static_cast<double>((std::chrono::duration_cast<std::chrono::milliseconds>(end_time - start_time)).count()) / 1000.0;
143+
qps = (duration_in_seconds == 0 ? num_query : num_query / duration_in_seconds);
144+
latency = (num_query == 0 ? 0 : duration_in_ms / num_query);
145+
std::cout << "Start " << num_threads << " thread to exec " << num_query << " searchKnn, cost "
146+
<< duration_in_seconds << " seconds, qps: " << qps << ", latency: " << latency << "ms" << std::endl;
147+
148+
delete alg_hnsw;
149+
}
150+
151+
} // namespace
152+
153+
int main() {
154+
155+
std::mt19937 rng;
156+
rng.seed(47);
157+
std::uniform_real_distribution<> distrib;
158+
159+
for (size_t i = 0; i < total_items * dimension; ++i) {
160+
data[i] = distrib(rng);
161+
}
162+
for (size_t i = 0; i < num_query * dimension; ++i) {
163+
query[i] = distrib(rng);
164+
}
165+
166+
test_compatibility(true, true);
167+
test_compatibility(false, false);
168+
test_compatibility(true, false);
169+
test_compatibility(false, true);
170+
171+
test_performace(true);
172+
test_performace(false);
173+
174+
return 0;
175+
}

tests/cpp/epsilon_search_test.cpp

+15-3
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,14 @@
44
typedef unsigned int docidtype;
55
typedef float dist_t;
66

7-
int main() {
7+
void test(bool use_small_blocks_memory) {
88
int dim = 16; // Dimension of the elements
99
int max_elements = 10000; // Maximum number of elements, should be known beforehand
1010
int M = 16; // Tightly connected with internal dimensionality of the data
1111
// strongly affects the memory consumption
1212
int ef_construction = 200; // Controls index search speed/build speed tradeoff
13+
size_t random_seed = 100;
14+
bool allow_replace_deleted = false;
1315

1416
int num_queries = 100;
1517
float epsilon2 = 1.0; // Squared distance to query
@@ -20,7 +22,8 @@ int main() {
2022
// Initing index
2123
hnswlib::L2Space space(dim);
2224
hnswlib::BruteforceSearch<dist_t>* alg_brute = new hnswlib::BruteforceSearch<dist_t>(&space, max_elements);
23-
hnswlib::HierarchicalNSW<dist_t>* alg_hnsw = new hnswlib::HierarchicalNSW<dist_t>(&space, max_elements, M, ef_construction);
25+
hnswlib::HierarchicalNSW<dist_t>* alg_hnsw = new hnswlib::HierarchicalNSW<dist_t>(&space, max_elements,
26+
M, ef_construction, random_seed, allow_replace_deleted, use_small_blocks_memory);
2427

2528
// Generate random data
2629
std::mt19937 rng;
@@ -110,5 +113,14 @@ int main() {
110113
delete[] data;
111114
delete alg_brute;
112115
delete alg_hnsw;
113-
return 0;
116+
}
117+
118+
int main() {
119+
std::cout << "Testing with use default memory allocator..." << std::endl;
120+
test(false);
121+
std::cout << "Test ok" << std::endl;
122+
123+
std::cout << "Testing with use block memory allocator..." << std::endl;
124+
test(true);
125+
std::cout << "Test ok" << std::endl;
114126
}

tests/cpp/multiThreadLoad_test.cpp

+18-3
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,13 @@
33
#include <chrono>
44

55

6-
int main() {
6+
void test(bool use_small_blocks_memory) {
77
std::cout << "Running multithread load test" << std::endl;
8+
size_t M = 16;
9+
size_t ef_construction = 200;
10+
size_t random_seed = 100;
11+
bool allow_replace_deleted = false;
12+
813
int d = 16;
914
int max_elements = 1000;
1015

@@ -13,7 +18,8 @@ int main() {
1318
std::uniform_real_distribution<> distrib_real;
1419

1520
hnswlib::L2Space space(d);
16-
hnswlib::HierarchicalNSW<float>* alg_hnsw = new hnswlib::HierarchicalNSW<float>(&space, 2 * max_elements);
21+
hnswlib::HierarchicalNSW<float>* alg_hnsw = new hnswlib::HierarchicalNSW<float>(&space, 2 * max_elements,
22+
M, ef_construction, random_seed, allow_replace_deleted, use_small_blocks_memory);
1723

1824
std::cout << "Building index" << std::endl;
1925
int num_threads = 40;
@@ -136,5 +142,14 @@ int main() {
136142
}
137143

138144
std::cout << "Finish" << std::endl;
139-
return 0;
145+
}
146+
147+
int main() {
148+
std::cout << "Testing with use default memory allocator..." << std::endl;
149+
test(false);
150+
std::cout << "Test ok" << std::endl;
151+
152+
std::cout << "Testing with use block memory allocator..." << std::endl;
153+
test(true);
154+
std::cout << "Test ok" << std::endl;
140155
}

tests/cpp/multiThread_replace_test.cpp

+19-3
Original file line numberDiff line numberDiff line change
@@ -58,13 +58,19 @@ inline void ParallelFor(size_t start, size_t end, size_t numThreads, Function fn
5858
}
5959

6060

61-
int main() {
61+
void test(bool use_small_blocks_memory) {
6262
std::cout << "Running multithread load test" << std::endl;
6363
int d = 16;
6464
int num_elements = 1000;
6565
int max_elements = 2 * num_elements;
6666
int num_threads = 50;
6767

68+
int M = 16; // Tightly connected with internal dimensionality of the data
69+
// strongly affects the memory consumption
70+
int ef_construction = 200; // Controls index search speed/build speed tradeoff
71+
size_t random_seed = 100;
72+
bool allow_replace_deleted = true;
73+
6874
std::mt19937 rng;
6975
rng.seed(47);
7076
std::uniform_real_distribution<> distrib_real;
@@ -90,7 +96,7 @@ int main() {
9096

9197
int iter = 0;
9298
while (iter < 200) {
93-
hnswlib::HierarchicalNSW<float>* alg_hnsw = new hnswlib::HierarchicalNSW<float>(&space, max_elements, 16, 200, 123, true);
99+
hnswlib::HierarchicalNSW<float>* alg_hnsw = new hnswlib::HierarchicalNSW<float>(&space, max_elements, M, ef_construction, 123, true, use_small_blocks_memory);
94100

95101
// add batch1 data
96102
ParallelFor(0, max_elements, num_threads, [&](size_t row, size_t threadId) {
@@ -117,5 +123,15 @@ int main() {
117123

118124
delete[] batch1;
119125
delete[] batch2;
120-
return 0;
126+
}
127+
128+
129+
int main() {
130+
std::cout << "Testing with use default memory allocator..." << std::endl;
131+
test(false);
132+
std::cout << "Test ok" << std::endl;
133+
134+
std::cout << "Testing with use block memory allocator..." << std::endl;
135+
test(true);
136+
std::cout << "Test ok" << std::endl;
121137
}

tests/cpp/multivector_search_test.cpp

+15-3
Original file line numberDiff line numberDiff line change
@@ -4,12 +4,14 @@
44
typedef unsigned int docidtype;
55
typedef float dist_t;
66

7-
int main() {
7+
void test(bool use_small_blocks_memory) {
88
int dim = 16; // Dimension of the elements
99
int max_elements = 1000; // Maximum number of elements, should be known beforehand
1010
int M = 16; // Tightly connected with internal dimensionality of the data
1111
// strongly affects the memory consumption
1212
int ef_construction = 200; // Controls index search speed/build speed tradeoff
13+
size_t random_seed = 100;
14+
bool allow_replace_deleted = false;
1315

1416
int num_queries = 100;
1517
int num_docs = 10; // Number of documents to search
@@ -21,7 +23,8 @@ int main() {
2123
// Initing index
2224
hnswlib::MultiVectorL2Space<docidtype> space(dim);
2325
hnswlib::BruteforceSearch<dist_t>* alg_brute = new hnswlib::BruteforceSearch<dist_t>(&space, max_elements);
24-
hnswlib::HierarchicalNSW<dist_t>* alg_hnsw = new hnswlib::HierarchicalNSW<dist_t>(&space, max_elements, M, ef_construction);
26+
hnswlib::HierarchicalNSW<dist_t>* alg_hnsw = new hnswlib::HierarchicalNSW<dist_t>(&space, max_elements,
27+
M, ef_construction, random_seed, allow_replace_deleted, use_small_blocks_memory);
2528

2629
// Generate random data
2730
std::mt19937 rng;
@@ -122,5 +125,14 @@ int main() {
122125
delete[] data;
123126
delete alg_brute;
124127
delete alg_hnsw;
125-
return 0;
128+
}
129+
130+
int main() {
131+
std::cout << "Testing with use default memory allocator..." << std::endl;
132+
test(false);
133+
std::cout << "Test ok" << std::endl;
134+
135+
std::cout << "Testing with use block memory allocator..." << std::endl;
136+
test(true);
137+
std::cout << "Test ok" << std::endl;
126138
}

tests/cpp/searchKnnCloserFirst_test.cpp

+14-4
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,12 @@ namespace {
1414

1515
using idx_t = hnswlib::labeltype;
1616

17-
void test() {
17+
void test(bool use_small_blocks_memory) {
18+
size_t M = 16;
19+
size_t ef_construction = 200;
20+
size_t random_seed = 100;
21+
bool allow_replace_deleted = false;
22+
1823
int d = 4;
1924
idx_t n = 100;
2025
idx_t nq = 10;
@@ -36,7 +41,8 @@ void test() {
3641

3742
hnswlib::L2Space space(d);
3843
hnswlib::AlgorithmInterface<float>* alg_brute = new hnswlib::BruteforceSearch<float>(&space, 2 * n);
39-
hnswlib::AlgorithmInterface<float>* alg_hnsw = new hnswlib::HierarchicalNSW<float>(&space, 2 * n);
44+
hnswlib::AlgorithmInterface<float>* alg_hnsw = new hnswlib::HierarchicalNSW<float>(&space, 2 * n,
45+
M, ef_construction, random_seed, allow_replace_deleted, use_small_blocks_memory);
4046

4147
for (size_t i = 0; i < n; ++i) {
4248
alg_brute->addPoint(data.data() + d * i, i);
@@ -74,8 +80,12 @@ void test() {
7480
} // namespace
7581

7682
int main() {
77-
std::cout << "Testing ..." << std::endl;
78-
test();
83+
std::cout << "Testing with use default memory allocator..." << std::endl;
84+
test(false);
85+
std::cout << "Test ok" << std::endl;
86+
87+
std::cout << "Testing with use block memory allocator..." << std::endl;
88+
test(true);
7989
std::cout << "Test ok" << std::endl;
8090

8191
return 0;

0 commit comments

Comments
 (0)