Skip to content

Commit

Permalink
[skip e2e] Optimize qps benchmark to show real recall (#1002)
Browse files Browse the repository at this point in the history 
Signed-off-by: Cai Yudong <[email protected]>
  • Loading branch information
@cydrain
cydrain authored Dec 23, 2024
1 parent 52c0303 commit 0341b62
Show file tree
Hide file tree
Showing 3 changed files with 117 additions and 41 deletions.
67 changes: 65 additions & 2 deletions benchmark/hdf5/benchmark_float.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -105,6 +105,39 @@ class Benchmark_float : public Benchmark_knowhere, public ::testing::Test {
printf("[%.3f s] Test '%s/%s' done\n\n", get_time_diff(), ann_test_name_.c_str(), index_type_.c_str());
}

template <typename T>
void
test_scann(const knowhere::Json& cfg) {
auto conf = cfg;

const auto reorder_k = conf[knowhere::indexparam::REORDER_K].get<int32_t>();
const auto with_raw_data = conf[knowhere::indexparam::WITH_RAW_DATA].get<bool>();
auto nlist = conf[knowhere::indexparam::NLIST].get<int32_t>();
std::string data_type_str = get_data_type_name<T>();

printf("\n[%0.3f s] %s | %s(%s) | nlist=%d, reorder_k=%d\n", get_time_diff(), ann_test_name_.c_str(),
index_type_.c_str(), data_type_str.c_str(), nlist, reorder_k);
printf("================================================================================\n");
for (auto nprobe : NPROBEs_) {
conf[knowhere::indexparam::NPROBE] = nprobe;
for (auto nq : NQs_) {
auto ds_ptr = knowhere::GenDataSet(nq, dim_, xq_);
auto query = knowhere::ConvertToDataTypeIfNeeded<T>(ds_ptr);
for (auto k : TOPKs_) {
conf[knowhere::meta::TOPK] = k;
CALC_TIME_SPAN(auto result = index_.value().Search(query, conf, nullptr));
auto ids = result.value()->GetIds();
float recall = CalcRecall(ids, nq, k);
printf(" nprobe = %4d, nq = %4d, k = %4d, elapse = %6.3fs, R@ = %.4f\n", nprobe, nq, k, TDIFF_,
recall);
std::fflush(stdout);
}
}
}
printf("================================================================================\n");
printf("[%.3f s] Test '%s/%s' done\n\n", get_time_diff(), ann_test_name_.c_str(), index_type_.c_str());
}

template <typename T>
void
test_hnsw(const knowhere::Json& cfg) {
Expand Down Expand Up @@ -233,6 +266,10 @@ class Benchmark_float : public Benchmark_knowhere, public ::testing::Test {
const std::vector<int32_t> Ms_ = {8, 16, 32};
const int32_t NBITS_ = 8;

// SCANN index params
const std::vector<int32_t> SCANN_REORDER_K = {256, 512, 1024};
const std::vector<bool> SCANN_WITH_RAW_DATA = {true};

// HNSW index params
const std::vector<int32_t> HNSW_Ms_ = {16};
const std::vector<int32_t> EFCONs_ = {200};
Expand Down Expand Up @@ -336,6 +373,32 @@ TEST_F(Benchmark_float, TEST_IVF_PQ) {
}
}

TEST_F(Benchmark_float, TEST_SCANN) {
index_type_ = knowhere::IndexEnum::INDEX_FAISS_SCANN;

#define TEST_SCANN(T, X) \
index_file_name = get_index_name<T>(X); \
create_index<T>(index_file_name, conf); \
test_scann<T>(conf);

std::string index_file_name;
knowhere::Json conf = cfg_;
for (auto reorder_k : SCANN_REORDER_K) {
conf[knowhere::indexparam::REORDER_K] = reorder_k;
for (auto nlist : NLISTs_) {
conf[knowhere::indexparam::NLIST] = nlist;
for (const auto with_raw_data : SCANN_WITH_RAW_DATA) {
conf[knowhere::indexparam::WITH_RAW_DATA] = with_raw_data;
std::vector<int32_t> params = {nlist, reorder_k, with_raw_data};

TEST_SCANN(knowhere::fp32, params);
TEST_SCANN(knowhere::fp16, params);
TEST_SCANN(knowhere::bf16, params);
}
}
}
}

TEST_F(Benchmark_float, TEST_HNSW) {
index_type_ = knowhere::IndexEnum::INDEX_HNSW;

Expand Down Expand Up @@ -386,8 +449,8 @@ TEST_F(Benchmark_float, TEST_DISKANN) {
index_type_, knowhere::Version::GetCurrentVersion().VersionNumber(), diskann_index_pack);
printf("[%.3f s] Building all on %d vectors\n", get_time_diff(), nb_);
knowhere::DataSetPtr ds_ptr = nullptr;
index_.value().Build(ds_ptr, conf);

CALC_TIME_SPAN(index_.value().Build(ds_ptr, conf));
printf("Build index %s time: %.3fs \n", index_.value().Type().c_str(), TDIFF_);
knowhere::BinarySet binset;
index_.value().Serialize(binset);
index_.value().Deserialize(binset, conf);
Expand Down
83 changes: 46 additions & 37 deletions benchmark/hdf5/benchmark_float_qps.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,7 +52,8 @@ class Benchmark_float_qps : public Benchmark_knowhere, public ::testing::Test {
auto nlist = conf[knowhere::indexparam::NLIST].get<int32_t>();
std::string data_type_str = get_data_type_name<T>();

auto find_smallest_nprobe = [&](float expected_recall) -> int32_t {
auto find_smallest_nprobe = [&](float expected_recall) -> std::tuple<int32_t, float> {
std::unordered_map<int32_t, float> recall_map;
conf[knowhere::meta::TOPK] = topk_;
auto ds_ptr = knowhere::GenDataSet(nq_, dim_, xq_);
auto query = knowhere::ConvertToDataTypeIfNeeded<T>(ds_ptr);
Expand All @@ -65,29 +66,29 @@ class Benchmark_float_qps : public Benchmark_knowhere, public ::testing::Test {

auto result = index_.value().Search(query, conf, nullptr);
recall = CalcRecall(result.value()->GetIds(), nq_, topk_);
recall_map[nprobe] = recall;
printf("[%0.3f s] iterate IVF param for recall %.4f: nlist=%d, nprobe=%4d, k=%d, R@=%.4f\n",
get_time_diff(), expected_recall, nlist, nprobe, topk_, recall);
std::fflush(stdout);
if (std::abs(recall - expected_recall) <= 0.0001) {
return nprobe;
return {nprobe, recall_map[nprobe]};
}
if (recall < expected_recall) {
left = nprobe + 1;
} else {
right = nprobe - 1;
}
}
return left;
return {left, recall_map[left]};
};

for (auto expected_recall : EXPECTED_RECALLs_) {
auto nprobe = find_smallest_nprobe(expected_recall);
auto [nprobe, recall] = find_smallest_nprobe(expected_recall);
conf[knowhere::indexparam::NPROBE] = nprobe;
conf[knowhere::meta::TOPK] = topk_;

printf("\n[%0.3f s] %s | %s(%s) | nlist=%d, nprobe=%d, k=%d, R@=%.4f\n", get_time_diff(),
ann_test_name_.c_str(), index_type_.c_str(), data_type_str.c_str(), nlist, nprobe, topk_,
expected_recall);
ann_test_name_.c_str(), index_type_.c_str(), data_type_str.c_str(), nlist, nprobe, topk_, recall);
printf("================================================================================\n");
for (auto thread_num : THREAD_NUMs_) {
CALC_TIME_SPAN(task<T>(conf, thread_num, nq_));
Expand All @@ -105,41 +106,44 @@ class Benchmark_float_qps : public Benchmark_knowhere, public ::testing::Test {
auto conf = cfg;
std::string data_type_str = get_data_type_name<T>();

auto find_smallest_max_iters = [&](float expected_recall) -> int32_t {
auto find_smallest_itopk_size = [&](float expected_recall) -> std::tuple<int32_t, float> {
std::unordered_map<int32_t, float> recall_map;
auto ds_ptr = knowhere::GenDataSet(nq_, dim_, xq_);
auto left = 32;
auto right = 256;
auto max_iterations = left;
auto itopk_size = left;

float recall;
while (left <= right) {
max_iterations = left + (right - left) / 2;
conf[knowhere::indexparam::MAX_ITERATIONS] = max_iterations;
itopk_size = left + (right - left) / 2;
conf[knowhere::indexparam::ITOPK_SIZE] = itopk_size;

auto result = index_.value().Search(ds_ptr, conf, nullptr);
recall = CalcRecall(result.value()->GetIds(), nq_, topk_);
printf("[%0.3f s] iterate CAGRA param for recall %.4f: max_iterations=%d, k=%d, R@=%.4f\n",
get_time_diff(), expected_recall, max_iterations, topk_, recall);
recall_map[itopk_size] = recall;
printf("[%0.3f s] iterate CAGRA param for recall %.4f: itopk_size=%d, k=%d, R@=%.4f\n", get_time_diff(),
expected_recall, itopk_size, topk_, recall);
std::fflush(stdout);
if (std::abs(recall - expected_recall) <= 0.0001) {
return max_iterations;
return {itopk_size, recall_map[itopk_size]};
}
if (recall < expected_recall) {
left = max_iterations + 1;
left = itopk_size + 1;
} else {
right = max_iterations - 1;
right = itopk_size - 1;
}
}
return left;
return {left, recall_map[itopk_size]};
};

for (auto expected_recall : EXPECTED_RECALLs_) {
auto [itopk_size, recall] = find_smallest_itopk_size(expected_recall);
conf[knowhere::indexparam::ITOPK_SIZE] = ((int{topk_} + 32 - 1) / 32) * 32;
conf[knowhere::meta::TOPK] = topk_;
conf[knowhere::indexparam::MAX_ITERATIONS] = find_smallest_max_iters(expected_recall);
conf[knowhere::indexparam::ITOPK_SIZE] = itopk_size;

printf("\n[%0.3f s] %s | %s(%s) | k=%d, R@=%.4f\n", get_time_diff(), ann_test_name_.c_str(),
index_type_.c_str(), data_type_str.c_str(), topk_, expected_recall);
index_type_.c_str(), data_type_str.c_str(), topk_, recall);
printf("================================================================================\n");
for (auto thread_num : THREAD_NUMs_) {
CALC_TIME_SPAN(task<T>(conf, thread_num, nq_));
Expand All @@ -159,7 +163,8 @@ class Benchmark_float_qps : public Benchmark_knowhere, public ::testing::Test {
auto efConstruction = conf[knowhere::indexparam::EFCONSTRUCTION].get<int32_t>();
std::string data_type_str = get_data_type_name<T>();

auto find_smallest_ef = [&](float expected_recall) -> int32_t {
auto find_smallest_ef = [&](float expected_recall) -> std::tuple<int32_t, float> {
std::unordered_map<int32_t, float> recall_map;
conf[knowhere::meta::TOPK] = topk_;
auto ds_ptr = knowhere::GenDataSet(nq_, dim_, xq_);
auto query = knowhere::ConvertToDataTypeIfNeeded<T>(ds_ptr);
Expand All @@ -170,31 +175,32 @@ class Benchmark_float_qps : public Benchmark_knowhere, public ::testing::Test {
ef = left + (right - left) / 2;
conf[knowhere::indexparam::EF] = ef;

auto result = index_.value().Search(ds_ptr, conf, nullptr);
auto result = index_.value().Search(query, conf, nullptr);
recall = CalcRecall(result.value()->GetIds(), nq_, topk_);
recall_map[ef] = recall;
printf("[%0.3f s] iterate HNSW param for expected recall %.4f: M=%d, efc=%d, ef=%4d, k=%d, R@=%.4f\n",
get_time_diff(), expected_recall, M, efConstruction, ef, topk_, recall);
std::fflush(stdout);
if (std::abs(recall - expected_recall) <= 0.0001) {
return ef;
return {ef, recall_map[ef]};
}
if (recall < expected_recall) {
left = ef + 1;
} else {
right = ef - 1;
}
}
return left;
return {left, recall_map[left]};
};

for (auto expected_recall : EXPECTED_RECALLs_) {
auto ef = find_smallest_ef(expected_recall);
auto [ef, recall] = find_smallest_ef(expected_recall);
conf[knowhere::indexparam::EF] = ef;
conf[knowhere::meta::TOPK] = topk_;

printf("\n[%0.3f s] %s | %s(%s) | M=%d | efConstruction=%d, ef=%d, k=%d, R@=%.4f\n", get_time_diff(),
ann_test_name_.c_str(), index_type_.c_str(), data_type_str.c_str(), M, efConstruction, ef, topk_,
expected_recall);
recall);
printf("================================================================================\n");
for (auto thread_num : THREAD_NUMs_) {
CALC_TIME_SPAN(task<T>(conf, thread_num, nq_));
Expand All @@ -216,7 +222,8 @@ class Benchmark_float_qps : public Benchmark_knowhere, public ::testing::Test {
auto nlist = conf[knowhere::indexparam::NLIST].get<int32_t>();
std::string data_type_str = get_data_type_name<T>();

auto find_smallest_nprobe = [&](float expected_recall) -> int32_t {
auto find_smallest_nprobe = [&](float expected_recall) -> std::tuple<int32_t, float> {
std::unordered_map<int32_t, float> recall_map;
conf[knowhere::meta::TOPK] = topk_;
auto ds_ptr = knowhere::GenDataSet(nq_, dim_, xq_);
auto query = knowhere::ConvertToDataTypeIfNeeded<T>(ds_ptr);
Expand All @@ -227,33 +234,34 @@ class Benchmark_float_qps : public Benchmark_knowhere, public ::testing::Test {
nprobe = left + (right - left) / 2;
conf[knowhere::indexparam::NPROBE] = nprobe;

auto result = index_.value().Search(ds_ptr, conf, nullptr);
auto result = index_.value().Search(query, conf, nullptr);
recall = CalcRecall(result.value()->GetIds(), nq_, topk_);
recall_map[nprobe] = recall;
printf(
"[%0.3f s] iterate scann param for recall %.4f: nlist=%d, nprobe=%4d, reorder_k=%d, "
"with_raw_data=%d, k=%d, R@=%.4f\n",
get_time_diff(), expected_recall, nlist, nprobe, reorder_k, with_raw_data ? 1 : 0, topk_, recall);
std::fflush(stdout);
if (std::abs(recall - expected_recall) <= 0.0001) {
return nprobe;
return {nprobe, recall_map[nprobe]};
}
if (recall < expected_recall) {
left = nprobe + 1;
} else {
right = nprobe - 1;
}
}
return left;
return {left, recall_map[left]};
};

for (auto expected_recall : EXPECTED_RECALLs_) {
auto nprobe = find_smallest_nprobe(expected_recall);
auto [nprobe, recall] = find_smallest_nprobe(expected_recall);
conf[knowhere::indexparam::NPROBE] = nprobe;
conf[knowhere::meta::TOPK] = topk_;

printf("\n[%0.3f s] %s | %s(%s) | nlist=%d, nprobe=%d, reorder_k=%d, with_raw_data=%d, k=%d, R@=%.4f\n",
get_time_diff(), ann_test_name_.c_str(), index_type_.c_str(), data_type_str.c_str(), nlist, nprobe,
reorder_k, with_raw_data ? 1 : 0, topk_, expected_recall);
reorder_k, with_raw_data ? 1 : 0, topk_, recall);
printf("================================================================================\n");
for (auto thread_num : THREAD_NUMs_) {
CALC_TIME_SPAN(task<T>(conf, thread_num, nq_));
Expand All @@ -272,7 +280,8 @@ class Benchmark_float_qps : public Benchmark_knowhere, public ::testing::Test {
auto conf = cfg;
std::string data_type_str = get_data_type_name<T>();

auto find_smallest_search_list_size = [&](float expected_recall) -> int32_t {
auto find_smallest_search_list_size = [&](float expected_recall) -> std::tuple<int32_t, float> {
std::unordered_map<int32_t, float> recall_map;
conf[knowhere::meta::TOPK] = topk_;
auto ds_ptr = knowhere::GenDataSet(nq_, dim_, xq_);
auto query = knowhere::ConvertToDataTypeIfNeeded<T>(ds_ptr);
Expand All @@ -283,32 +292,32 @@ class Benchmark_float_qps : public Benchmark_knowhere, public ::testing::Test {
search_list_size = left + (right - left) / 2;
conf[knowhere::indexparam::SEARCH_LIST_SIZE] = search_list_size;

auto result = index_.value().Search(ds_ptr, conf, nullptr);
auto result = index_.value().Search(query, conf, nullptr);
recall = CalcRecall(result.value()->GetIds(), nq_, topk_);
recall_map[search_list_size] = recall;
printf(
"[%0.3f s] iterate DISKANN param for expected recall %.4f: search_list_size=%4d, k=%d, R@=%.4f\n",
get_time_diff(), expected_recall, search_list_size, topk_, recall);
std::fflush(stdout);
if (std::abs(recall - expected_recall) <= 0.0001) {
return search_list_size;
return {search_list_size, recall_map[search_list_size]};
}
if (recall < expected_recall) {
left = search_list_size + 1;
} else {
right = search_list_size - 1;
}
}
return left;
return {left, recall_map[left]};
};

for (auto expected_recall : EXPECTED_RECALLs_) {
auto search_list_size = find_smallest_search_list_size(expected_recall);
auto [search_list_size, recall] = find_smallest_search_list_size(expected_recall);
conf[knowhere::indexparam::SEARCH_LIST_SIZE] = search_list_size;
conf[knowhere::meta::TOPK] = topk_;

printf("\n[%0.3f s] %s | %s(%s) | search_list_size=%d, k=%d, R@=%.4f\n", get_time_diff(),
ann_test_name_.c_str(), index_type_.c_str(), data_type_str.c_str(), search_list_size, topk_,
expected_recall);
ann_test_name_.c_str(), index_type_.c_str(), data_type_str.c_str(), search_list_size, topk_, recall);
printf("================================================================================\n");
for (auto thread_num : THREAD_NUMs_) {
CALC_TIME_SPAN(task<T>(conf, thread_num, nq_));
Expand Down
Loading

0 comments on commit 0341b62

Please sign in to comment.