SNMG ANN

.gitignore

@@ -66,3 +66,6 @@ _text
 # clang tooling
 compile_commands.json
 .clangd/
+
+datasets/
+index/

build.sh

@@ -78,7 +78,7 @@ INSTALL_TARGET=install
 BUILD_REPORT_METRICS=""
 BUILD_REPORT_INCL_CACHE_STATS=OFF
 
-TEST_TARGETS="CLUSTER_TEST;CORE_TEST;DISTANCE_TEST;LABEL_TEST;LINALG_TEST;MATRIX_TEST;NEIGHBORS_TEST;NEIGHBORS_ANN_BRUTE_FORCE_TEST;NEIGHBORS_ANN_CAGRA_TEST;NEIGHBORS_ANN_NN_DESCENT_TEST;NEIGHBORS_ANN_IVF_TEST;RANDOM_TEST;SOLVERS_TEST;SPARSE_TEST;SPARSE_DIST_TEST;SPARSE_NEIGHBORS_TEST;STATS_TEST;UTILS_TEST"
+TEST_TARGETS="CLUSTER_TEST;CORE_TEST;DISTANCE_TEST;LABEL_TEST;LINALG_TEST;MATRIX_TEST;NEIGHBORS_TEST;NEIGHBORS_ANN_BRUTE_FORCE_TEST;NEIGHBORS_ANN_CAGRA_TEST;NEIGHBORS_ANN_NN_DESCENT_TEST;NEIGHBORS_ANN_IVF_TEST;NEIGHBORS_ANN_MG_TEST;RANDOM_TEST;SOLVERS_TEST;SPARSE_TEST;SPARSE_DIST_TEST;SPARSE_NEIGHBORS_TEST;STATS_TEST;UTILS_TEST"
 BENCH_TARGETS="CLUSTER_BENCH;CORE_BENCH;NEIGHBORS_BENCH;DISTANCE_BENCH;LINALG_BENCH;MATRIX_BENCH;SPARSE_BENCH;RANDOM_BENCH"
 
 CACHE_ARGS=""
@@ -326,6 +326,7 @@ if hasArg tests || (( ${NUMARGS} == 0 )); then
           $CMAKE_TARGET == *"NEIGHBORS_ANN_BRUTE_FORCE_TEST"* || \
           $CMAKE_TARGET == *"NEIGHBORS_ANN_CAGRA_TEST"* || \
           $CMAKE_TARGET == *"NEIGHBORS_ANN_IVF_TEST"* || \
+          $CMAKE_TARGET == *"NEIGHBORS_ANN_MG_TEST"* || \
           $CMAKE_TARGET == *"NEIGHBORS_ANN_NN_DESCENT_TEST"* || \
           $CMAKE_TARGET == *"NEIGHBORS_TEST"* || \
           $CMAKE_TARGET == *"SPARSE_DIST_TEST" || \

cpp/bench/ann/CMakeLists.txt

