Support filtered search in IVF-RaBitQ

cpp/CMakeLists.txt

@@ -683,6 +683,20 @@ if(NOT BUILD_CPU_ONLY)
     KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
   )
   set(ivf_rabitq_ns "cuvs::neighbors::ivf_rabitq::detail")
+  generate_jit_lto_kernels(
+    jit_lto_files
+    NAME_FORMAT "ivf_rabitq_sample_filter_@filter_name@"
+    MATRIX_JSON_FILE
+      "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/sample_filter_matrix.json"
+    KERNEL_INPUT_FILE
+      "${CMAKE_CURRENT_SOURCE_DIR}/src/neighbors/ivf_rabitq/jit_lto_kernels/sample_filter_kernel.cu.in"
+    FRAGMENT_TAG_FORMAT
+      "${ivf_rabitq_ns}::fragment_tag_sample_filter<${neighbors_ns}::tag_filter_@filter_name@>"
+    FRAGMENT_TAG_HEADER_FILES "<cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp>"
+                              "<cuvs/detail/jit_lto/common_fragments.hpp>"
+    OUTPUT_DIRECTORY "${CMAKE_CURRENT_BINARY_DIR}/generated_kernels/ivf_rabitq/sample_filter"
+    KERNEL_LINK_LIBRARIES jit_lto_kernel_usage_requirements
+  )
   generate_jit_lto_kernels(
     jit_lto_files
     NAME_FORMAT "ivf_rabitq_compute_inner_products_with_lut"

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

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -78,6 +78,8 @@ struct algo_property {
   MemoryType dataset_memory_type;
   // neighbors/distances should have same memory type as queries
   MemoryType query_memory_type;
+  // GPU filters should be device-resident by default, independently of the dataset.
+  MemoryType filter_memory_type = MemoryType::kDevice;
 };
 
 class algo_base {

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

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -102,6 +102,9 @@ inline auto parse_algo_property(algo_property prop, const nlohmann::json& conf)
   if (conf.contains("query_memory_type")) {
     prop.query_memory_type = parse_memory_type(conf.at("query_memory_type"));
   }
+  if (conf.contains("filter_memory_type")) {
+    prop.filter_memory_type = parse_memory_type(conf.at("filter_memory_type"));
+  }
   return prop;
 };
 
@@ -263,7 +266,7 @@ void bench_search(::benchmark::State& state,
     }
     try {
       a->set_search_param(*search_param,
-                          dataset->filter_bitset(current_algo_props->dataset_memory_type));
+                          dataset->filter_bitset(current_algo_props->filter_memory_type));
     } catch (const std::exception& ex) {
       state.SkipWithError("An error occurred setting search parameters: " + std::string(ex.what()));
       return;
@@ -537,15 +540,16 @@ void dispatch_benchmark(std::string cmdline,
   auto base_file     = dataset_conf.base_file;
   auto query_file    = dataset_conf.query_file;
   auto gt_file       = dataset_conf.groundtruth_neighbors_file;
-  auto dataset =
-    std::make_shared<bench::dataset<T>>(dataset_conf.name,
-                                        base_file,
-                                        dataset_conf.subset_first_row,
-                                        dataset_conf.subset_size,
-                                        query_file,
-                                        dataset_conf.distance,
-                                        gt_file,
-                                        search_mode ? dataset_conf.filtering_rate : std::nullopt);
+  auto dataset       = std::make_shared<bench::dataset<T>>(
+    dataset_conf.name,
+    base_file,
+    dataset_conf.subset_first_row,
+    dataset_conf.subset_size,
+    query_file,
+    dataset_conf.distance,
+    gt_file,
+    search_mode ? dataset_conf.filtering_rate : std::nullopt,
+    search_mode ? dataset_conf.filter_bitset_file : std::nullopt);
   ::benchmark::AddCustomContext("dataset", dataset_conf.name);
   ::benchmark::AddCustomContext("distance", dataset_conf.distance);
   std::vector<configuration::index>& indices = conf.get_indices();

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

@@ -8,6 +8,7 @@
 
 #include <iostream>
 #include <optional>
+#include <stdexcept>
 #include <string>
 #include <unordered_set>
 #include <vector>
@@ -44,6 +45,7 @@ class configuration {
     std::string dtype;
 
     std::optional<double> filtering_rate{std::nullopt};
+    std::optional<std::string> filter_bitset_file{std::nullopt};
   };
 
   [[nodiscard]] inline auto get_dataset_conf() const -> const dataset_conf&
@@ -89,6 +91,14 @@ class configuration {
     if (conf.contains("filtering_rate")) {
       dataset_conf_.filtering_rate.emplace(conf.at("filtering_rate"));
     }
+    if (conf.contains("filter_bitset_file")) {
+      dataset_conf_.filter_bitset_file = combine_path(data_prefix, conf.at("filter_bitset_file"));
+    }
+    if (dataset_conf_.filtering_rate.has_value() && dataset_conf_.filter_bitset_file.has_value()) {
+      throw std::invalid_argument(
+        "dataset config must set at most one of 'filtering_rate' or "
+        "'filter_bitset_file'; got both");
+    }
 
     if (conf.contains("groundtruth_neighbors_file")) {
       dataset_conf_.groundtruth_neighbors_file =

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

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -41,7 +41,8 @@ struct ground_truth_map {
 
   explicit ground_truth_map(std::string file_name,
                             uint32_t n_queries,
-                            std::optional<blob<bitset_carrier_type>>& filter_bitset)
+                            std::optional<blob<bitset_carrier_type>>& filter_bitset,
+                            bool gt_is_prefiltered = false)
     : gt_maps_(n_queries)
   {
     // Eagerly iterate over and optionally filter the ground truth set to build gt_maps_ for up to
@@ -59,12 +60,19 @@ struct ground_truth_map {
     */
     auto ground_truth_set = blob<T>(file_name);
     max_k_                = ground_truth_set.n_cols();
-    auto filter           = [&](T i) -> bool {
+    // Counts ground-truth ids that the bitset rejects. When the GT was
+    // produced with a matching pre-filter (gt_is_prefiltered), every id
+    // is expected to pass, so any nonzero count signals a bitset/GT
+    // mismatch and is reported as a warning after the workers join.
+    std::atomic<size_t> filtered_out_count{0};
+    auto filter = [&](T i) -> bool {
       if (!filter_bitset.has_value()) { return true; }
       // bitset is `32 = bitset_carrier_type * 8` times more dense than the data
       // use bitwise arithmetic to get the `row_id` and correct bit pos in the `word`
-      auto word = filter_bitset->data(MemoryType::kHostMmap)[i >> 5];
-      return word & (1 << (i & 31));
+      auto word   = filter_bitset->data(MemoryType::kHostMmap)[i >> 5];
+      bool passes = word & (1 << (i & 31));
+      if (!passes) { filtered_out_count.fetch_add(1, std::memory_order_relaxed); }
+      return passes;
     };
     // Avoid CPU oversubscription when parallelizing recall calculation loop
     int num_map_building_worker_threads =
@@ -113,6 +121,18 @@ struct ground_truth_map {
     for (auto& worker : gt_map_building_workers) {
       worker.join();
     }
+    if (gt_is_prefiltered) {
+      auto drops = filtered_out_count.load(std::memory_order_relaxed);
+      if (drops > 0) {
+        log_warn(
+          "ground truth file '%s' was declared pre-filtered, but %zu of its "
+          "neighbor ids are rejected by the supplied filter bitset. This "
+          "indicates the bitset and ground truth files are mismatched; "
+          "recall numbers will be unreliable.",
+          file_name.c_str(),
+          drops);
+      }
+    }
   }
 
   [[nodiscard]] auto max_k() const -> uint32_t { return max_k_; }
@@ -179,32 +199,61 @@ struct dataset {
           std::string query_file,
           std::string distance,
           std::optional<std::string> groundtruth_neighbors_file,
-          std::optional<double> filtering_rate = std::nullopt)
+          std::optional<double> filtering_rate          = std::nullopt,
+          std::optional<std::string> filter_bitset_file = std::nullopt)
     : name_{std::move(name)},
       distance_{std::move(distance)},
       base_set_{base_file, subset_first_row, subset_size},
       query_set_{query_file}
   {
-    if (filtering_rate.has_value()) {
-      // Generate a random bitset for filtering
+    if (filtering_rate.has_value() && filter_bitset_file.has_value()) {
+      throw std::invalid_argument(
+        "dataset: at most one of 'filtering_rate' or 'filter_bitset_file' may be set");
+    }
+
+    if (filtering_rate.has_value() || filter_bitset_file.has_value()) {
       auto n_rows = static_cast<size_t>(subset_size) + static_cast<size_t>(subset_first_row);
       if (subset_size == 0) {
         // Read the base set size as a last resort only - for better laziness
         n_rows = base_set_size();
       }
       auto bitset_size = (n_rows - 1) / kBitsPerCarrierValue + 1;
-      blob_file<bitset_carrier_type> bitset_blob_file{static_cast<uint32_t>(bitset_size), 1};
-      blob_mmap<bitset_carrier_type> bitset_blob{
-        std::move(bitset_blob_file), false, HugePages::kDisable};
-      generate_bernoulli(const_cast<bitset_carrier_type*>(bitset_blob.data()),
-                         bitset_size,
-                         1.0 - filtering_rate.value());
-      filter_bitset_.emplace(std::move(bitset_blob));
+
+      if (filter_bitset_file.has_value()) {
+        // Load a pre-generated bitset from disk. The producer
+        // (cuvs_bench.generate_groundtruth) writes it in cuVS .bin format with
+        // shape (ceil(n_rows / 32), 1), matching the in-memory layout below.
+        filter_bitset_.emplace(blob<bitset_carrier_type>{filter_bitset_file.value()});
+        // Reject a bitset that does not cover the whole dataset up front: the
+        // filter is indexed by absolute row id, so a short file would be read
+        // out of bounds (past the end of the mmap) rather than fail cleanly.
+        auto loaded_words = static_cast<size_t>(filter_bitset_->n_rows()) *
+                            static_cast<size_t>(filter_bitset_->n_cols());
+        if (loaded_words < bitset_size) {
+          throw std::invalid_argument("dataset: filter bitset file '" + filter_bitset_file.value() +
+                                      "' holds " + std::to_string(loaded_words) + " words, but " +
+                                      std::to_string(bitset_size) + " are needed to cover " +
+                                      std::to_string(n_rows) +
+                                      " rows; the bitset and dataset do not match");
+        }
+      } else {
+        // Generate a random bitset for filtering
+        blob_file<bitset_carrier_type> bitset_blob_file{static_cast<uint32_t>(bitset_size), 1};
+        blob_mmap<bitset_carrier_type> bitset_blob{
+          std::move(bitset_blob_file), false, HugePages::kDisable};
+        generate_bernoulli(const_cast<bitset_carrier_type*>(bitset_blob.data()),
+                           bitset_size,
+                           1.0 - filtering_rate.value());
+        filter_bitset_.emplace(std::move(bitset_blob));
+      }
     }
 
     if (groundtruth_neighbors_file.has_value()) {
-      ground_truth_map_.emplace(ground_truth_map<IdxT>{
-        groundtruth_neighbors_file.value(), query_set_.n_rows(), filter_bitset_});
+      ground_truth_map_.emplace(
+        ground_truth_map<IdxT>{groundtruth_neighbors_file.value(),
+                               query_set_.n_rows(),
+                               filter_bitset_,
+                               /*gt_is_prefiltered=*/filter_bitset_file.has_value()});
     }
   }
 

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

@@ -81,6 +81,7 @@ class cuvs_ivf_rabitq : public algo<T>, public algo_gpu {
   int dimension_;
   float refine_ratio_ = 1.0;
   raft::device_matrix_view<const T, IdxT> dataset_;
+  std::shared_ptr<cuvs::neighbors::filtering::base_filter> filter_;
 };
 
 template <typename T, typename IdxT>
@@ -118,8 +119,10 @@ std::unique_ptr<algo<T>> cuvs_ivf_rabitq<T, IdxT>::copy()
 }
 
 template <typename T, typename IdxT>
-void cuvs_ivf_rabitq<T, IdxT>::set_search_param(const search_param_base& param, const void*)
+void cuvs_ivf_rabitq<T, IdxT>::set_search_param(const search_param_base& param,
+                                                const void* filter_bitset)
 {
+  filter_        = make_cuvs_filter(filter_bitset, index_->size());
   auto sp        = dynamic_cast<const search_param&>(param);
   search_params_ = sp.rabitq_param;
   refine_ratio_  = sp.refine_ratio;
@@ -154,7 +157,7 @@ void cuvs_ivf_rabitq<T, IdxT>::search(
   auto distances_view = raft::make_device_matrix_view<float, int64_t>(distances, batch_size, k);
 
   cuvs::neighbors::ivf_rabitq::search(
-    handle_, search_params_, *index_, queries_view, neighbors_view, distances_view);
+    handle_, search_params_, *index_, queries_view, neighbors_view, distances_view, *filter_);
 
   if constexpr (sizeof(IdxT) != sizeof(algo_base::index_type)) {
     raft::linalg::unaryOp(neighbors,

cpp/include/cuvs/detail/jit_lto/ivf_rabitq/ivf_rabitq_fragments.hpp

@@ -48,4 +48,7 @@ struct fragment_tag_compute_lut_ip_for_vec {};
 template <int NumBits>
 struct fragment_tag_compute_bitwise_quantized_ip_for_vec {};
 
+template <typename FilterTag>
+struct fragment_tag_sample_filter {};
+
 }  // namespace cuvs::neighbors::ivf_rabitq::detail

cpp/include/cuvs/neighbors/ivf_rabitq.hpp

@@ -18,6 +18,7 @@
 #include <raft/core/resources.hpp>
 #include <raft/util/integer_utils.hpp>
 
+#include <limits>
 #include <optional>
 #include <tuple>
 #include <variant>
@@ -100,6 +101,13 @@ enum class search_mode {
   QUANT8 = 3,
 };
 
+/**
+ * Default value returned by `search` when fewer than k samples pass the filter or the combined
+ * size of the probed clusters is smaller than k.
+ */
+template <typename IdxT>
+constexpr static IdxT kOutOfBoundsRecord = std::numeric_limits<IdxT>::max();
+
 struct search_params : cuvs::neighbors::search_params {
   /** The number of clusters to search. */
   uint32_t n_probes = 20;
@@ -240,13 +248,17 @@ auto build(raft::resources const& handle,
  * [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. Supports `none_sample_filter` and `bitset_filter<uint32_t, int64_t>`.
  */
 void search(raft::resources const& handle,
             const cuvs::neighbors::ivf_rabitq::search_params& search_params,
             cuvs::neighbors::ivf_rabitq::index<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,
+            const cuvs::neighbors::filtering::base_filter& sample_filter =
+              cuvs::neighbors::filtering::none_sample_filter{});
 
 /**
  * @}