diff --git a/cpp/bench/ann/src/common/ann_types.hpp b/cpp/bench/ann/src/common/ann_types.hpp index bd669dff78..2562ef58a3 100644 --- a/cpp/bench/ann/src/common/ann_types.hpp +++ b/cpp/bench/ann/src/common/ann_types.hpp @@ -9,6 +9,7 @@ #include #include +#include #include #include #include @@ -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 filter_memory_type{std::nullopt}; }; class algo_base { diff --git a/cpp/bench/ann/src/common/benchmark.hpp b/cpp/bench/ann/src/common/benchmark.hpp index a588b1e2a6..5bbddcbcbc 100644 --- a/cpp/bench/ann/src/common/benchmark.hpp +++ b/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(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 recall_calculation_workers; - recall_calculation_workers.reserve(num_recall_calculation_worker_threads); - std::vector local_match_count(num_recall_calculation_worker_threads + 1); - std::vector 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(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; } diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h index 87111e4761..125f10535f 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h @@ -13,7 +13,7 @@ #include #include #include -#include +#include #include #include #include @@ -27,7 +27,6 @@ #include #include -#include #include #include #include @@ -36,37 +35,11 @@ #include #include #include -#include #include #include 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(std::thread::hardware_concurrency())); - const unsigned bn = static_cast(std::max(0, benchmark_n_threads)); - if (bn >= 2u * hc) { return; } - - static std::atomic 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, 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, 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::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 @@ -482,21 +453,33 @@ void cuvs_cagra::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*> 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>> + 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(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(n_streams)); - cuvs::neighbors::composite::composite_index 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); } } }