adapt for raft ivf

Signed-off-by: Yusheng.Ma <[email protected]>

benchmark/CMakeLists.txt

@@ -12,15 +12,14 @@
 include_directories(${CMAKE_SOURCE_DIR})
 include_directories(${CMAKE_SOURCE_DIR}/include)
 
-include_directories(/usr/local/hdf5/include)
-link_directories(/usr/local/hdf5/lib)
-
+find_package(HDF5 REQUIRED)
+include_directories(${HDF5_INCLUDE_DIRS})
 set(unittest_libs
         gtest gmock gtest_main gmock_main)
 
 set(depend_libs
         knowhere
-        hdf5
+        ${HDF5_LIBRARIES}
         ${OpenBLAS_LIBRARIES}
         ${LAPACK_LIBRARIES}
         )

benchmark/hdf5/benchmark_knowhere_float_qps.cpp

@@ -18,8 +18,9 @@
 #include "knowhere/comp/knowhere_config.h"
 #include "knowhere/dataset.h"
 
-const int32_t GPU_DEVICE_ID = 0;
-const int32_t CLIENT_NUM = 4;
+constexpr int32_t GPU_DEVICE_ID = 0;
+constexpr int32_t CLIENT_NUM = 1;
+constexpr int32_t THREAD_NUM = 8;
 
 class Benchmark_knowhere_float_qps : public Benchmark_knowhere, public ::testing::Test {
  public:
@@ -65,6 +66,49 @@ class Benchmark_knowhere_float_qps : public Benchmark_knowhere, public ::testing
         }
     }
 
+    void
+    test_raft_ivf(const knowhere::Json& cfg) {
+        auto conf = cfg;
+        auto nlist = conf[knowhere::indexparam::NLIST].get<int32_t>();
+
+        auto find_smallest_nprobe = [&](float expected_recall) -> int32_t {
+            int32_t golden_nq = 10000, golden_topk = 100;
+            int32_t nprobe = 1;
+            float recall;
+            while (nprobe <= NLIST_) {
+                conf[knowhere::indexparam::NPROBE] = nprobe;
+                conf[knowhere::meta::TOPK] = golden_topk;
+                auto ds_ptr = knowhere::GenDataSet(golden_nq, dim_, xq_);
+
+                auto result = index_.Search(*ds_ptr, conf, nullptr);
+                recall = CalcRecall(result.value()->GetIds(), golden_nq, golden_topk);
+                printf("\n[%0.3f s] iterate IVF param for recall %.4f: nlist=%d, nprobe=%d, k=%d, R@=%.4f\n",
+                       get_time_diff(), expected_recall, nlist, nprobe, golden_topk, recall);
+                if (recall >= expected_recall) {
+                    break;
+                }
+                nprobe *= 2;
+            }
+            return std::min(nprobe, NLIST_);
+        };
+
+        for (auto expected_recall : EXPECTED_RECALLs_) {
+            auto nprobe = find_smallest_nprobe(expected_recall);
+            conf[knowhere::indexparam::NPROBE] = nprobe;
+            conf[knowhere::meta::TOPK] = topk_;
+
+            printf("\n[%0.3f s] %s | %s | nlist=%d, nprobe=%d, k=%d, R@=%.4f\n", get_time_diff(),
+                   ann_test_name_.c_str(), index_type_.c_str(), nlist, nprobe, topk_, expected_recall);
+            printf("================================================================================\n");
+            CALC_TIME_SPAN(task(conf, CLIENT_NUM * THREAD_NUM, nq_));
+            printf("  client_num = %d, elapse = %6.3fs, QPS = %.3f\n", CLIENT_NUM, t_diff,
+                   nq_ * CLIENT_NUM * THREAD_NUM / t_diff);
+            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());
+        }
+    }
+
     void
     test_hnsw(const knowhere::Json& cfg) {
         auto conf = cfg;
@@ -120,7 +164,7 @@ class Benchmark_knowhere_float_qps : public Benchmark_knowhere, public ::testing
 
         std::vector<std::thread> thread_vector(worker_num);
         for (int32_t i = 0; i < worker_num; i++) {
-            thread_vector[i] = std::thread(worker, i, nq);
+            thread_vector[i] = std::thread(worker, i % CLIENT_NUM, nq);
         }
         for (int32_t i = 0; i < worker_num; i++) {
             thread_vector[i].join();
@@ -140,7 +184,7 @@ class Benchmark_knowhere_float_qps : public Benchmark_knowhere, public ::testing
         cfg_[knowhere::meta::METRIC_TYPE] = metric_type_;
         knowhere::KnowhereConfig::SetSimdType(knowhere::KnowhereConfig::SimdType::AUTO);
 #ifdef USE_CUDA
-        knowhere::KnowhereConfig::InitGPUResource(GPU_DEVICE_ID, CLIENT_NUM);
+        // knowhere::KnowhereConfig::InitGPUResource(GPU_DEVICE_ID, CLIENT_NUM);
         cfg_[knowhere::meta::DEVICE_ID] = GPU_DEVICE_ID;
 #endif
     }
@@ -227,3 +271,35 @@ TEST_F(Benchmark_knowhere_float_qps, TEST_HNSW) {
     binary_set_.clear();
     test_hnsw(conf);
 }
+
+TEST_F(Benchmark_knowhere_float_qps, TEST_RAFT_IVF_FLAT) {
+    index_type_ = knowhere::IndexEnum::INDEX_RAFT_IVFFLAT;
+
+    knowhere::Json conf = cfg_;
+    conf[knowhere::indexparam::NLIST] = NLIST_;
+
+    std::string index_file_name = get_index_name({NLIST_});
+
+    for (int i = 0; i < CLIENT_NUM; i++) {
+        indices_.emplace_back(create_index(index_file_name, conf));
+        indices_.back().Deserialize(binary_set_);
+    }
+    binary_set_.clear();
+    test_raft_ivf(conf);
+}
+
+TEST_F(Benchmark_knowhere_float_qps, TEST_RAFT_IVF_PQ) {
+    index_type_ = knowhere::IndexEnum::INDEX_RAFT_IVFPQ;
+
+    knowhere::Json conf = cfg_;
+    conf[knowhere::indexparam::NLIST] = NLIST_;
+
+    std::string index_file_name = get_index_name({NLIST_});
+
+    for (int i = 0; i < CLIENT_NUM; i++) {
+        indices_.emplace_back(create_index(index_file_name, conf));
+        indices_.back().Deserialize(binary_set_);
+    }
+    binary_set_.clear();
+    test_raft_ivf(conf);
+}

benchmark/prepare.sh

@@ -3,19 +3,6 @@
 SCRIPTPATH="$( cd -- "$(dirname "$0")" >/dev/null 2>&1 ; pwd -P )"
 
 ROOT="$(dirname "$(dirname "$SCRIPTPATH")")"
-HDF5_DIR=$ROOT"/hdf5-hdf5-1_13_2"
-wget https://github.com/HDFGroup/hdf5/archive/refs/tags/hdf5-1_13_2.tar.gz
-tar xvfz hdf5-1_13_2.tar.gz
-rm hdf5-1_13_2.tar.gz
-cd $HDF5_DIR
-./configure --prefix=/usr/local/hdf5 --enable-fortran
-make -j8
-make install
-cd $ROOT
-rm -r hdf5-hdf5-1_13_2
-
-./build.sh  -u -t Release -b
-
 SIFT_FILE=$ROOT"/output/unittest/sift-128-euclidean.hdf5"
 
 wget -P $SIFT_FILE http://ann-benchmarks.com/sift-128-euclidean.hdf5

include/knowhere/comp/index_param.h

@@ -38,6 +38,8 @@ constexpr const char* INDEX_ANNOY = "ANNOY";
 constexpr const char* INDEX_HNSW = "HNSW";
 
 constexpr const char* INDEX_DISKANN = "DISKANN";
+constexpr const char* INDEX_RAFT_IVFFLAT = "RAFT_IVF_FLAT";
+constexpr const char* INDEX_RAFT_IVFPQ = "RAFT_IVF_PQ";
 
 }  // namespace IndexEnum
 