@@ -31,6 +31,9 @@ option(RAFT_ANN_BENCH_USE_RAFT_IVF_PQ "Include raft's ivf pq algorithm in benchm
 option(RAFT_ANN_BENCH_USE_RAFT_CAGRA "Include raft's CAGRA in benchmark" ON)
 option(RAFT_ANN_BENCH_USE_RAFT_BRUTE_FORCE "Include raft's brute force knn in benchmark" ON)
 option(RAFT_ANN_BENCH_USE_RAFT_CAGRA_HNSWLIB "Include raft's CAGRA in benchmark" ON)
+option(RAFT_ANN_BENCH_USE_RAFT_ANN_MG_IVF_FLAT "Include raft's MG ANN IVF-FLAT in benchmark" ON)
+option(RAFT_ANN_BENCH_USE_RAFT_ANN_MG_IVF_PQ "Include raft's MG ANN IVF-PQ in benchmark" ON)
+option(RAFT_ANN_BENCH_USE_RAFT_ANN_MG_CAGRA "Include raft's MG ANN CAGRA in benchmark" ON)
 option(RAFT_ANN_BENCH_USE_HNSWLIB "Include hnsw algorithm in benchmark" ON)
 option(RAFT_ANN_BENCH_USE_GGNN "Include ggnn algorithm in benchmark" ON)
 option(RAFT_ANN_BENCH_SINGLE_EXE
@@ -57,6 +60,9 @@ if(BUILD_CPU_ONLY)
   set(RAFT_ANN_BENCH_USE_RAFT_CAGRA OFF)
   set(RAFT_ANN_BENCH_USE_RAFT_BRUTE_FORCE OFF)
   set(RAFT_ANN_BENCH_USE_RAFT_CAGRA_HNSWLIB OFF)
+  set(RAFT_ANN_BENCH_USE_RAFT_ANN_MG_IVF_FLAT OFF)
+  set(RAFT_ANN_BENCH_USE_RAFT_ANN_MG_IVF_PQ OFF)
+  set(RAFT_ANN_BENCH_USE_RAFT_ANN_MG_CAGRA OFF)
   set(RAFT_ANN_BENCH_USE_GGNN OFF)
 endif()
 
@@ -66,6 +72,9 @@ if(RAFT_ANN_BENCH_USE_RAFT_IVF_PQ
    OR RAFT_ANN_BENCH_USE_RAFT_IVF_FLAT
    OR RAFT_ANN_BENCH_USE_RAFT_CAGRA
    OR RAFT_ANN_BENCH_USE_RAFT_CAGRA_HNSWLIB
+   OR RAFT_ANN_BENCH_USE_RAFT_ANN_MG_IVF_FLAT
+   OR RAFT_ANN_BENCH_USE_RAFT_ANN_MG_IVF_PQ
+   OR RAFT_ANN_BENCH_USE_RAFT_ANN_MG_CAGRA
 )
   set(RAFT_ANN_BENCH_USE_RAFT ON)
 endif()
@@ -252,6 +261,50 @@ if(RAFT_ANN_BENCH_USE_RAFT_CAGRA_HNSWLIB)
   )
 endif()
 
+if(RAFT_ANN_BENCH_USE_RAFT_ANN_MG_IVF_FLAT)
+  ConfigureAnnBench(
+    NAME
+    RAFT_ANN_MG_IVF_FLAT
+    PATH
+    src/raft/raft_benchmark.cu
+    src/raft/raft_ann_mg_ivf_flat.cu
+    LINKS
+    raft::compiled
+    ucp ucs ucxx nccl
+  )
+endif()
+
+if(RAFT_ANN_BENCH_USE_RAFT_ANN_MG_IVF_PQ)
+  ConfigureAnnBench(
+    NAME
+    RAFT_ANN_MG_IVF_PQ
+    PATH
+    src/raft/raft_benchmark.cu
+    src/raft/raft_ann_mg_ivf_pq.cu
+    LINKS
+    raft::compiled
+    ucp ucs ucxx nccl
+  )
+endif()
+
+if(RAFT_ANN_BENCH_USE_RAFT_ANN_MG_CAGRA)
+  ConfigureAnnBench(
+    NAME
+    RAFT_ANN_MG_CAGRA
+    PATH
+    src/raft/raft_benchmark.cu
+    src/raft/raft_ann_mg_cagra.cu
+    LINKS
+    raft::compiled
+    ucp ucs ucxx nccl
+  )
+endif()
+
+set(RAFT_FAISS_TARGETS faiss::faiss)
+if(TARGET faiss::faiss_avx2)
+  set(RAFT_FAISS_TARGETS faiss::faiss_avx2)
+endif()
+
 message("RAFT_FAISS_TARGETS: ${RAFT_FAISS_TARGETS}")
 message("CUDAToolkit_LIBRARY_DIR: ${CUDAToolkit_LIBRARY_DIR}")
 if(RAFT_ANN_BENCH_USE_FAISS_CPU_FLAT)

cpp/bench/ann/src/common/benchmark.hpp

@@ -36,6 +36,7 @@
 #include <numeric>
 #include <sstream>
 #include <string>
+#include <thread>
 #include <vector>
 
 namespace raft::bench::ann {

cpp/bench/ann/src/raft/raft_ann_bench_param_parser.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -30,15 +30,16 @@ extern template class raft::bench::ann::RaftIvfFlatGpu<uint8_t, int64_t>;
 extern template class raft::bench::ann::RaftIvfFlatGpu<int8_t, int64_t>;
 #endif
 #if defined(RAFT_ANN_BENCH_USE_RAFT_IVF_PQ) || defined(RAFT_ANN_BENCH_USE_RAFT_CAGRA) || \
-  defined(RAFT_ANN_BENCH_USE_RAFT_CAGRA_HNSWLIB)
+  defined(RAFT_ANN_BENCH_USE_RAFT_CAGRA_HNSWLIB) || defined(RAFT_ANN_BENCH_USE_RAFT_ANN_MG_CAGRA)
 #include "raft_ivf_pq_wrapper.h"
 #endif
 #ifdef RAFT_ANN_BENCH_USE_RAFT_IVF_PQ
 extern template class raft::bench::ann::RaftIvfPQ<float, int64_t>;
 extern template class raft::bench::ann::RaftIvfPQ<uint8_t, int64_t>;
 extern template class raft::bench::ann::RaftIvfPQ<int8_t, int64_t>;
 #endif
-#if defined(RAFT_ANN_BENCH_USE_RAFT_CAGRA) || defined(RAFT_ANN_BENCH_USE_RAFT_CAGRA_HNSWLIB)
+#if defined(RAFT_ANN_BENCH_USE_RAFT_CAGRA) || defined(RAFT_ANN_BENCH_USE_RAFT_CAGRA_HNSWLIB) || \
+  defined(RAFT_ANN_BENCH_USE_RAFT_ANN_MG_CAGRA)
 #include "raft_cagra_wrapper.h"
 #endif
 #ifdef RAFT_ANN_BENCH_USE_RAFT_CAGRA
@@ -47,11 +48,33 @@ extern template class raft::bench::ann::RaftCagra<half, uint32_t>;
 extern template class raft::bench::ann::RaftCagra<uint8_t, uint32_t>;
 extern template class raft::bench::ann::RaftCagra<int8_t, uint32_t>;
 #endif
+#ifdef RAFT_ANN_BENCH_USE_RAFT_ANN_MG_IVF_FLAT
+#include "raft_ann_mg_ivf_flat_wrapper.hpp"
+extern template class raft::bench::ann::RaftAnnMG_IvfFlat<float, int64_t>;
+extern template class raft::bench::ann::RaftAnnMG_IvfFlat<uint8_t, int64_t>;
+extern template class raft::bench::ann::RaftAnnMG_IvfFlat<int8_t, int64_t>;
+#endif
+#ifdef RAFT_ANN_BENCH_USE_RAFT_ANN_MG_IVF_PQ
+#include "raft_ann_mg_ivf_pq_wrapper.hpp"
+extern template class raft::bench::ann::RaftAnnMG_IvfPq<float, int64_t>;
+extern template class raft::bench::ann::RaftAnnMG_IvfPq<uint8_t, int64_t>;
+extern template class raft::bench::ann::RaftAnnMG_IvfPq<int8_t, int64_t>;
+#endif
+#ifdef RAFT_ANN_BENCH_USE_RAFT_ANN_MG_CAGRA
+#include "raft_ann_mg_cagra_wrapper.hpp"
+extern template class raft::bench::ann::RaftAnnMG_Cagra<float, uint32_t>;
+extern template class raft::bench::ann::RaftAnnMG_Cagra<uint8_t, uint32_t>;
+extern template class raft::bench::ann::RaftAnnMG_Cagra<int8_t, uint32_t>;
+#endif
 
-#ifdef RAFT_ANN_BENCH_USE_RAFT_IVF_FLAT
+#if defined(RAFT_ANN_BENCH_USE_RAFT_IVF_FLAT) || defined(RAFT_ANN_BENCH_USE_RAFT_ANN_MG_IVF_FLAT)
 template <typename T, typename IdxT>
 void parse_build_param(const nlohmann::json& conf,
+#ifdef RAFT_ANN_BENCH_USE_RAFT_IVF_FLAT
                        typename raft::bench::ann::RaftIvfFlatGpu<T, IdxT>::BuildParam& param)
+#else
+                       typename raft::bench::ann::RaftAnnMG_IvfFlat<T, IdxT>::BuildParam& param)
+#endif
 {
   param.n_lists = conf.at("nlist");
   if (conf.contains("niter")) { param.kmeans_n_iters = conf.at("niter"); }
@@ -60,17 +83,26 @@ void parse_build_param(const nlohmann::json& conf,
 
 template <typename T, typename IdxT>
 void parse_search_param(const nlohmann::json& conf,
+#ifdef RAFT_ANN_BENCH_USE_RAFT_IVF_FLAT
                         typename raft::bench::ann::RaftIvfFlatGpu<T, IdxT>::SearchParam& param)
+#else
+                        typename raft::bench::ann::RaftAnnMG_IvfFlat<T, IdxT>::SearchParam& param)
+#endif
 {
   param.ivf_flat_params.n_probes = conf.at("nprobe");
 }
 #endif
 
 #if defined(RAFT_ANN_BENCH_USE_RAFT_IVF_PQ) || defined(RAFT_ANN_BENCH_USE_RAFT_CAGRA) || \
-  defined(RAFT_ANN_BENCH_USE_RAFT_CAGRA_HNSWLIB)
+  defined(RAFT_ANN_BENCH_USE_RAFT_CAGRA_HNSWLIB) ||                                      \
+  defined(RAFT_ANN_BENCH_USE_RAFT_ANN_MG_IVF_PQ) || defined(RAFT_ANN_BENCH_USE_RAFT_ANN_MG_CAGRA)
 template <typename T, typename IdxT>
 void parse_build_param(const nlohmann::json& conf,
+#ifdef RAFT_ANN_BENCH_USE_RAFT_ANN_MG_IVF_PQ
+                       typename raft::bench::ann::RaftAnnMG_IvfPq<T, IdxT>::BuildParam& param)
+#else
                        typename raft::bench::ann::RaftIvfPQ<T, IdxT>::BuildParam& param)
+#endif
 {
   if (conf.contains("nlist")) { param.n_lists = conf.at("nlist"); }
   if (conf.contains("niter")) { param.kmeans_n_iters = conf.at("niter"); }
@@ -92,7 +124,11 @@ void parse_build_param(const nlohmann::json& conf,
 
 template <typename T, typename IdxT>
 void parse_search_param(const nlohmann::json& conf,
+#ifdef RAFT_ANN_BENCH_USE_RAFT_ANN_MG_IVF_PQ
+                        typename raft::bench::ann::RaftAnnMG_IvfPq<T, IdxT>::SearchParam& param)
+#else
                         typename raft::bench::ann::RaftIvfPQ<T, IdxT>::SearchParam& param)
+#endif
 {
   if (conf.contains("nprobe")) { param.pq_param.n_probes = conf.at("nprobe"); }
   if (conf.contains("internalDistanceDtype")) {
@@ -133,7 +169,8 @@ void parse_search_param(const nlohmann::json& conf,
 }
 #endif
 
-#if defined(RAFT_ANN_BENCH_USE_RAFT_CAGRA) || defined(RAFT_ANN_BENCH_USE_RAFT_CAGRA_HNSWLIB)
+#if defined(RAFT_ANN_BENCH_USE_RAFT_CAGRA) || defined(RAFT_ANN_BENCH_USE_RAFT_CAGRA_HNSWLIB) || \
+  defined(RAFT_ANN_BENCH_USE_RAFT_ANN_MG_CAGRA)
 template <typename T, typename IdxT>
 void parse_build_param(const nlohmann::json& conf,
                        raft::neighbors::experimental::nn_descent::index_params& param)
@@ -180,7 +217,11 @@ nlohmann::json collect_conf_with_prefix(const nlohmann::json& conf,
 
 template <typename T, typename IdxT>
 void parse_build_param(const nlohmann::json& conf,
+#ifdef RAFT_ANN_BENCH_USE_RAFT_ANN_MG_CAGRA
+                       typename raft::bench::ann::RaftAnnMG_Cagra<T, IdxT>::BuildParam& param)
+#else
                        typename raft::bench::ann::RaftCagra<T, IdxT>::BuildParam& param)
+#endif
 {
   if (conf.contains("graph_degree")) {
     param.cagra_params.graph_degree              = conf.at("graph_degree");
@@ -244,7 +285,11 @@ raft::bench::ann::AllocatorType parse_allocator(std::string mem_type)
 
 template <typename T, typename IdxT>
 void parse_search_param(const nlohmann::json& conf,
+#ifdef RAFT_ANN_BENCH_USE_RAFT_ANN_MG_CAGRA
+                        typename raft::bench::ann::RaftAnnMG_Cagra<T, IdxT>::SearchParam& param)
+#else
                         typename raft::bench::ann::RaftCagra<T, IdxT>::SearchParam& param)
+#endif
 {
   if (conf.contains("itopk")) { param.p.itopk_size = conf.at("itopk"); }
   if (conf.contains("search_width")) { param.p.search_width = conf.at("search_width"); }
@@ -263,12 +308,15 @@ void parse_search_param(const nlohmann::json& conf,
       THROW("Invalid value for algo: %s", tmp.c_str());
     }
   }
+
+#if defined(RAFT_ANN_BENCH_USE_RAFT_CAGRA) || defined(RAFT_ANN_BENCH_USE_RAFT_CAGRA_HNSWLIB)
   if (conf.contains("graph_memory_type")) {
     param.graph_mem = parse_allocator(conf.at("graph_memory_type"));
   }
   if (conf.contains("internal_dataset_memory_type")) {
     param.dataset_mem = parse_allocator(conf.at("internal_dataset_memory_type"));
   }
+#endif
   // Same ratio as in IVF-PQ
   param.refine_ratio = conf.value("refine_ratio", 1.0f);
 }

cpp/bench/ann/src/raft/raft_ann_mg_cagra.cu

@@ -0,0 +1,28 @@
+/*
+ * Copyright (c) 2024, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "raft_ann_mg_cagra_wrapper.hpp"
+
+#include <raft/comms/std_comms.hpp>
+#include <raft/neighbors/ann_mg_helpers.cuh>
+
+namespace raft::bench::ann {
+
+template class RaftAnnMG_Cagra<float, uint32_t>;
+template class RaftAnnMG_Cagra<uint8_t, uint32_t>;
+template class RaftAnnMG_Cagra<int8_t, uint32_t>;
+
+}  // namespace raft::bench::ann