Add FAISS with RAFT enabled Benchmarking to raft-ann-bench

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

@@ -18,6 +18,7 @@
 #include "ann_types.hpp"
 #include "conf.hpp"
 #include "dataset.hpp"
+#include "raft/util/cudart_utils.hpp"
 #include "util.hpp"
 
 #include <benchmark/benchmark.h>
@@ -317,6 +318,7 @@ void bench_search(::benchmark::State& state,
   }
   auto end      = std::chrono::high_resolution_clock::now();
   auto duration = std::chrono::duration_cast<std::chrono::duration<double>>(end - start).count();
+  // std::cout << "duration" << duration << std::endl;
   if (state.thread_index() == 0) { state.counters.insert({{"end_to_end", duration}}); }
   state.counters.insert(
     {"Latency", {duration / double(state.iterations()), benchmark::Counter::kAvgThreads}});

cpp/bench/ann/src/faiss/faiss_gpu_benchmark.cu

@@ -43,6 +43,11 @@ void parse_build_param(const nlohmann::json& conf,
                        typename raft::bench::ann::FaissGpuIVFFlat<T>::BuildParam& param)
 {
   parse_base_build_param<T>(conf, param);
+  if (conf.contains("use_raft")) {
+    param.use_raft = conf.at("use_raft");
+  } else {
+    param.use_raft = false;
+  }
 }
 
 template <typename T>
@@ -61,6 +66,16 @@ void parse_build_param(const nlohmann::json& conf,
   } else {
     param.useFloat16 = false;
   }
+  if (conf.contains("use_raft")) {
+    param.use_raft = conf.at("use_raft");
+  } else {
+    param.use_raft = false;
+  }
+  if (conf.contains("bitsPerCode")) {
+    param.bitsPerCode = conf.at("bitsPerCode");
+  } else {
+    param.bitsPerCode = 8;
+  }
 }
 
 template <typename T>
