Change snmg index to use updated multi gpu resource API

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

@@ -18,7 +18,7 @@
 #include "cuvs_ann_bench_utils.h"
 #include "cuvs_cagra_wrapper.h"
 #include <cuvs/neighbors/cagra.hpp>
-#include <raft/core/device_resources_snmg.hpp>
+#include <raft/core/device_resources_snmg_nccl.hpp>
 
 namespace cuvs::bench {
 using namespace cuvs::neighbors;

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

@@ -19,7 +19,7 @@
 #include "cuvs_ann_bench_utils.h"
 #include "cuvs_ivf_flat_wrapper.h"
 #include <cuvs/neighbors/ivf_flat.hpp>
-#include <raft/core/device_resources_snmg.hpp>
+#include <raft/core/device_resources_snmg_nccl.hpp>
 
 namespace cuvs::bench {
 using namespace cuvs::neighbors;

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

@@ -19,7 +19,7 @@
 #include "cuvs_ann_bench_utils.h"
 #include "cuvs_ivf_pq_wrapper.h"
 #include <cuvs/neighbors/ivf_pq.hpp>
-#include <raft/core/device_resources_snmg.hpp>
+#include <raft/core/device_resources_snmg_nccl.hpp>
 
 namespace cuvs::bench {
 using namespace cuvs::neighbors;

cpp/cmake/thirdparty/get_raft.cmake

@@ -45,8 +45,8 @@ function(find_and_configure_raft)
             COMPONENTS          ${RAFT_COMPONENTS}
             CPM_ARGS
               EXCLUDE_FROM_ALL TRUE
-              GIT_REPOSITORY        https://github.com/${PKG_FORK}/raft.git
-              GIT_TAG               ${PKG_PINNED_TAG}
+              GIT_REPOSITORY        https://github.com/jinsolp/raft.git
+              GIT_TAG               multi-gpu-resource
               SOURCE_SUBDIR         cpp
               OPTIONS
               "BUILD_TESTS OFF"

cpp/src/neighbors/mg/snmg.cuh

@@ -17,7 +17,8 @@
 #pragma once
 
 #include "../detail/knn_merge_parts.cuh"
-#include <raft/core/resource/nccl_clique.hpp>
+#include <raft/core/resource/multi_gpu.hpp>
+#include <raft/core/resource/nccl_comm.hpp>
 #include <raft/core/serialize.hpp>
 #include <raft/linalg/add.cuh>
 #include <raft/util/cuda_dev_essentials.cuh>
@@ -75,10 +76,10 @@ void deserialize(const raft::resources& clique,
   index.mode_        = (cuvs::neighbors::distribution_mode)deserialize_scalar<int>(handle, is);
   index.num_ranks_   = deserialize_scalar<int>(handle, is);
 
-  if (index.num_ranks_ != raft::resource::get_nccl_num_ranks(clique)) {
+  if (index.num_ranks_ != raft::resource::get_num_ranks(clique)) {
     RAFT_FAIL("Serialized index has %d ranks whereas NCCL clique has %d ranks",
               index.num_ranks_,
-              raft::resource::get_nccl_num_ranks(clique));
+              raft::resource::get_num_ranks(clique));
   }
 
   for (int rank = 0; rank < index.num_ranks_; rank++) {
@@ -215,8 +216,8 @@ void sharded_search_with_direct_merge(const raft::resources& clique,
       const raft::resources& dev_res = raft::resource::set_current_device_to_rank(clique, rank);
       auto& ann_if                   = index.ann_interfaces_[rank];
 
-      if (rank == raft::resource::get_nccl_clique_root_rank(clique)) {  // root rank
-        uint64_t batch_offset = raft::resource::get_nccl_clique_root_rank(clique) * part_size;
+      if (rank == raft::resource::get_root_rank(clique)) {  // root rank
+        uint64_t batch_offset = raft::resource::get_root_rank(clique) * part_size;
         auto d_neighbors      = raft::make_device_matrix_view<IdxT, int64_t, row_major>(
           in_neighbors.data_handle() + batch_offset, n_rows_of_current_batch, n_neighbors);
         auto d_distances = raft::make_device_matrix_view<float, int64_t, row_major>(
@@ -227,7 +228,7 @@ void sharded_search_with_direct_merge(const raft::resources& clique,
         // wait for other ranks
         ncclGroupStart();
         for (int from_rank = 0; from_rank < index.num_ranks_; from_rank++) {
-          if (from_rank == raft::resource::get_nccl_clique_root_rank(clique)) continue;
+          if (from_rank == raft::resource::get_root_rank(clique)) continue;
 
           batch_offset = from_rank * part_size;
           ncclRecv(in_neighbors.data_handle() + batch_offset,
@@ -258,13 +259,13 @@ void sharded_search_with_direct_merge(const raft::resources& clique,
         ncclSend(d_neighbors.data_handle(),
                  part_size * sizeof(IdxT),
                  ncclUint8,
-                 raft::resource::get_nccl_clique_root_rank(clique),
+                 raft::resource::get_root_rank(clique),
                  raft::resource::get_nccl_comm(dev_res),
                  raft::resource::get_cuda_stream(dev_res));
         ncclSend(d_distances.data_handle(),
                  part_size * sizeof(float),
                  ncclUint8,
-                 raft::resource::get_nccl_clique_root_rank(clique),
+                 raft::resource::get_root_rank(clique),
                  raft::resource::get_nccl_comm(dev_res),
                  raft::resource::get_cuda_stream(dev_res));
         ncclGroupEnd();
@@ -655,7 +656,7 @@ template <typename AnnIndexType, typename T, typename IdxT>
 mg_index<AnnIndexType, T, IdxT>::mg_index(const raft::resources& clique, distribution_mode mode)
   : mode_(mode), round_robin_counter_(std::make_shared<std::atomic<int64_t>>(0))
 {
-  num_ranks_ = raft::resource::get_nccl_num_ranks(clique);
+  num_ranks_ = raft::resource::get_num_ranks(clique);
 }
 
 template <typename AnnIndexType, typename T, typename IdxT>

cpp/tests/neighbors/mg.cuh

@@ -22,7 +22,7 @@
 #include <cuvs/neighbors/cagra.hpp>
 #include <cuvs/neighbors/ivf_flat.hpp>
 #include <cuvs/neighbors/ivf_pq.hpp>
-#include <raft/core/device_resources_snmg.hpp>
+#include <raft/core/device_resources_snmg_nccl.hpp>
 
 namespace cuvs::neighbors::mg {
 
@@ -645,7 +645,7 @@ class AnnMGTest : public ::testing::TestWithParam<AnnMGInputs> {
   void TearDown() override {}
 
  private:
-  raft::device_resources_snmg clique_;
+  raft::device_resources_snmg_nccl clique_;
   AnnMGInputs ps;
   std::vector<DataT> h_index_dataset;
   std::vector<DataT> h_queries;