Expose search function with pre-filter for ANN

cpp/CMakeLists.txt

@@ -215,6 +215,7 @@ add_library(
   src/cluster/kmeans_balanced_predict_int8.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
@@ -405,9 +406,6 @@ add_library(
   src/neighbors/ivf_pq/detail/ivf_pq_search_float_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_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_int8.cu

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/cagra.hpp

@@ -951,9 +951,6 @@ void extend(
  *
  * See the [cagra::build](#cagra::build) documentation for a usage example.
  *
- * @tparam T data element type
- * @tparam IdxT type of the indices
- *
  * @param[in] res raft resources
  * @param[in] params configure the search
  * @param[in] idx cagra index
@@ -962,23 +959,24 @@ 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 bitmap filter function that greenlights samples for a
+ * given query
  */
 
 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,
+            std::optional<cuvs::neighbors::filtering::bitset_filter<uint32_t, int64_t>>
+              sample_filter = std::nullopt);
 
 /**
  * @brief Search ANN using the constructed index.
  *
  * See the [cagra::build](#cagra::build) documentation for a usage example.
  *
- * @tparam T data element type
- * @tparam IdxT type of the indices
- *
  * @param[in] res raft resources
  * @param[in] params configure the search
  * @param[in] index cagra index
@@ -987,22 +985,23 @@ 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 bitmap filter function that greenlights samples for a
+ * given query
  */
 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,
+            std::optional<cuvs::neighbors::filtering::bitset_filter<uint32_t, int64_t>>
+              sample_filter = std::nullopt);
 
 /**
  * @brief Search ANN using the constructed index.
  *
  * See the [cagra::build](#cagra::build) documentation for a usage example.
  *
- * @tparam T data element type
- * @tparam IdxT type of the indices
- *
  * @param[in] res raft resources
  * @param[in] params configure the search
  * @param[in] index cagra index
@@ -1011,13 +1010,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 bitmap filter function that greenlights samples for a
+ * given query
  */
 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,
+            std::optional<cuvs::neighbors::filtering::bitset_filter<uint32_t, int64_t>>
+              sample_filter = std::nullopt);
+
 /**
  * @}
  */

cpp/include/cuvs/neighbors/ivf_flat.hpp

@@ -19,6 +19,7 @@
 #include "common.hpp"
 #include <cstdint>
 #include <cuvs/neighbors/common.hpp>
+#include <optional>
 #include <raft/core/host_mdarray.hpp>
 #include <raft/core/host_mdspan.hpp>
 
@@ -1163,13 +1164,17 @@ void extend(raft::resources const& handle,
  * dataset [n_queries, k]
  * @param[out] distances raft::device_matrix_view to the distances to the selected neighbors
  * [n_queries, k]
+ * @param[in] sample_filter a device bitset filter function that greenlights samples for a given
+ * query.
  */
 void search(raft::resources const& handle,
             const cuvs::neighbors::ivf_flat::search_params& params,
             cuvs::neighbors::ivf_flat::index<float, int64_t>& index,
             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);
+            raft::device_matrix_view<float, int64_t, raft::row_major> distances,
+            std::optional<cuvs::neighbors::filtering::bitset_filter<uint32_t, int64_t>>
+              sample_filter = std::nullopt);
 
 /**
  * @brief Search ANN using the constructed index.
@@ -1200,13 +1205,17 @@ void search(raft::resources const& handle,
  * dataset [n_queries, k]
  * @param[out] distances raft::device_matrix_view to the distances to the selected neighbors
  * [n_queries, k]
+ * @param[in] sample_filter a device bitset filter function that greenlights samples for a given
+ * query.
  */
 void search(raft::resources const& handle,
             const cuvs::neighbors::ivf_flat::search_params& params,
             cuvs::neighbors::ivf_flat::index<int8_t, int64_t>& index,
             raft::device_matrix_view<const int8_t, 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);
+            raft::device_matrix_view<float, int64_t, raft::row_major> distances,
+            std::optional<cuvs::neighbors::filtering::bitset_filter<uint32_t, int64_t>>
+              sample_filter = std::nullopt);
 
 /**
  * @brief Search ANN using the constructed index.
@@ -1237,112 +1246,18 @@ void search(raft::resources const& handle,
  * dataset [n_queries, k]
  * @param[out] distances raft::device_matrix_view to the distances to the selected neighbors
  * [n_queries, k]
+ * @param[in] sample_filter a device bitset filter function that greenlights samples for a given
+ * query.
  */
 void search(raft::resources const& handle,
             const cuvs::neighbors::ivf_flat::search_params& params,
             cuvs::neighbors::ivf_flat::index<uint8_t, int64_t>& index,
             raft::device_matrix_view<const uint8_t, 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);
