Expose search function with pre-filter for ANN (#302)

Related to #136

This PR:
- Expose the `search()` function with optional filters in the public API of ANNs (CAGRA, IVF-PQ, IVF-Flat)
- Add compiled bitset objects
- Restore tests for filters in IVF-Flat
- Add tests for filter in CAGRA

Authors:
  - Micka (https://github.com/lowener)
  - Artem M. Chirkin (https://github.com/achirkin)

Approvers:
  - Corey J. Nolet (https://github.com/cjnolet)

URL: #302

cpp/CMakeLists.txt

@@ -308,6 +308,7 @@ add_library(
   src/cluster/kmeans_transform_double.cu
   src/cluster/kmeans_transform_float.cu
   src/cluster/single_linkage_float.cu
+  src/core/bitset.cu
   src/distance/detail/pairwise_matrix/dispatch_canberra_float_float_float_int.cu
   src/distance/detail/pairwise_matrix/dispatch_canberra_half_float_float_int.cu
   src/distance/detail/pairwise_matrix/dispatch_canberra_double_double_double_int.cu
@@ -406,10 +407,6 @@ add_library(
   src/neighbors/ivf_pq/detail/ivf_pq_search_half_int64_t.cu
   src/neighbors/ivf_pq/detail/ivf_pq_search_int8_t_int64_t.cu
   src/neighbors/ivf_pq/detail/ivf_pq_search_uint8_t_int64_t.cu
-  src/neighbors/ivf_pq/detail/ivf_pq_search_with_filter_float_int64_t.cu
-  src/neighbors/ivf_pq/detail/ivf_pq_search_with_filter_half_int64_t.cu
-  src/neighbors/ivf_pq/detail/ivf_pq_search_with_filter_int8_t_int64_t.cu
-  src/neighbors/ivf_pq/detail/ivf_pq_search_with_filter_uint8_t_int64_t.cu
   src/neighbors/nn_descent.cu
   src/neighbors/nn_descent_float.cu
   src/neighbors/nn_descent_half.cu

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

@@ -134,7 +134,7 @@ class BruteForceKNNBenchmark {
             search_queries.data(), params_.num_queries, params_.dim),
           indices,
           distances,
-          std::nullopt);
+          cuvs::neighbors::filtering::none_sample_filter{});
         flush_l2_cache();
         raft::resource::sync_stream(handle_, stream_);
       }
@@ -158,7 +158,7 @@ class BruteForceKNNBenchmark {
           search_queries.data(), params_.num_queries, params_.dim),
         indices,
         distances,
-        std::nullopt);
+        cuvs::neighbors::filtering::none_sample_filter{});
       raft::resource::sync_stream(handle_, stream_);
       end        = std::chrono::high_resolution_clock::now();
       search_dur = end - start;
@@ -178,7 +178,7 @@ class BruteForceKNNBenchmark {
             search_queries.data(), params_.num_queries, params_.dim),
           indices,
           distances,
-          std::nullopt);
+          cuvs::neighbors::filtering::none_sample_filter{});
         flush_l2_cache();
         raft::resource::sync_stream(handle_, stream_);
       }
@@ -202,7 +202,7 @@ class BruteForceKNNBenchmark {
           search_queries.data(), params_.num_queries, params_.dim),
         indices,
         distances,
-        std::nullopt);
+        cuvs::neighbors::filtering::none_sample_filter{});
       raft::resource::sync_stream(handle_, stream_);
       end        = std::chrono::high_resolution_clock::now();
       search_dur = end - start;

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

@@ -155,8 +155,12 @@ void cuvs_gpu<T>::search(
     raft::make_device_matrix_view<algo_base::index_type, int64_t>(neighbors, batch_size, k);
   auto distances_view = raft::make_device_matrix_view<float, int64_t>(distances, batch_size, k);
 
-  cuvs::neighbors::brute_force::search(
-    handle_, *index_, queries_view, neighbors_view, distances_view, std::nullopt);
+  cuvs::neighbors::brute_force::search(handle_,
+                                       *index_,
+                                       queries_view,
+                                       neighbors_view,
+                                       distances_view,
+                                       cuvs::neighbors::filtering::none_sample_filter{});
 }
 
 template <typename T>

cpp/include/cuvs/core/bitset.hpp