@@ -77,6 +92,12 @@ void parse_search_param(const nlohmann::json& conf,
 {
   param.nprobe = conf.at("nprobe");
   if (conf.contains("refine_ratio")) { param.refine_ratio = conf.at("refine_ratio"); }
+  if (conf.contains("raft_refinement")) {
+    RAFT_LOG_INFO("found raft_refinement");
+    param.raft_refinement = conf.at("raft_refinement");
+  } else {
+    param.raft_refinement = false;
+  }
 }
 
 template <typename T, template <typename> class Algo>
@@ -158,5 +179,15 @@ REGISTER_ALGO_INSTANCE(std::uint8_t);
 
 #ifdef ANN_BENCH_BUILD_MAIN
 #include "../common/benchmark.hpp"
-int main(int argc, char** argv) { return raft::bench::ann::run_main(argc, argv); }
+int main(int argc, char** argv)
+{
+  rmm::mr::cuda_memory_resource cuda_mr;
+  // Construct a resource that uses a coalescing best-fit pool allocator
+  rmm::mr::pool_memory_resource<rmm::mr::cuda_memory_resource> pool_mr{&cuda_mr};
+  rmm::mr::set_current_device_resource(
+    &pool_mr);  // Updates the current device resource pointer to `pool_mr`
+  rmm::mr::device_memory_resource* mr =
+    rmm::mr::get_current_device_resource();  // Points to `pool_mr`
+  return raft::bench::ann::run_main(argc, argv);
+}
 #endif

cpp/bench/ann/src/faiss/faiss_gpu_wrapper.h

@@ -17,10 +17,15 @@
 #define FAISS_WRAPPER_H_
 
 #include "../common/ann_types.hpp"
+#include <raft/core/device_mdarray.hpp>
+#include <raft/core/host_mdarray.hpp>
+#include <raft/core/host_mdspan.hpp>
+#include <raft/distance/distance_types.hpp>
 
 #include <raft/core/logger.hpp>
 #include <raft/util/cudart_utils.hpp>
 
+#include <faiss/MetricType.h>
 #include <faiss/IndexFlat.h>
 #include <faiss/IndexIVFFlat.h>
 #include <faiss/IndexIVFPQ.h>
@@ -37,6 +42,10 @@
 
 #include <raft/core/device_resources.hpp>
 #include <raft/core/resource/stream_view.hpp>
+#include <raft_runtime/neighbors/refine.hpp>
+#include <rmm/cuda_device.hpp>
+#include <rmm/mr/device/device_memory_resource.hpp>
+#include <rmm/mr/device/per_device_resource.hpp>
 
 #include <cassert>
 #include <memory>
@@ -99,7 +108,8 @@ class FaissGpu : public ANN<T> {
   using typename ANN<T>::AnnSearchParam;
   struct SearchParam : public AnnSearchParam {
     int nprobe;
-    float refine_ratio = 1.0;
+    float refine_ratio   = 1.0;
+    bool raft_refinement = false;
     auto needs_dataset() const -> bool override { return refine_ratio > 1.0f; }
   };
 
@@ -143,6 +153,8 @@ class FaissGpu : public ANN<T> {
     return property;
   }
 
+  auto metric_faiss_to_raft(faiss::MetricType metric) const -> raft::distance::DistanceType;
+
  protected:
   template <typename GpuIndex, typename CpuIndex>
   void save_(const std::string& file) const;
@@ -181,13 +193,27 @@ class FaissGpu : public ANN<T> {
   copyable_event sync_{};
   double training_sample_fraction_;
   std::shared_ptr<faiss::SearchParameters> search_params_;
+  std::shared_ptr<faiss::IndexRefineSearchParameters> refine_search_params_{nullptr};
   const T* dataset_;
-  float refine_ratio_ = 1.0;
+  float refine_ratio_   = 1.0;
+  bool raft_refinement_ = false;
 };
 
+template <typename T>
+auto FaissGpu<T>::metric_faiss_to_raft(faiss::MetricType metric) const
+  -> raft::distance::DistanceType
+{
+  switch (metric) {
+    case faiss::MetricType::METRIC_L2: return raft::distance::DistanceType::L2Expanded;
+    case faiss::MetricType::METRIC_INNER_PRODUCT:
+    default: throw std::runtime_error("FAISS supports only metric type of inner product and L2");
+  }
+}
+
 template <typename T>
 void FaissGpu<T>::build(const T* dataset, size_t nrow, cudaStream_t stream)
 {
+  // raft::print_host_vector("faiss dataset", dataset, 100, std::cout);
   OmpSingleThreadScope omp_single_thread;
   auto index_ivf = dynamic_cast<faiss::gpu::GpuIndexIVF*>(index_.get());
   if (index_ivf != nullptr) {
@@ -208,7 +234,7 @@ void FaissGpu<T>::build(const T* dataset, size_t nrow, cudaStream_t stream)
         nlist_,
         index_ivf->cp.min_points_per_centroid);
     }
-    index_ivf->cp.max_points_per_centroid = max_ppc;
+    index_ivf->cp.max_points_per_centroid = 300;
     index_ivf->cp.min_points_per_centroid = min_ppc;
   }
   index_->train(nrow, dataset);  // faiss::gpu::GpuIndexFlat::train() will do nothing
@@ -225,19 +251,79 @@ void FaissGpu<T>::search(const T* queries,
                          float* distances,
                          cudaStream_t stream) const
 {
+  using IdxT = faiss::idx_t;
   static_assert(sizeof(size_t) == sizeof(faiss::idx_t),
                 "sizes of size_t and faiss::idx_t are different");
 
-  if (this->refine_ratio_ > 1.0) {
-    // TODO: FAISS changed their search APIs to accept the search parameters as a struct object
-    // but their refine API doesn't allow the struct to be passed in. Once this is fixed, we
-    // need to re-enable refinement below
-    // index_refine_->search(batch_size, queries, k, distances,
-    // reinterpret_cast<faiss::idx_t*>(neighbors), this->search_params_.get()); Related FAISS issue:
-    // https://github.com/facebookresearch/faiss/issues/3118
-    throw std::runtime_error(
-      "FAISS doesn't support refinement in their new APIs so this feature is disabled in the "
-      "benchmarks for the time being.");
+  if (refine_ratio_ > 1.0) {
+    if (raft_refinement_) {
+      uint32_t k0        = static_cast<uint32_t>(refine_ratio_ * k);
+      // auto distances_tmp = raft::make_host_matrix<float, IdxT>(batch_size, k0);
+      // auto candidates    = raft::make_host_matrix<IdxT, IdxT>(batch_size, k0);
+      auto distances_tmp = raft::make_device_matrix<float, IdxT>(gpu_resource_->getRaftHandle(device_), batch_size, k0);
+      auto candidates    = raft::make_device_matrix<IdxT, IdxT>(gpu_resource_->getRaftHandle(device_), batch_size, k0);
+      index_->search(batch_size,
+                     queries,
+                     k0,
+                     distances_tmp.data_handle(),
+                     candidates.data_handle(),
+                     this->search_params_.get());
+      // auto queries_v = raft::make_host_matrix_view<const T, IdxT>(queries, batch_size, index_->d);
+
+
+      // auto dataset_v = raft::make_host_matrix_view<const T, faiss::idx_t>(
+      //   this->dataset_, index_->ntotal, index_->d);
+
+      // auto neighbors_v =
+      //   raft::make_host_matrix_view<IdxT, IdxT>(reinterpret_cast<IdxT*>(neighbors), batch_size, k);
+      // auto distances_v = raft::make_host_matrix_view<float, IdxT>(distances, batch_size, k);
+
+      // raft::runtime::neighbors::refine(gpu_resource_->getRaftHandle(device_),
+      //                                  dataset_v,
+      //                                  queries_v,
+      //                                  candidates.view(),
+      //                                  neighbors_v,
+      //                                  distances_v,
+      //                                  metric_faiss_to_raft(index_->metric_type));
+
+      auto queries_host    = raft::make_host_matrix<T, IdxT>(batch_size, index_->d);
+      auto candidates_host = raft::make_host_matrix<IdxT, IdxT>(batch_size, k0);
+      auto neighbors_host  = raft::make_host_matrix<IdxT, IdxT>(batch_size, k);
+      auto distances_host  = raft::make_host_matrix<float, IdxT>(batch_size, k);
+      auto dataset_v = raft::make_host_matrix_view<const T, faiss::idx_t>(
+        this->dataset_, index_->ntotal, index_->d);
+
+      auto handle_ = gpu_resource_->getRaftHandle(device_);
+
+      raft::copy(queries_host.data_handle(), queries, queries_host.size(), stream);
+      raft::copy(candidates_host.data_handle(),
+                 candidates.data_handle(),
+                 candidates_host.size(),
+                 resource::get_cuda_stream(handle_));
+
+      // wait for the queries to copy to host in 'stream` and for IVF-PQ::search to finish
+      // RAFT_CUDA_TRY(cudaEventRecord(handle_.get_sync_event(), resource::get_cuda_stream(handle_)));
+      // RAFT_CUDA_TRY(cudaEventRecord(handle_.get_sync_event(), stream));
+      // RAFT_CUDA_TRY(cudaEventSynchronize(handle_.get_sync_event()));
+      handle_.sync_stream();
+      raft::runtime::neighbors::refine(handle_,
+                                       dataset_v,
+                                       queries_host.view(),
+                                       candidates_host.view(),
+                                       neighbors_host.view(),
+                                       distances_host.view(),
+                                       metric_faiss_to_raft(index_->metric_type));
+
+      raft::copy(neighbors, (size_t*)neighbors_host.data_handle(), neighbors_host.size(), stream);
+      raft::copy(distances, distances_host.data_handle(), distances_host.size(), stream);
+    } else {
+      index_refine_->search(batch_size,
+                            queries,
+                            k,
+                            distances,
+                            reinterpret_cast<faiss::idx_t*>(neighbors),
+                            this->refine_search_params_.get());
+    }
   } else {
     index_->search(batch_size,
                    queries,
@@ -280,13 +366,16 @@ void FaissGpu<T>::load_(const std::string& file)
 template <typename T>
 class FaissGpuIVFFlat : public FaissGpu<T> {
  public:
-  using typename FaissGpu<T>::BuildParam;
+  struct BuildParam : public FaissGpu<T>::BuildParam {
+    bool use_raft;
+  };
 
   FaissGpuIVFFlat(Metric metric, int dim, const BuildParam& param) : FaissGpu<T>(metric, dim, param)
   {
     faiss::gpu::GpuIndexIVFFlatConfig config;
-    config.device = this->device_;
-    this->index_  = std::make_shared<faiss::gpu::GpuIndexIVFFlat>(
+    config.device   = this->device_;
+    config.use_raft = param.use_raft;
+    this->index_    = std::make_shared<faiss::gpu::GpuIndexIVFFlat>(
       this->gpu_resource_.get(), dim, param.nlist, this->metric_type_, config);
   }
 
@@ -320,21 +409,25 @@ class FaissGpuIVFPQ : public FaissGpu<T> {
     int M;
     bool useFloat16;
     bool usePrecomputed;
+    bool use_raft;
+    int bitsPerCode;
   };
 
   FaissGpuIVFPQ(Metric metric, int dim, const BuildParam& param) : FaissGpu<T>(metric, dim, param)
   {
     faiss::gpu::GpuIndexIVFPQConfig config;
     config.useFloat16LookupTables = param.useFloat16;
     config.usePrecomputedTables   = param.usePrecomputed;
+    config.use_raft               = param.use_raft;
+    config.interleavedLayout      = param.use_raft;
     config.device                 = this->device_;
 
     this->index_ =
       std::make_shared<faiss::gpu::GpuIndexIVFPQ>(this->gpu_resource_.get(),
                                                   dim,
                                                   param.nlist,
                                                   param.M,
-                                                  8,  // FAISS only supports bitsPerCode=8
+                                                  param.bitsPerCode,
                                                   this->metric_type_,
                                                   config);
   }
@@ -354,7 +447,14 @@ class FaissGpuIVFPQ : public FaissGpu<T> {
       this->index_refine_ =
         std::make_shared<faiss::IndexRefineFlat>(this->index_.get(), this->dataset_);
       this->index_refine_.get()->k_factor = search_param.refine_ratio;
+      faiss::IndexRefineSearchParameters faiss_refine_search_params;
+      faiss_refine_search_params.k_factor          = this->index_refine_.get()->k_factor;
+      faiss_refine_search_params.base_index_params = this->search_params_.get();
+      this->refine_search_params_ =
+        std::make_unique<faiss::IndexRefineSearchParameters>(faiss_refine_search_params);
     }
+    this->raft_refinement_     = search_param.raft_refinement;
+    RAFT_LOG_INFO("refine_ratio %f raft_refinement %d", this->refine_ratio_, this->raft_refinement_);
   }
 
   void save(const std::string& file) const override
@@ -410,6 +510,11 @@ class FaissGpuIVFSQ : public FaissGpu<T> {
       this->index_refine_ =
         std::make_shared<faiss::IndexRefineFlat>(this->index_.get(), this->dataset_);
       this->index_refine_.get()->k_factor = search_param.refine_ratio;
+      faiss::IndexRefineSearchParameters faiss_refine_search_params;
+      faiss_refine_search_params.k_factor          = this->index_refine_.get()->k_factor;
+      faiss_refine_search_params.base_index_params = this->search_params_.get();
+      this->refine_search_params_ =
+        std::make_unique<faiss::IndexRefineSearchParameters>(faiss_refine_search_params);
     }
   }
 

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

@@ -78,7 +78,7 @@ class RaftIvfPQ : public ANN<T> {
   {
     AlgoProperty property;
     property.dataset_memory_type = MemoryType::Host;
-    property.query_memory_type   = MemoryType::Device;
+    property.query_memory_type   = MemoryType::Host;
     return property;
   }
   void save(const std::string& file) const override;
@@ -209,7 +209,7 @@ void RaftIvfPQ<T, IdxT>::search(const T* queries,
       raft::make_device_matrix_view<const T, IdxT>(queries, batch_size, index_->dim());
     auto neighbors_v = raft::make_device_matrix_view<IdxT, IdxT>((IdxT*)neighbors, batch_size, k);
     auto distances_v = raft::make_device_matrix_view<float, IdxT>(distances, batch_size, k);
-
+    
     raft::runtime::neighbors::ivf_pq::search(
       handle_, search_params_, *index_, queries_v, neighbors_v, distances_v);
     handle_.stream_wait(stream);  // RAFT stream -> bench stream

cpp/cmake/thirdparty/get_faiss.cmake

@@ -50,6 +50,7 @@ function(find_and_configure_faiss)
                 EXCLUDE_FROM_ALL ${PKG_EXCLUDE_FROM_ALL}
                 OPTIONS
                 "FAISS_ENABLE_GPU ${PKG_ENABLE_GPU}"
+                "FAISS_ENABLE_RAFT ON"
                 "FAISS_ENABLE_PYTHON OFF"
                 "FAISS_OPT_LEVEL ${RAFT_FAISS_OPT_LEVEL}"
                 "FAISS_USE_CUDA_TOOLKIT_STATIC ${CUDA_STATIC_RUNTIME}"
@@ -90,14 +91,16 @@ endfunction()
 if(NOT RAFT_FAISS_GIT_TAG)
     # TODO: Remove this once faiss supports FAISS_USE_CUDA_TOOLKIT_STATIC
     # (https://github.com/facebookresearch/faiss/pull/2446)
-    set(RAFT_FAISS_GIT_TAG fea/statically-link-ctk)
+    set(RAFT_FAISS_GIT_TAG rmm-pool-alloc)
+    # set(RAFT_FAISS_GIT_TAG fea/statically-link-ctk)
     # set(RAFT_FAISS_GIT_TAG bde7c0027191f29c9dadafe4f6e68ca0ee31fb30)
 endif()
 
 if(NOT RAFT_FAISS_GIT_REPOSITORY)
     # TODO: Remove this once faiss supports FAISS_USE_CUDA_TOOLKIT_STATIC
     # (https://github.com/facebookresearch/faiss/pull/2446)
-    set(RAFT_FAISS_GIT_REPOSITORY https://github.com/cjnolet/faiss.git)
+    set(RAFT_FAISS_GIT_REPOSITORY https://github.com/tarang-jain/faiss.git)
+    # set(RAFT_FAISS_GIT_REPOSITORY https://github.com/cjnolet/faiss.git)
     # set(RAFT_FAISS_GIT_REPOSITORY https://github.com/facebookresearch/faiss.git)
 endif()