include/knowhere/gpu/gpu_res_mgr.h

@@ -82,18 +82,6 @@ class GPUResMgr {
         LOG_KNOWHERE_DEBUG_ << "InitDevice gpu_id " << gpu_id_ << ", resource count " << gpu_params_.res_num_
                             << ", tmp_mem_sz " << gpu_params_.tmp_mem_sz_ / MB << "MB, pin_mem_sz "
                             << gpu_params_.pin_mem_sz_ / MB << "MB";
-#ifdef KNOWHERE_WITH_RAFT
-        if (gpu_id >= std::numeric_limits<int>::min() && gpu_id <= std::numeric_limits<int>::max()) {
-            auto rmm_id = rmm::cuda_device_id{int(gpu_id)};
-            rmm_memory_resources_.push_back(
-                std::make_unique<rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource>>(
-                    rmm::mr::get_per_device_resource(rmm_id)));
-            rmm::mr::set_per_device_resource(rmm_id, rmm_memory_resources_.back().get());
-        } else {
-            LOG_KNOWHERE_WARNING_ << "Could not init pool memory resource on GPU " << gpu_id_
-                                  << ". ID is outside expected range.";
-        }
-#endif
     }
 
     void
@@ -125,11 +113,6 @@ class GPUResMgr {
             res_bq_.Take();
         }
         init_ = false;
-#ifdef KNOWHERE_WITH_RAFT
-        for (auto&& rmm_res : rmm_memory_resources_) {
-            rmm_res.release();
-        }
-#endif
     }
 
     ResPtr
@@ -156,9 +139,6 @@ class GPUResMgr {
     int64_t gpu_id_ = 0;
     GPUParams gpu_params_;
     ResBQ res_bq_;
-#ifdef KNOWHERE_WITH_RAFT
-    std::vector<std::unique_ptr<rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource>>> rmm_memory_resources_;
-#endif
 };
 
 class ResScope {