@@ -18,6 +18,12 @@
 
 #include <raft/core/bitset.hpp>
 
+extern template struct raft::core::bitset<uint8_t, uint32_t>;
+extern template struct raft::core::bitset<uint16_t, uint32_t>;
+extern template struct raft::core::bitset<uint32_t, uint32_t>;
+extern template struct raft::core::bitset<uint32_t, int64_t>;
+extern template struct raft::core::bitset<uint64_t, int64_t>;
+
 namespace cuvs::core {
 /* To use bitset functions containing CUDA code, include <raft/core/bitset.cuh> */
 

cpp/include/cuvs/neighbors/brute_force.hpp

@@ -291,7 +291,8 @@ void search(raft::resources const& handle,
             raft::device_matrix_view<const float, int64_t, raft::row_major> queries,
             raft::device_matrix_view<int64_t, int64_t, raft::row_major> neighbors,
             raft::device_matrix_view<float, int64_t, raft::row_major> distances,
-            std::optional<cuvs::core::bitmap_view<const uint32_t, int64_t>> sample_filter);
+            const cuvs::neighbors::filtering::base_filter& sample_filter =
+              cuvs::neighbors::filtering::none_sample_filter{});
 
 /**
  * @brief Search ANN using the constructed index.
@@ -326,7 +327,8 @@ void search(raft::resources const& handle,
             raft::device_matrix_view<const half, int64_t, raft::row_major> queries,
             raft::device_matrix_view<int64_t, int64_t, raft::row_major> neighbors,
             raft::device_matrix_view<float, int64_t, raft::row_major> distances,
-            std::optional<cuvs::core::bitmap_view<const uint32_t, int64_t>> sample_filter);
+            const cuvs::neighbors::filtering::base_filter& sample_filter =
+              cuvs::neighbors::filtering::none_sample_filter{});
 /**
  * @brief Search ANN using the constructed index.
  *
@@ -346,7 +348,8 @@ void search(raft::resources const& handle,
             raft::device_matrix_view<const float, int64_t, raft::col_major> queries,
             raft::device_matrix_view<int64_t, int64_t, raft::row_major> neighbors,
             raft::device_matrix_view<float, int64_t, raft::row_major> distances,
-            std::optional<cuvs::core::bitmap_view<const uint32_t, int64_t>> sample_filter);
+            const cuvs::neighbors::filtering::base_filter& sample_filter =
+              cuvs::neighbors::filtering::none_sample_filter{});
 /**
  * @brief Search ANN using the constructed index.
  *
@@ -366,7 +369,8 @@ void search(raft::resources const& handle,
             raft::device_matrix_view<const half, int64_t, raft::col_major> queries,
             raft::device_matrix_view<int64_t, int64_t, raft::row_major> neighbors,
             raft::device_matrix_view<float, int64_t, raft::row_major> distances,
-            std::optional<cuvs::core::bitmap_view<const uint32_t, int64_t>> sample_filter);
+            const cuvs::neighbors::filtering::base_filter& sample_filter =
+              cuvs::neighbors::filtering::none_sample_filter{});
 /**
  * @}
  */

cpp/include/cuvs/neighbors/cagra.hpp

@@ -1055,14 +1055,18 @@ void extend(
  * [n_queries, k]
  * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries,
  * k]
