Moving MG functions into unified API + raft::device_resources_snmg as device resource type for MG functions

cpp/CMakeLists.txt

@@ -306,7 +306,6 @@ if(BUILD_SHARED_LIBS)
         src/neighbors/mg/mg_cagra_int8_t_uint32_t.cu
         src/neighbors/mg/mg_cagra_uint8_t_uint32_t.cu
         src/neighbors/mg/omp_checks.cpp
-        src/neighbors/mg/nccl_comm.cpp
     )
   endif()
 
@@ -609,6 +608,7 @@ if(BUILD_SHARED_LIBS)
   if(BUILD_MG_ALGOS)
     target_compile_definitions(cuvs PUBLIC CUVS_BUILD_MG_ALGOS)
     target_compile_definitions(cuvs_objs PUBLIC CUVS_BUILD_MG_ALGOS)
+    target_compile_definitions(cuvs-cagra-search PUBLIC CUVS_BUILD_MG_ALGOS)
   endif()
 
   if(BUILD_CAGRA_HNSWLIB)

cpp/bench/ann/src/cuvs/cuvs_benchmark.cu

@@ -30,38 +30,38 @@
 namespace cuvs::bench {
 
 #ifdef CUVS_ANN_BENCH_USE_CUVS_MG
-void add_distribution_mode(cuvs::neighbors::mg::distribution_mode* dist_mode,
+void add_distribution_mode(cuvs::neighbors::distribution_mode* dist_mode,
                            const nlohmann::json& conf)
 {
   if (conf.contains("distribution_mode")) {
     std::string distribution_mode = conf.at("distribution_mode");
     if (distribution_mode == "replicated") {
-      *dist_mode = cuvs::neighbors::mg::distribution_mode::REPLICATED;
+      *dist_mode = cuvs::neighbors::distribution_mode::REPLICATED;
     } else if (distribution_mode == "sharded") {
-      *dist_mode = cuvs::neighbors::mg::distribution_mode::SHARDED;
+      *dist_mode = cuvs::neighbors::distribution_mode::SHARDED;
     } else {
       throw std::runtime_error("invalid value for distribution_mode");
     }
   } else {
     // default
-    *dist_mode = cuvs::neighbors::mg::distribution_mode::SHARDED;
+    *dist_mode = cuvs::neighbors::distribution_mode::SHARDED;
   }
 };
 
-void add_merge_mode(cuvs::neighbors::mg::sharded_merge_mode* merge_mode, const nlohmann::json& conf)
+void add_merge_mode(cuvs::neighbors::sharded_merge_mode* merge_mode, const nlohmann::json& conf)
 {
   if (conf.contains("merge_mode")) {
     std::string sharded_merge_mode = conf.at("merge_mode");
     if (sharded_merge_mode == "tree_merge") {
-      *merge_mode = cuvs::neighbors::mg::sharded_merge_mode::TREE_MERGE;
+      *merge_mode = cuvs::neighbors::sharded_merge_mode::TREE_MERGE;
     } else if (sharded_merge_mode == "merge_on_root_rank") {
-      *merge_mode = cuvs::neighbors::mg::sharded_merge_mode::MERGE_ON_ROOT_RANK;
+      *merge_mode = cuvs::neighbors::sharded_merge_mode::MERGE_ON_ROOT_RANK;
     } else {
       throw std::runtime_error("invalid value for merge_mode");
     }
   } else {
     // default
-    *merge_mode = cuvs::neighbors::mg::sharded_merge_mode::TREE_MERGE;
+    *merge_mode = cuvs::neighbors::sharded_merge_mode::TREE_MERGE;
   }
 };
 #endif

cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h

@@ -17,8 +17,8 @@
 
 #include "cuvs_ann_bench_utils.h"
 #include "cuvs_cagra_wrapper.h"
-#include <cuvs/neighbors/mg.hpp>
-#include <raft/core/resource/nccl_clique.hpp>
+#include <cuvs/neighbors/cagra.hpp>
+#include <raft/core/device_resources_snmg.hpp>
 
 namespace cuvs::bench {
 using namespace cuvs::neighbors;
@@ -33,21 +33,20 @@ class cuvs_mg_cagra : public algo<T>, public algo_gpu {
   using algo<T>::dim_;
 
   struct build_param : public cuvs::bench::cuvs_cagra<T, IdxT>::build_param {
-    cuvs::neighbors::mg::distribution_mode mode;
+    cuvs::neighbors::distribution_mode mode;
   };
 
   struct search_param : public cuvs::bench::cuvs_cagra<T, IdxT>::search_param {
-    cuvs::neighbors::mg::sharded_merge_mode merge_mode;
+    cuvs::neighbors::sharded_merge_mode merge_mode;
   };
 
   cuvs_mg_cagra(Metric metric, int dim, const build_param& param, int concurrent_searches = 1)
-    : algo<T>(metric, dim), index_params_(param)
+    : algo<T>(metric, dim), index_params_(param), clique_()
   {
     index_params_.cagra_params.metric         = parse_metric_type(metric);
     index_params_.ivf_pq_build_params->metric = parse_metric_type(metric);
 
-    // init nccl clique outside as to not affect benchmark
-    const raft::comms::nccl_clique& clique = raft::resource::get_nccl_clique(handle_);
+    clique_.set_memory_pool(80);
   }
 
   void build(const T* dataset, size_t nrow) final;
@@ -69,7 +68,7 @@ class cuvs_mg_cagra : public algo<T>, public algo_gpu {
 
   [[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override
   {
-    auto stream = raft::resource::get_cuda_stream(handle_);
+    auto stream = raft::resource::get_cuda_stream(clique_);
     return stream;
   }
 
@@ -87,11 +86,11 @@ class cuvs_mg_cagra : public algo<T>, public algo_gpu {
   std::unique_ptr<algo<T>> copy() override;
 
  private:
-  raft::device_resources handle_;
+  raft::device_resources_snmg clique_;
   float refine_ratio_;
   build_param index_params_;
-  cuvs::neighbors::mg::search_params<cagra::search_params> search_params_;
-  std::shared_ptr<cuvs::neighbors::mg::index<cuvs::neighbors::cagra::index<T, IdxT>, T, IdxT>>
+  cuvs::neighbors::mg_search_params<cagra::search_params> search_params_;
+  std::shared_ptr<cuvs::neighbors::mg_index<cuvs::neighbors::cagra::index<T, IdxT>, T, IdxT>>
     index_;
 };
 
@@ -100,14 +99,14 @@ void cuvs_mg_cagra<T, IdxT>::build(const T* dataset, size_t nrow)
 {
   auto dataset_extents = raft::make_extents<IdxT>(nrow, dim_);
   index_params_.prepare_build_params(dataset_extents);
-  cuvs::neighbors::mg::index_params<cagra::index_params> build_params = index_params_.cagra_params;
-  build_params.mode                                                   = index_params_.mode;
+  cuvs::neighbors::mg_index_params<cagra::index_params> build_params = index_params_.cagra_params;
+  build_params.mode                                                  = index_params_.mode;
 
   auto dataset_view =
     raft::make_host_matrix_view<const T, int64_t, raft::row_major>(dataset, nrow, dim_);
-  auto idx = cuvs::neighbors::mg::build(handle_, build_params, dataset_view);
+  auto idx = cuvs::neighbors::cagra::build(clique_, build_params, dataset_view);
   index_ =
-    std::make_shared<cuvs::neighbors::mg::index<cuvs::neighbors::cagra::index<T, IdxT>, T, IdxT>>(
+    std::make_shared<cuvs::neighbors::mg_index<cuvs::neighbors::cagra::index<T, IdxT>, T, IdxT>>(
       std::move(idx));
 }
 
@@ -118,8 +117,7 @@ void cuvs_mg_cagra<T, IdxT>::set_search_param(const search_param_base& param,
                                               const void* filter_bitset)
 {
   if (filter_bitset != nullptr) { throw std::runtime_error("Filtering is not supported yet."); }
-  auto sp = dynamic_cast<const search_param&>(param);
-  // search_params_ = static_cast<mg::search_params<cagra::search_params>>(sp.p);
+  auto sp                                  = dynamic_cast<const search_param&>(param);
   cagra::search_params* search_params_ptr_ = static_cast<cagra::search_params*>(&search_params_);
   *search_params_ptr_                      = sp.p;
   search_params_.merge_mode                = sp.merge_mode;
@@ -134,15 +132,15 @@ void cuvs_mg_cagra<T, IdxT>::set_search_dataset(const T* dataset, size_t nrow)
 template <typename T, typename IdxT>
 void cuvs_mg_cagra<T, IdxT>::save(const std::string& file) const
 {
-  cuvs::neighbors::mg::serialize(handle_, *index_, file);
+  cuvs::neighbors::cagra::serialize(clique_, *index_, file);
 }
 
 template <typename T, typename IdxT>
 void cuvs_mg_cagra<T, IdxT>::load(const std::string& file)
 {
   index_ =
-    std::make_shared<cuvs::neighbors::mg::index<cuvs::neighbors::cagra::index<T, IdxT>, T, IdxT>>(
-      std::move(cuvs::neighbors::mg::deserialize_cagra<T, IdxT>(handle_, file)));
+    std::make_shared<cuvs::neighbors::mg_index<cuvs::neighbors::cagra::index<T, IdxT>, T, IdxT>>(
+      std::move(cuvs::neighbors::cagra::deserialize<T, IdxT>(clique_, file)));
 }
 
 template <typename T, typename IdxT>
@@ -165,8 +163,8 @@ void cuvs_mg_cagra<T, IdxT>::search_base(
   auto distances_view =
     raft::make_host_matrix_view<float, int64_t, raft::row_major>(distances, batch_size, k);
 
-  cuvs::neighbors::mg::search(
-    handle_, *index_, search_params_, queries_view, neighbors_view, distances_view);
+  cuvs::neighbors::cagra::search(
+    clique_, *index_, search_params_, queries_view, neighbors_view, distances_view);
 }
 
 template <typename T, typename IdxT>

cpp/bench/ann/src/cuvs/cuvs_mg_ivf_flat_wrapper.h

@@ -18,8 +18,8 @@
 
 #include "cuvs_ann_bench_utils.h"
 #include "cuvs_ivf_flat_wrapper.h"
-#include <cuvs/neighbors/mg.hpp>
-#include <raft/core/resource/nccl_clique.hpp>
+#include <cuvs/neighbors/ivf_flat.hpp>
+#include <raft/core/device_resources_snmg.hpp>
 
 namespace cuvs::bench {
 using namespace cuvs::neighbors;
@@ -30,18 +30,18 @@ class cuvs_mg_ivf_flat : public algo<T>, public algo_gpu {
   using search_param_base = typename algo<T>::search_param;
   using algo<T>::dim_;
 
-  using build_param = cuvs::neighbors::mg::index_params<ivf_flat::index_params>;
+  using build_param = cuvs::neighbors::mg_index_params<ivf_flat::index_params>;
 
   struct search_param : public cuvs::bench::cuvs_ivf_flat<T, IdxT>::search_param {
-    cuvs::neighbors::mg::sharded_merge_mode merge_mode;
+    cuvs::neighbors::sharded_merge_mode merge_mode;
   };
 
   cuvs_mg_ivf_flat(Metric metric, int dim, const build_param& param)
-    : algo<T>(metric, dim), index_params_(param)
+    : algo<T>(metric, dim), index_params_(param), clique_()
   {
     index_params_.metric = parse_metric_type(metric);
-    // init nccl clique outside as to not affect benchmark
-    const raft::comms::nccl_clique& clique = raft::resource::get_nccl_clique(handle_);
+
+    clique_.set_memory_pool(80);
   }
 
   void build(const T* dataset, size_t nrow) final;
@@ -62,7 +62,7 @@ class cuvs_mg_ivf_flat : public algo<T>, public algo_gpu {
 
   [[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override
   {
-    auto stream = raft::resource::get_cuda_stream(handle_);
+    auto stream = raft::resource::get_cuda_stream(clique_);
     return stream;
   }
 
@@ -73,10 +73,10 @@ class cuvs_mg_ivf_flat : public algo<T>, public algo_gpu {
   std::unique_ptr<algo<T>> copy() override;
 
  private:
-  raft::device_resources handle_;
+  raft::device_resources_snmg clique_;
   build_param index_params_;
-  cuvs::neighbors::mg::search_params<ivf_flat::search_params> search_params_;
-  std::shared_ptr<cuvs::neighbors::mg::index<cuvs::neighbors::ivf_flat::index<T, IdxT>, T, IdxT>>
+  cuvs::neighbors::mg_search_params<ivf_flat::search_params> search_params_;
+  std::shared_ptr<cuvs::neighbors::mg_index<cuvs::neighbors::ivf_flat::index<T, IdxT>, T, IdxT>>
     index_;
 };
 
@@ -85,9 +85,10 @@ void cuvs_mg_ivf_flat<T, IdxT>::build(const T* dataset, size_t nrow)
 {
   auto dataset_view =
     raft::make_host_matrix_view<const T, int64_t, raft::row_major>(dataset, IdxT(nrow), IdxT(dim_));
-  auto idx = cuvs::neighbors::mg::build(handle_, index_params_, dataset_view);
-  index_   = std::make_shared<
-    cuvs::neighbors::mg::index<cuvs::neighbors::ivf_flat::index<T, IdxT>, T, IdxT>>(std::move(idx));
+  auto idx = cuvs::neighbors::ivf_flat::build(clique_, index_params_, dataset_view);
+  index_ =
+    std::make_shared<cuvs::neighbors::mg_index<cuvs::neighbors::ivf_flat::index<T, IdxT>, T, IdxT>>(
+      std::move(idx));
 }
 
 template <typename T, typename IdxT>
@@ -107,15 +108,15 @@ void cuvs_mg_ivf_flat<T, IdxT>::set_search_param(const search_param_base& param,
 template <typename T, typename IdxT>
 void cuvs_mg_ivf_flat<T, IdxT>::save(const std::string& file) const
 {
-  cuvs::neighbors::mg::serialize(handle_, *index_, file);
+  cuvs::neighbors::ivf_flat::serialize(clique_, *index_, file);
 }
 
 template <typename T, typename IdxT>
 void cuvs_mg_ivf_flat<T, IdxT>::load(const std::string& file)
 {
-  index_ = std::make_shared<
-    cuvs::neighbors::mg::index<cuvs::neighbors::ivf_flat::index<T, IdxT>, T, IdxT>>(
-    std::move(cuvs::neighbors::mg::deserialize_flat<T, IdxT>(handle_, file)));
+  index_ =
+    std::make_shared<cuvs::neighbors::mg_index<cuvs::neighbors::ivf_flat::index<T, IdxT>, T, IdxT>>(
+      std::move(cuvs::neighbors::ivf_flat::deserialize<T, IdxT>(clique_, file)));
 }
 
 template <typename T, typename IdxT>
@@ -135,8 +136,8 @@ void cuvs_mg_ivf_flat<T, IdxT>::search(
   auto distances_view = raft::make_host_matrix_view<float, int64_t, raft::row_major>(
     distances, IdxT(batch_size), IdxT(k));
 
-  cuvs::neighbors::mg::search(
-    handle_, *index_, search_params_, queries_view, neighbors_view, distances_view);
+  cuvs::neighbors::ivf_flat::search(
+    clique_, *index_, search_params_, queries_view, neighbors_view, distances_view);
 }
 
 }  // namespace cuvs::bench

cpp/bench/ann/src/cuvs/cuvs_mg_ivf_pq_wrapper.h

@@ -18,8 +18,8 @@
 
 #include "cuvs_ann_bench_utils.h"
 #include "cuvs_ivf_pq_wrapper.h"
-#include <cuvs/neighbors/mg.hpp>
-#include <raft/core/resource/nccl_clique.hpp>
+#include <cuvs/neighbors/ivf_pq.hpp>
+#include <raft/core/device_resources_snmg.hpp>
 
 namespace cuvs::bench {
 using namespace cuvs::neighbors;
@@ -30,18 +30,18 @@ class cuvs_mg_ivf_pq : public algo<T>, public algo_gpu {
   using search_param_base = typename algo<T>::search_param;
   using algo<T>::dim_;
 
-  using build_param = cuvs::neighbors::mg::index_params<ivf_pq::index_params>;
+  using build_param = cuvs::neighbors::mg_index_params<ivf_pq::index_params>;
 
   struct search_param : public cuvs::bench::cuvs_ivf_pq<T, IdxT>::search_param {
-    cuvs::neighbors::mg::sharded_merge_mode merge_mode;
+    cuvs::neighbors::sharded_merge_mode merge_mode;
   };
 
   cuvs_mg_ivf_pq(Metric metric, int dim, const build_param& param)
-    : algo<T>(metric, dim), index_params_(param)
+    : algo<T>(metric, dim), index_params_(param), clique_()
   {
     index_params_.metric = parse_metric_type(metric);
-    // init nccl clique outside as to not affect benchmark
-    const raft::comms::nccl_clique& clique = raft::resource::get_nccl_clique(handle_);
+
+    clique_.set_memory_pool(80);
   }
 
   void build(const T* dataset, size_t nrow) final;
@@ -62,7 +62,7 @@ class cuvs_mg_ivf_pq : public algo<T>, public algo_gpu {
 
   [[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override
   {
-    auto stream = raft::resource::get_cuda_stream(handle_);
+    auto stream = raft::resource::get_cuda_stream(clique_);
     return stream;
   }
 
@@ -73,20 +73,20 @@ class cuvs_mg_ivf_pq : public algo<T>, public algo_gpu {
   std::unique_ptr<algo<T>> copy() override;
 
  private:
-  raft::device_resources handle_;
+  raft::device_resources_snmg clique_;
   build_param index_params_;
-  cuvs::neighbors::mg::search_params<ivf_pq::search_params> search_params_;
-  std::shared_ptr<cuvs::neighbors::mg::index<cuvs::neighbors::ivf_pq::index<IdxT>, T, IdxT>> index_;
+  cuvs::neighbors::mg_search_params<ivf_pq::search_params> search_params_;
+  std::shared_ptr<cuvs::neighbors::mg_index<cuvs::neighbors::ivf_pq::index<IdxT>, T, IdxT>> index_;
 };
 
 template <typename T, typename IdxT>
 void cuvs_mg_ivf_pq<T, IdxT>::build(const T* dataset, size_t nrow)
 {
   auto dataset_view =
     raft::make_host_matrix_view<const T, int64_t, raft::row_major>(dataset, IdxT(nrow), IdxT(dim_));
-  auto idx = cuvs::neighbors::mg::build(handle_, index_params_, dataset_view);
+  auto idx = cuvs::neighbors::ivf_pq::build(clique_, index_params_, dataset_view);
   index_ =
-    std::make_shared<cuvs::neighbors::mg::index<cuvs::neighbors::ivf_pq::index<IdxT>, T, IdxT>>(
+    std::make_shared<cuvs::neighbors::mg_index<cuvs::neighbors::ivf_pq::index<IdxT>, T, IdxT>>(
       std::move(idx));
 }
 
@@ -95,8 +95,7 @@ void cuvs_mg_ivf_pq<T, IdxT>::set_search_param(const search_param_base& param,
                                                const void* filter_bitset)
 {
   if (filter_bitset != nullptr) { throw std::runtime_error("Filtering is not supported yet."); }
-  auto sp = dynamic_cast<const search_param&>(param);
-  // search_params_ = static_cast<mg::search_params<ivf_pq::search_params>>(sp.pq_param);
+  auto sp                                   = dynamic_cast<const search_param&>(param);
   ivf_pq::search_params* search_params_ptr_ = static_cast<ivf_pq::search_params*>(&search_params_);
   *search_params_ptr_                       = sp.pq_param;
   search_params_.merge_mode                 = sp.merge_mode;
@@ -106,15 +105,15 @@ void cuvs_mg_ivf_pq<T, IdxT>::set_search_param(const search_param_base& param,
 template <typename T, typename IdxT>
 void cuvs_mg_ivf_pq<T, IdxT>::save(const std::string& file) const
 {
-  cuvs::neighbors::mg::serialize(handle_, *index_, file);
+  cuvs::neighbors::ivf_pq::serialize(clique_, *index_, file);
 }
 
 template <typename T, typename IdxT>
 void cuvs_mg_ivf_pq<T, IdxT>::load(const std::string& file)
 {
   index_ =
-    std::make_shared<cuvs::neighbors::mg::index<cuvs::neighbors::ivf_pq::index<IdxT>, T, IdxT>>(
-      std::move(cuvs::neighbors::mg::deserialize_pq<T, IdxT>(handle_, file)));
+    std::make_shared<cuvs::neighbors::mg_index<cuvs::neighbors::ivf_pq::index<IdxT>, T, IdxT>>(
+      std::move(cuvs::neighbors::ivf_pq::deserialize<T, IdxT>(clique_, file)));
 }
 
 template <typename T, typename IdxT>
@@ -134,8 +133,8 @@ void cuvs_mg_ivf_pq<T, IdxT>::search(
   auto distances_view = raft::make_host_matrix_view<float, int64_t, raft::row_major>(
     distances, IdxT(batch_size), IdxT(k));
 
-  cuvs::neighbors::mg::search(
-    handle_, *index_, search_params_, queries_view, neighbors_view, distances_view);
+  cuvs::neighbors::ivf_pq::search(
+    clique_, *index_, search_params_, queries_view, neighbors_view, distances_view);
 }
 
 }  // namespace cuvs::bench