Fix filtered CAGRA search returning incorrect recall due to kHostMmap

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

@@ -9,6 +9,7 @@
 
 #include <memory>
 #include <nlohmann/json.hpp>
+#include <optional>
 #include <stdexcept>
 #include <string>
 #include <vector>
@@ -78,6 +79,8 @@ struct algo_property {
   MemoryType dataset_memory_type;
   // neighbors/distances should have same memory type as queries
   MemoryType query_memory_type;
+  // filter bitset memory type; defaults to dataset_memory_type if not explicitly set
+  std::optional<MemoryType> filter_memory_type{std::nullopt};
 };
 
 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-2025, NVIDIA CORPORATION.
  * 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,9 @@ 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.value_or(
+                              current_algo_props->dataset_memory_type)));
     } catch (const std::exception& ex) {
       state.SkipWithError("An error occurred setting search parameters: " + std::string(ex.what()));
       return;
@@ -351,9 +356,15 @@ void bench_search(::benchmark::State& state,
 
   // Each thread calculates recall on their partition of queries.
   // evaluate recall
-  if (dataset->max_k() >= k && dataset->gt_maps().has_value()) {
-    // gt_maps[i] is a hash map of {id, neighbor_rank} for query i
-    const auto& gt_maps = dataset->gt_maps();
+  if (dataset->max_k() >= k) {
+    const std::int32_t* gt             = dataset->gt_set();
+    const std::uint32_t* filter_bitset = dataset->filter_bitset(MemoryType::kHostMmap);
+    auto filter                        = [filter_bitset](std::int32_t i) -> bool {
+      if (filter_bitset == nullptr) { return true; }
+      auto word = filter_bitset[i >> 5];
+      return word & (1 << (i & 31));
+    };
+    const std::uint32_t max_k = dataset->max_k();
     result_buf.transfer_data(MemoryType::kHost, current_algo_props->query_memory_type);
     auto* neighbors_host    = reinterpret_cast<index_type*>(result_buf.data(MemoryType::kHost));
     std::size_t rows        = std::min(queries_processed, query_set_size);
@@ -363,49 +374,39 @@ void bench_search(::benchmark::State& state,
     // We go through the groundtruth with same stride as the benchmark loop.
     size_t out_offset   = 0;
     size_t batch_offset = (state.thread_index() * n_queries) % query_set_size;
-    // Avoid CPU oversubscription when parallelizing recall calculation loop
-    int num_recall_calculation_worker_threads =
-      std::thread::hardware_concurrency() / benchmark_n_threads - 1;  // -1 for the main thread
-    // ensure non-negative number of workers (possible if hardware_concurrency()
-    // does not return an expected value) by clamping to 0
-    if (num_recall_calculation_worker_threads < 0) { num_recall_calculation_worker_threads = 0; }
     while (out_offset < rows) {
-      std::vector<std::thread> recall_calculation_workers;
-      recall_calculation_workers.reserve(num_recall_calculation_worker_threads);
-      std::vector<std::size_t> local_match_count(num_recall_calculation_worker_threads + 1);
-      std::vector<std::size_t> local_total_count(num_recall_calculation_worker_threads + 1);
-      int chunk_size =
-        n_queries / (num_recall_calculation_worker_threads + 1);  // +1 for the main thread
-      int remainder           = n_queries % (num_recall_calculation_worker_threads + 1);
-      auto recall_calculation = [&](int start, int end, int tid) -> void {
-        for (int i = start; i < end; ++i) {
-          size_t i_orig_idx = batch_offset + i;
-          size_t i_out_idx  = out_offset + i;
-          if (i_out_idx < rows) {
-            auto* candidates       = neighbors_host + i_out_idx * k;
-            auto [matching, total] = gt_maps->count_matches(i_orig_idx, candidates, k);
-            local_match_count[tid] += matching;
-            local_total_count[tid] += total;
+      for (std::size_t i = 0; i < n_queries; i++) {
+        size_t i_orig_idx = batch_offset + i;
+        size_t i_out_idx  = out_offset + i;
+        if (i_out_idx < rows) {
+          /* NOTE: recall correctness & filtering
+
+          In the loop below, we filter the ground truth values on-the-fly.
+          We need enough ground truth values to compute recall correctly though.
+          But the ground truth file only contains `max_k` values per row; if there are less valid
+          values than k among them, we overestimate the recall. Essentially, we compare the first
+          `filter_pass_count` values of the algorithm output, and this counter can be less than `k`.
+          In the extreme case of very high filtering rate, we may be bypassing entire rows of
+          results. However, this is still better than no recall estimate at all.
+
+          TODO: consider generating the filtered ground truth on-the-fly
+          */
+          uint32_t filter_pass_count = 0;
+          for (std::uint32_t l = 0; l < max_k && filter_pass_count < k; l++) {
+            auto exp_idx = gt[i_orig_idx * max_k + l];
+            if (!filter(exp_idx)) { continue; }
+            filter_pass_count++;
+            for (std::uint32_t j = 0; j < k; j++) {
+              auto act_idx = static_cast<std::int32_t>(neighbors_host[i_out_idx * k + j]);
+              if (act_idx == exp_idx) {
+                match_count++;
+                break;
+              }
+            }
           }
+          total_count += filter_pass_count;
         }
-      };
-      // launch worker threads
-      int start = 0;
-      for (int tid = 0; tid < num_recall_calculation_worker_threads; tid++) {
-        int end = start + chunk_size;
-        if (tid < remainder) { ++end; }
-        recall_calculation_workers.emplace_back(recall_calculation, start, end, tid);
-        start = end;
-      }
-      // main thread works on last chunk
-      recall_calculation(start, n_queries, num_recall_calculation_worker_threads);
-      // join all worker threads
-      for (auto& worker : recall_calculation_workers) {
-        worker.join();
       }
-      match_count += std::accumulate(local_match_count.begin(), local_match_count.end(), 0);
-      total_count += std::accumulate(local_total_count.begin(), local_total_count.end(), 0);
-
       out_offset += n_queries;
       batch_offset = (batch_offset + queries_stride) % query_set_size;
     }

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

@@ -13,7 +13,7 @@
 #include <cuvs/distance/distance.hpp>
 #include <cuvs/neighbors/cagra.hpp>
 #include <cuvs/neighbors/common.hpp>
-#include <cuvs/neighbors/composite/index.hpp>
+#include <cuvs/neighbors/composite/merge.hpp>
 #include <cuvs/neighbors/dynamic_batching.hpp>
 #include <cuvs/neighbors/ivf_pq.hpp>
 #include <cuvs/neighbors/nn_descent.hpp>
@@ -27,7 +27,6 @@
 #include <rmm/device_uvector.hpp>
 #include <rmm/resource_ref.hpp>
 
-#include <atomic>
 #include <cassert>
 #include <fstream>
 #include <iostream>
@@ -36,37 +35,11 @@
 #include <raft/util/integer_utils.hpp>
 #include <stdexcept>
 #include <string>
-#include <thread>
 #include <type_traits>
 #include <vector>
 
 namespace cuvs::bench {
 
-namespace detail {
-
-/** If persistent CAGRA search uses few benchmark threads, log a throughput hint once per process.
- */
-inline void maybe_log_cagra_persistent_concurrency_hint(bool persistent_search)
-{
-  if (!persistent_search) { return; }
-
-  const unsigned hc = std::max(1u, static_cast<unsigned>(std::thread::hardware_concurrency()));
-  const unsigned bn = static_cast<unsigned>(std::max(0, benchmark_n_threads));
-  if (bn >= 2u * hc) { return; }
-
-  static std::atomic<bool> logged{false};
-  bool expected = false;
-  if (!logged.compare_exchange_strong(expected, true)) { return; }
-
-  const unsigned threads_rec = 16u * hc;
-  RAFT_LOG_INFO(
-    "CAGRA persistent search benefits from high client concurrency (try `--mode=throughput "
-    "--threads=1:%u`).",
-    threads_rec);
-}
-
-}  // namespace detail
-
 enum class AllocatorType { kHostPinned, kHostHugePage, kDevice };
 enum class CagraBuildAlgo { kAuto, kIvfPq, kNnDescent };
 enum class CagraMergeType { kPhysical, kLogical };
@@ -149,12 +122,12 @@ class cuvs_cagra : public algo<T>, public algo_gpu {
     return !search_params_.persistent;
   }
 
-  // to enable dataset access from GPU memory
   [[nodiscard]] auto get_preference() const -> algo_property override
   {
     algo_property property;
     property.dataset_memory_type = MemoryType::kHostMmap;
     property.query_memory_type   = MemoryType::kDevice;
+    property.filter_memory_type  = MemoryType::kDevice;
     return property;
   }
   void save(const std::string& file) const override;
@@ -193,9 +166,9 @@ class cuvs_cagra : public algo<T>, public algo_gpu {
   inline rmm::device_async_resource_ref get_mr(AllocatorType mem_type)
   {
     switch (mem_type) {
-      case (AllocatorType::kHostPinned): return mr_pinned_;
-      case (AllocatorType::kHostHugePage): return mr_huge_page_;
-      default: return rmm::mr::get_current_device_resource_ref();
+      case (AllocatorType::kHostPinned): return &mr_pinned_;
+      case (AllocatorType::kHostHugePage): return &mr_huge_page_;
+      default: return rmm::mr::get_current_device_resource();
     }
   }
 };
@@ -352,8 +325,6 @@ void cuvs_cagra<T, IdxT>::set_search_param(const search_param_base& param,
   } else {
     if (dynamic_batcher_) { dynamic_batcher_.reset(); }
   }
-
-  detail::maybe_log_cagra_persistent_concurrency_hint(search_params_.persistent);
 }
 
 template <typename T, typename IdxT>
@@ -482,21 +453,33 @@ void cuvs_cagra<T, IdxT>::search_base(
     } else {
       if (index_params_.merge_type == CagraMergeType::kLogical) {
         // TODO: index merge must happen outside of search, otherwise what are we benchmarking?
-        std::vector<cuvs::neighbors::cagra::index<T, IdxT>*> cagra_indices;
-        cagra_indices.reserve(sub_indices_.size());
+        cuvs::neighbors::cagra::merge_params merge_params{cuvs::neighbors::cagra::index_params{}};
+        merge_params.merge_strategy = cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL;
+
+        // Create wrapped indices for composite merge
+        std::vector<std::shared_ptr<cuvs::neighbors::IndexWrapper<T, IdxT, algo_base::index_type>>>
+          wrapped_indices;
+        wrapped_indices.reserve(sub_indices_.size());
         for (auto& ptr : sub_indices_) {
-          cagra_indices.push_back(ptr.get());
+          auto index_wrapper =
+            cuvs::neighbors::cagra::make_index_wrapper<T, IdxT, algo_base::index_type>(ptr.get());
+          wrapped_indices.push_back(index_wrapper);
         }
 
         raft::resources composite_handle(handle_);
-        size_t n_streams = cagra_indices.size();
+        size_t n_streams = wrapped_indices.size();
         raft::resource::set_cuda_stream_pool(composite_handle,
                                              std::make_shared<rmm::cuda_stream_pool>(n_streams));
 
-        cuvs::neighbors::composite::composite_index<T, IdxT, algo_base::index_type> composite(
-          cagra_indices);
-        composite.search(
-          composite_handle, search_params_, queries_view, neighbors_view, distances_view);
+        auto merged_index =
+          cuvs::neighbors::composite::merge(composite_handle, merge_params, wrapped_indices);
+        cuvs::neighbors::filtering::none_sample_filter empty_filter;
+        merged_index->search(composite_handle,
+                             search_params_,
+                             queries_view,
+                             neighbors_view,
+                             distances_view,
+                             empty_filter);
       }
     }
   }