+ * @param[in] sample_filter an optional device filter function object that greenlights samples
+ * for a given query. (none_sample_filter for no filtering)
  */
 
 void search(raft::resources const& res,
             cuvs::neighbors::cagra::search_params const& params,
             const cuvs::neighbors::cagra::index<float, uint32_t>& index,
             raft::device_matrix_view<const float, int64_t, raft::row_major> queries,
             raft::device_matrix_view<uint32_t, int64_t, raft::row_major> neighbors,
-            raft::device_matrix_view<float, int64_t, raft::row_major> distances);
+            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{});
 
 /**
  * @brief Search ANN using the constructed index.
@@ -1077,13 +1081,17 @@ void search(raft::resources const& res,
  * [n_queries, k]
  * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries,
  * k]
+ * @param[in] sample_filter an optional device filter function object that greenlights samples
+ * for a given query. (none_sample_filter for no filtering)
  */
 void search(raft::resources const& res,
             cuvs::neighbors::cagra::search_params const& params,
             const cuvs::neighbors::cagra::index<half, uint32_t>& index,
             raft::device_matrix_view<const half, int64_t, raft::row_major> queries,
             raft::device_matrix_view<uint32_t, int64_t, raft::row_major> neighbors,
-            raft::device_matrix_view<float, int64_t, raft::row_major> distances);
+            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{});
 
 /**
  * @brief Search ANN using the constructed index.
@@ -1098,13 +1106,17 @@ void search(raft::resources const& res,
  * [n_queries, k]
  * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries,
  * k]
+ * @param[in] sample_filter an optional device filter function object that greenlights samples
+ * for a given query. (none_sample_filter for no filtering)
  */
 void search(raft::resources const& res,
             cuvs::neighbors::cagra::search_params const& params,
             const cuvs::neighbors::cagra::index<int8_t, uint32_t>& index,
             raft::device_matrix_view<const int8_t, int64_t, raft::row_major> queries,
             raft::device_matrix_view<uint32_t, int64_t, raft::row_major> neighbors,
-            raft::device_matrix_view<float, int64_t, raft::row_major> distances);
+            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{});
 
 /**
  * @brief Search ANN using the constructed index.
@@ -1119,13 +1131,18 @@ void search(raft::resources const& res,
  * [n_queries, k]
  * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries,
  * k]
+ * @param[in] sample_filter an optional device filter function object that greenlights samples
+ * for a given query. (none_sample_filter for no filtering)
  */
 void search(raft::resources const& res,
             cuvs::neighbors::cagra::search_params const& params,
             const cuvs::neighbors::cagra::index<uint8_t, uint32_t>& index,
             raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> queries,
             raft::device_matrix_view<uint32_t, int64_t, raft::row_major> neighbors,
-            raft::device_matrix_view<float, int64_t, raft::row_major> distances);
+            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{});
+
 /**
  * @}
  */

cpp/include/cuvs/neighbors/common.h

@@ -44,7 +44,7 @@ enum cuvsFilterType {
 };
 
 /**
- * @brief Struct to hold address of cuvs::neighbor::prefilter and its type
+ * @brief Struct to hold address of cuvs::neighbors::prefilter and its type
  *
  */
 typedef struct {

cpp/include/cuvs/neighbors/common.hpp

@@ -383,19 +383,20 @@ inline constexpr bool is_vpq_dataset_v = is_vpq_dataset<DatasetT>::value;
 
 namespace filtering {
 
+struct base_filter {
+  virtual ~base_filter() = default;
+};
+
 /* A filter that filters nothing. This is the default behavior. */
-struct none_ivf_sample_filter {
+struct none_sample_filter : public base_filter {
   inline _RAFT_HOST_DEVICE bool operator()(
     // query index
     const uint32_t query_ix,
     // the current inverted list index
     const uint32_t cluster_ix,
     // the index of the current sample inside the current inverted list
     const uint32_t sample_ix) const;
-};
 
-/* A filter that filters nothing. This is the default behavior. */
-struct none_cagra_sample_filter {
   inline _RAFT_HOST_DEVICE bool operator()(
     // query index
     const uint32_t query_ix,
@@ -431,13 +432,33 @@ struct ivf_to_sample_filter {
     const uint32_t sample_ix) const;
 };
 
+/**
+ * @brief Filter an index with a bitmap
+ *
+ * @tparam bitmap_t Data type of the bitmap
+ * @tparam index_t Indexing type
+ */
+template <typename bitmap_t, typename index_t>
+struct bitmap_filter : public base_filter {
+  // View of the bitset to use as a filter
+  const cuvs::core::bitmap_view<bitmap_t, index_t> bitmap_view_;
+
+  bitmap_filter(const cuvs::core::bitmap_view<bitmap_t, index_t> bitmap_for_filtering);
+  inline _RAFT_HOST_DEVICE bool operator()(
+    // query index
+    const uint32_t query_ix,
+    // the index of the current sample
+    const uint32_t sample_ix) const;
+};
+
 /**
  * @brief Filter an index with a bitset
  *
+ * @tparam bitset_t Data type of the bitset
  * @tparam index_t Indexing type
  */
 template <typename bitset_t, typename index_t>
-struct bitset_filter {
+struct bitset_filter : public base_filter {
   // View of the bitset to use as a filter
   const cuvs::core::bitset_view<bitset_t, index_t> bitset_view_;