Fold multi-GPU CAGRA build into train(), delete trainMultiGpu (#5500)

CMakeLists.txt

@@ -125,6 +125,10 @@ if(FAISS_ENABLE_CUVS AND NOT TARGET cuvs::cuvs)
    find_package(cuvs)
  endif()
 
+if(FAISS_ENABLE_CUVS AND NOT TARGET rmm::rmm)
+  find_package(rmm REQUIRED)
+endif()
+
 add_subdirectory(faiss)
 
 if(FAISS_ENABLE_GPU)

faiss/IndexHNSW.cpp

@@ -17,7 +17,6 @@
 #include <limits>
 #include <memory>
 #include <queue>
-#include <random>
 
 #include <cstdint>
 #include "faiss/Index.h"
@@ -1433,13 +1432,11 @@ void IndexHNSWCagra::search(
                 // first real candidate will always be strictly better.
                 nearest_d[i] = C::neutral();
 
-                std::random_device rd;
-                std::mt19937 gen(rd());
-                std::uniform_int_distribution<idx_t> distrib(
-                        0, this->ntotal - 1);
+                // Seeded per query so entrypoints are reproducible.
+                SplitMix64RandomGenerator gen(i);
 
                 for (idx_t j = 0; j < num_base_level_search_entrypoints; j++) {
-                    auto idx = distrib(gen);
+                    idx_t idx = gen.rand_int64() % this->ntotal;
                     auto distance = (*dis)(idx);
                     if (C::cmp(nearest_d[i], distance)) {
                         nearest[i] = static_cast<storage_idx_t>(idx);
@@ -1498,12 +1495,11 @@ void IndexHNSWCagra::range_search(
             // real candidate will always be strictly better.
             float nearest_d = C::neutral();
 
-            std::random_device rd;
-            std::mt19937 gen(rd());
-            std::uniform_int_distribution<idx_t> distrib(0, ntotal - 1);
+            // For reproducible entrypoint.
+            SplitMix64RandomGenerator gen(i);
 
             for (idx_t j = 0; j < num_base_level_search_entrypoints; j++) {
-                auto idx = distrib(gen);
+                idx_t idx = gen.rand_int64() % ntotal;
                 auto distance = (*dis)(idx);
                 // C::cmp(nearest_d, distance) is true iff distance is
                 // strictly better than the current nearest_d.

faiss/IndexHNSW.h

@@ -253,7 +253,7 @@ struct IndexHNSWCagra : IndexHNSW {
     /// searches only the base level knn graph of the HNSW index.
     /// This parameter selects the entry point by randomly selecting
     /// some points and using the best one.
-    int num_base_level_search_entrypoints = 32;
+    int num_base_level_search_entrypoints = 256;
 
     void add(idx_t n, const float* x) override;
 

faiss/gpu/CMakeLists.txt

@@ -363,7 +363,7 @@ else()
 
 
   find_package(CUDAToolkit REQUIRED)
-  target_link_libraries(faiss_gpu_objs PRIVATE ${CUDA_LIBS} $<$<BOOL:${FAISS_ENABLE_CUVS}>:cuvs::cuvs> $<$<BOOL:${FAISS_ENABLE_CUVS}>:OpenMP::OpenMP_CXX>)
+  target_link_libraries(faiss_gpu_objs PRIVATE ${CUDA_LIBS} $<$<BOOL:${FAISS_ENABLE_CUVS}>:cuvs::cuvs> $<$<BOOL:${FAISS_ENABLE_CUVS}>:rmm::rmm> $<$<BOOL:${FAISS_ENABLE_CUVS}>:OpenMP::OpenMP_CXX>)
   target_compile_options(faiss_gpu_objs PRIVATE
     $<$<COMPILE_LANGUAGE:CUDA>:-Xfatbin=-compress-all
     --expt-extended-lambda --expt-relaxed-constexpr

faiss/gpu/GpuIndexBinaryCagra.cu

@@ -78,6 +78,13 @@ std::shared_ptr<GpuResources> GpuIndexBinaryCagra::getResources() {
 }
 
 void GpuIndexBinaryCagra::train(idx_t n, const uint8_t* x) {
+    // The config is shared with the float index; there is no binary
+    // multi-GPU build, so reject rather than silently ignoring the request.
+    FAISS_THROW_IF_MSG(
+            cagraConfig_.devices.size() > 1,
+            "binary CAGRA has no multi-GPU build; "
+            "GpuIndexCagraConfig::devices must name at most one device");
+
     DeviceScope scope(cagraConfig_.device);
     if (this->is_trained) {
         FAISS_ASSERT(index_);