[Feat] Support bitset filter for Brute Force (#560)

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

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

URL: #560

Author: rhdong

cpp/include/cuvs/neighbors/brute_force.hpp

@@ -332,15 +332,28 @@ auto build(raft::resources const& handle,
  * Note, this function requires a temporary buffer to store intermediate results between cuda kernel
  * calls, which may lead to undesirable allocations and slowdown. To alleviate the problem, you can
  * pass a pool memory resource or a large enough pre-allocated memory resource to reduce or
- * eliminate entirely allocations happening within `search`:
+ * eliminate entirely allocations happening within `search`.
+ *
+ * Usage example:
  * @code{.cpp}
- *   ...
- *   // Use the same allocator across multiple searches to reduce the number of
- *   // cuda memory allocations
- *   brute_force::search(handle, index, queries1, out_inds1, out_dists1);
- *   brute_force::search(handle, index, queries2, out_inds2, out_dists2);
- *   brute_force::search(handle, index, queries3, out_inds3, out_dists3);
- *   ...
+ *   using namespace cuvs::neighbors;
+ *
+ *   // use default index parameters
+ *   brute_force::index_params index_params;
+ *   // create and fill the index from a [N, D] dataset
+ *   brute_force::index_params index_params;
+ *   auto index = brute_force::build(handle, index_params, dataset);
+ *   // use default search parameters
+ *   brute_force::search_params search_params;
+ *   // create a bitset to filter the search
+ *   auto removed_indices = raft::make_device_vector<int64_t, int64_t>(res, n_removed_indices);
+ *   raft::core::bitset<std::uint32_t, int64_t> removed_indices_bitset(
+ *     res, removed_indices.view(), dataset.extent(0));
+ *   // search K nearest neighbours according to a bitset
+ *   auto neighbors = raft::make_device_matrix<uint32_t>(res, n_queries, k);
+ *   auto distances = raft::make_device_matrix<float>(res, n_queries, k);
+ *   auto filter    = filtering::bitset_filter(removed_indices_bitset.view());
+ *   brute_force::search(res, search_params, index, queries, neighbors, distances, filter);
  * @endcode
  *
  * @param[in] handle
@@ -350,9 +363,17 @@ auto build(raft::resources const& handle,
  * @param[out] neighbors a device pointer to the indices of the neighbors in the source dataset
  * [n_queries, k]
  * @param[out] distances a device pointer to the distances to the selected neighbors [n_queries, k]
- * @param[in] sample_filter An optional device bitmap filter function with a `row-major` layout and
- * the shape of [n_queries, index->size()], which means the filter will use the first
- * `index->size()` bits to indicate whether queries[0] should compute the distance with dataset.
+ * @param[in] sample_filter An optional device filter that restricts which dataset elements should
+ * be considered for each query.
+ *
+ * - Supports two types of filters:
+ *   1. **Bitset Filter**: A shared filter where each bit corresponds to a dataset element.
+ *      All queries share the same filter, with a logical shape of `[1, index->size()]`.
+ *   2. **Bitmap Filter**: A per-query filter with a logical shape of `[n_queries, index->size()]`,
+ *      where each bit indicates whether a specific dataset element should be considered for a
+ *      particular query. (1 for inclusion, 0 for exclusion).
+ *
+ * - The default value is `none_sample_filter`, which applies no filtering.
  */
 void search(raft::resources const& handle,
             const cuvs::neighbors::brute_force::search_params& params,
@@ -379,15 +400,28 @@ void search(raft::resources const& handle,
  * Note, this function requires a temporary buffer to store intermediate results between cuda kernel
  * calls, which may lead to undesirable allocations and slowdown. To alleviate the problem, you can
  * pass a pool memory resource or a large enough pre-allocated memory resource to reduce or
- * eliminate entirely allocations happening within `search`:
+ * eliminate entirely allocations happening within `search`.
+ *
+ * Usage example:
  * @code{.cpp}
- *   ...
- *   // Use the same allocator across multiple searches to reduce the number of
- *   // cuda memory allocations
- *   brute_force::search(handle, index, queries1, out_inds1, out_dists1);
- *   brute_force::search(handle, index, queries2, out_inds2, out_dists2);
- *   brute_force::search(handle, index, queries3, out_inds3, out_dists3);
- *   ...
+ *   using namespace cuvs::neighbors;
+ *
+ *   // use default index parameters
+ *   brute_force::index_params index_params;
+ *   // create and fill the index from a [N, D] dataset
+ *   brute_force::index_params index_params;
+ *   auto index = brute_force::build(handle, index_params, dataset);
+ *   // use default search parameters
+ *   brute_force::search_params search_params;
+ *   // create a bitset to filter the search
+ *   auto removed_indices = raft::make_device_vector<int64_t, int64_t>(res, n_removed_indices);
+ *   raft::core::bitset<std::uint32_t, int64_t> removed_indices_bitset(
+ *     res, removed_indices.view(), dataset.extent(0));
+ *   // search K nearest neighbours according to a bitset
+ *   auto neighbors = raft::make_device_matrix<uint32_t>(res, n_queries, k);
+ *   auto distances = raft::make_device_matrix<half>(res, n_queries, k);
+ *   auto filter    = filtering::bitset_filter(removed_indices_bitset.view());
+ *   brute_force::search(res, search_params, index, queries, neighbors, distances, filter);
  * @endcode
  *
  * @param[in] handle
@@ -397,8 +431,17 @@ void search(raft::resources const& handle,
  * @param[out] neighbors a device pointer to the indices of the neighbors in the source dataset
  * [n_queries, k]
  * @param[out] distances a device pointer to the distances to the selected neighbors [n_queries, k]
- * @param[in] sample_filter a optional device bitmap filter function that greenlights samples for a
- * given
+ * @param[in] sample_filter An optional device filter that restricts which dataset elements should
+ * be considered for each query.
+ *
+ * - Supports two types of filters:
+ *   1. **Bitset Filter**: A shared filter where each bit corresponds to a dataset element.
+ *      All queries share the same filter, with a logical shape of `[1, index->size()]`.
+ *   2. **Bitmap Filter**: A per-query filter with a logical shape of `[n_queries, index->size()]`,
+ *      where each bit indicates whether a specific dataset element should be considered for a
+ *      particular query. (1 for inclusion, 0 for exclusion).
+ *
+ * - The default value is `none_sample_filter`, which applies no filtering.
  */
 void search(raft::resources const& handle,
             const cuvs::neighbors::brute_force::search_params& params,
@@ -421,15 +464,51 @@ void search(raft::resources const& handle,
  *
  * See the [brute_force::build](#brute_force::build) documentation for a usage example.
  *
+ * Note, this function requires a temporary buffer to store intermediate results between cuda kernel
+ * calls, which may lead to undesirable allocations and slowdown. To alleviate the problem, you can
+ * pass a pool memory resource or a large enough pre-allocated memory resource to reduce or
+ * eliminate entirely allocations happening within `search`.
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   using namespace cuvs::neighbors;
+ *
+ *   // use default index parameters
+ *   brute_force::index_params index_params;
+ *   // create and fill the index from a [N, D] dataset
+ *   brute_force::index_params index_params;
+ *   auto index = brute_force::build(handle, index_params, dataset);
+ *   // use default search parameters
+ *   brute_force::search_params search_params;
+ *   // create a bitset to filter the search
+ *   auto removed_indices = raft::make_device_vector<int64_t, int64_t>(res, n_removed_indices);
+ *   raft::core::bitset<std::uint32_t, int64_t> removed_indices_bitset(
+ *     res, removed_indices.view(), dataset.extent(0));
+ *   // search K nearest neighbours according to a bitset
+ *   auto neighbors = raft::make_device_matrix<uint32_t>(res, n_queries, k);
+ *   auto distances = raft::make_device_matrix<float>(res, n_queries, k);
+ *   auto filter    = filtering::bitset_filter(removed_indices_bitset.view());
+ *   brute_force::search(res, search_params, index, queries, neighbors, distances, filter);
+ * @endcode
+ *
  * @param[in] handle
  * @param[in] params parameters configuring the search
  * @param[in] index bruteforce constructed index
  * @param[in] queries a device pointer to a col-major matrix [n_queries, index->dim()]
  * @param[out] neighbors a device pointer to the indices of the neighbors in the source dataset
  * [n_queries, k]
  * @param[out] distances a device pointer to the distances to the selected neighbors [n_queries, k]
- * @param[in] sample_filter an optional device bitmap filter function that greenlights samples for a
- * given query
+ * @param[in] sample_filter An optional device filter that restricts which dataset elements should
+ * be considered for each query.
+ *
+ * - Supports two types of filters:
+ *   1. **Bitset Filter**: A shared filter where each bit corresponds to a dataset element.
+ *      All queries share the same filter, with a logical shape of `[1, index->size()]`.
+ *   2. **Bitmap Filter**: A per-query filter with a logical shape of `[n_queries, index->size()]`,
+ *      where each bit indicates whether a specific dataset element should be considered for a
+ *      particular query. (1 for inclusion, 0 for exclusion).
+ *
+ * - The default value is `none_sample_filter`, which applies no filtering.
  */
 void search(raft::resources const& handle,
             const cuvs::neighbors::brute_force::search_params& params,
@@ -452,15 +531,51 @@ void search(raft::resources const& handle,
  *
  * See the [brute_force::build](#brute_force::build) documentation for a usage example.
  *
+ * Note, this function requires a temporary buffer to store intermediate results between cuda kernel
+ * calls, which may lead to undesirable allocations and slowdown. To alleviate the problem, you can
+ * pass a pool memory resource or a large enough pre-allocated memory resource to reduce or
+ * eliminate entirely allocations happening within `search`.
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   using namespace cuvs::neighbors;
+ *
+ *   // use default index parameters
+ *   brute_force::index_params index_params;
+ *   // create and fill the index from a [N, D] dataset
+ *   brute_force::index_params index_params;
+ *   auto index = brute_force::build(handle, index_params, dataset);
+ *   // use default search parameters
+ *   brute_force::search_params search_params;
+ *   // create a bitset to filter the search
+ *   auto removed_indices = raft::make_device_vector<int64_t, int64_t>(res, n_removed_indices);
+ *   raft::core::bitset<std::uint32_t, int64_t> removed_indices_bitset(
+ *     res, removed_indices.view(), dataset.extent(0));
+ *   // search K nearest neighbours according to a bitset
+ *   auto neighbors = raft::make_device_matrix<uint32_t>(res, n_queries, k);
+ *   auto distances = raft::make_device_matrix<half>(res, n_queries, k);
+ *   auto filter    = filtering::bitset_filter(removed_indices_bitset.view());
+ *   brute_force::search(res, search_params, index, queries, neighbors, distances, filter);
+ * @endcode
+ *
  * @param[in] handle
  * @param[in] params parameters configuring the search
  * @param[in] index bruteforce constructed index
  * @param[in] queries a device pointer to a col-major matrix [n_queries, index->dim()]
  * @param[out] neighbors a device pointer to the indices of the neighbors in the source dataset
  * [n_queries, k]
  * @param[out] distances a device pointer to the distances to the selected neighbors [n_queries, k]
- * @param[in] sample_filter an optional device bitmap filter function that greenlights samples for a
- * given query
+ * @param[in] sample_filter An optional device filter that restricts which dataset elements should
+ * be considered for each query.
+ *
+ * - Supports two types of filters:
+ *   1. **Bitset Filter**: A shared filter where each bit corresponds to a dataset element.
+ *      All queries share the same filter, with a logical shape of `[1, index->size()]`.
+ *   2. **Bitmap Filter**: A per-query filter with a logical shape of `[n_queries, index->size()]`,
+ *      where each bit indicates whether a specific dataset element should be considered for a
+ *      particular query. (1 for inclusion, 0 for exclusion).
+ *
+ * - The default value is `none_sample_filter`, which applies no filtering.
  */
 void search(raft::resources const& handle,
             const cuvs::neighbors::brute_force::search_params& params,

cpp/include/cuvs/neighbors/common.hpp

@@ -18,6 +18,7 @@
 
 #include <cstdint>
 #include <cuvs/distance/distance.hpp>
+#include <raft/core/device_csr_matrix.hpp>
 #include <raft/core/device_mdarray.hpp>
 #include <raft/core/device_resources.hpp>
 #include <raft/core/host_mdspan.hpp>
@@ -456,8 +457,11 @@ inline constexpr bool is_vpq_dataset_v = is_vpq_dataset<DatasetT>::value;
 
 namespace filtering {
 
+enum class FilterType { None, Bitmap, Bitset };
+
 struct base_filter {
-  virtual ~base_filter() = default;
+  virtual ~base_filter()                     = default;
+  virtual FilterType get_filter_type() const = 0;
 };
 
 /* A filter that filters nothing. This is the default behavior. */
@@ -475,6 +479,8 @@ struct none_sample_filter : public base_filter {
     const uint32_t query_ix,
     // the index of the current sample
     const uint32_t sample_ix) const;
+
+  FilterType get_filter_type() const override { return FilterType::None; }
 };
 
 /**
@@ -513,15 +519,24 @@ struct ivf_to_sample_filter {
  */
 template <typename bitmap_t, typename index_t>
 struct bitmap_filter : public base_filter {
+  using view_t = cuvs::core::bitmap_view<bitmap_t, index_t>;
+
   // View of the bitset to use as a filter
-  const cuvs::core::bitmap_view<bitmap_t, index_t> bitmap_view_;
+  const view_t bitmap_view_;
 
-  bitmap_filter(const cuvs::core::bitmap_view<bitmap_t, index_t> bitmap_for_filtering);
+  bitmap_filter(const view_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;
+
+  FilterType get_filter_type() const override { return FilterType::Bitmap; }
+
+  view_t view() const { return bitmap_view_; }
+
+  template <typename csr_matrix_t>
+  void to_csr(raft::resources const& handle, csr_matrix_t& csr);
 };
 
 /**
@@ -532,15 +547,24 @@ struct bitmap_filter : public base_filter {
  */
 template <typename bitset_t, typename index_t>
 struct bitset_filter : public base_filter {
+  using view_t = cuvs::core::bitset_view<bitset_t, index_t>;
+
   // View of the bitset to use as a filter
-  const cuvs::core::bitset_view<bitset_t, index_t> bitset_view_;
+  const view_t bitset_view_;
 
-  bitset_filter(const cuvs::core::bitset_view<bitset_t, index_t> bitset_for_filtering);
+  bitset_filter(const view_t bitset_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;
+
+  FilterType get_filter_type() const override { return FilterType::Bitset; }
+
+  view_t view() const { return bitset_view_; }
+
+  template <typename csr_matrix_t>
+  void to_csr(raft::resources const& handle, csr_matrix_t& csr);
 };
 
 /**

cpp/src/neighbors/brute_force_c.cpp

@@ -67,8 +67,8 @@ void _search(cuvsResources_t res,
   using queries_mdspan_type   = raft::device_matrix_view<T const, int64_t, QueriesLayoutT>;
   using neighbors_mdspan_type = raft::device_matrix_view<int64_t, int64_t, raft::row_major>;
   using distances_mdspan_type = raft::device_matrix_view<float, int64_t, raft::row_major>;
-  using prefilter_mds_type    = raft::device_vector_view<const uint32_t, int64_t>;
-  using prefilter_bmp_type    = cuvs::core::bitmap_view<const uint32_t, int64_t>;
+  using prefilter_mds_type    = raft::device_vector_view<uint32_t, int64_t>;
+  using prefilter_bmp_type    = cuvs::core::bitmap_view<uint32_t, int64_t>;
 
   auto queries_mds   = cuvs::core::from_dlpack<queries_mdspan_type>(queries_tensor);
   auto neighbors_mds = cuvs::core::from_dlpack<neighbors_mdspan_type>(neighbors_tensor);
@@ -85,14 +85,14 @@ void _search(cuvsResources_t res,
                                          distances_mds,
                                          cuvs::neighbors::filtering::none_sample_filter{});
   } else if (prefilter.type == BITMAP) {
-    auto prefilter_ptr  = reinterpret_cast<DLManagedTensor*>(prefilter.addr);
-    auto prefilter_mds  = cuvs::core::from_dlpack<prefilter_mds_type>(prefilter_ptr);
-    auto prefilter_view = cuvs::neighbors::filtering::bitmap_filter(
-      prefilter_bmp_type((const uint32_t*)prefilter_mds.data_handle(),
+    auto prefilter_ptr   = reinterpret_cast<DLManagedTensor*>(prefilter.addr);
+    auto prefilter_mds   = cuvs::core::from_dlpack<prefilter_mds_type>(prefilter_ptr);
+    const auto prefilter = cuvs::neighbors::filtering::bitmap_filter(
+      prefilter_bmp_type((uint32_t*)prefilter_mds.data_handle(),
                          queries_mds.extent(0),
                          index_ptr->dataset().extent(0)));
     cuvs::neighbors::brute_force::search(
-      *res_ptr, params, *index_ptr, queries_mds, neighbors_mds, distances_mds, prefilter_view);
+      *res_ptr, params, *index_ptr, queries_mds, neighbors_mds, distances_mds, prefilter);
   } else {
     RAFT_FAIL("Unsupported prefilter type: BITSET");
   }