Repository navigation
Expand file tree
/
Copy pathtask2.cpp
More file actions
334 lines (317 loc) · 19.1 KB
/
Copy pathtask2.cpp
File metadata and controls
334 lines (317 loc) · 19.1 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
/**
* @file task2.cpp
* @brief EVP Benchmark Task 2 — Entrypoint to dispatch HDF5-based MIPS ANN benchmarks for LLM attention workloads.
*
* Overview
* --------
* This executable implements the baseline benchmark mode (mode1) exercising FP32 graph construction
* and FP32 search.
* It accepts a single HDF5 input file (llama-dev.h5 ONLY — no other HDF5 files are supported for task2) containing at least the "train" dataset and
* the "test/knns" ground-truth for recall evaluation.
*
* Input Format
* ------------
* The input HDF5 file must contain:
* - "train" : N × D dataset of FP32 (float) feature vectors (database).
* - "test/queries" : Q × D query vectors (separate from database).
* - "test/knns" : Q × K ground-truth nearest-neighbor indices (0-based, int32),
* required when running in recall mode (the default).
*
* Output Modes
* ------------
* Two mutually exclusive output behaviours are supported:
*
* 1. Recall mode (default):
* Loads the "test/knns" ground truth, runs the benchmark, and prints Recall@K.
*
* 2. Save mode (--no-recall --output <path>):
* Skips ground-truth loading, runs the search once, and writes the retrieved
* neighbor indices to a binary ivecs file (one row per query; each row:
* uint32_t count followed by count × uint32_t label values).
*
* Benchmark Modes
* ---------------
* mode1 (baseline | fp32-build-fp32-explore)
* Builds a SizeBoundedGraph using FP32 features (Metric::InnerProduct) and
* explores via FP32 inner-product similarity. Serves as the high-quality baseline.
*
* mode2 (fp16-build-fp16-explore)
* Builds the graph with FP16 features and explores via FP16 inner-product similarity.
*
* mode3 (baseline-fp16 | fp32-build-fp16-explore)
* Builds the graph with FP32 features and explores via FP16 inner-product similarity.
*
* mode4 (l2-converted | fp32-build-l2-explore)
* Builds the graph with FP32 features converted to L2 space (dimension d+1)
* and explores via FP32 L2 similarity.
*
* mode5 (l2-fp16-ip | l2-build-fp16-ip-explore)
* Builds the graph with FP32 features converted to L2 space (dimension d+1)
* and explores via FP16 Inner Product similarity.
*
* mode6 (l2-fp16-l2 | l2-build-fp16-l2-explore)
* Builds the graph with FP32 features converted to L2 space (dimension d+1)
* and explores via FP16 L2 similarity.
*
* mode7 (l2-fp16-d2 | l2-build-fp16-d2-explore)
* Builds the graph with FP32 features converted to L2-converted space (dimension d+2)
* and explores via FP16 L2 similarity.
*
* FLAS Settings
* -------------
* --flas Enable FLAS pre-sort (default: off). Works with any mode.
* --flas-metric <l2|ip> FLAS distance metric (default: l2).
* --flas-radius-decay <f> FLAS swap radius decay factor (default: 0.93).
*
* CLI Usage
* ---------
* deglib_sisap_task2.exe <hdf5_file> <mode> [options...]
*
* <mode> accepts:
* - "baseline"
* - "fp32-build-fp32-explore"
* - "mode1"
*
* Key Options
* -----------
* --threads <n> Parallel thread count (default: 6).
* --k-top <n> Number of nearest neighbors to retrieve (default: 30).
* --k-graph <n> Graph degree / edges per vertex (default: 32).
* --k-ext <n> Builder extension degree (default: 32).
* --eps-ext <f> Builder entry-search expansion coefficient (default: 0.001).
* --max-dist <n> Maximum distance computations per query during exploration (default: 200).
* --no-recall Skip recall computation; requires --output.
* --output <path> Write results to a binary ivecs file.
* --graph <path> File path to save the pre-built graph to, or load a pre-built graph from.
* --prune-worst <n> Number of worst (least similar) neighbors per vertex to replace with self-loops.
*
* Examples
* --------
* # FP32 baseline (mode1)
* .\deglib_sisap_task2.exe llama-dev.h5 mode1 --max-dist 100
* # Or via combined entry point
* .\deglib_sisap.exe task2 llama-dev.h5 mode1 --max-dist 100
*/
#if defined(_WIN32)
#define _CRT_SECURE_NO_WARNINGS
#endif
#include <cstdio>
#include <string>
#include <filesystem>
#include <iostream>
#include "sisap_common.h"
#include "task2/mode1.h"
#include "task2/mode2.h"
#include "task2/mode3.h"
#include "task2/mode4.h"
#include "task2/mode5.h"
#include "task2/mode6.h"
#include "task2/mode7.h"
#include "flas/fast_linear_assignment_sorter.hpp"
#ifdef BUILD_COMBINED_SISAP
int run_task2(int argc, char* argv[]) {
#else
int main(int argc, char* argv[]) {
#endif
try {
std::string data_path;
std::string mode;
uint32_t threads = 8;
uint32_t build_threads = 1;
uint32_t k_top = 30;
uint8_t k_graph = 30;
uint8_t k_ext = 30;
float eps_ext = 0.001f;
bool run_recall = true;
std::string output_path;
std::string graph_path;
std::string max_dist_str = "20000";
uint32_t prune_worst = 0;
std::string eps_search_str = "0.3";
uint32_t num_runs = 1;
bool use_flas = false;
float goal_recall = 0.8f;
std::string flas_metric_str = "l2";
float flas_radius_decay = 0.93f;
deglib::builder::OptimizationTarget opt_target = deglib::builder::OptimizationTarget::LowLID;
uint64_t opt_iterations = 0;
if (argc < 3) {
std::fprintf(stderr, "Usage: %s <hdf5_file_path> <mode_name> [options...]\n", argv[0]);
std::fprintf(stderr, "Modes:\n");
std::fprintf(stderr, " baseline | fp32-build-fp32-explore | mode1 : FP32 build + FP32 explore\n");
std::fprintf(stderr, " fp16-build-fp16-explore | mode2 : FP16 build + FP16 explore\n");
std::fprintf(stderr, " baseline-fp16 | fp32-build-fp16-explore | mode3 : FP32 build + FP16 explore\n");
std::fprintf(stderr, " l2-converted | fp32-build-l2-explore | mode4 : FP32 build (L2 converted) + FP32 L2 search\n");
std::fprintf(stderr, " l2-fp16-ip | l2-build-fp16-ip-explore | mode5 : FP32 build (L2 converted) + FP16 InnerProduct search\n");
std::fprintf(stderr, " l2-fp16-l2 | l2-build-fp16-l2-explore | mode6 : FP32 build (L2 converted) + FP16 L2 search\n");
std::fprintf(stderr, " l2-fp16-d2 | l2-build-fp16-d2-explore | mode7 : FP32 build (L2-converted d+2) + FP16 L2 search\n");
std::fprintf(stderr, " Use --flas with any FP32 mode to enable FLAS pre-sort.\n\n");
std::fprintf(stderr, " --threads <n> Number of CPU worker threads used for query exploration (default: 6).\n");
std::fprintf(stderr, " --build-threads <n> Number of CPU worker threads used for graph construction (default: same as --threads).\n");
std::fprintf(stderr, " --k-top <n> The final number of nearest neighbors (top-K) retrieved per query\n");
std::fprintf(stderr, " and evaluated for recall or written to the output file (default: 30).\n");
std::fprintf(stderr, " --k-graph <n> The degree of the regular graph. Specifies the exact number of edges\n");
std::fprintf(stderr, " (neighbors) per vertex. Must be an even integer >= 4 (default: 32).\n");
std::fprintf(stderr, " --k-ext <n> The search size (k-top parameter) used during graph construction. Decides\n");
std::fprintf(stderr, " how many good neighbors are shown to each newly added node, from which it\n");
std::fprintf(stderr, " selects nodes to connect with up to k-graph (default: 32).\n");
std::fprintf(stderr, " --eps-ext <f> Search expansion parameter used together with k-ext during graph construction.\n");
std::fprintf(stderr, " Decides if nodes whose distance is slightly worse (e.g. 0.01 = 1%%) than the\n");
std::fprintf(stderr, " current worst in the search list should be explored (default: 0.001).\n");
std::fprintf(stderr, " --no-recall Disables loading ground-truth datasets and calculating Recall@K metrics.\n");
std::fprintf(stderr, " Required when exporting search results to an output file.\n");
std::fprintf(stderr, " --output <path> Path to a binary `.ivecs` file where the retrieved nearest-neighbor indices\n");
std::fprintf(stderr, " will be saved (one row per query; uint32_t count followed by indices).\n");
std::fprintf(stderr, " --max-dist <list> Exploration search budget or comma-separated list of budgets. Specifies the maximum number of\n");
std::fprintf(stderr, " distance computations allowed per query. Main parameter to trade search speed for recall (default: 20000).\n");
std::fprintf(stderr, " --graph <path> File path to save the pre-built graph to, or load a pre-built graph from\n");
std::fprintf(stderr, " to bypass the construction phase.\n");
std::fprintf(stderr, " --prune-worst <n> Number of worst (least similar) neighbors per vertex to replace with self-loops (default: 0).\n");
std::fprintf(stderr, " --flas Enable FLAS pre-sort of training vectors before graph construction (default: off).\n");
std::fprintf(stderr, " --flas-metric <l2|ip> FLAS distance metric: l2 (Euclidean) or ip (Inner Product) (default: l2).\n");
std::fprintf(stderr, " --flas-radius-decay <f> FLAS swap radius decay factor per iteration (default: 0.93).\n");
std::fprintf(stderr, " --num-runs <n> Number of query explorations to perform and average (default: 1).\n");
std::fprintf(stderr, " --opt-target <str> Optimization target for graph builder (LowLID, HighLID, StreamingData_SchemeA, ..., default: LowLID).\n");
std::fprintf(stderr, " --goal-recall <f> Recall threshold. Sweeps will select configuration minimizing search time\n");
std::fprintf(stderr, " among those with recall >= threshold (default: 0.8).\n");
std::fprintf(stderr, " --opt-iterations <n> Number of graph optimization iterations to perform after building (default: 0).\n");
return 1;
}
data_path = argv[1];
mode = argv[2];
for (int i = 3; i < argc; ++i) {
std::string arg = argv[i];
if (arg == "--threads" && i + 1 < argc) {
threads = std::stoul(argv[++i]);
} else if (arg == "--build-threads" && i + 1 < argc) {
build_threads = std::stoul(argv[++i]);
} else if (arg == "--k-top" && i + 1 < argc) {
k_top = std::stoul(argv[++i]);
} else if (arg == "--k-graph" && i + 1 < argc) {
k_graph = static_cast<uint8_t>(std::stoul(argv[++i]));
} else if (arg == "--k-ext" && i + 1 < argc) {
k_ext = static_cast<uint8_t>(std::stoul(argv[++i]));
} else if (arg == "--eps-ext" && i + 1 < argc) {
eps_ext = std::stof(argv[++i]);
} else if (arg == "--no-recall") {
run_recall = false;
} else if (arg == "--output" && i + 1 < argc) {
output_path = argv[++i];
} else if (arg == "--max-dist" && i + 1 < argc) {
max_dist_str = argv[++i];
} else if (arg == "--graph" && i + 1 < argc) {
graph_path = argv[++i];
} else if (arg == "--prune-worst" && i + 1 < argc) {
prune_worst = std::stoul(argv[++i]);
} else if (arg == "--eps-search" && i + 1 < argc) {
eps_search_str = argv[++i];
} else if (arg == "--flas") {
use_flas = true;
} else if (arg == "--flas-metric" && i + 1 < argc) {
flas_metric_str = argv[++i];
} else if (arg == "--flas-radius-decay" && i + 1 < argc) {
flas_radius_decay = std::stof(argv[++i]);
} else if (arg == "--num-runs" && i + 1 < argc) {
num_runs = std::stoul(argv[++i]);
} else if (arg == "--goal-recall" && i + 1 < argc) {
goal_recall = std::stof(argv[++i]);
} else if (arg == "--opt-iterations" && i + 1 < argc) {
opt_iterations = std::stoull(argv[++i]);
} else if (arg == "--opt-target" && i + 1 < argc) {
std::string target_str = argv[++i];
if (target_str == "LowLID") {
opt_target = deglib::builder::OptimizationTarget::LowLID;
} else if (target_str == "HighLID") {
opt_target = deglib::builder::OptimizationTarget::HighLID;
} else if (target_str == "StreamingData_SchemeA") {
opt_target = deglib::builder::OptimizationTarget::StreamingData_SchemeA;
} else if (target_str == "StreamingData_SchemeB") {
opt_target = deglib::builder::OptimizationTarget::StreamingData_SchemeB;
} else if (target_str == "StreamingData_SchemeC") {
opt_target = deglib::builder::OptimizationTarget::StreamingData_SchemeC;
} else if (target_str == "StreamingData_SchemeD") {
opt_target = deglib::builder::OptimizationTarget::StreamingData_SchemeD;
} else if (target_str == "SchemeA") {
opt_target = deglib::builder::OptimizationTarget::SchemeA;
} else if (target_str == "SchemeB") {
opt_target = deglib::builder::OptimizationTarget::SchemeB;
} else {
std::fprintf(stderr, "Error: Unknown optimization target '%s'\n", target_str.c_str());
return 1;
}
} else {
std::fprintf(stderr, "Warning: Unknown or malformed option '%s'\n", arg.c_str());
}
}
if (build_threads == 0) {
build_threads = threads;
}
std::vector<uint32_t> max_dist_list = sisap_common::parse_list(max_dist_str);
std::vector<float> eps_search_list = sisap_common::parse_list_f(eps_search_str);
// Validate --prune-worst
if (prune_worst >= k_graph) {
std::fprintf(stderr, "Error: --prune-worst (%u) must be less than --k-graph (%u)\n", prune_worst, (uint32_t)k_graph);
return 1;
}
// Validate --no-recall restrictions
if (!run_recall) {
if (output_path.empty()) {
std::fprintf(stderr, "Error: --output path must be specified when --no-recall is set.\n");
return 1;
}
if (max_dist_list.size() > 1) {
std::fprintf(stderr, "Error: Multiple parameters for --max-dist are not allowed when --no-recall is set.\n");
return 1;
}
}
std::setvbuf(stdout, NULL, _IONBF, 0);
std::setvbuf(stderr, NULL, _IONBF, 0);
// Validate that data_path has an HDF5 extension
std::filesystem::path path(data_path);
std::string ext = path.extension().string();
if (ext != ".h5" && ext != ".hdf5") {
std::fprintf(stderr, "Error: Only HDF5 (.h5 or .hdf5) input files are accepted by task2.\n");
return 1;
}
// Resolve FLAS metric string to enum
FlasMetric flas_metric = FlasMetric::L2;
if (use_flas) {
if (flas_metric_str == "ip" || flas_metric_str == "innerproduct") {
flas_metric = FlasMetric::InnerProduct;
} else if (flas_metric_str != "l2") {
std::fprintf(stderr, "Warning: Unknown --flas-metric '%s', using l2\n", flas_metric_str.c_str());
}
}
// Print compiled-in SIMD instruction set info
#ifdef USE_AVX512
std::fprintf(stderr, "SIMD: AVX-512, AVX2, SSE\n");
#elif defined(USE_AVX)
std::fprintf(stderr, "SIMD: AVX2, SSE\n");
#elif defined(USE_SSE)
std::fprintf(stderr, "SIMD: SSE\n");
#else
std::fprintf(stderr, "SIMD: none (scalar)\n");
#endif
if (mode == "fp32-build-fp32-explore" || mode == "baseline" || mode == "mode1") {
return task2::mode_baseline::run(path, threads, build_threads, use_flas, flas_metric, flas_radius_decay, k_graph, k_ext, eps_ext, opt_target, opt_iterations, prune_worst, k_top, num_runs, max_dist_list, eps_search_list, run_recall, goal_recall, output_path, graph_path);
} else if (mode == "fp16-build-fp16-explore" || mode == "mode2") {
return task2::mode_fp16_build::run(path, threads, build_threads, use_flas, flas_metric, flas_radius_decay, k_graph, k_ext, eps_ext, opt_target, opt_iterations, prune_worst, k_top, num_runs, max_dist_list, eps_search_list, run_recall, goal_recall, output_path, graph_path);
} else if (mode == "fp32-build-fp16-explore" || mode == "baseline-fp16" || mode == "mode3") {
return task2::mode_fp16::run(path, threads, build_threads, use_flas, flas_metric, flas_radius_decay, k_graph, k_ext, eps_ext, opt_target, opt_iterations, prune_worst, k_top, num_runs, max_dist_list, eps_search_list, run_recall, goal_recall, output_path, graph_path);
} else if (mode == "l2-converted" || mode == "fp32-build-l2-explore" || mode == "mode4") {
return task2::mode_l2_converted::run(path, threads, build_threads, use_flas, flas_metric, flas_radius_decay, k_graph, k_ext, eps_ext, opt_target, opt_iterations, prune_worst, k_top, num_runs, max_dist_list, eps_search_list, run_recall, goal_recall, output_path, graph_path);
} else if (mode == "l2-fp16-ip" || mode == "l2-build-fp16-ip-explore" || mode == "mode5") {
return task2::mode_l2_fp16_ip::run(path, threads, build_threads, use_flas, flas_metric, flas_radius_decay, k_graph, k_ext, eps_ext, opt_target, opt_iterations, prune_worst, k_top, num_runs, max_dist_list, eps_search_list, run_recall, goal_recall, output_path, graph_path);
} else if (mode == "l2-fp16-l2" || mode == "l2-build-fp16-l2-explore" || mode == "mode6") {
return task2::mode_l2_fp16::run(path, threads, build_threads, use_flas, flas_metric, flas_radius_decay, k_graph, k_ext, eps_ext, opt_target, opt_iterations, prune_worst, k_top, num_runs, max_dist_list, eps_search_list, run_recall, goal_recall, output_path, graph_path);
} else if (mode == "l2-fp16-d2" || mode == "l2-build-fp16-d2-explore" || mode == "mode7") {
return task2::mode_l2_fp16_d2::run(path, threads, build_threads, use_flas, flas_metric, flas_radius_decay, k_graph, k_ext, eps_ext, opt_target, opt_iterations, prune_worst, k_top, num_runs, max_dist_list, eps_search_list, run_recall, goal_recall, output_path, graph_path);
} else {
std::fprintf(stderr, "Error: Unknown mode '%s'. Supported modes: mode1, mode2, mode3, mode4, mode5, mode6, mode7\n", mode.c_str());
return 1;
}
} catch (const std::exception& ex) {
std::fprintf(stderr, "Fatal error: %s\n", ex.what());
return 1;
}
}