-
-/**
- * @brief Search ANN using the constructed index with the given filter.
- *
- * See the [ivf_flat::build](#ivf_flat::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`.
- * The exact size of the temporary buffer depends on multiple factors and is an implementation
- * detail. However, you can safely specify a small initial size for the memory pool, so that only a
- * few allocations happen to grow it during the first invocations of the `search`.
- *
- * @param[in] handle
- * @param[in] params configure the search
- * @param[in] idx ivf-flat constructed index
- * @param[in] queries a device matrix view to a row-major matrix [n_queries, index->dim()]
- * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset
- * [n_queries, k]
- * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries,
- * k]
- * @param[in] sample_filter a device bitset filter function that greenlights samples for a given
- * query.
- */
-void search_with_filtering(
-  raft::resources const& handle,
-  const search_params& params,
-  index<float, int64_t>& idx,
-  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,
-  cuvs::neighbors::filtering::bitset_filter<uint32_t, int64_t> sample_filter);
+            raft::device_matrix_view<float, int64_t, raft::row_major> distances,
+            std::optional<cuvs::neighbors::filtering::bitset_filter<uint32_t, int64_t>>
+              sample_filter = std::nullopt);
 
-/**
- * @brief Search ANN using the constructed index with the given filter.
- *
- * See the [ivf_flat::build](#ivf_flat::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`.
- * The exact size of the temporary buffer depends on multiple factors and is an implementation
- * detail. However, you can safely specify a small initial size for the memory pool, so that only a
- * few allocations happen to grow it during the first invocations of the `search`.
- *
- * @param[in] handle
- * @param[in] params configure the search
- * @param[in] idx ivf-flat constructed index
- * @param[in] queries a device matrix view to a row-major matrix [n_queries, index->dim()]
- * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset
- * [n_queries, k]
- * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries,
- * k]
- * @param[in] sample_filter a device bitset filter function that greenlights samples for a given
- * query.
- */
-void search_with_filtering(
-  raft::resources const& handle,
-  const search_params& params,
-  index<int8_t, int64_t>& idx,
-  raft::device_matrix_view<const int8_t, 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,
-  cuvs::neighbors::filtering::bitset_filter<uint32_t, int64_t> sample_filter);
-
-/**
- * @brief Search ANN using the constructed index with the given filter.
- *
- * See the [ivf_flat::build](#ivf_flat::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`.
- * The exact size of the temporary buffer depends on multiple factors and is an implementation
- * detail. However, you can safely specify a small initial size for the memory pool, so that only a
- * few allocations happen to grow it during the first invocations of the `search`.
- *
- * @param[in] handle
- * @param[in] params configure the search
- * @param[in] idx ivf-flat constructed index
- * @param[in] queries a device matrix view to a row-major matrix [n_queries, index->dim()]
- * @param[out] neighbors a device matrix view to the indices of the neighbors in the source dataset
- * [n_queries, k]
- * @param[out] distances a device matrix view to the distances to the selected neighbors [n_queries,
- * k]
- * @param[in] sample_filter a device bitset filter function that greenlights samples for a given
- * query.
- */
-void search_with_filtering(
-  raft::resources const& handle,
-  const search_params& params,
-  index<uint8_t, int64_t>& idx,
-  raft::device_matrix_view<const uint8_t, 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,
-  cuvs::neighbors::filtering::bitset_filter<uint32_t, int64_t> sample_filter);
 /**
  * @}
  */
