Update diskann benchmark (#984)

Signed-off-by: Yudong Cai <yudong.cai@zilliz.com>

Author: cydrain

CMakeLists.txt

@@ -97,6 +97,7 @@ knowhere_file_glob(GLOB_RECURSE KNOWHERE_SRCS src/common/*.cc src/index/*.cc
 set(KNOWHERE_LINKER_LIBS "")
 
 if(WITH_DISKANN)
+  add_definitions(-DKNOWHERE_WITH_DISKANN)
   include(cmake/libs/libdiskann.cmake)
 else()
   knowhere_file_glob(GLOB_RECURSE KNOWHERE_DISKANN_SRCS src/index/diskann/*.cc)

benchmark/hdf5/benchmark_float.cpp

@@ -27,12 +27,6 @@ std::string kIPIndexDir = kDir + "/ip_index";
 std::string kL2IndexPrefix = kL2IndexDir + "/l2";
 std::string kIPIndexPrefix = kIPIndexDir + "/ip";
 
-constexpr uint32_t kNumRows = 10000;
-constexpr uint32_t kNumQueries = 100;
-constexpr uint32_t kDim = 128;
-constexpr uint32_t kK = 10;
-constexpr float kL2KnnRecall = 0.8;
-
 void
 WriteRawDataToDisk(const std::string data_path, const float* raw_data, const uint32_t num, const uint32_t dim) {
     std::ofstream writer(data_path.c_str(), std::ios::binary);
@@ -119,14 +113,17 @@ 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());
     }
 
+#ifdef KNOWHERE_WITH_DISKANN
     void
     test_diskann(const knowhere::Json& cfg) {
         auto conf = cfg;
         conf["index_prefix"] = (metric_type_ == knowhere::metric::L2 ? kL2IndexPrefix : kIPIndexPrefix);
-        conf["num_threads"] = 8;
         conf["search_cache_budget_gb"] = 0;
         conf["beamwidth"] = 8;
 
+        knowhere::BinarySet binset;
+        index_.Deserialize(binset, conf);
+
         printf("\n[%0.3f s] %s | %s \n", get_time_diff(), ann_test_name_.c_str(), index_type_.c_str());
         printf("================================================================================\n");
         for (auto nq : NQs_) {
@@ -144,6 +141,7 @@ class Benchmark_float : public Benchmark_knowhere, public ::testing::Test {
         printf("================================================================================\n");
         printf("[%.3f s] Test '%s/%s' done\n\n", get_time_diff(), ann_test_name_.c_str(), index_type_.c_str());
     }
+#endif
 
  protected:
     void
@@ -255,6 +253,7 @@ TEST_F(Benchmark_float, TEST_HNSW) {
     }
 }
 
+#ifdef KNOWHERE_WITH_DISKANN
 TEST_F(Benchmark_float, TEST_DISKANN) {
     index_type_ = knowhere::IndexEnum::INDEX_DISKANN;
 
@@ -264,9 +263,8 @@ TEST_F(Benchmark_float, TEST_DISKANN) {
     conf["data_path"] = kRawDataPath;
     conf["max_degree"] = 56;
     conf["search_list_size"] = 128;
-    conf["pq_code_budget_gb"] = sizeof(float) * kDim * kNumRows * 0.125 / (1024 * 1024 * 1024);
+    conf["pq_code_budget_gb"] = sizeof(float) * dim_ * nb_ * 0.125 / (1024 * 1024 * 1024);
     conf["build_dram_budget_gb"] = 32.0;
-    conf["num_threads"] = 8;
 
     fs::create_directory(kDir);
     fs::create_directory(kL2IndexDir);
@@ -283,3 +281,4 @@ TEST_F(Benchmark_float, TEST_DISKANN) {
     index_.Build(*ds_ptr, conf);
     test_diskann(conf);
 }
+#endif

benchmark/hdf5/ref_logs/Makefile

@@ -22,7 +22,7 @@ test_binary_range_hnsw:
 
 ###################################################################################################
 # Test Knowhere float index
-test_float: test_float_idmap test_float_ivf_flat test_float_ivf_sq8 test_float_ivf_pq test_float_hnsw
+test_float: test_float_idmap test_float_ivf_flat test_float_ivf_sq8 test_float_ivf_pq test_float_hnsw test_float_diskann
 test_float_gpu: test_float_ivf_flat test_float_ivf_pq
 
 test_float_idmap:
@@ -35,6 +35,8 @@ test_float_ivf_pq:
 	./benchmark_float --gtest_filter="Benchmark_float.TEST_IVF_PQ" | tee test_float_ivf_pq.log
 test_float_hnsw:
 	./benchmark_float --gtest_filter="Benchmark_float.TEST_HNSW" | tee test_float_hnsw.log
+test_float_diskann:
+	./benchmark_float --gtest_filter="Benchmark_float.TEST_DISKANN" | tee test_float_diskann.log
 
 ###################################################################################################
 # Test Knowhere float index bitset

src/common/comp/knowhere_config.cc

@@ -13,9 +13,11 @@
 
 #include <string>
 
+#ifdef KNOWHERE_WITH_DISKANN
+#include "diskann/aio_context_pool.h"
+#endif
 #include "faiss/Clustering.h"
 #include "faiss/utils/distances.h"
-#include "faiss/utils/utils.h"
 #include "knowhere/log.h"
 #ifdef KNOWHERE_WITH_GPU
 #include "index/gpu/gpu_res_mgr.h"

tests/ut/test_search.cc

@@ -428,6 +428,7 @@ TEST_CASE("Test Mem Index With Binary Vector", "[bool metrics]") {
         }
     }
 
+#if 0
     SECTION("Test Range Search") {
         using std::make_tuple;
         auto [name, gen] = GENERATE_REF(table<std::string, std::function<knowhere::Json()>>({
@@ -449,4 +450,5 @@ TEST_CASE("Test Mem Index With Binary Vector", "[bool metrics]") {
         auto results = idx.RangeSearch(*query_ds, json, nullptr);
         REQUIRE(results.error() == knowhere::Status::faiss_inner_error);
     }
+#endif
 }

tests/ut/test_simd.cc

@@ -15,7 +15,6 @@
 #include "knowhere/comp/brute_force.h"
 #include "knowhere/comp/index_param.h"
 #include "knowhere/comp/knowhere_config.h"
-#include "knowhere/utils.h"
 #include "utils.h"
 
 TEST_CASE("Test BruteForce Search SIMD", "[bf]") {
@@ -54,9 +53,9 @@ TEST_CASE("Test BruteForce Search SIMD", "[bf]") {
         }
     };
 
-    for (auto simd_type : {knowhere::KnowhereConfig::SimdType::AUTO, knowhere::KnowhereConfig::SimdType::AVX512,
-                           knowhere::KnowhereConfig::SimdType::AVX2, knowhere::KnowhereConfig::SimdType::SSE4_2,
-                           knowhere::KnowhereConfig::SimdType::GENERIC}) {
+    for (auto simd_type : {knowhere::KnowhereConfig::SimdType::AVX512, knowhere::KnowhereConfig::SimdType::AVX2,
+                           knowhere::KnowhereConfig::SimdType::SSE4_2, knowhere::KnowhereConfig::SimdType::GENERIC,
+                           knowhere::KnowhereConfig::SimdType::AUTO}) {
         test_search_with_simd(simd_type);
     }
 }
@@ -105,7 +104,7 @@ TEST_CASE("Test PQ Search SIMD", "[pq]") {
         REQUIRE(recall > 0.2);
     };
 
-    for (auto simd_type : {knowhere::KnowhereConfig::SimdType::SSE4_2, knowhere::KnowhereConfig::SimdType::GENERIC}) {
+    for (auto simd_type : {knowhere::KnowhereConfig::SimdType::GENERIC, knowhere::KnowhereConfig::SimdType::AUTO}) {
         for (int64_t m : {8, 16, 32, 64, 128}) {
             test_search_with_simd(m, simd_type);
         }