From 0341b623afebdafd1922b078bd0625e84d7f47b6 Mon Sep 17 00:00:00 2001 From: Cai Yudong Date: Mon, 23 Dec 2024 17:08:42 +0800 Subject: [PATCH] [skip e2e] Optimize qps benchmark to show real recall (#1002) Signed-off-by: Cai Yudong --- benchmark/hdf5/benchmark_float.cpp | 67 ++++++++++++++++++++- benchmark/hdf5/benchmark_float_qps.cpp | 83 ++++++++++++++------------ benchmark/hdf5/ref_logs/Makefile | 8 ++- 3 files changed, 117 insertions(+), 41 deletions(-) diff --git a/benchmark/hdf5/benchmark_float.cpp b/benchmark/hdf5/benchmark_float.cpp index d73596d41..7687e2a26 100644 --- a/benchmark/hdf5/benchmark_float.cpp +++ b/benchmark/hdf5/benchmark_float.cpp @@ -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 + void + test_scann(const knowhere::Json& cfg) { + auto conf = cfg; + + const auto reorder_k = conf[knowhere::indexparam::REORDER_K].get(); + const auto with_raw_data = conf[knowhere::indexparam::WITH_RAW_DATA].get(); + auto nlist = conf[knowhere::indexparam::NLIST].get(); + std::string data_type_str = get_data_type_name(); + + 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(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 void test_hnsw(const knowhere::Json& cfg) { @@ -233,6 +266,10 @@ class Benchmark_float : public Benchmark_knowhere, public ::testing::Test { const std::vector Ms_ = {8, 16, 32}; const int32_t NBITS_ = 8; + // SCANN index params + const std::vector SCANN_REORDER_K = {256, 512, 1024}; + const std::vector SCANN_WITH_RAW_DATA = {true}; + // HNSW index params const std::vector HNSW_Ms_ = {16}; const std::vector EFCONs_ = {200}; @@ -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(X); \ + create_index(index_file_name, conf); \ + test_scann(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 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; @@ -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); diff --git a/benchmark/hdf5/benchmark_float_qps.cpp b/benchmark/hdf5/benchmark_float_qps.cpp index ea3b23788..e6c99945f 100644 --- a/benchmark/hdf5/benchmark_float_qps.cpp +++ b/benchmark/hdf5/benchmark_float_qps.cpp @@ -52,7 +52,8 @@ class Benchmark_float_qps : public Benchmark_knowhere, public ::testing::Test { auto nlist = conf[knowhere::indexparam::NLIST].get(); std::string data_type_str = get_data_type_name(); - auto find_smallest_nprobe = [&](float expected_recall) -> int32_t { + auto find_smallest_nprobe = [&](float expected_recall) -> std::tuple { + std::unordered_map recall_map; conf[knowhere::meta::TOPK] = topk_; auto ds_ptr = knowhere::GenDataSet(nq_, dim_, xq_); auto query = knowhere::ConvertToDataTypeIfNeeded(ds_ptr); @@ -65,11 +66,12 @@ 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; @@ -77,17 +79,16 @@ class Benchmark_float_qps : public Benchmark_knowhere, public ::testing::Test { 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(conf, thread_num, nq_)); @@ -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(); - auto find_smallest_max_iters = [&](float expected_recall) -> int32_t { + auto find_smallest_itopk_size = [&](float expected_recall) -> std::tuple { + std::unordered_map 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(conf, thread_num, nq_)); @@ -159,7 +163,8 @@ class Benchmark_float_qps : public Benchmark_knowhere, public ::testing::Test { auto efConstruction = conf[knowhere::indexparam::EFCONSTRUCTION].get(); std::string data_type_str = get_data_type_name(); - auto find_smallest_ef = [&](float expected_recall) -> int32_t { + auto find_smallest_ef = [&](float expected_recall) -> std::tuple { + std::unordered_map recall_map; conf[knowhere::meta::TOPK] = topk_; auto ds_ptr = knowhere::GenDataSet(nq_, dim_, xq_); auto query = knowhere::ConvertToDataTypeIfNeeded(ds_ptr); @@ -170,13 +175,14 @@ 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; @@ -184,17 +190,17 @@ class Benchmark_float_qps : public Benchmark_knowhere, public ::testing::Test { 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(conf, thread_num, nq_)); @@ -216,7 +222,8 @@ class Benchmark_float_qps : public Benchmark_knowhere, public ::testing::Test { auto nlist = conf[knowhere::indexparam::NLIST].get(); std::string data_type_str = get_data_type_name(); - auto find_smallest_nprobe = [&](float expected_recall) -> int32_t { + auto find_smallest_nprobe = [&](float expected_recall) -> std::tuple { + std::unordered_map recall_map; conf[knowhere::meta::TOPK] = topk_; auto ds_ptr = knowhere::GenDataSet(nq_, dim_, xq_); auto query = knowhere::ConvertToDataTypeIfNeeded(ds_ptr); @@ -227,15 +234,16 @@ 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; @@ -243,17 +251,17 @@ class Benchmark_float_qps : public Benchmark_knowhere, public ::testing::Test { 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(conf, thread_num, nq_)); @@ -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(); - auto find_smallest_search_list_size = [&](float expected_recall) -> int32_t { + auto find_smallest_search_list_size = [&](float expected_recall) -> std::tuple { + std::unordered_map recall_map; conf[knowhere::meta::TOPK] = topk_; auto ds_ptr = knowhere::GenDataSet(nq_, dim_, xq_); auto query = knowhere::ConvertToDataTypeIfNeeded(ds_ptr); @@ -283,14 +292,15 @@ 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; @@ -298,17 +308,16 @@ class Benchmark_float_qps : public Benchmark_knowhere, public ::testing::Test { 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(conf, thread_num, nq_)); diff --git a/benchmark/hdf5/ref_logs/Makefile b/benchmark/hdf5/ref_logs/Makefile index b73904b6b..bc7bc2c58 100644 --- a/benchmark/hdf5/ref_logs/Makefile +++ b/benchmark/hdf5/ref_logs/Makefile @@ -22,7 +22,7 @@ test_binary_range_hnsw: ################################################################################################### # Test Knowhere float index -test_float: test_float_brute_force test_float_idmap test_float_ivf_flat test_float_ivf_sq8 test_float_ivf_pq test_float_hnsw test_float_diskann +test_float: test_float_brute_force test_float_idmap test_float_ivf_flat test_float_ivf_sq8 test_float_ivf_pq test_float_scann test_float_hnsw test_float_diskann test_float_raft: test_float_raft_brute_force test_float_raft_ivf_flat test_float_raft_ivf_pq test_float_raft_cagra test_float_ivf: test_float_ivf_flat test_float_ivf_pq @@ -36,6 +36,8 @@ test_float_ivf_sq8: ./benchmark_float --gtest_filter="Benchmark_float.TEST_IVF_SQ8" | tee test_float_ivf_sq8.log test_float_ivf_pq: ./benchmark_float --gtest_filter="Benchmark_float.TEST_IVF_PQ" | tee test_float_ivf_pq.log +test_float_scann: + ./benchmark_float --gtest_filter="Benchmark_float.TEST_SCANN" | tee test_float_scann.log test_float_hnsw: ./benchmark_float --gtest_filter="Benchmark_float.TEST_HNSW" | tee test_float_hnsw.log test_float_diskann: @@ -115,7 +117,7 @@ test_float_range_multi_hnsw: ################################################################################################### # Test Knowhere float index qps -test_float_qps: test_float_qps_idmap test_float_qps_ivf_flat test_float_qps_ivf_sq8 test_float_qps_ivf_pq test_float_qps_hnsw test_float_qps_scann +test_float_qps: test_float_qps_idmap test_float_qps_ivf_flat test_float_qps_ivf_sq8 test_float_qps_ivf_pq test_float_qps_hnsw test_float_qps_scann test_float_qps_diskann test_float_qps_raft: test_float_qps_raft_brute_force test_float_qps_raft_ivf_flat test_float_qps_raft_ivf_pq test_float_qps_raft_cagra test_float_qps_idmap: @@ -130,6 +132,8 @@ test_float_qps_hnsw: ./benchmark_float_qps --gtest_filter="Benchmark_float_qps.TEST_HNSW" | tee test_float_qps_hnsw.log test_float_qps_scann: ./benchmark_float_qps --gtest_filter="Benchmark_float_qps.TEST_SCANN" | tee test_float_qps_scann.log +test_float_qps_diskann: + ./benchmark_float_qps --gtest_filter="Benchmark_float_qps.TEST_DISKANN" | tee test_float_qps_diskann.log test_float_qps_raft_brute_force: ./benchmark_float_qps --gtest_filter="Benchmark_float_qps.TEST_RAFT_BRUTE_FORCE" | tee test_float_qps_raft_brute_force.log