@@ -2039,18 +1954,18 @@ void reset_index(const raft::resources& res, index<uint8_t, int64_t>* index);
  *   using namespace cuvs::neighbors;
  *   raft::resources res;
  *   // use default index parameters
- *   ivf_pq::index_params index_params;
+ *   ivf_flat::index_params index_params;
  *   // initialize an empty index
- *   ivf_pq::index<int64_t> index(res, index_params, D);
- *   ivf_pq::helpers::reset_index(res, &index);
+ *   ivf_flat::index<uint8_t, int64_t> index(res, index_params, D);
+ *   ivf_flat::helpers::reset_index(res, &index);
  *   // resize the first IVF list to hold 5 records
- *   auto spec = list_spec<uint32_t, int64_t>{
- *     index->pq_bits(), index->pq_dim(), index->conservative_memory_allocation()};
+ *   auto spec = list_spec<uint32_t, uint8_t, int64_t>{
+ *     index->dim(), index->conservative_memory_allocation()};
  *   uint32_t new_size = 5;
  *   ivf::resize_list(res, list, spec, new_size, 0);
  *   raft::update_device(index.list_sizes(), &new_size, 1, stream);
  *   // recompute the internal state of the index
- *   ivf_pq::helpers::recompute_internal_state(res, index);
+ *   ivf_flat::helpers::recompute_internal_state(res, index);
  * @endcode
  *
  * @param[in] res raft resource
@@ -2067,18 +1982,18 @@ void recompute_internal_state(const raft::resources& res, index<float, int64_t>*
  *   using namespace cuvs::neighbors;
  *   raft::resources res;
  *   // use default index parameters
- *   ivf_pq::index_params index_params;
+ *   ivf_flat::index_params index_params;
  *   // initialize an empty index
- *   ivf_pq::index<int64_t> index(res, index_params, D);
- *   ivf_pq::helpers::reset_index(res, &index);
+ *   ivf_flat::index<uint8_t, int64_t> index(res, index_params, D);
+ *   ivf_flat::helpers::reset_index(res, &index);
  *   // resize the first IVF list to hold 5 records
- *   auto spec = list_spec<uint32_t, int64_t>{
- *     index->pq_bits(), index->pq_dim(), index->conservative_memory_allocation()};
+ *   auto spec = list_spec<uint32_t, uint8_t, int64_t>{
+ *     index->dim(), index->conservative_memory_allocation()};
  *   uint32_t new_size = 5;
  *   ivf::resize_list(res, list, spec, new_size, 0);
  *   raft::update_device(index.list_sizes(), &new_size, 1, stream);
  *   // recompute the internal state of the index
- *   ivf_pq::helpers::recompute_internal_state(res, index);
+ *   ivf_flat::helpers::recompute_internal_state(res, index);
  * @endcode
  *
  * @param[in] res raft resource
@@ -2095,18 +2010,18 @@ void recompute_internal_state(const raft::resources& res, index<int8_t, int64_t>
  *   using namespace cuvs::neighbors;
  *   raft::resources res;
  *   // use default index parameters
- *   ivf_pq::index_params index_params;
+ *   ivf_flat::index_params index_params;
  *   // initialize an empty index
- *   ivf_pq::index<int64_t> index(res, index_params, D);
- *   ivf_pq::helpers::reset_index(res, &index);
+ *   ivf_flat::index<uint8_t, int64_t> index(res, index_params, D);
+ *   ivf_flat::helpers::reset_index(res, &index);
  *   // resize the first IVF list to hold 5 records
- *   auto spec = list_spec<uint32_t, int64_t>{
- *     index->pq_bits(), index->pq_dim(), index->conservative_memory_allocation()};
+ *   auto spec = list_spec<uint32_t, uint8_t, int64_t>{
+ *     index->dim(), index->conservative_memory_allocation()};
  *   uint32_t new_size = 5;
  *   ivf::resize_list(res, list, spec, new_size, 0);
  *   raft::update_device(index.list_sizes(), &new_size, 1, stream);
  *   // recompute the internal state of the index
- *   ivf_pq::helpers::recompute_internal_state(res, index);
+ *   ivf_flat::helpers::recompute_internal_state(res, index);
  * @endcode
  *
  * @param[in] res raft resource