[Feat] CAGRA filtering with BFKNN when sparsity matching threshold

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

@@ -43,6 +43,7 @@
 #include <iostream>
 #include <memory>
 #include <optional>
+#include <random>
 #include <stdexcept>
 #include <string>
 #include <type_traits>
@@ -52,10 +53,13 @@ namespace cuvs::bench {
 enum class AllocatorType { kHostPinned, kHostHugePage, kDevice };
 enum class CagraBuildAlgo { kAuto, kIvfPq, kNnDescent };
 
+constexpr double sparsity = 0.0f;
+
 template <typename T, typename IdxT>
 class cuvs_cagra : public algo<T>, public algo_gpu {
  public:
   using search_param_base = typename algo<T>::search_param;
+  // TODO: Move to arguments
 
   struct search_param : public search_param_base {
     cuvs::neighbors::cagra::search_params p;
@@ -91,6 +95,40 @@ class cuvs_cagra : public algo<T>, public algo_gpu {
     }
   };
 
+  int64_t create_sparse_bitset(int64_t total, float sparsity, std::vector<uint32_t>& bitset) const
+  {
+    int64_t num_ones = static_cast<int64_t>((total * 1.0f) * (1.0f - sparsity));
+    int64_t res      = num_ones;
+
+    for (auto& item : bitset) {
+      item = static_cast<uint32_t>(0);
+    }
+
+    if (sparsity == 0.0) {
+      for (auto& item : bitset) {
+        item = static_cast<uint32_t>(0xffffffff);
+      }
+      return total;
+    }
+
+    std::random_device rd;
+    std::mt19937 gen(rd());
+    std::uniform_int_distribution<int64_t> dis(0, total - 1);
+
+    while (num_ones > 0) {
+      int64_t index = dis(gen);
+
+      uint32_t& element    = bitset[index / (8 * sizeof(uint32_t))];
+      int64_t bit_position = index % (8 * sizeof(uint32_t));
+
+      if (((element >> bit_position) & 1) == 0) {
+        element |= (static_cast<uint32_t>(1) << bit_position);
+        num_ones--;
+      }
+    }
+    return res;
+  }
+
   cuvs_cagra(Metric metric, int dim, const build_param& param, int concurrent_searches = 1)
     : algo<T>(metric, dim),
       index_params_(param),
@@ -102,8 +140,9 @@ class cuvs_cagra : public algo<T>, public algo_gpu {
         std::move(raft::make_device_matrix<IdxT, int64_t>(handle_, 0, 0)))),
       input_dataset_v_(
         std::make_shared<raft::device_matrix_view<const T, int64_t, raft::row_major>>(
-          nullptr, 0, 0))
-
+          nullptr, 0, 0)),
+      bitset_filter_(std::make_shared<cuvs::core::bitset<std::uint32_t, int64_t>>(
+        std::move(cuvs::core::bitset<std::uint32_t, int64_t>(handle_, 0, false))))
   {
     index_params_.cagra_params.metric         = parse_metric_type(metric);
     index_params_.ivf_pq_build_params->metric = parse_metric_type(metric);
@@ -171,6 +210,9 @@ class cuvs_cagra : public algo<T>, public algo_gpu {
   std::shared_ptr<raft::device_matrix<T, int64_t, raft::row_major>> dataset_;
   std::shared_ptr<raft::device_matrix_view<const T, int64_t, raft::row_major>> input_dataset_v_;
 
+  //  std::shared_ptr<raft::device_vector<std::uint32_t, int64_t>> bitset_filter_;
+  std::shared_ptr<cuvs::core::bitset<std::uint32_t, int64_t>> bitset_filter_;
+
   inline rmm::device_async_resource_ref get_mr(AllocatorType mem_type)
   {
     switch (mem_type) {
@@ -256,6 +298,15 @@ void cuvs_cagra<T, IdxT>::set_search_param(const search_param_base& param)
 
     need_dataset_update_ = false;
   }
+
+  {  // create bitset filter in advance.
+    auto stream_             = raft::resource::get_cuda_stream(handle_);
+    size_t filter_n_elements = size_t((input_dataset_v_->extent(0) + 31) / 32);
+    bitset_filter_->resize(handle_, input_dataset_v_->extent(0), false);
+    std::vector<std::uint32_t> bitset_cpu(filter_n_elements);
+    create_sparse_bitset(input_dataset_v_->extent(0), sparsity, bitset_cpu);
+    raft::copy(bitset_filter_->data(), bitset_cpu.data(), filter_n_elements, stream_);
+  }
 }
 
 template <typename T, typename IdxT>
@@ -328,8 +379,15 @@ void cuvs_cagra<T, IdxT>::search_base(const T* queries,
     raft::make_device_matrix_view<IdxT, int64_t>(neighbors_idx_t, batch_size, k);
   auto distances_view = raft::make_device_matrix_view<float, int64_t>(distances, batch_size, k);
 
-  cuvs::neighbors::cagra::search(
-    handle_, search_params_, *index_, queries_view, neighbors_view, distances_view);
+  if constexpr ((std::is_same_v<T, float> || std::is_same_v<T, half>)&&sparsity >= 0.0f) {
+    auto filter = cuvs::neighbors::filtering::bitset_filter(bitset_filter_->view());
+    cuvs::neighbors::cagra::search(
+      handle_, search_params_, *index_, queries_view, neighbors_view, distances_view, filter);
+
+  } else {
+    cuvs::neighbors::cagra::search(
+      handle_, search_params_, *index_, queries_view, neighbors_view, distances_view);
+  }
 
   if constexpr (sizeof(IdxT) != sizeof(algo_base::index_type)) {
     if (raft::get_device_for_address(neighbors) < 0 &&

cpp/include/cuvs/neighbors/cagra.hpp

@@ -1057,6 +1057,8 @@ void extend(
  * k]
  * @param[in] sample_filter an optional device filter function object that greenlights samples
  * for a given query. (none_sample_filter for no filtering)
+ * @param[in] threshold_to_bf A sparsity threshold; brute force is used when sparsity exceeds this
+ * threshold, in the range [0, 1]
  */
 
 void search(raft::resources const& res,
@@ -1066,7 +1068,8 @@ void search(raft::resources const& res,
             raft::device_matrix_view<uint32_t, int64_t, raft::row_major> neighbors,
             raft::device_matrix_view<float, int64_t, raft::row_major> distances,
             const cuvs::neighbors::filtering::base_filter& sample_filter =
-              cuvs::neighbors::filtering::none_sample_filter{});
+              cuvs::neighbors::filtering::none_sample_filter{},
+            double threshold_to_bf = 0.9f);
 
 /**
  * @brief Search ANN using the constructed index.
@@ -1083,6 +1086,8 @@ void search(raft::resources const& res,
  * k]
  * @param[in] sample_filter an optional device filter function object that greenlights samples
  * for a given query. (none_sample_filter for no filtering)
+ * @param[in] threshold_to_bf A sparsity threshold; brute force is used when sparsity exceeds this
+ * threshold, in the range [0, 1]
  */
 void search(raft::resources const& res,
             cuvs::neighbors::cagra::search_params const& params,
@@ -1091,7 +1096,8 @@ void search(raft::resources const& res,
             raft::device_matrix_view<uint32_t, int64_t, raft::row_major> neighbors,
             raft::device_matrix_view<float, int64_t, raft::row_major> distances,
             const cuvs::neighbors::filtering::base_filter& sample_filter =
-              cuvs::neighbors::filtering::none_sample_filter{});
+              cuvs::neighbors::filtering::none_sample_filter{},
+            double threshold_to_bf = 0.9f);
 
 /**
  * @brief Search ANN using the constructed index.
@@ -1108,6 +1114,8 @@ void search(raft::resources const& res,
  * k]
  * @param[in] sample_filter an optional device filter function object that greenlights samples
  * for a given query. (none_sample_filter for no filtering)
+ * @param[in] threshold_to_bf A sparsity threshold; brute force is used when sparsity exceeds this
+ * threshold, in the range [0, 1]
  */
 void search(raft::resources const& res,
             cuvs::neighbors::cagra::search_params const& params,
@@ -1116,7 +1124,8 @@ void search(raft::resources const& res,
             raft::device_matrix_view<uint32_t, int64_t, raft::row_major> neighbors,
             raft::device_matrix_view<float, int64_t, raft::row_major> distances,
             const cuvs::neighbors::filtering::base_filter& sample_filter =
-              cuvs::neighbors::filtering::none_sample_filter{});
+              cuvs::neighbors::filtering::none_sample_filter{},
+            double threshold_to_bf = 0.9f);
 
 /**
  * @brief Search ANN using the constructed index.
@@ -1133,6 +1142,8 @@ void search(raft::resources const& res,
  * k]
  * @param[in] sample_filter an optional device filter function object that greenlights samples
  * for a given query. (none_sample_filter for no filtering)
+ * @param[in] threshold_to_bf A sparsity threshold; brute force is used when sparsity exceeds this
+ * threshold, in the range [0, 1]
  */
 void search(raft::resources const& res,
             cuvs::neighbors::cagra::search_params const& params,
@@ -1141,7 +1152,8 @@ void search(raft::resources const& res,
             raft::device_matrix_view<uint32_t, int64_t, raft::row_major> neighbors,
             raft::device_matrix_view<float, int64_t, raft::row_major> distances,
             const cuvs::neighbors::filtering::base_filter& sample_filter =
-              cuvs::neighbors::filtering::none_sample_filter{});
+              cuvs::neighbors::filtering::none_sample_filter{},
+            double threshold_to_bf = 0.9f);
 
 /**
  * @}

cpp/src/neighbors/cagra.cuh

@@ -293,6 +293,9 @@ index<T, IdxT> build(
  * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries,
  * k]
  * @param[in] sample_filter a device filter function that greenlights samples for a given query
+ * @param[in] threshold_to_bf A sparsity threshold; brute force is used when sparsity exceeds this
+ * threshold, in the range [0, 1]
+ *
  */
 template <typename T, typename IdxT, typename CagraSampleFilterT>
 void search_with_filtering(raft::resources const& res,
@@ -301,7 +304,8 @@ void search_with_filtering(raft::resources const& res,
                            raft::device_matrix_view<const T, int64_t, raft::row_major> queries,
                            raft::device_matrix_view<IdxT, int64_t, raft::row_major> neighbors,
                            raft::device_matrix_view<float, int64_t, raft::row_major> distances,
-                           CagraSampleFilterT sample_filter = CagraSampleFilterT())
+                           CagraSampleFilterT sample_filter = CagraSampleFilterT(),
+                           double threshold_to_bf           = 0.9)
 {
   RAFT_EXPECTS(
     queries.extent(0) == neighbors.extent(0) && queries.extent(0) == distances.extent(0),
@@ -322,8 +326,14 @@ void search_with_filtering(raft::resources const& res,
   auto distances_internal = raft::make_device_matrix_view<float, int64_t, raft::row_major>(
     distances.data_handle(), distances.extent(0), distances.extent(1));
 
-  return cagra::detail::search_main<T, internal_IdxT, CagraSampleFilterT, IdxT>(
-    res, params, idx, queries_internal, neighbors_internal, distances_internal, sample_filter);
+  return cagra::detail::search_main<T, internal_IdxT, CagraSampleFilterT, IdxT>(res,
+                                                                                params,
+                                                                                idx,
+                                                                                queries_internal,
+                                                                                neighbors_internal,
+                                                                                distances_internal,
+                                                                                sample_filter,
+                                                                                threshold_to_bf);
 }
 
 template <typename T, typename IdxT>
@@ -333,14 +343,15 @@ void search(raft::resources const& res,
             raft::device_matrix_view<const T, int64_t, raft::row_major> queries,
             raft::device_matrix_view<IdxT, int64_t, raft::row_major> neighbors,
             raft::device_matrix_view<float, int64_t, raft::row_major> distances,
-            const cuvs::neighbors::filtering::base_filter& sample_filter_ref)
+            const cuvs::neighbors::filtering::base_filter& sample_filter_ref,
+            double threshold_to_bf = 0.9)
 {
   try {
     using none_filter_type  = cuvs::neighbors::filtering::none_sample_filter;
     auto& sample_filter     = dynamic_cast<const none_filter_type&>(sample_filter_ref);
     auto sample_filter_copy = sample_filter;
     return search_with_filtering<T, IdxT, none_filter_type>(
-      res, params, idx, queries, neighbors, distances, sample_filter_copy);
+      res, params, idx, queries, neighbors, distances, sample_filter_copy, threshold_to_bf);
     return;
   } catch (const std::bad_cast&) {
   }
@@ -351,7 +362,7 @@ void search(raft::resources const& res,
         sample_filter_ref);
     auto sample_filter_copy = sample_filter;
     return search_with_filtering<T, IdxT, decltype(sample_filter_copy)>(
-      res, params, idx, queries, neighbors, distances, sample_filter_copy);
+      res, params, idx, queries, neighbors, distances, sample_filter_copy, threshold_to_bf);
   } catch (const std::bad_cast&) {
     RAFT_FAIL("Unsupported sample filter type");
   }

cpp/src/neighbors/cagra_search_float.cu

@@ -19,17 +19,18 @@
 
 namespace cuvs::neighbors::cagra {
 
-#define CUVS_INST_CAGRA_SEARCH(T, IdxT)                                            \
-  void search(raft::resources const& handle,                                       \
-              cuvs::neighbors::cagra::search_params const& params,                 \
-              const cuvs::neighbors::cagra::index<T, IdxT>& index,                 \
-              raft::device_matrix_view<const T, int64_t, raft::row_major> queries, \
-              raft::device_matrix_view<IdxT, int64_t, raft::row_major> neighbors,  \
-              raft::device_matrix_view<float, int64_t, raft::row_major> distances, \
-              const cuvs::neighbors::filtering::base_filter& sample_filter)        \
-  {                                                                                \
-    cuvs::neighbors::cagra::search<T, IdxT>(                                       \
-      handle, params, index, queries, neighbors, distances, sample_filter);        \
+#define CUVS_INST_CAGRA_SEARCH(T, IdxT)                                                      \
+  void search(raft::resources const& handle,                                                 \
+              cuvs::neighbors::cagra::search_params const& params,                           \
+              const cuvs::neighbors::cagra::index<T, IdxT>& index,                           \
+              raft::device_matrix_view<const T, int64_t, raft::row_major> queries,           \
+              raft::device_matrix_view<IdxT, int64_t, raft::row_major> neighbors,            \
+              raft::device_matrix_view<float, int64_t, raft::row_major> distances,           \
+              const cuvs::neighbors::filtering::base_filter& sample_filter,                  \
+              double threshold_to_bf)                                                        \
+  {                                                                                          \
+    cuvs::neighbors::cagra::search<T, IdxT>(                                                 \
+      handle, params, index, queries, neighbors, distances, sample_filter, threshold_to_bf); \
   }
 
 CUVS_INST_CAGRA_SEARCH(float, uint32_t);

cpp/src/neighbors/cagra_search_half.cu

@@ -19,17 +19,18 @@
 
 namespace cuvs::neighbors::cagra {
 
-#define CUVS_INST_CAGRA_SEARCH(T, IdxT)                                            \
-  void search(raft::resources const& handle,                                       \
-              cuvs::neighbors::cagra::search_params const& params,                 \
-              const cuvs::neighbors::cagra::index<T, IdxT>& index,                 \
-              raft::device_matrix_view<const T, int64_t, raft::row_major> queries, \
-              raft::device_matrix_view<IdxT, int64_t, raft::row_major> neighbors,  \
-              raft::device_matrix_view<float, int64_t, raft::row_major> distances, \
-              const cuvs::neighbors::filtering::base_filter& sample_filter)        \
-  {                                                                                \
-    cuvs::neighbors::cagra::search<T, IdxT>(                                       \
-      handle, params, index, queries, neighbors, distances, sample_filter);        \
+#define CUVS_INST_CAGRA_SEARCH(T, IdxT)                                                      \
+  void search(raft::resources const& handle,                                                 \
+              cuvs::neighbors::cagra::search_params const& params,                           \
+              const cuvs::neighbors::cagra::index<T, IdxT>& index,                           \
+              raft::device_matrix_view<const T, int64_t, raft::row_major> queries,           \
+              raft::device_matrix_view<IdxT, int64_t, raft::row_major> neighbors,            \
+              raft::device_matrix_view<float, int64_t, raft::row_major> distances,           \
+              const cuvs::neighbors::filtering::base_filter& sample_filter,                  \
+              double threshold_to_bf)                                                        \
+  {                                                                                          \
+    cuvs::neighbors::cagra::search<T, IdxT>(                                                 \
+      handle, params, index, queries, neighbors, distances, sample_filter, threshold_to_bf); \
   }
 
 CUVS_INST_CAGRA_SEARCH(half, uint32_t);

cpp/src/neighbors/cagra_search_int8.cu

@@ -18,17 +18,18 @@
 #include <cuvs/neighbors/cagra.hpp>
 namespace cuvs::neighbors::cagra {
 
-#define CUVS_INST_CAGRA_SEARCH(T, IdxT)                                            \
-  void search(raft::resources const& handle,                                       \
-              cuvs::neighbors::cagra::search_params const& params,                 \
-              const cuvs::neighbors::cagra::index<T, IdxT>& index,                 \
-              raft::device_matrix_view<const T, int64_t, raft::row_major> queries, \
-              raft::device_matrix_view<IdxT, int64_t, raft::row_major> neighbors,  \
-              raft::device_matrix_view<float, int64_t, raft::row_major> distances, \
-              const cuvs::neighbors::filtering::base_filter& sample_filter)        \
-  {                                                                                \
-    cuvs::neighbors::cagra::search<T, IdxT>(                                       \
-      handle, params, index, queries, neighbors, distances, sample_filter);        \
+#define CUVS_INST_CAGRA_SEARCH(T, IdxT)                                                      \
+  void search(raft::resources const& handle,                                                 \
+              cuvs::neighbors::cagra::search_params const& params,                           \
+              const cuvs::neighbors::cagra::index<T, IdxT>& index,                           \
+              raft::device_matrix_view<const T, int64_t, raft::row_major> queries,           \
+              raft::device_matrix_view<IdxT, int64_t, raft::row_major> neighbors,            \
+              raft::device_matrix_view<float, int64_t, raft::row_major> distances,           \
+              const cuvs::neighbors::filtering::base_filter& sample_filter,                  \
+              double threshold_to_bf)                                                        \
+  {                                                                                          \
+    cuvs::neighbors::cagra::search<T, IdxT>(                                                 \
+      handle, params, index, queries, neighbors, distances, sample_filter, threshold_to_bf); \
   }
 
 CUVS_INST_CAGRA_SEARCH(int8_t, uint32_t);