Optimize recall calculation and enhance CAGRA search features

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

@@ -360,55 +360,56 @@ void bench_search(::benchmark::State& state,
     std::size_t match_count = 0;
     std::size_t total_count = 0;
 
-    // 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;
-          }
-        }
-      };
-      // 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();
+    // Map result-buffer row -> original query index using the same stride as the
+    // timed search loop. Parallelize once over all `rows` (not once per search
+    // batch): when n_queries==1 the old per-batch spawn/join dominated wall time.
+    const size_t start_batch_offset = (state.thread_index() * n_queries) % query_set_size;
+    auto orig_query_idx             = [&](size_t i_out_idx) -> size_t {
+      const size_t batch_num  = i_out_idx / n_queries;
+      const size_t i_in_batch = i_out_idx % n_queries;
+      const size_t batch_offset =
+        (start_batch_offset + batch_num * queries_stride) % query_set_size;
+      return batch_offset + i_in_batch;
+    };
+
+    // Avoid CPU oversubscription when parallelizing recall calculation
+    int num_workers =
+      static_cast<int>(std::thread::hardware_concurrency()) / benchmark_n_threads;  // includes main
+    if (num_workers < 1) { num_workers = 1; }
+    // No benefit from more workers than rows
+    num_workers                  = std::min(num_workers, std::max(1, static_cast<int>(rows)));
+    const int num_helper_threads = num_workers - 1;
+
+    std::vector<std::thread> recall_workers;
+    recall_workers.reserve(num_helper_threads);
+    std::vector<std::size_t> local_match_count(num_workers, 0);
+    std::vector<std::size_t> local_total_count(num_workers, 0);
+
+    auto recall_range = [&](size_t start, size_t end, int tid) {
+      for (size_t i_out_idx = start; i_out_idx < end; ++i_out_idx) {
+        auto* candidates       = neighbors_host + i_out_idx * k;
+        auto [matching, total] = gt_maps->count_matches(orig_query_idx(i_out_idx), candidates, k);
+        local_match_count[tid] += matching;
+        local_total_count[tid] += total;
       }
-      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;
+    };
+
+    const size_t chunk_size = rows / static_cast<size_t>(num_workers);
+    const size_t remainder  = rows % static_cast<size_t>(num_workers);
+    size_t start            = 0;
+    for (int tid = 0; tid < num_helper_threads; ++tid) {
+      size_t end = start + chunk_size + (static_cast<size_t>(tid) < remainder ? 1 : 0);
+      recall_workers.emplace_back(recall_range, start, end, tid);
+      start = end;
     }
+    // main thread works on last chunk
+    recall_range(start, rows, num_helper_threads);
+    for (auto& worker : recall_workers) {
+      worker.join();
+    }
+    match_count = std::accumulate(local_match_count.begin(), local_match_count.end(), size_t{0});
+    total_count = std::accumulate(local_total_count.begin(), local_total_count.end(), size_t{0});
+
     double actual_recall = static_cast<double>(match_count) / static_cast<double>(total_count);
     /* NOTE: recall in the throughput mode & filtering
 

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

@@ -489,6 +489,9 @@ void parse_search_param(const nlohmann::json& conf,
   if (conf.contains("thread_block_size")) {
     param.p.thread_block_size = conf.at("thread_block_size");
   }
+  if (conf.contains("hashmap_min_bitlen")) {
+    param.p.hashmap_min_bitlen = conf.at("hashmap_min_bitlen");
+  }
   if (conf.contains("algo")) {
     if (conf.at("algo") == "single_cta") {
       param.p.algo = cuvs::neighbors::cagra::search_algo::SINGLE_CTA;

cpp/include/cuvs/neighbors/cagra.hpp

@@ -356,8 +356,9 @@ struct search_params : cuvs::neighbors::search_params {
   /** Number of threads used to calculate a single distance. 4, 8, 16, or 32. */
   size_t team_size = 0;
 
-  /** Number of graph nodes to select as the starting point for the search in each iteration. aka
-   * search width?*/
+  /** Number of graph nodes to select as the starting point for the search in each iteration.
+   *  Auto select as ceil(itopk_size / graph_degree) when 0.
+   */
   size_t search_width = 1;
   /** Lower limit of search iterations. */
   size_t min_iterations = 0;
