adapt for raft ivf (#734)

Signed-off-by: Yusheng.Ma <[email protected]>
milvus-io · Mar 16, 2023 · a926d86 · a926d86
1 parent 218a655
commit a926d86
Show file tree

Hide file tree

Showing 7 changed files with 245 additions and 88 deletions.
diff --git a/benchmark/CMakeLists.txt b/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}
         )

diff --git a/benchmark/prepare.sh b/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

diff --git a/include/knowhere/comp/index_param.h b/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
 

diff --git a/include/knowhere/gpu/gpu_res_mgr.h b/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 {

diff --git a/src/index/ivf_raft/ivf_pq_raft.cu b/src/index/ivf_raft/ivf_pq_raft.cu
diff --git a/src/index/ivf_raft/ivf_flat_raft.cu → src/index/ivf_raft/ivf_raft.cu b/src/index/ivf_raft/ivf_flat_raft.cu → src/index/ivf_raft/ivf_raft.cu
@@ -23,4 +23,10 @@ KNOWHERE_REGISTER_GLOBAL(RAFT_IVF_FLAT, [](const Object& object) {
     return Index<IndexNodeThreadPoolWrapper>::Create(
         std::make_unique<RaftIvfIndexNode<detail::raft_ivf_flat_index>>(object));
 });
+
+KNOWHERE_REGISTER_GLOBAL(RAFT_IVF_PQ, [](const Object& object) {
+    return Index<IndexNodeThreadPoolWrapper>::Create(
+        std::make_unique<RaftIvfIndexNode<detail::raft_ivf_pq_index>>(object));
+});
+
 }  // namespace knowhere