@@ -366,7 +367,7 @@ struct search_params : cuvs::neighbors::search_params {
   size_t thread_block_size = 0;
   /** Hashmap type. Auto selection when AUTO. */
   hash_mode hashmap_mode = hash_mode::AUTO;
-  /** Lower limit of hashmap bit length. More than 8. */
+  /** Lower limit of hashmap bit length. 0 selects the default; otherwise, 8 to 20. */
   size_t hashmap_min_bitlen = 0;
   /** Upper limit of hashmap fill rate. More than 0.1, less than 0.9.*/
   float hashmap_max_fill_rate = 0.5;

cpp/src/neighbors/detail/cagra/search_multi_cta.cuh

@@ -117,13 +117,14 @@ struct search
   void set_params(raft::resources const& res, const search_params& params)
   {
     size_t global_itopk_size                = itopk_size;
+    size_t global_search_width              = search_width;
     constexpr unsigned multi_cta_itopk_size = 32;
     this->itopk_size                        = multi_cta_itopk_size;
     search_width                            = 1;
     RAFT_LOG_DEBUG("params.itopk_size: %lu", (uint64_t)params.itopk_size);
     RAFT_LOG_DEBUG("global_itopk_size: %lu", (uint64_t)global_itopk_size);
     num_cta_per_query =
-      max(params.search_width, raft::ceildiv(global_itopk_size, (size_t)multi_cta_itopk_size));
+      max(global_search_width, raft::ceildiv(global_itopk_size, (size_t)multi_cta_itopk_size));
     result_buffer_size = itopk_size + (search_width * graph_degree);
     typedef raft::Pow2<32> AlignBytes;
     unsigned result_buffer_size_32 = AlignBytes::roundUp(result_buffer_size);

cpp/src/neighbors/detail/cagra/search_plan.cuh

@@ -10,6 +10,7 @@
 #include <cuvs/neighbors/common.hpp>
 #include <neighbors/detail/cagra/compute_distance-ext.cuh>
 #include <raft/core/resource/cuda_stream.hpp>
+#include <raft/core/resource/device_properties.hpp>
 // #include "topk_for_cagra/topk.h"
 
 #include <raft/core/device_mdspan.hpp>
@@ -198,12 +199,22 @@ struct search_plan_impl : public search_plan_impl_base {
 
   void adjust_search_params()
   {
+    if (search_width == 0) {
+      search_width = raft::ceildiv(itopk_size, static_cast<size_t>(graph_degree));
+      RAFT_LOG_DEBUG("# search_width is auto-selected as %lu.", search_width);
+    }
+
     uint32_t _max_iterations = max_iterations;
     if (max_iterations == 0) {
       if (algo == search_algo::MULTI_CTA) {
-        constexpr uint32_t mc_itopk_size   = 32;
-        constexpr uint32_t mc_search_width = 1;
-        _max_iterations                    = mc_itopk_size / mc_search_width;
+        constexpr size_t mc_itopk_size  = 32;
+        constexpr size_t search_quality = 3;
+        const size_t minimum_depth      = 8 * (search_quality - 1);
+        const auto effective_itopk_size = raft::ceildiv(itopk_size, mc_itopk_size) * mc_itopk_size;
+        const auto num_ctas = max(search_width, raft::ceildiv(effective_itopk_size, mc_itopk_size));
+
+        _max_iterations = minimum_depth + raft::ceildiv(mc_itopk_size - minimum_depth, num_ctas);
+        _max_iterations += raft::ceildiv(static_cast<size_t>(topk), mc_itopk_size) - 1;
       } else {
         _max_iterations = itopk_size / search_width;
       }
@@ -287,6 +298,17 @@ struct search_plan_impl : public search_plan_impl_base {
       }
       RAFT_EXPECTS(hash_bitlen <= 25, "hash_bitlen cannot be largen than 25 (32M)");
     } else {
+      constexpr size_t kib = 1024;
+      const auto shared_mem_per_sm =
+        raft::resource::get_device_properties(res).sharedMemPerMultiprocessor;
+      unsigned default_min_bitlen = 8;
+      if (shared_mem_per_sm >= 196 * kib) {
+        default_min_bitlen = 11;
+      } else if (shared_mem_per_sm >= 128 * kib) {
+        default_min_bitlen = 10;
+      } else if (shared_mem_per_sm >= 64 * kib) {
+        default_min_bitlen = 9;
+      }
       while (hashmap_mode == hash_mode::AUTO || hashmap_mode == hash_mode::SMALL) {
         //
         // The small-hash reduces hash table size by initializing the hash table
@@ -296,10 +318,9 @@ struct search_plan_impl : public search_plan_impl_base {
         // visited per iteration.
         //
         const auto max_visited_nodes = itopk_size + (search_width * graph_degree * 1);
-        unsigned min_bitlen          = 8;   // 256
+        unsigned min_bitlen = hashmap_min_bitlen == 0 ? default_min_bitlen : hashmap_min_bitlen;
         unsigned max_bitlen          = 13;  // 8K
-        if (min_bitlen < hashmap_min_bitlen) { min_bitlen = hashmap_min_bitlen; }
-        hash_bitlen = min_bitlen;
+        hash_bitlen                  = min_bitlen;
         while (max_visited_nodes > hashmap::get_size(hash_bitlen) * max_fill_rate) {
           hash_bitlen += 1;
         }

fern/pages/cuvs_bench/param_tuning.md

@@ -89,7 +89,7 @@ CAGRA uses a graph-based index, which creates an intermediate, approximate kNN g
 | `use_disk` | `build` | N | Boolean | `false` | Whether to use disk-based storage for ACE build. When true, forces ACE to use disk-based storage even if the graph fits in host and GPU memory. When false, ACE will use in-memory storage if the graph fits in host and GPU memory and disk-based storage otherwise. |
 | `query_memory_type` | `search` | N | [`device`, `host`, `mmap`] | `device` | Where should the queries reside? |
 | `itopk` | `search` | N | Positive integer >0 | 64 | Number of intermediate search results retained during the search. Higher values improve search accuracy at the cost of speed |
-| `search_width` | `search` | N | Positive integer >0 | 1 | Number of graph nodes to select as the starting point for the search in each iteration. |
+| `search_width` | `search` | N | Non-negative integer | 1 | Number of graph nodes to select as the starting point for the search in each iteration. A value of 0 selects `ceil(itopk_size / graph_degree)`. |
 | `max_iterations` | `search` | N | Positive integer >=0 | 0 | Upper limit of search iterations. Auto select when 0 |
 | `algo` | `search` | N | [`auto`, `single_cta`, `multi_cta`, `multi_kernel`] | `auto` | Algorithm to use for search. It's usually best to leave this to `auto`. |
 | `persistent` | `search` | N | Boolean | `false` | Enables the persistent CAGRA search kernel for high-throughput search with many concurrent client threads. Persistent search currently requires `single_cta`; `auto` selects it when persistent mode is enabled. |