From 21e5f20f6c5d9b9b67af11d95b37c5e6ce2e3484 Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Wed, 10 Jun 2026 23:33:07 +0000 Subject: [PATCH 01/11] bloom filter --- cpp/CMakeLists.txt | 4 +- .../cuvs/detail/jit_lto/common_fragments.hpp | 1 + cpp/include/cuvs/neighbors/common.hpp | 28 ++- cpp/src/neighbors/cagra.cuh | 19 ++ .../detail/cagra/cagra_filter_payload.hpp | 10 +- .../jit_lto_kernels/sample_filter_impl.cuh | 13 ++ .../jit_lto_kernels/sample_filter_matrix.json | 2 +- .../detail/cagra/search_multi_cta_inst.cu.in | 3 + .../detail/cagra/search_single_cta_inst.cu.in | 3 + .../search_single_cta_kernel_launcher_jit.cuh | 2 + .../detail/cagra/shared_launcher_jit.hpp | 4 + .../neighbors/detail/sample_filter_data.cuh | 10 + examples/cpp/CMakeLists.txt | 6 + .../cpp/src/cagra_bloom_filter_example.cu | 171 ++++++++++++++++++ 14 files changed, 272 insertions(+), 4 deletions(-) create mode 100644 examples/cpp/src/cagra_bloom_filter_example.cu diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 227c2906cc..b3f6e23ca9 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -367,7 +367,9 @@ if(NOT BUILD_CPU_ONLY) "$<$:${CUVS_CUDA_FLAGS}>" ) target_compile_features(jit_lto_kernel_usage_requirements INTERFACE cuda_std_20) - target_link_libraries(jit_lto_kernel_usage_requirements INTERFACE rmm::rmm raft::raft CCCL::CCCL) + target_link_libraries( + jit_lto_kernel_usage_requirements INTERFACE rmm::rmm raft::raft CCCL::CCCL cuco::cuco + ) block(PROPAGATE jit_lto_files) set(jit_lto_files) diff --git a/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp index ef2a8e6002..3ae94a4b13 100644 --- a/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp +++ b/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp @@ -14,6 +14,7 @@ struct tag_i8 {}; struct tag_u8 {}; struct tag_filter_none {}; struct tag_filter_bitset {}; +struct tag_filter_bloom_filter {}; struct tag_filter_udf {}; struct tag_bitset_u32 {}; diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 2fd804f115..da9b812f60 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -497,7 +497,7 @@ namespace filtering { * @{ */ -enum class FilterType { None, Bitmap, Bitset, UDF }; +enum class FilterType { None, Bitmap, Bitset, Bloom, UDF }; struct base_filter { ~base_filter() = default; @@ -617,6 +617,32 @@ struct bitset_filter : public base_filter { void to_csr(raft::resources const& handle, csr_matrix_t& csr); }; +/** + * @brief Filter CAGRA candidates with a global @c cuco bloom filter over the index. + * + * Build the filter once on the host with bulk @c add() over the allowed dataset row ids, obtain a + * @c ref() from the owning @c cuco::bloom_filter, copy that ref to device memory, and pass the + * device pointer as @c filter_data. The linked JIT-LTO fragment probes the same filter for every + * query and candidate, similar to @ref bitset_filter but with probabilistic membership tests. + * + * Bloom filters have no false negatives: if a row was inserted, @c contains returns @c true. False + * positives are possible, so highly selective predicates may still need a bitset or UDF for exact + * filtering. + */ +struct bloom_filter : public base_filter { + void* filter_data{nullptr}; + float filtering_rate{-1.0f}; + + bloom_filter() = default; + + explicit bloom_filter(void* filter_data, float filtering_rate = -1.0f) + : filter_data(filter_data), filtering_rate(filtering_rate) + { + } + + FilterType get_filter_type() const override { return FilterType::Bloom; } +}; + /** * @brief JIT-LTO user-defined filter predicate. * diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index ee87c2c0ab..af8331a9f1 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -385,6 +385,25 @@ void search(raft::resources const& res, } catch (const std::bad_cast&) { } + try { + auto& sample_filter = + dynamic_cast(sample_filter_ref); + search_params params_copy = params; + if (params.filtering_rate < 0.0) { + const float min_filtering_rate = 0.0f; + const float max_filtering_rate = 0.999f; + params_copy.filtering_rate = + sample_filter.filtering_rate < 0.0f + ? 0.0f + : std::min(std::max(sample_filter.filtering_rate, min_filtering_rate), + max_filtering_rate); + } + auto sample_filter_copy = sample_filter; + return search_with_filtering( + res, params_copy, idx, queries, neighbors, distances, sample_filter_copy); + } catch (const std::bad_cast&) { + } + try { auto& sample_filter = dynamic_cast(sample_filter_ref); diff --git a/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp b/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp index 88f30f7745..aa927215ef 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp +++ b/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp @@ -139,6 +139,12 @@ template struct is_bitset_filter<::cuvs::neighbors::filtering::bitset_filter> : std::true_type {}; +template +struct is_bloom_filter : std::false_type {}; + +template <> +struct is_bloom_filter<::cuvs::neighbors::filtering::bloom_filter> : std::true_type {}; + template struct is_udf_filter : std::false_type {}; @@ -177,6 +183,8 @@ void fill_cagra_sample_filter(cagra_sample_filter& out, using DecayedFilter = std::decay_t; if constexpr (is_bitset_filter::value) { out.filter_data = make_cagra_bitset_filter_payload(filter, stream); + } else if constexpr (is_bloom_filter::value) { + out.filter_data = filter.filter_data; } else if constexpr (is_udf_filter::value) { out.filter_data = filter.filter_data; } @@ -199,7 +207,7 @@ template void* cagra_filter_data_ptr(const FilterT& filter) { using DecayedFilter = std::decay_t; - if constexpr (is_udf_filter::value) { + if constexpr (is_bloom_filter::value || is_udf_filter::value) { return filter.filter_data; } else if constexpr (requires { filter.filter; }) { return cagra_filter_data_ptr(filter.filter); diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_impl.cuh b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_impl.cuh index d01f58166d..1b3b3825f2 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_impl.cuh +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_impl.cuh @@ -9,6 +9,8 @@ #include "../../sample_filter_data.cuh" +#include + #include #include @@ -38,4 +40,15 @@ __device__ bool sample_filter_bitset_impl(uint32_t /*query_id*/, return view.test(node_id); } +template +__device__ bool sample_filter_bloom_filter_impl(uint32_t /*query_id*/, + SourceIndexT node_id, + void* filter_data) +{ + if (filter_data == nullptr) { return true; } + + auto* data = static_cast*>(filter_data); + return data->filter.contains(static_cast(node_id)); +} + } // namespace cuvs::neighbors::detail diff --git a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_matrix.json b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_matrix.json index 0136587b48..b58f56ceb6 100644 --- a/cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_matrix.json +++ b/cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_matrix.json @@ -1,5 +1,5 @@ { - "filter_name": ["none", "bitset"], + "filter_name": ["none", "bitset", "bloom_filter"], "_bitset": [ { "bitset_type": "uint32_t", diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in b/cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in index 7c642fe406..e60af91046 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in @@ -11,6 +11,8 @@ namespace { using data_t = @data_type@; using bitset_filter_t = cuvs::neighbors::cagra::detail::CagraSampleFilterWithQueryIdOffset< cuvs::neighbors::filtering::bitset_filter>; +using bloom_filter_t = cuvs::neighbors::cagra::detail::CagraSampleFilterWithQueryIdOffset< + cuvs::neighbors::filtering::bloom_filter>; using udf_filter_t = cuvs::neighbors::cagra::detail::CagraSampleFilterWithQueryIdOffset< cuvs::neighbors::filtering::udf_filter>; @@ -22,6 +24,7 @@ instantiate_kernel_selection(data_t, float, cuvs::neighbors::filtering::none_sample_filter); instantiate_kernel_selection(data_t, uint32_t, float, bitset_filter_t); +instantiate_kernel_selection(data_t, uint32_t, float, bloom_filter_t); instantiate_kernel_selection(data_t, uint32_t, float, udf_filter_t); } // namespace cuvs::neighbors::cagra::detail::multi_cta_search diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in b/cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in index 4616a9652b..c9eac33d44 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in @@ -11,6 +11,8 @@ namespace { using data_t = @data_type@; using bitset_filter_t = cuvs::neighbors::cagra::detail::CagraSampleFilterWithQueryIdOffset< cuvs::neighbors::filtering::bitset_filter>; +using bloom_filter_t = cuvs::neighbors::cagra::detail::CagraSampleFilterWithQueryIdOffset< + cuvs::neighbors::filtering::bloom_filter>; using udf_filter_t = cuvs::neighbors::cagra::detail::CagraSampleFilterWithQueryIdOffset< cuvs::neighbors::filtering::udf_filter>; @@ -22,6 +24,7 @@ instantiate_kernel_selection(data_t, float, cuvs::neighbors::filtering::none_sample_filter); instantiate_kernel_selection(data_t, uint32_t, float, bitset_filter_t); +instantiate_kernel_selection(data_t, uint32_t, float, bloom_filter_t); instantiate_kernel_selection(data_t, uint32_t, float, udf_filter_t); } // namespace cuvs::neighbors::cagra::detail::single_cta_search diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel_launcher_jit.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel_launcher_jit.cuh index a8731d06b5..28d660e2e1 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel_launcher_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel_launcher_jit.cuh @@ -80,6 +80,8 @@ std::uint64_t cagra_sample_filter_type_id(const SampleFilterT& sample_filter) { using DecayedFilter = std::decay_t; if constexpr (is_udf_filter::value) { + return 3; + } else if constexpr (is_bloom_filter::value) { return 2; } else if constexpr (is_bitset_filter::value) { return 1; diff --git a/cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp b/cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp index e5157ffa6a..dd118757e4 100644 --- a/cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp +++ b/cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp @@ -100,6 +100,8 @@ struct sample_filter_jit_tag { using namespace cuvs::neighbors::filtering; if constexpr (std::is_same_v) { return cuvs::neighbors::detail::tag_filter_none{}; + } else if constexpr (is_bloom_filter::value) { + return cuvs::neighbors::detail::tag_filter_bloom_filter{}; } else if constexpr (is_udf_filter::value) { return cuvs::neighbors::detail::tag_filter_udf{}; } else if constexpr (requires { std::declval().filter; }) { @@ -109,6 +111,8 @@ struct sample_filter_jit_tag { std::is_same_v, bitset_filter>) { return cuvs::neighbors::detail::tag_filter_bitset{}; + } else if constexpr (is_bloom_filter>::value) { + return cuvs::neighbors::detail::tag_filter_bloom_filter{}; } else if constexpr (is_udf_filter>::value) { return cuvs::neighbors::detail::tag_filter_udf{}; } else { diff --git a/cpp/src/neighbors/detail/sample_filter_data.cuh b/cpp/src/neighbors/detail/sample_filter_data.cuh index 4c99ca1e3a..3f2e412681 100644 --- a/cpp/src/neighbors/detail/sample_filter_data.cuh +++ b/cpp/src/neighbors/detail/sample_filter_data.cuh @@ -5,6 +5,8 @@ #pragma once +#include + #include #include @@ -21,4 +23,12 @@ struct bitset_filter_data_t { SourceIndexT original_nbits{}; }; +/// Global cuco bloom filter ref for linked @c sample_filter in CAGRA JIT LTO. +template +struct bloom_filter_data_t { + using ref_type = typename cuco::bloom_filter::ref_type<>; + + ref_type filter{}; +}; + } // namespace cuvs::neighbors::detail diff --git a/examples/cpp/CMakeLists.txt b/examples/cpp/CMakeLists.txt index d63ddbdb71..b1ba8eecb6 100644 --- a/examples/cpp/CMakeLists.txt +++ b/examples/cpp/CMakeLists.txt @@ -27,11 +27,14 @@ find_package(Threads) rapids_cpm_init() set(BUILD_CUVS_C_LIBRARY OFF) include(../cmake/thirdparty/get_cuvs.cmake) +include(${rapids-cmake-dir}/cpm/cuco.cmake) +rapids_cpm_cuco() # -------------- compile tasks ----------------- # add_executable(BRUTE_FORCE_EXAMPLE src/brute_force_bitmap.cu) add_executable(CAGRA_EXAMPLE src/cagra_example.cu) add_executable(CAGRA_FILTER_UDF_EXAMPLE src/cagra_filter_udf_example.cu) +add_executable(CAGRA_BLOOM_FILTER_EXAMPLE src/cagra_bloom_filter_example.cu) add_executable(CAGRA_HNSW_ACE_EXAMPLE src/cagra_hnsw_ace_example.cu) add_executable(CAGRA_PERSISTENT_EXAMPLE src/cagra_persistent_example.cu) add_executable(DYNAMIC_BATCHING_EXAMPLE src/dynamic_batching_example.cu) @@ -48,6 +51,9 @@ target_link_libraries(CAGRA_EXAMPLE PRIVATE cuvs::cuvs $ ) +target_link_libraries( + CAGRA_BLOOM_FILTER_EXAMPLE PRIVATE cuvs::cuvs cuco::cuco $ +) target_link_libraries(CAGRA_HNSW_ACE_EXAMPLE PRIVATE cuvs::cuvs $) target_link_libraries( CAGRA_PERSISTENT_EXAMPLE PRIVATE cuvs::cuvs $ Threads::Threads diff --git a/examples/cpp/src/cagra_bloom_filter_example.cu b/examples/cpp/src/cagra_bloom_filter_example.cu new file mode 100644 index 0000000000..7d16989bbb --- /dev/null +++ b/examples/cpp/src/cagra_bloom_filter_example.cu @@ -0,0 +1,171 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include + +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include + +namespace { + +constexpr int64_t n_rows = 4096; +constexpr int64_t n_dim = 32; +constexpr int64_t n_queries = 4; +constexpr int64_t k = 8; +constexpr int sub_filters = 256; + +using key_type = std::uint32_t; +using filter_type = cuco::bloom_filter; +using ref_type = filter_type::ref_type<>; + +// Layout must match cuvs::neighbors::detail::bloom_filter_data_t in the JIT fragment. +struct bloom_payload { + ref_type filter; +}; + +// Global index filter: even row ids are valid candidates (same rule for every query). +bool is_valid_row(key_type source_id) { return (source_id % 2) == 0; } + +std::vector copy_neighbors_to_host( + raft::device_resources const& res, + raft::device_matrix_view neighbors) +{ + std::vector host(neighbors.size()); + raft::copy( + host.data(), neighbors.data_handle(), host.size(), raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + return host; +} + +} // namespace + +int main() +{ + raft::device_resources res; + auto stream = raft::resource::get_cuda_stream(res); + + rmm::mr::pool_memory_resource pool_mr(rmm::mr::get_current_device_resource_ref(), + 1024 * 1024 * 1024ull); + rmm::mr::set_current_device_resource(pool_mr); + + auto dataset = raft::make_device_matrix(res, n_rows, n_dim); + auto queries = raft::make_device_matrix(res, n_queries, n_dim); + + raft::random::RngState rng(1234ULL); + raft::random::uniform(res, rng, dataset.data_handle(), dataset.size(), -1.0f, 1.0f); + raft::random::uniform(res, rng, queries.data_handle(), queries.size(), -1.0f, 1.0f); + + cuvs::neighbors::cagra::index_params index_params; + index_params.metric = cuvs::distance::DistanceType::L2Expanded; + index_params.graph_degree = 32; + index_params.intermediate_graph_degree = 64; + index_params.graph_build_params = cuvs::neighbors::cagra::graph_build_params::nn_descent_params( + index_params.intermediate_graph_degree); + + std::cout << "Building CAGRA index" << std::endl; + auto index = + cuvs::neighbors::cagra::build(res, index_params, raft::make_const_mdspan(dataset.view())); + + // Build one global bloom filter over the index: bulk-insert every valid row id once. + std::vector valid_ids_host; + valid_ids_host.reserve(static_cast(n_rows / 2)); + for (int64_t i = 0; i < n_rows; ++i) { + if (is_valid_row(static_cast(i))) { + valid_ids_host.push_back(static_cast(i)); + } + } + + rmm::device_uvector valid_ids_device(valid_ids_host.size(), stream); + raft::copy(valid_ids_device.data(), valid_ids_host.data(), valid_ids_host.size(), stream); + + filter_type allowed_rows{sub_filters}; + allowed_rows.add_async( + valid_ids_device.data(), valid_ids_device.data() + valid_ids_device.size(), stream); + raft::resource::sync_stream(res); + + std::cout << "Inserted " << valid_ids_host.size() + << " valid row ids into global bloom filter via bulk add_async" << std::endl; + + // Copy the owning filter's device ref into a payload the JIT fragment can probe. + auto payload_device = raft::make_device_vector(res, 1); + bloom_payload host_payload{allowed_rows.ref()}; + raft::copy(payload_device.data_handle(), &host_payload, 1, stream); + raft::resource::sync_stream(res); + + auto neighbors = raft::make_device_matrix(res, n_queries, k); + auto distances = raft::make_device_matrix(res, n_queries, k); + + cuvs::neighbors::cagra::search_params search_params; + search_params.algo = cuvs::neighbors::cagra::search_algo::MULTI_CTA; + search_params.itopk_size = 128; + search_params.max_queries = n_queries; + search_params.thread_block_size = 256; + + // ~50% of rows are rejected by the global even-id predicate. + auto filter = cuvs::neighbors::filtering::bloom_filter(payload_device.data_handle(), 0.5f); + + cuvs::neighbors::cagra::search(res, + search_params, + index, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view(), + filter); + + auto host_neighbors = copy_neighbors_to_host(res, neighbors.view()); + + std::cout << "bloom_filter first query neighbors:"; + for (int64_t i = 0; i < k; ++i) { + std::cout << " " << host_neighbors[static_cast(i)]; + } + std::cout << std::endl; + + // Validate with cuco's bulk contains API over the returned neighbors. + rmm::device_uvector neighbor_ids_device(host_neighbors.size(), stream); + rmm::device_uvector bloom_hits_device(host_neighbors.size(), stream); + raft::copy(neighbor_ids_device.data(), host_neighbors.data(), host_neighbors.size(), stream); + allowed_rows.contains_async(neighbor_ids_device.data(), + neighbor_ids_device.data() + neighbor_ids_device.size(), + bloom_hits_device.data(), + stream); + raft::resource::sync_stream(res); + + std::vector bloom_hits_host(bloom_hits_device.size()); + raft::copy(bloom_hits_host.data(), bloom_hits_device.data(), bloom_hits_host.size(), stream); + raft::resource::sync_stream(res); + + for (size_t i = 0; i < host_neighbors.size(); ++i) { + auto source_id = host_neighbors[i]; + if (source_id >= static_cast(n_rows)) { + std::cerr << "bloom_filter produced out-of-range source_id=" << source_id << std::endl; + return 1; + } + if (bloom_hits_host[i] == 0) { + std::cerr << "bloom_filter rejected source_id=" << source_id + << " but global bloom filter bulk contains says absent" << std::endl; + return 1; + } + if (!is_valid_row(source_id)) { + std::cerr << "bloom_filter allowed invalid source_id=" << source_id + << " (unexpected bloom false positive)" << std::endl; + return 1; + } + } + + std::cout << "CAGRA bloom filter example produced valid filtered neighbors." << std::endl; + return 0; +} From 16946a4d5ed3f92041c3b8f917128de1408a6c06 Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Mon, 15 Jun 2026 17:10:05 +0000 Subject: [PATCH 02/11] benchmarks --- examples/cpp/CMakeLists.txt | 4 + examples/cpp/plot_filter_benchmark.py | 457 ++++++++++++++++++ examples/cpp/src/cagra_filter_benchmark.cu | 537 +++++++++++++++++++++ 3 files changed, 998 insertions(+) create mode 100755 examples/cpp/plot_filter_benchmark.py create mode 100644 examples/cpp/src/cagra_filter_benchmark.cu diff --git a/examples/cpp/CMakeLists.txt b/examples/cpp/CMakeLists.txt index b1ba8eecb6..07573d40dd 100644 --- a/examples/cpp/CMakeLists.txt +++ b/examples/cpp/CMakeLists.txt @@ -35,6 +35,7 @@ add_executable(BRUTE_FORCE_EXAMPLE src/brute_force_bitmap.cu) add_executable(CAGRA_EXAMPLE src/cagra_example.cu) add_executable(CAGRA_FILTER_UDF_EXAMPLE src/cagra_filter_udf_example.cu) add_executable(CAGRA_BLOOM_FILTER_EXAMPLE src/cagra_bloom_filter_example.cu) +add_executable(CAGRA_FILTER_BENCHMARK src/cagra_filter_benchmark.cu) add_executable(CAGRA_HNSW_ACE_EXAMPLE src/cagra_hnsw_ace_example.cu) add_executable(CAGRA_PERSISTENT_EXAMPLE src/cagra_persistent_example.cu) add_executable(DYNAMIC_BATCHING_EXAMPLE src/dynamic_batching_example.cu) @@ -54,6 +55,9 @@ target_link_libraries( target_link_libraries( CAGRA_BLOOM_FILTER_EXAMPLE PRIVATE cuvs::cuvs cuco::cuco $ ) +target_link_libraries( + CAGRA_FILTER_BENCHMARK PRIVATE cuvs::cuvs cuco::cuco $ +) target_link_libraries(CAGRA_HNSW_ACE_EXAMPLE PRIVATE cuvs::cuvs $) target_link_libraries( CAGRA_PERSISTENT_EXAMPLE PRIVATE cuvs::cuvs $ Threads::Threads diff --git a/examples/cpp/plot_filter_benchmark.py b/examples/cpp/plot_filter_benchmark.py new file mode 100755 index 0000000000..8840c5e9fc --- /dev/null +++ b/examples/cpp/plot_filter_benchmark.py @@ -0,0 +1,457 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +"""Plot CAGRA filter benchmark CSV (bitset vs bloom_filter line charts).""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import seaborn as sns + +FILTER_LABELS = { + "bitset": "Bitset", + "bloom_filter": "Bloom", +} + +FILTER_ORDER = ["Bitset", "Bloom"] +FILTER_COLORS = {"Bitset": "#0173b2", "Bloom": "#de8f05"} +FILTER_MARKERS = {"Bitset": "o", "Bloom": "s"} + +SEARCH_LABELS = { + 10000: "10k queries", + 25000: "25k queries", +} + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Generate line charts from cagra_filter_benchmark CSV output.", + ) + parser.add_argument( + "csv", + nargs="?", + default="filter_bench.csv", + help="Path to benchmark CSV (default: filter_bench.csv)", + ) + parser.add_argument( + "-o", + "--output-dir", + default="filter_bench_plots", + help="Directory for PNG output (default: filter_bench_plots)", + ) + parser.add_argument( + "--with-recall", + action="store_true", + help="Also plot recall charts when non-NaN values are present", + ) + return parser.parse_args() + + +def load_csv(path: Path) -> pd.DataFrame: + df = pd.read_csv(path) + df["recall"] = pd.to_numeric(df["recall"], errors="coerce") + df["filter_label"] = ( + df["filter_type"].map(FILTER_LABELS).fillna(df["filter_type"]) + ) + df["search_label"] = ( + df["search_n_rows"] + .map(SEARCH_LABELS) + .fillna(df["search_n_rows"].astype(str) + " queries") + ) + return df + + +def save_figure(fig: plt.Figure, path: Path) -> None: + fig.subplots_adjust(left=0.11, bottom=0.08, right=0.98, top=0.94) + fig.savefig(path, dpi=150, bbox_inches="tight") + plt.close(fig) + + +def plot_panel_bars( + ax: plt.Axes, + panel: pd.DataFrame, + x_col: str, + y_col: str, +) -> list: + """Grouped bars for two filters at each x value — visible even when latencies are close.""" + x_values = sorted(panel[x_col].unique()) + x_idx = np.arange(len(x_values)) + bar_width = 0.36 + handles = [] + + for series_idx, label in enumerate(FILTER_ORDER): + heights = [] + for x_val in x_values: + row = panel[ + (panel["filter_label"] == label) & (panel[x_col] == x_val) + ] + heights.append(row[y_col].iloc[0] if not row.empty else 0.0) + offset = (series_idx - 0.5) * bar_width + bars = ax.bar( + x_idx + offset, + heights, + width=bar_width, + label=label, + color=FILTER_COLORS[label], + edgecolor="black", + linewidth=0.6, + alpha=0.9, + ) + handles.append(bars[0]) + + ax.set_xticks(x_idx) + if x_col == "build_n_rows": + ax.set_xticklabels([f"{int(v):,}" for v in x_values]) + else: + ax.set_xticklabels( + [ + str(int(v)) if float(v).is_integer() else str(v) + for v in x_values + ] + ) + return handles + + +def plot_panel_lines( + ax: plt.Axes, + panel: pd.DataFrame, + x_col: str, + y_col: str, + log_x: bool = False, +) -> list: + """Draw bitset vs bloom as two explicit series (no seaborn hue/style mashup).""" + handles = [] + x_values = sorted(panel[x_col].unique()) + if log_x: + ax.set_xscale("log") + ax.set_xticks(x_values) + ax.set_xticklabels([f"{int(v):,}" for v in x_values]) + ax.minorticks_off() + + for label in FILTER_ORDER: + series = panel[panel["filter_label"] == label].sort_values(x_col) + if series.empty: + continue + (line,) = ax.plot( + series[x_col], + series[y_col], + label=label, + color=FILTER_COLORS[label], + marker=FILTER_MARKERS[label], + markersize=7, + linewidth=2.0, + markerfacecolor="white", + markeredgewidth=1.5, + ) + handles.append(line) + + return handles + + +def add_row_band_label( + fig: plt.Figure, axes: np.ndarray, row_idx: int, text: str +) -> None: + ax = axes[row_idx, 0] + pos = ax.get_position() + fig.text( + 0.02, + (pos.y0 + pos.y1) / 2, + text, + ha="left", + va="center", + rotation=90, + fontsize=10, + fontweight="bold", + ) + + +def plot_latency_grid_for_valid_pct( + df: pd.DataFrame, + out_dir: Path, + valid_pct: int, +) -> None: + """One PNG per valid_pct: 6×3 grid. + + Row bands (top to bottom): 10k queries × dims 128/512/1024, then 25k × dims. + Columns: k = 64 / 256 / 1024. + Each panel has only two lines (bitset vs bloom) over index size. + """ + sub = df[df["valid_pct"] == valid_pct].copy() + if sub.empty: + print(f"valid_pct={valid_pct}%: no rows, skipping") + return + + k_values = sorted(sub["k"].unique()) + col_values = sorted(sub["build_n_cols"].unique()) + search_values = sorted(sub["search_n_rows"].unique()) + + row_specs = [ + (search, dims) for search in search_values for dims in col_values + ] + n_rows = len(row_specs) + n_cols = len(k_values) + + fig, axes = plt.subplots( + n_rows, + n_cols, + figsize=(4.0 * n_cols + 0.8, 2.4 * n_rows), + squeeze=False, + ) + + legend_handles = None + for row_idx, (search_n_rows, build_n_cols) in enumerate(row_specs): + for col_idx, k in enumerate(k_values): + ax = axes[row_idx, col_idx] + panel = sub[ + (sub["build_n_cols"] == build_n_cols) + & (sub["k"] == k) + & (sub["search_n_rows"] == search_n_rows) + ] + if panel.empty: + ax.set_visible(False) + continue + + handles = plot_panel_bars( + ax, + panel, + x_col="build_n_rows", + y_col="avg_latency_per_query_ms", + ) + if legend_handles is None and handles: + legend_handles = handles + + ax.set_title(f"k={k}", fontsize=10) + ax.set_xlabel("index rows") + if col_idx == 0: + ax.set_ylabel(f"dims={build_n_cols}\nlatency / query (ms)") + else: + ax.set_ylabel("") + + for search_idx, search_n_rows in enumerate(search_values): + band_row = search_idx * len(col_values) + add_row_band_label( + fig, + axes, + band_row, + SEARCH_LABELS.get(search_n_rows, f"{search_n_rows:,} queries"), + ) + + if legend_handles: + fig.legend( + legend_handles, + FILTER_ORDER, + loc="lower center", + bbox_to_anchor=(0.5, -0.02), + ncol=2, + frameon=False, + fontsize=11, + ) + + fig.suptitle( + f"Per-query search latency — {valid_pct}% rows valid", + fontsize=13, + y=1.01, + ) + path = out_dir / f"latency_valid_{valid_pct}pct.png" + save_figure(fig, path) + print(f"wrote {path}") + + +def plot_all_valid_pct_grids(df: pd.DataFrame, out_dir: Path) -> None: + for valid_pct in sorted(df["valid_pct"].unique()): + plot_latency_grid_for_valid_pct(df, out_dir, int(valid_pct)) + + +def plot_valid_pct_overview( + df: pd.DataFrame, + out_dir: Path, + build_n_rows: int, + search_n_rows: int, +) -> None: + """valid_pct on x-axis at one build/search point; 3×3 grid (dims × k).""" + sub = df[ + (df["build_n_rows"] == build_n_rows) + & (df["search_n_rows"] == search_n_rows) + ].copy() + if sub.empty: + print( + "valid_pct overview: no rows for selected build/search slice, skipping" + ) + return + + k_values = sorted(sub["k"].unique()) + col_values = sorted(sub["build_n_cols"].unique()) + + fig, axes = plt.subplots( + len(col_values), + len(k_values), + figsize=(4.0 * len(k_values), 2.8 * len(col_values)), + squeeze=False, + ) + + legend_handles = None + for row_idx, build_n_cols in enumerate(col_values): + for col_idx, k in enumerate(k_values): + ax = axes[row_idx, col_idx] + panel = sub[ + (sub["build_n_cols"] == build_n_cols) & (sub["k"] == k) + ] + if panel.empty: + ax.set_visible(False) + continue + + handles = plot_panel_bars( + ax, + panel, + x_col="valid_pct", + y_col="avg_latency_per_query_ms", + ) + if legend_handles is None and handles: + legend_handles = handles + + ax.set_title(f"k={k}", fontsize=10) + ax.set_xlabel("valid rows (%)") + if col_idx == 0: + ax.set_ylabel(f"dims={build_n_cols}\nlatency / query (ms)") + else: + ax.set_ylabel("") + + if legend_handles: + fig.legend( + legend_handles, + FILTER_ORDER, + loc="lower center", + bbox_to_anchor=(0.5, -0.02), + ncol=2, + frameon=False, + fontsize=11, + ) + + fig.suptitle( + f"Latency vs filter selectivity — " + f"build={build_n_rows:,} rows search={search_n_rows:,} queries", + fontsize=12, + y=1.01, + ) + path = out_dir / "overview_valid_pct_sweep.png" + save_figure(fig, path) + print(f"wrote {path}") + + +def plot_recall_grid_for_valid_pct( + df: pd.DataFrame, + out_dir: Path, + valid_pct: int, +) -> None: + sub = df[(df["valid_pct"] == valid_pct) & df["recall"].notna()].copy() + if sub.empty: + return + + k_values = sorted(sub["k"].unique()) + col_values = sorted(sub["build_n_cols"].unique()) + search_values = sorted(sub["search_n_rows"].unique()) + row_specs = [ + (search, dims) for search in search_values for dims in col_values + ] + + fig, axes = plt.subplots( + len(row_specs), + len(k_values), + figsize=(4.0 * len(k_values) + 0.8, 2.4 * len(row_specs)), + squeeze=False, + ) + + legend_handles = None + for row_idx, (search_n_rows, build_n_cols) in enumerate(row_specs): + for col_idx, k in enumerate(k_values): + ax = axes[row_idx, col_idx] + panel = sub[ + (sub["build_n_cols"] == build_n_cols) + & (sub["k"] == k) + & (sub["search_n_rows"] == search_n_rows) + ] + if panel.empty: + ax.set_visible(False) + continue + + handles = plot_panel_bars( + ax, + panel, + x_col="build_n_rows", + y_col="recall", + ) + if legend_handles is None and handles: + legend_handles = handles + + ax.set_ylim(0.0, 1.0) + ax.set_title(f"k={k}", fontsize=10) + ax.set_xlabel("index rows") + if col_idx == 0: + ax.set_ylabel(f"dims={build_n_cols}\nrecall@k") + else: + ax.set_ylabel("") + + for search_idx, search_n_rows in enumerate(search_values): + band_row = search_idx * len(col_values) + add_row_band_label( + fig, + axes, + band_row, + SEARCH_LABELS.get(search_n_rows, f"{search_n_rows:,} queries"), + ) + + if legend_handles: + fig.legend( + legend_handles, + FILTER_ORDER, + loc="lower center", + bbox_to_anchor=(0.5, -0.02), + ncol=2, + frameon=False, + fontsize=11, + ) + + fig.suptitle(f"Recall@k — {valid_pct}% rows valid", fontsize=13, y=1.01) + path = out_dir / f"recall_valid_{valid_pct}pct.png" + save_figure(fig, path) + print(f"wrote {path}") + + +def main() -> None: + args = parse_args() + csv_path = Path(args.csv) + out_dir = Path(args.output_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + if not csv_path.is_file(): + raise SystemExit(f"CSV not found: {csv_path}") + + df = load_csv(csv_path) + print(f"loaded {len(df)} rows from {csv_path}") + + sns.set_theme(style="whitegrid", context="notebook") + + plot_all_valid_pct_grids(df, out_dir) + + default_build = int(df["build_n_rows"].min()) + default_search = int(df["search_n_rows"].min()) + plot_valid_pct_overview(df, out_dir, default_build, default_search) + + if args.with_recall and df["recall"].notna().any(): + for valid_pct in sorted(df["valid_pct"].unique()): + plot_recall_grid_for_valid_pct(df, out_dir, int(valid_pct)) + elif args.with_recall: + print( + "recall charts skipped: CSV has no recall values (run benchmark with --ground-truth)" + ) + else: + print("recall charts skipped (default; pass --with-recall to enable)") + + +if __name__ == "__main__": + main() diff --git a/examples/cpp/src/cagra_filter_benchmark.cu b/examples/cpp/src/cagra_filter_benchmark.cu new file mode 100644 index 0000000000..2c72c22ddd --- /dev/null +++ b/examples/cpp/src/cagra_filter_benchmark.cu @@ -0,0 +1,537 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +constexpr int k_warmup_runs = 1; +constexpr int k_timed_runs = 3; + +using key_type = std::uint32_t; +using filter_type = cuco::bloom_filter; +using ref_type = filter_type::ref_type<>; + +struct bloom_payload { + ref_type filter; +}; + +constexpr std::array k_build_rows{100'000, 500'000, 1'000'000}; +constexpr std::array k_build_cols{128, 512, 1024}; +constexpr std::array k_search_rows{10'000, 25'000}; +constexpr std::array k_search_cols{128, 512, 1024}; +constexpr std::array k_valid_pcts{10, 50, 90}; +constexpr std::array k_values{64, 256, 1024}; + +bool is_valid_row(key_type row_id, int valid_pct) +{ + // Deterministic ~valid_pct% membership independent of dataset size. + return (static_cast(row_id) * 2654435761ULL) % 100ULL < + static_cast(valid_pct); +} + +std::size_t bloom_num_blocks(std::size_t num_valid_rows) +{ + // Scale filter size with the number of inserted keys; keep a reasonable minimum. + std::size_t blocks = std::max(256, num_valid_rows / 8); + blocks = std::min(blocks, static_cast(1 << 20)); + return blocks; +} + +double compute_recall(std::vector const& expected, + std::vector const& actual, + int64_t n_queries, + int64_t k, + int64_t n_rows) +{ + std::size_t match_count = 0; + std::size_t total_count = static_cast(n_queries) * static_cast(k); + for (int64_t q = 0; q < n_queries; ++q) { + for (int64_t ki = 0; ki < k; ++ki) { + auto const act = actual[static_cast(q * k + ki)]; + if (act >= static_cast(n_rows)) { continue; } + for (int64_t kj = 0; kj < k; ++kj) { + if (expected[static_cast(q * k + kj)] == act) { + ++match_count; + break; + } + } + } + } + return total_count == 0 ? 0.0 + : static_cast(match_count) / static_cast(total_count); +} + +std::vector copy_neighbors_to_host(raft::device_resources const& res, + raft::device_matrix_view neighbors) +{ + std::vector host(neighbors.size()); + auto stream = raft::resource::get_cuda_stream(res); + raft::copy(host.data(), neighbors.data_handle(), host.size(), stream); + raft::resource::sync_stream(res); + return host; +} + +struct filter_assets { + cuvs::core::bitset removed_bitset; + cuvs::neighbors::filtering::bitset_filter bitset_filter; + filter_type bloom; + rmm::device_uvector bloom_payload; + cuvs::neighbors::filtering::bloom_filter bloom_filter; + float filtering_rate{0.0f}; +}; + +filter_assets make_filters(raft::device_resources const& res, + int64_t n_rows, + int valid_pct, + rmm::cuda_stream_view stream) +{ + std::vector valid_ids_host; + std::vector removed_ids_host; + valid_ids_host.reserve(static_cast(n_rows)); + removed_ids_host.reserve(static_cast(n_rows)); + + for (int64_t i = 0; i < n_rows; ++i) { + auto const row = static_cast(i); + if (is_valid_row(row, valid_pct)) { + valid_ids_host.push_back(row); + } else { + removed_ids_host.push_back(i); + } + } + + auto removed_ids = + raft::make_device_vector(res, static_cast(removed_ids_host.size())); + if (!removed_ids_host.empty()) { + raft::copy(removed_ids.data_handle(), removed_ids_host.data(), removed_ids_host.size(), stream); + } + + auto removed_bitset = cuvs::core::bitset(res, removed_ids.view(), n_rows); + auto bitset_filter = + cuvs::neighbors::filtering::bitset_filter(removed_bitset.view()); + auto bloom = filter_type{bloom_num_blocks(valid_ids_host.size()), {}, {}, {}, stream}; + auto payload_device = rmm::device_uvector{1, stream}; + float const filtering_rate = static_cast(100 - valid_pct) / 100.0f; + + if (!valid_ids_host.empty()) { + rmm::device_uvector valid_ids_device(valid_ids_host.size(), stream); + raft::copy(valid_ids_device.data(), valid_ids_host.data(), valid_ids_host.size(), stream); + bloom.add_async( + valid_ids_device.data(), valid_ids_device.data() + valid_ids_device.size(), stream); + } + + bloom_payload host_payload{bloom.ref()}; + raft::copy(payload_device.data(), &host_payload, 1, stream); + auto bloom_filter_obj = + cuvs::neighbors::filtering::bloom_filter(payload_device.data(), filtering_rate); + + raft::resource::sync_stream(res); + return filter_assets{std::move(removed_bitset), + std::move(bitset_filter), + std::move(bloom), + std::move(payload_device), + std::move(bloom_filter_obj), + filtering_rate}; +} + +struct benchmark_case { + int64_t build_n_rows; + int64_t build_n_cols; + int64_t search_n_rows; + int64_t search_n_cols; + int valid_pct; + int64_t k; +}; + +struct csv_row { + benchmark_case config; + std::string filter_name; + double build_time_ms; + double avg_search_latency_ms; + double avg_latency_per_query_ms; + double recall; +}; + +void write_csv_header(std::ostream& os) +{ + os << "build_n_rows,build_n_cols,search_n_rows,search_n_cols,valid_pct,filter_type," + "build_time_ms,avg_search_latency_ms,avg_latency_per_query_ms,recall,k,warmup_runs," + "timed_runs\n"; +} + +void write_csv_row(std::ostream& os, csv_row const& row) +{ + os << row.config.build_n_rows << ',' << row.config.build_n_cols << ',' << row.config.search_n_rows + << ',' << row.config.search_n_cols << ',' << row.config.valid_pct << ',' << row.filter_name + << ',' << row.build_time_ms << ',' << row.avg_search_latency_ms << ',' + << row.avg_latency_per_query_ms << ',' << row.recall << ',' << row.config.k << ',' + << k_warmup_runs << ',' << k_timed_runs << '\n'; +} + +template +double time_cuda_ms(raft::device_resources const& res, int runs, Fn&& fn) +{ + auto stream = raft::resource::get_cuda_stream(res); + cudaEvent_t start{}; + cudaEvent_t stop{}; + RAFT_CUDA_TRY(cudaEventCreate(&start)); + RAFT_CUDA_TRY(cudaEventCreate(&stop)); + + RAFT_CUDA_TRY(cudaEventRecord(start, stream)); + for (int i = 0; i < runs; ++i) { + fn(); + } + RAFT_CUDA_TRY(cudaEventRecord(stop, stream)); + RAFT_CUDA_TRY(cudaEventSynchronize(stop)); + + float elapsed_ms = 0.0f; + RAFT_CUDA_TRY(cudaEventElapsedTime(&elapsed_ms, start, stop)); + RAFT_CUDA_TRY(cudaEventDestroy(start)); + RAFT_CUDA_TRY(cudaEventDestroy(stop)); + return static_cast(elapsed_ms) / static_cast(runs); +} + +void append_cases(std::vector& cases, + std::vector const& build_rows, + std::vector const& build_cols, + std::vector const& search_rows, + std::vector const& search_cols, + std::vector const& valid_pcts, + std::vector const& k_sweep) +{ + for (auto build_n_rows : build_rows) { + for (auto build_n_cols : build_cols) { + for (auto search_n_rows : search_rows) { + for (auto search_n_cols : search_cols) { + if (search_n_cols != build_n_cols) { continue; } + for (auto valid_pct : valid_pcts) { + for (auto k : k_sweep) { + cases.push_back(benchmark_case{ + build_n_rows, build_n_cols, search_n_rows, search_n_cols, valid_pct, k}); + } + } + } + } + } + } +} + +std::vector make_cases(bool quick) +{ + std::vector cases; + if (quick) { + append_cases(cases, {100'000}, {128}, {10'000}, {128}, {1, 50}, {64}); + } else { + append_cases(cases, + {k_build_rows.begin(), k_build_rows.end()}, + {k_build_cols.begin(), k_build_cols.end()}, + {k_search_rows.begin(), k_search_rows.end()}, + {k_search_cols.begin(), k_search_cols.end()}, + {k_valid_pcts.begin(), k_valid_pcts.end()}, + {k_values.begin(), k_values.end()}); + } + return cases; +} + +constexpr std::size_t k_max_bf_bytes = 20ULL << 30; // skip BF recall above this estimate +constexpr std::size_t k_max_bf_chunk_bytes = 2ULL << 30; // cap each BF chunk when computing recall + +std::size_t estimate_bf_distance_matrix_bytes(int64_t n_queries, int64_t n_dataset) +{ + return static_cast(n_queries) * static_cast(n_dataset) * sizeof(float); +} + +bool should_compute_bf_recall(int64_t n_queries, int64_t n_dataset) +{ + return estimate_bf_distance_matrix_bytes(n_queries, n_dataset) <= k_max_bf_bytes; +} + +int64_t choose_gt_chunk_queries(int64_t n_dataset) +{ + int64_t chunk = 256; + while (chunk > 1 && estimate_bf_distance_matrix_bytes(chunk, n_dataset) > k_max_bf_chunk_bytes) { + chunk /= 2; + } + return chunk; +} + +std::vector brute_force_ground_truth( + raft::device_resources const& res, + cuvs::neighbors::brute_force::index& bf_index, + cuvs::neighbors::brute_force::search_params const& bf_search_params, + raft::device_matrix_view queries, + cuvs::neighbors::filtering::bitset_filter const& bitset_filter, + int64_t k, + int64_t gt_chunk_queries) +{ + int64_t const n_queries = queries.extent(0); + std::vector gt_host(static_cast(n_queries * k)); + auto stream = raft::resource::get_cuda_stream(res); + + for (int64_t query_offset = 0; query_offset < n_queries; query_offset += gt_chunk_queries) { + int64_t const chunk_queries = std::min(gt_chunk_queries, n_queries - query_offset); + auto query_chunk = raft::make_device_matrix_view( + queries.data_handle() + query_offset * queries.extent(1), chunk_queries, queries.extent(1)); + auto gt_neighbors = raft::make_device_matrix(res, chunk_queries, k); + auto gt_distances = raft::make_device_matrix(res, chunk_queries, k); + + cuvs::neighbors::brute_force::search(res, + bf_search_params, + bf_index, + raft::make_const_mdspan(query_chunk), + gt_neighbors.view(), + gt_distances.view(), + bitset_filter); + raft::resource::sync_stream(res); + + std::vector chunk_host(static_cast(chunk_queries * k)); + raft::copy(chunk_host.data(), gt_neighbors.data_handle(), chunk_host.size(), stream); + raft::resource::sync_stream(res); + + for (int64_t q = 0; q < chunk_queries; ++q) { + for (int64_t ki = 0; ki < k; ++ki) { + gt_host[static_cast((query_offset + q) * k + ki)] = + static_cast(chunk_host[static_cast(q * k + ki)]); + } + } + } + + return gt_host; +} + +} // namespace + +int main(int argc, char** argv) +{ + std::string output_path = "cagra_filter_benchmark_results.csv"; + bool quick = false; + bool compute_ground_truth = false; + for (int i = 1; i < argc; ++i) { + std::string arg = argv[i]; + if (arg == "--quick") { + quick = true; + } else if (arg == "--ground-truth") { + compute_ground_truth = true; + } else if (arg == "--skip-ground-truth") { + compute_ground_truth = false; + } else if (arg == "--output" && i + 1 < argc) { + output_path = argv[++i]; + } else if (arg == "--help" || arg == "-h") { + std::cout << "Usage: " << argv[0] + << " [--quick] [--ground-truth] [--skip-ground-truth] [--output path.csv]\n" + << "\n" + << "Brute-force recall is skipped by default. Pass --ground-truth to compute it.\n"; + return 0; + } else { + output_path = arg; + } + } + + auto cases = make_cases(quick); + std::cout << "Running " << cases.size() << " benchmark configurations" + << (quick ? " (quick mode)" : "") + << (compute_ground_truth ? " (with ground-truth recall)" : " (ground-truth skipped)") + << std::endl; + + std::ofstream csv(output_path); + if (!csv) { + std::cerr << "Failed to open output file: " << output_path << std::endl; + return 1; + } + write_csv_header(csv); + + raft::device_resources res; + auto stream = raft::resource::get_cuda_stream(res); + + // Large enough for the biggest benchmark configuration (1M x 1024 dataset + index overhead). + rmm::mr::pool_memory_resource pool_mr(rmm::mr::get_current_device_resource_ref(), 16ULL << 30); + rmm::mr::set_current_device_resource(pool_mr); + + int64_t prev_build_rows = -1; + int64_t prev_build_cols = -1; + double last_build_time_ms = 0.0; + + std::optional> index; + std::optional> dataset; + std::optional> bf_index; + + cuvs::neighbors::cagra::index_params index_params; + index_params.metric = cuvs::distance::DistanceType::L2Expanded; + index_params.graph_degree = 32; + index_params.intermediate_graph_degree = 64; + index_params.graph_build_params = cuvs::neighbors::cagra::graph_build_params::nn_descent_params( + index_params.intermediate_graph_degree); + + cuvs::neighbors::cagra::search_params search_params; + search_params.algo = cuvs::neighbors::cagra::search_algo::MULTI_CTA; + search_params.itopk_size = 128; + search_params.thread_block_size = 256; + + cuvs::neighbors::brute_force::index_params bf_index_params; + cuvs::neighbors::brute_force::search_params bf_search_params; + + for (std::size_t case_idx = 0; case_idx < cases.size(); ++case_idx) { + auto const& cfg = cases[case_idx]; + + if (cfg.build_n_rows != prev_build_rows || cfg.build_n_cols != prev_build_cols) { + std::cout << "Building CAGRA index: n_rows=" << cfg.build_n_rows + << " n_cols=" << cfg.build_n_cols << std::endl; + + dataset.emplace( + raft::make_device_matrix(res, cfg.build_n_rows, cfg.build_n_cols)); + raft::random::RngState rng( + static_cast(cfg.build_n_rows * 17 + cfg.build_n_cols)); + raft::random::uniform(res, rng, dataset->data_handle(), dataset->size(), -1.0f, 1.0f); + + auto build_start = std::chrono::steady_clock::now(); + index.emplace( + cuvs::neighbors::cagra::build(res, index_params, raft::make_const_mdspan(dataset->view()))); + bf_index.reset(); + raft::resource::sync_stream(res); + auto build_end = std::chrono::steady_clock::now(); + last_build_time_ms = + std::chrono::duration(build_end - build_start).count(); + + prev_build_rows = cfg.build_n_rows; + prev_build_cols = cfg.build_n_cols; + } + + int64_t const total_queries = cfg.search_n_rows; + search_params.max_queries = total_queries; + search_params.itopk_size = static_cast(cfg.k); + + std::cout << "Case " << (case_idx + 1) << '/' << cases.size() + << ": build_n_rows=" << cfg.build_n_rows << " search_n_rows=" << total_queries + << " k=" << cfg.k << " valid_pct=" << cfg.valid_pct << '%' << std::endl; + + try { + auto queries = + raft::make_device_matrix(res, total_queries, cfg.search_n_cols); + raft::random::RngState query_rng(static_cast( + cfg.build_n_rows * 31 + cfg.search_n_rows * 17 + cfg.search_n_cols + cfg.valid_pct)); + raft::random::uniform(res, query_rng, queries.data_handle(), queries.size(), -1.0f, 1.0f); + + auto neighbors = raft::make_device_matrix(res, total_queries, cfg.k); + auto distances = raft::make_device_matrix(res, total_queries, cfg.k); + + auto filters = make_filters(res, cfg.build_n_rows, cfg.valid_pct, stream); + + bool const run_bf_recall = + compute_ground_truth && should_compute_bf_recall(total_queries, cfg.build_n_rows); + std::optional> gt_host; + if (run_bf_recall) { + if (!bf_index.has_value()) { + bf_index.emplace(cuvs::neighbors::brute_force::build( + res, bf_index_params, raft::make_const_mdspan(dataset->view()))); + } + gt_host.emplace(brute_force_ground_truth(res, + *bf_index, + bf_search_params, + queries.view(), + filters.bitset_filter, + cfg.k, + choose_gt_chunk_queries(cfg.build_n_rows))); + } else if (compute_ground_truth) { + auto const est_gib = + static_cast(estimate_bf_distance_matrix_bytes(total_queries, cfg.build_n_rows)) / + static_cast(1ULL << 30); + std::cout << " skipping brute-force recall (estimated " << est_gib + << " GiB distance matrix > 20 GiB limit)" << std::endl; + } + + auto run_cagra_search = [&](cuvs::neighbors::filtering::base_filter const& filter) { + cuvs::neighbors::cagra::search(res, + search_params, + *index, + raft::make_const_mdspan(queries.view()), + neighbors.view(), + distances.view(), + filter); + raft::resource::sync_stream(res); + }; + + struct filter_run { + std::string name; + cuvs::neighbors::filtering::base_filter const* filter; + }; + std::vector filter_runs{ + {"bitset", &filters.bitset_filter}, + {"bloom_filter", &filters.bloom_filter}, + }; + + for (auto const& fr : filter_runs) { + for (int w = 0; w < k_warmup_runs; ++w) { + run_cagra_search(*fr.filter); + } + + double const avg_search_ms = + time_cuda_ms(res, k_timed_runs, [&] { run_cagra_search(*fr.filter); }); + double const avg_per_query_ms = avg_search_ms / static_cast(total_queries); + + double const recall = [&]() { + if (!gt_host.has_value()) { return std::numeric_limits::quiet_NaN(); } + auto result_host = copy_neighbors_to_host(res, neighbors.view()); + return compute_recall(*gt_host, result_host, total_queries, cfg.k, cfg.build_n_rows); + }(); + + write_csv_row( + csv, csv_row{cfg, fr.name, last_build_time_ms, avg_search_ms, avg_per_query_ms, recall}); + csv.flush(); + + std::cout << " " << fr.name << ": search_ms=" << avg_search_ms + << " per_query_ms=" << avg_per_query_ms << " recall="; + if (gt_host.has_value()) { + std::cout << recall; + } else { + std::cout << "n/a"; + } + std::cout << std::endl; + } + } catch (std::exception const& ex) { + std::cerr << " case failed: " << ex.what() << std::endl; + for (auto const* filter_name : {"bitset", "bloom_filter"}) { + write_csv_row(csv, + csv_row{cfg, + filter_name, + last_build_time_ms, + std::numeric_limits::quiet_NaN(), + std::numeric_limits::quiet_NaN(), + std::numeric_limits::quiet_NaN()}); + csv.flush(); + } + } + } + + std::cout << "Wrote results to " << output_path << std::endl; + return 0; +} From a3d7b96a7a6f707c12690a6d6ac2c00e9683b458 Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Wed, 8 Jul 2026 21:43:34 +0000 Subject: [PATCH 03/11] add api wrapper --- cpp/CMakeLists.txt | 1 + cpp/include/cuvs/core/bloom_filter.hpp | 93 +++ cpp/include/cuvs/neighbors/common.hpp | 8 +- cpp/src/core/bloom_filter.cu | 170 ++++++ cpp/tests/neighbors/ann_cagra.cuh | 74 +++ examples/cpp/CMakeLists.txt | 5 +- examples/cpp/plot_filter_benchmark.py | 457 --------------- .../cpp/src/cagra_bloom_filter_example.cu | 45 +- examples/cpp/src/cagra_filter_benchmark.cu | 537 ------------------ 9 files changed, 361 insertions(+), 1029 deletions(-) create mode 100644 cpp/include/cuvs/core/bloom_filter.hpp create mode 100644 cpp/src/core/bloom_filter.cu delete mode 100755 examples/cpp/plot_filter_benchmark.py delete mode 100644 examples/cpp/src/cagra_filter_benchmark.cu diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index b3f6e23ca9..d81b66ba3c 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -1144,6 +1144,7 @@ if(NOT BUILD_CPU_ONLY) src/cluster/single_linkage_float.cu src/cluster/spectral.cu src/core/bitset.cu + src/core/bloom_filter.cu src/core/omp_wrapper.cpp src/util/file_io.cpp src/util/host_memory.cpp diff --git a/cpp/include/cuvs/core/bloom_filter.hpp b/cpp/include/cuvs/core/bloom_filter.hpp new file mode 100644 index 0000000000..e9d7de56df --- /dev/null +++ b/cpp/include/cuvs/core/bloom_filter.hpp @@ -0,0 +1,93 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include + +#include +#include +#include + +namespace CUVS_EXPORT cuvs { +namespace core { + +/** + * @brief cuVS-owned Bloom filter wrapper with opaque implementation. + * + * This class intentionally hides cuCollections types from the cuVS public API. + * The wrapper supports the expected bulk host APIs used by ANN workflows. + */ +class CUVS_EXPORT bloom_filter { + public: + using key_type = std::uint32_t; + + /** + * @brief Construct a Bloom filter with user-facing quality knobs. + * + * The first add/add_async call sets `dataset_rows` from `keys.size()` and then computes a target + * filter size from the first two parameters. The filter keeps at least `num_blocks` blocks, and + * may grow above that floor to satisfy the requested false-positive rate. + * + * The primary tuning knobs are: + * - @p expected_valid_rate: expected fraction of dataset rows that will be inserted as valid ids. + * - @p target_false_positive_rate: desired Bloom filter false-positive probability. + * + * Sizing math used internally: + * - `expected_insertions = ceil(dataset_rows * expected_valid_rate)` + * - `required_bits = -expected_insertions * ln(target_false_positive_rate) / (ln(2)^2)` + * - `required_blocks = ceil(required_bits / 256)` (default cuco policy uses 256-bit blocks) + * - `final_blocks = max(num_blocks, required_blocks)` + * + * Practical knob behavior: + * - Lower @p target_false_positive_rate -> larger filter, fewer false positives, typically higher + * filtered-search recall. + * - Higher @p expected_valid_rate -> larger filter for the same target false-positive rate. + * - @p num_blocks is an expert floor; keep default unless you need a hard minimum memory budget. + */ + bloom_filter(raft::resources const& res, + float expected_valid_rate = 1.0f, + float target_false_positive_rate = 0.01f, + std::size_t num_blocks = 256); + ~bloom_filter(); + + bloom_filter(bloom_filter const&) = delete; + bloom_filter& operator=(bloom_filter const&) = delete; + bloom_filter(bloom_filter&&) noexcept; + bloom_filter& operator=(bloom_filter&&) noexcept; + + void clear(raft::resources const& res); + void clear_async(raft::resources const& res); + + void add(raft::resources const& res, raft::device_vector_view keys); + void add_async(raft::resources const& res, + raft::device_vector_view keys); + + void contains(raft::resources const& res, + raft::device_vector_view keys, + raft::device_vector_view output) const; + void contains_async(raft::resources const& res, + raft::device_vector_view keys, + raft::device_vector_view output) const; + + [[nodiscard]] std::size_t num_blocks() const noexcept; + + /** + * @brief Device pointer to the CAGRA JIT sample-filter payload. + * + * The pointed object is device memory owned by this wrapper and remains valid while this object + * is alive. + */ + [[nodiscard]] void* filter_data() const noexcept; + + private: + struct impl; + std::unique_ptr impl_; +}; + +} // namespace core +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index da9b812f60..403091f3fb 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -618,12 +618,12 @@ struct bitset_filter : public base_filter { }; /** - * @brief Filter CAGRA candidates with a global @c cuco bloom filter over the index. + * @brief Filter CAGRA candidates with a global @c cuvs::core::bloom_filter over the index. * * Build the filter once on the host with bulk @c add() over the allowed dataset row ids, obtain a - * @c ref() from the owning @c cuco::bloom_filter, copy that ref to device memory, and pass the - * device pointer as @c filter_data. The linked JIT-LTO fragment probes the same filter for every - * query and candidate, similar to @ref bitset_filter but with probabilistic membership tests. + * @c ref() from the owning @c cuvs::core::bloom_filter, copy that ref to device memory, and pass + * the device pointer as @c filter_data. The linked JIT-LTO fragment probes the same filter for + * every query and candidate, similar to @ref bitset_filter but with probabilistic membership tests. * * Bloom filters have no false negatives: if a row was inserted, @c contains returns @c true. False * positives are possible, so highly selective predicates may still need a bitset or UDF for exact diff --git a/cpp/src/core/bloom_filter.cu b/cpp/src/core/bloom_filter.cu new file mode 100644 index 0000000000..cb493c7fd1 --- /dev/null +++ b/cpp/src/core/bloom_filter.cu @@ -0,0 +1,170 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../neighbors/detail/sample_filter_data.cuh" +#include + +#include + +#include +#include +#include + +#include + +#include +#include +#include +#include + +namespace cuvs::core { + +std::size_t compute_num_blocks_from_rates(std::size_t dataset_rows, + float expected_valid_rate, + float target_false_positive_rate) +{ + RAFT_EXPECTS(dataset_rows > 0, + "dataset_rows must be greater than zero when deriving bloom size."); + RAFT_EXPECTS(expected_valid_rate > 0.0f && expected_valid_rate <= 1.0f, + "expected_valid_rate must be in (0, 1]."); + RAFT_EXPECTS(target_false_positive_rate > 0.0f && target_false_positive_rate < 1.0f, + "target_false_positive_rate must be in (0, 1)."); + + // Bloom sizing: m = -n * ln(p) / (ln(2)^2), then blocks = ceil(m / 256 bits-per-block). + constexpr double kBitsPerBlock = 256.0; + constexpr double kLn2 = 0.6931471805599453; + constexpr double kLn2Sq = kLn2 * kLn2; + + auto expected_insertions = std::max( + 1, + static_cast(std::ceil(static_cast(dataset_rows) * expected_valid_rate))); + auto required_bits = -static_cast(expected_insertions) * + std::log(static_cast(target_false_positive_rate)) / kLn2Sq; + return std::max(1, + static_cast(std::ceil(required_bits / kBitsPerBlock))); +} + +struct bloom_filter::impl { + using key_type = bloom_filter::key_type; + using cuco_filter_type = cuco::bloom_filter; + using sample_filter_payload = cuvs::neighbors::detail::bloom_filter_data_t; + + cuco_filter_type filter; + rmm::device_uvector payload; + std::optional configured_dataset_rows; + float expected_valid_rate; + float target_false_positive_rate; + + impl(raft::resources const& res, + std::size_t num_blocks, + float expected_valid_rate_, + float target_false_positive_rate_) + : filter(num_blocks, {}, {}, {}, raft::resource::get_cuda_stream(res)), + payload(1, raft::resource::get_cuda_stream(res)), + expected_valid_rate(expected_valid_rate_), + target_false_positive_rate(target_false_positive_rate_) + { + auto stream = raft::resource::get_cuda_stream(res); + sample_filter_payload host_payload{filter.ref()}; + raft::copy(payload.data(), &host_payload, 1, stream); + raft::resource::sync_stream(res); + } + + void rebuild_filter(std::size_t num_blocks, cudaStream_t stream) + { + filter = cuco_filter_type(num_blocks, {}, {}, {}, stream); + sample_filter_payload host_payload{filter.ref()}; + raft::copy(payload.data(), &host_payload, 1, stream); + } + + void configure_or_validate_dataset_rows(raft::device_vector_view keys, + cudaStream_t stream) + { + if (keys.extent(0) == 0) { return; } + + auto inferred_dataset_rows = static_cast(keys.extent(0)); + if (!configured_dataset_rows.has_value()) { + configured_dataset_rows = inferred_dataset_rows; + auto required_num_blocks = compute_num_blocks_from_rates( + inferred_dataset_rows, expected_valid_rate, target_false_positive_rate); + auto target_num_blocks = std::max(filter.block_extent(), required_num_blocks); + if (target_num_blocks != filter.block_extent()) { rebuild_filter(target_num_blocks, stream); } + return; + } + + RAFT_EXPECTS( + inferred_dataset_rows == *configured_dataset_rows, + "keys.size() must match dataset_rows established by the first add/add_async call."); + } +}; + +bloom_filter::bloom_filter(raft::resources const& res, + float expected_valid_rate, + float target_false_positive_rate, + std::size_t num_blocks) + : impl_(std::make_unique( + res, std::max(1, num_blocks), expected_valid_rate, target_false_positive_rate)) +{ + RAFT_EXPECTS(expected_valid_rate > 0.0f && expected_valid_rate <= 1.0f, + "expected_valid_rate must be in (0, 1]."); + RAFT_EXPECTS(target_false_positive_rate > 0.0f && target_false_positive_rate < 1.0f, + "target_false_positive_rate must be in (0, 1)."); +} + +bloom_filter::~bloom_filter() = default; +bloom_filter::bloom_filter(bloom_filter&&) noexcept = default; +bloom_filter& bloom_filter::operator=(bloom_filter&&) noexcept = default; + +void bloom_filter::clear(raft::resources const& res) +{ + impl_->filter.clear(raft::resource::get_cuda_stream(res)); +} + +void bloom_filter::clear_async(raft::resources const& res) +{ + impl_->filter.clear_async(raft::resource::get_cuda_stream(res)); +} + +void bloom_filter::add(raft::resources const& res, + raft::device_vector_view keys) +{ + auto stream = raft::resource::get_cuda_stream(res); + impl_->configure_or_validate_dataset_rows(keys, stream); + impl_->filter.add(keys.data_handle(), keys.data_handle() + keys.extent(0), stream); +} + +void bloom_filter::add_async(raft::resources const& res, + raft::device_vector_view keys) +{ + auto stream = raft::resource::get_cuda_stream(res); + impl_->configure_or_validate_dataset_rows(keys, stream); + impl_->filter.add_async(keys.data_handle(), keys.data_handle() + keys.extent(0), stream); +} + +void bloom_filter::contains(raft::resources const& res, + raft::device_vector_view keys, + raft::device_vector_view output) const +{ + impl_->filter.contains(keys.data_handle(), + keys.data_handle() + keys.extent(0), + output.data_handle(), + raft::resource::get_cuda_stream(res)); +} + +void bloom_filter::contains_async(raft::resources const& res, + raft::device_vector_view keys, + raft::device_vector_view output) const +{ + impl_->filter.contains_async(keys.data_handle(), + keys.data_handle() + keys.extent(0), + output.data_handle(), + raft::resource::get_cuda_stream(res)); +} + +std::size_t bloom_filter::num_blocks() const noexcept { return impl_->filter.block_extent(); } + +void* bloom_filter::filter_data() const noexcept { return impl_->payload.data(); } + +} // namespace cuvs::core diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index a6704f892a..e7e4cc8356 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -10,6 +10,7 @@ #include "naive_knn.cuh" +#include #include #include #include @@ -34,6 +35,7 @@ #include +#include #include #include #include @@ -893,6 +895,78 @@ class AnnCagraFilterTest : public ::testing::TestWithParam { raft::update_host(distances_Cagra.data(), distances_dev.data(), queries_size, stream_); raft::update_host(indices_Cagra.data(), indices_dev.data(), queries_size, stream_); raft::resource::sync_stream(handle_); + + std::vector valid_ids_host; + valid_ids_host.reserve(ps.n_rows - test_cagra_sample_filter::offset); + for (std::uint32_t i = test_cagra_sample_filter::offset; + i < static_cast(ps.n_rows); + ++i) { + valid_ids_host.push_back(i); + } + + rmm::device_uvector valid_ids_device(valid_ids_host.size(), stream_); + raft::copy(valid_ids_device.data(), valid_ids_host.data(), valid_ids_host.size(), stream_); + + auto bloom_num_blocks = std::max(4096, static_cast(ps.n_rows)); + auto valid_ids_view = raft::make_device_vector_view( + valid_ids_device.data(), static_cast(valid_ids_device.size())); + auto candidate_fprs = std::vector{0.01f}; + if (ps.n_rows == 1000 && ps.dim == 8 && ps.k == 16 && + ps.metric == cuvs::distance::DistanceType::L2Expanded && + ps.algo == search_algo::SINGLE_CTA && !ps.compression.has_value()) { + // Keep this sweep narrow to avoid exploding test time while still validating the knob. + candidate_fprs = {0.25f, 0.05f, 0.01f}; + } + + std::vector bloom_recalls; + bloom_recalls.reserve(candidate_fprs.size()); + for (auto target_false_positive_rate : candidate_fprs) { + cuvs::core::bloom_filter global_bloom_filter( + handle_, 1.0f, target_false_positive_rate, bloom_num_blocks); + global_bloom_filter.add_async(handle_, valid_ids_view); + raft::resource::sync_stream(handle_); + + auto bloom_filter_obj = cuvs::neighbors::filtering::bloom_filter( + global_bloom_filter.filter_data(), + static_cast(test_cagra_sample_filter::offset) / static_cast(ps.n_rows)); + + cagra::search(handle_, + search_params, + index, + search_queries_view, + indices_out_view, + dists_out_view, + bloom_filter_obj); + + std::vector bloom_indices_host(queries_size); + std::vector bloom_distances_host(queries_size); + raft::update_host(bloom_indices_host.data(), indices_dev.data(), queries_size, stream_); + raft::update_host( + bloom_distances_host.data(), distances_dev.data(), queries_size, stream_); + raft::resource::sync_stream(handle_); + + auto [bloom_recall, bloom_match_count, bloom_total_count] = + calc_recall(indices_naive, + bloom_indices_host, + distances_naive, + bloom_distances_host, + ps.n_queries, + ps.k, + 0.003); + RAFT_LOG_INFO("Bloom filter recall = %f (%zu/%zu), target_false_positive_rate = %f", + bloom_recall, + bloom_match_count, + bloom_total_count, + target_false_positive_rate); + bloom_recalls.push_back(bloom_recall); + } + if (bloom_recalls.size() > 1) { + // As target_false_positive_rate decreases, recall should not regress in a meaningful way. + for (size_t i = 1; i < bloom_recalls.size(); ++i) { + EXPECT_GE(bloom_recalls[i] + 0.02, bloom_recalls[i - 1]); + } + EXPECT_GE(bloom_recalls.back() + 1e-3, bloom_recalls.front()); + } } // Test search results for nodes marked as filtered diff --git a/examples/cpp/CMakeLists.txt b/examples/cpp/CMakeLists.txt index 07573d40dd..5c0013eaf4 100644 --- a/examples/cpp/CMakeLists.txt +++ b/examples/cpp/CMakeLists.txt @@ -35,7 +35,6 @@ add_executable(BRUTE_FORCE_EXAMPLE src/brute_force_bitmap.cu) add_executable(CAGRA_EXAMPLE src/cagra_example.cu) add_executable(CAGRA_FILTER_UDF_EXAMPLE src/cagra_filter_udf_example.cu) add_executable(CAGRA_BLOOM_FILTER_EXAMPLE src/cagra_bloom_filter_example.cu) -add_executable(CAGRA_FILTER_BENCHMARK src/cagra_filter_benchmark.cu) add_executable(CAGRA_HNSW_ACE_EXAMPLE src/cagra_hnsw_ace_example.cu) add_executable(CAGRA_PERSISTENT_EXAMPLE src/cagra_persistent_example.cu) add_executable(DYNAMIC_BATCHING_EXAMPLE src/dynamic_batching_example.cu) @@ -55,8 +54,8 @@ target_link_libraries( target_link_libraries( CAGRA_BLOOM_FILTER_EXAMPLE PRIVATE cuvs::cuvs cuco::cuco $ ) -target_link_libraries( - CAGRA_FILTER_BENCHMARK PRIVATE cuvs::cuvs cuco::cuco $ +target_include_directories( + CAGRA_BLOOM_FILTER_EXAMPLE PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/../../cpp/include" ) target_link_libraries(CAGRA_HNSW_ACE_EXAMPLE PRIVATE cuvs::cuvs $) target_link_libraries( diff --git a/examples/cpp/plot_filter_benchmark.py b/examples/cpp/plot_filter_benchmark.py deleted file mode 100755 index 8840c5e9fc..0000000000 --- a/examples/cpp/plot_filter_benchmark.py +++ /dev/null @@ -1,457 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. -# SPDX-License-Identifier: Apache-2.0 -"""Plot CAGRA filter benchmark CSV (bitset vs bloom_filter line charts).""" - -from __future__ import annotations - -import argparse -from pathlib import Path - -import matplotlib.pyplot as plt -import numpy as np -import pandas as pd -import seaborn as sns - -FILTER_LABELS = { - "bitset": "Bitset", - "bloom_filter": "Bloom", -} - -FILTER_ORDER = ["Bitset", "Bloom"] -FILTER_COLORS = {"Bitset": "#0173b2", "Bloom": "#de8f05"} -FILTER_MARKERS = {"Bitset": "o", "Bloom": "s"} - -SEARCH_LABELS = { - 10000: "10k queries", - 25000: "25k queries", -} - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Generate line charts from cagra_filter_benchmark CSV output.", - ) - parser.add_argument( - "csv", - nargs="?", - default="filter_bench.csv", - help="Path to benchmark CSV (default: filter_bench.csv)", - ) - parser.add_argument( - "-o", - "--output-dir", - default="filter_bench_plots", - help="Directory for PNG output (default: filter_bench_plots)", - ) - parser.add_argument( - "--with-recall", - action="store_true", - help="Also plot recall charts when non-NaN values are present", - ) - return parser.parse_args() - - -def load_csv(path: Path) -> pd.DataFrame: - df = pd.read_csv(path) - df["recall"] = pd.to_numeric(df["recall"], errors="coerce") - df["filter_label"] = ( - df["filter_type"].map(FILTER_LABELS).fillna(df["filter_type"]) - ) - df["search_label"] = ( - df["search_n_rows"] - .map(SEARCH_LABELS) - .fillna(df["search_n_rows"].astype(str) + " queries") - ) - return df - - -def save_figure(fig: plt.Figure, path: Path) -> None: - fig.subplots_adjust(left=0.11, bottom=0.08, right=0.98, top=0.94) - fig.savefig(path, dpi=150, bbox_inches="tight") - plt.close(fig) - - -def plot_panel_bars( - ax: plt.Axes, - panel: pd.DataFrame, - x_col: str, - y_col: str, -) -> list: - """Grouped bars for two filters at each x value — visible even when latencies are close.""" - x_values = sorted(panel[x_col].unique()) - x_idx = np.arange(len(x_values)) - bar_width = 0.36 - handles = [] - - for series_idx, label in enumerate(FILTER_ORDER): - heights = [] - for x_val in x_values: - row = panel[ - (panel["filter_label"] == label) & (panel[x_col] == x_val) - ] - heights.append(row[y_col].iloc[0] if not row.empty else 0.0) - offset = (series_idx - 0.5) * bar_width - bars = ax.bar( - x_idx + offset, - heights, - width=bar_width, - label=label, - color=FILTER_COLORS[label], - edgecolor="black", - linewidth=0.6, - alpha=0.9, - ) - handles.append(bars[0]) - - ax.set_xticks(x_idx) - if x_col == "build_n_rows": - ax.set_xticklabels([f"{int(v):,}" for v in x_values]) - else: - ax.set_xticklabels( - [ - str(int(v)) if float(v).is_integer() else str(v) - for v in x_values - ] - ) - return handles - - -def plot_panel_lines( - ax: plt.Axes, - panel: pd.DataFrame, - x_col: str, - y_col: str, - log_x: bool = False, -) -> list: - """Draw bitset vs bloom as two explicit series (no seaborn hue/style mashup).""" - handles = [] - x_values = sorted(panel[x_col].unique()) - if log_x: - ax.set_xscale("log") - ax.set_xticks(x_values) - ax.set_xticklabels([f"{int(v):,}" for v in x_values]) - ax.minorticks_off() - - for label in FILTER_ORDER: - series = panel[panel["filter_label"] == label].sort_values(x_col) - if series.empty: - continue - (line,) = ax.plot( - series[x_col], - series[y_col], - label=label, - color=FILTER_COLORS[label], - marker=FILTER_MARKERS[label], - markersize=7, - linewidth=2.0, - markerfacecolor="white", - markeredgewidth=1.5, - ) - handles.append(line) - - return handles - - -def add_row_band_label( - fig: plt.Figure, axes: np.ndarray, row_idx: int, text: str -) -> None: - ax = axes[row_idx, 0] - pos = ax.get_position() - fig.text( - 0.02, - (pos.y0 + pos.y1) / 2, - text, - ha="left", - va="center", - rotation=90, - fontsize=10, - fontweight="bold", - ) - - -def plot_latency_grid_for_valid_pct( - df: pd.DataFrame, - out_dir: Path, - valid_pct: int, -) -> None: - """One PNG per valid_pct: 6×3 grid. - - Row bands (top to bottom): 10k queries × dims 128/512/1024, then 25k × dims. - Columns: k = 64 / 256 / 1024. - Each panel has only two lines (bitset vs bloom) over index size. - """ - sub = df[df["valid_pct"] == valid_pct].copy() - if sub.empty: - print(f"valid_pct={valid_pct}%: no rows, skipping") - return - - k_values = sorted(sub["k"].unique()) - col_values = sorted(sub["build_n_cols"].unique()) - search_values = sorted(sub["search_n_rows"].unique()) - - row_specs = [ - (search, dims) for search in search_values for dims in col_values - ] - n_rows = len(row_specs) - n_cols = len(k_values) - - fig, axes = plt.subplots( - n_rows, - n_cols, - figsize=(4.0 * n_cols + 0.8, 2.4 * n_rows), - squeeze=False, - ) - - legend_handles = None - for row_idx, (search_n_rows, build_n_cols) in enumerate(row_specs): - for col_idx, k in enumerate(k_values): - ax = axes[row_idx, col_idx] - panel = sub[ - (sub["build_n_cols"] == build_n_cols) - & (sub["k"] == k) - & (sub["search_n_rows"] == search_n_rows) - ] - if panel.empty: - ax.set_visible(False) - continue - - handles = plot_panel_bars( - ax, - panel, - x_col="build_n_rows", - y_col="avg_latency_per_query_ms", - ) - if legend_handles is None and handles: - legend_handles = handles - - ax.set_title(f"k={k}", fontsize=10) - ax.set_xlabel("index rows") - if col_idx == 0: - ax.set_ylabel(f"dims={build_n_cols}\nlatency / query (ms)") - else: - ax.set_ylabel("") - - for search_idx, search_n_rows in enumerate(search_values): - band_row = search_idx * len(col_values) - add_row_band_label( - fig, - axes, - band_row, - SEARCH_LABELS.get(search_n_rows, f"{search_n_rows:,} queries"), - ) - - if legend_handles: - fig.legend( - legend_handles, - FILTER_ORDER, - loc="lower center", - bbox_to_anchor=(0.5, -0.02), - ncol=2, - frameon=False, - fontsize=11, - ) - - fig.suptitle( - f"Per-query search latency — {valid_pct}% rows valid", - fontsize=13, - y=1.01, - ) - path = out_dir / f"latency_valid_{valid_pct}pct.png" - save_figure(fig, path) - print(f"wrote {path}") - - -def plot_all_valid_pct_grids(df: pd.DataFrame, out_dir: Path) -> None: - for valid_pct in sorted(df["valid_pct"].unique()): - plot_latency_grid_for_valid_pct(df, out_dir, int(valid_pct)) - - -def plot_valid_pct_overview( - df: pd.DataFrame, - out_dir: Path, - build_n_rows: int, - search_n_rows: int, -) -> None: - """valid_pct on x-axis at one build/search point; 3×3 grid (dims × k).""" - sub = df[ - (df["build_n_rows"] == build_n_rows) - & (df["search_n_rows"] == search_n_rows) - ].copy() - if sub.empty: - print( - "valid_pct overview: no rows for selected build/search slice, skipping" - ) - return - - k_values = sorted(sub["k"].unique()) - col_values = sorted(sub["build_n_cols"].unique()) - - fig, axes = plt.subplots( - len(col_values), - len(k_values), - figsize=(4.0 * len(k_values), 2.8 * len(col_values)), - squeeze=False, - ) - - legend_handles = None - for row_idx, build_n_cols in enumerate(col_values): - for col_idx, k in enumerate(k_values): - ax = axes[row_idx, col_idx] - panel = sub[ - (sub["build_n_cols"] == build_n_cols) & (sub["k"] == k) - ] - if panel.empty: - ax.set_visible(False) - continue - - handles = plot_panel_bars( - ax, - panel, - x_col="valid_pct", - y_col="avg_latency_per_query_ms", - ) - if legend_handles is None and handles: - legend_handles = handles - - ax.set_title(f"k={k}", fontsize=10) - ax.set_xlabel("valid rows (%)") - if col_idx == 0: - ax.set_ylabel(f"dims={build_n_cols}\nlatency / query (ms)") - else: - ax.set_ylabel("") - - if legend_handles: - fig.legend( - legend_handles, - FILTER_ORDER, - loc="lower center", - bbox_to_anchor=(0.5, -0.02), - ncol=2, - frameon=False, - fontsize=11, - ) - - fig.suptitle( - f"Latency vs filter selectivity — " - f"build={build_n_rows:,} rows search={search_n_rows:,} queries", - fontsize=12, - y=1.01, - ) - path = out_dir / "overview_valid_pct_sweep.png" - save_figure(fig, path) - print(f"wrote {path}") - - -def plot_recall_grid_for_valid_pct( - df: pd.DataFrame, - out_dir: Path, - valid_pct: int, -) -> None: - sub = df[(df["valid_pct"] == valid_pct) & df["recall"].notna()].copy() - if sub.empty: - return - - k_values = sorted(sub["k"].unique()) - col_values = sorted(sub["build_n_cols"].unique()) - search_values = sorted(sub["search_n_rows"].unique()) - row_specs = [ - (search, dims) for search in search_values for dims in col_values - ] - - fig, axes = plt.subplots( - len(row_specs), - len(k_values), - figsize=(4.0 * len(k_values) + 0.8, 2.4 * len(row_specs)), - squeeze=False, - ) - - legend_handles = None - for row_idx, (search_n_rows, build_n_cols) in enumerate(row_specs): - for col_idx, k in enumerate(k_values): - ax = axes[row_idx, col_idx] - panel = sub[ - (sub["build_n_cols"] == build_n_cols) - & (sub["k"] == k) - & (sub["search_n_rows"] == search_n_rows) - ] - if panel.empty: - ax.set_visible(False) - continue - - handles = plot_panel_bars( - ax, - panel, - x_col="build_n_rows", - y_col="recall", - ) - if legend_handles is None and handles: - legend_handles = handles - - ax.set_ylim(0.0, 1.0) - ax.set_title(f"k={k}", fontsize=10) - ax.set_xlabel("index rows") - if col_idx == 0: - ax.set_ylabel(f"dims={build_n_cols}\nrecall@k") - else: - ax.set_ylabel("") - - for search_idx, search_n_rows in enumerate(search_values): - band_row = search_idx * len(col_values) - add_row_band_label( - fig, - axes, - band_row, - SEARCH_LABELS.get(search_n_rows, f"{search_n_rows:,} queries"), - ) - - if legend_handles: - fig.legend( - legend_handles, - FILTER_ORDER, - loc="lower center", - bbox_to_anchor=(0.5, -0.02), - ncol=2, - frameon=False, - fontsize=11, - ) - - fig.suptitle(f"Recall@k — {valid_pct}% rows valid", fontsize=13, y=1.01) - path = out_dir / f"recall_valid_{valid_pct}pct.png" - save_figure(fig, path) - print(f"wrote {path}") - - -def main() -> None: - args = parse_args() - csv_path = Path(args.csv) - out_dir = Path(args.output_dir) - out_dir.mkdir(parents=True, exist_ok=True) - - if not csv_path.is_file(): - raise SystemExit(f"CSV not found: {csv_path}") - - df = load_csv(csv_path) - print(f"loaded {len(df)} rows from {csv_path}") - - sns.set_theme(style="whitegrid", context="notebook") - - plot_all_valid_pct_grids(df, out_dir) - - default_build = int(df["build_n_rows"].min()) - default_search = int(df["search_n_rows"].min()) - plot_valid_pct_overview(df, out_dir, default_build, default_search) - - if args.with_recall and df["recall"].notna().any(): - for valid_pct in sorted(df["valid_pct"].unique()): - plot_recall_grid_for_valid_pct(df, out_dir, int(valid_pct)) - elif args.with_recall: - print( - "recall charts skipped: CSV has no recall values (run benchmark with --ground-truth)" - ) - else: - print("recall charts skipped (default; pass --with-recall to enable)") - - -if __name__ == "__main__": - main() diff --git a/examples/cpp/src/cagra_bloom_filter_example.cu b/examples/cpp/src/cagra_bloom_filter_example.cu index 7d16989bbb..2d2f6ce9e9 100644 --- a/examples/cpp/src/cagra_bloom_filter_example.cu +++ b/examples/cpp/src/cagra_bloom_filter_example.cu @@ -3,7 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -#include +#include #include #include @@ -22,20 +22,13 @@ namespace { -constexpr int64_t n_rows = 4096; -constexpr int64_t n_dim = 32; -constexpr int64_t n_queries = 4; -constexpr int64_t k = 8; -constexpr int sub_filters = 256; +constexpr int64_t n_rows = 4096; +constexpr int64_t n_dim = 32; +constexpr int64_t n_queries = 4; +constexpr int64_t k = 8; +constexpr std::size_t sub_filters = 256; -using key_type = std::uint32_t; -using filter_type = cuco::bloom_filter; -using ref_type = filter_type::ref_type<>; - -// Layout must match cuvs::neighbors::detail::bloom_filter_data_t in the JIT fragment. -struct bloom_payload { - ref_type filter; -}; +using key_type = std::uint32_t; // Global index filter: even row ids are valid candidates (same rule for every query). bool is_valid_row(key_type source_id) { return (source_id % 2) == 0; } @@ -92,20 +85,15 @@ int main() rmm::device_uvector valid_ids_device(valid_ids_host.size(), stream); raft::copy(valid_ids_device.data(), valid_ids_host.data(), valid_ids_host.size(), stream); - filter_type allowed_rows{sub_filters}; - allowed_rows.add_async( - valid_ids_device.data(), valid_ids_device.data() + valid_ids_device.size(), stream); + auto valid_ids_view = raft::make_device_vector_view( + valid_ids_device.data(), static_cast(valid_ids_device.size())); + cuvs::core::bloom_filter allowed_rows(res, 1.0f, 0.01f, sub_filters); + allowed_rows.add_async(res, valid_ids_view); raft::resource::sync_stream(res); std::cout << "Inserted " << valid_ids_host.size() << " valid row ids into global bloom filter via bulk add_async" << std::endl; - // Copy the owning filter's device ref into a payload the JIT fragment can probe. - auto payload_device = raft::make_device_vector(res, 1); - bloom_payload host_payload{allowed_rows.ref()}; - raft::copy(payload_device.data_handle(), &host_payload, 1, stream); - raft::resource::sync_stream(res); - auto neighbors = raft::make_device_matrix(res, n_queries, k); auto distances = raft::make_device_matrix(res, n_queries, k); @@ -116,7 +104,7 @@ int main() search_params.thread_block_size = 256; // ~50% of rows are rejected by the global even-id predicate. - auto filter = cuvs::neighbors::filtering::bloom_filter(payload_device.data_handle(), 0.5f); + auto filter = cuvs::neighbors::filtering::bloom_filter(allowed_rows.filter_data(), 0.5f); cuvs::neighbors::cagra::search(res, search_params, @@ -138,10 +126,11 @@ int main() rmm::device_uvector neighbor_ids_device(host_neighbors.size(), stream); rmm::device_uvector bloom_hits_device(host_neighbors.size(), stream); raft::copy(neighbor_ids_device.data(), host_neighbors.data(), host_neighbors.size(), stream); - allowed_rows.contains_async(neighbor_ids_device.data(), - neighbor_ids_device.data() + neighbor_ids_device.size(), - bloom_hits_device.data(), - stream); + auto neighbor_ids_view = raft::make_device_vector_view( + neighbor_ids_device.data(), static_cast(neighbor_ids_device.size())); + auto bloom_hits_view = raft::make_device_vector_view( + bloom_hits_device.data(), static_cast(bloom_hits_device.size())); + allowed_rows.contains_async(res, neighbor_ids_view, bloom_hits_view); raft::resource::sync_stream(res); std::vector bloom_hits_host(bloom_hits_device.size()); diff --git a/examples/cpp/src/cagra_filter_benchmark.cu b/examples/cpp/src/cagra_filter_benchmark.cu deleted file mode 100644 index 2c72c22ddd..0000000000 --- a/examples/cpp/src/cagra_filter_benchmark.cu +++ /dev/null @@ -1,537 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#include -#include -#include -#include - -#include -#include -#include -#include -#include -#include - -#include -#include - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace { - -constexpr int k_warmup_runs = 1; -constexpr int k_timed_runs = 3; - -using key_type = std::uint32_t; -using filter_type = cuco::bloom_filter; -using ref_type = filter_type::ref_type<>; - -struct bloom_payload { - ref_type filter; -}; - -constexpr std::array k_build_rows{100'000, 500'000, 1'000'000}; -constexpr std::array k_build_cols{128, 512, 1024}; -constexpr std::array k_search_rows{10'000, 25'000}; -constexpr std::array k_search_cols{128, 512, 1024}; -constexpr std::array k_valid_pcts{10, 50, 90}; -constexpr std::array k_values{64, 256, 1024}; - -bool is_valid_row(key_type row_id, int valid_pct) -{ - // Deterministic ~valid_pct% membership independent of dataset size. - return (static_cast(row_id) * 2654435761ULL) % 100ULL < - static_cast(valid_pct); -} - -std::size_t bloom_num_blocks(std::size_t num_valid_rows) -{ - // Scale filter size with the number of inserted keys; keep a reasonable minimum. - std::size_t blocks = std::max(256, num_valid_rows / 8); - blocks = std::min(blocks, static_cast(1 << 20)); - return blocks; -} - -double compute_recall(std::vector const& expected, - std::vector const& actual, - int64_t n_queries, - int64_t k, - int64_t n_rows) -{ - std::size_t match_count = 0; - std::size_t total_count = static_cast(n_queries) * static_cast(k); - for (int64_t q = 0; q < n_queries; ++q) { - for (int64_t ki = 0; ki < k; ++ki) { - auto const act = actual[static_cast(q * k + ki)]; - if (act >= static_cast(n_rows)) { continue; } - for (int64_t kj = 0; kj < k; ++kj) { - if (expected[static_cast(q * k + kj)] == act) { - ++match_count; - break; - } - } - } - } - return total_count == 0 ? 0.0 - : static_cast(match_count) / static_cast(total_count); -} - -std::vector copy_neighbors_to_host(raft::device_resources const& res, - raft::device_matrix_view neighbors) -{ - std::vector host(neighbors.size()); - auto stream = raft::resource::get_cuda_stream(res); - raft::copy(host.data(), neighbors.data_handle(), host.size(), stream); - raft::resource::sync_stream(res); - return host; -} - -struct filter_assets { - cuvs::core::bitset removed_bitset; - cuvs::neighbors::filtering::bitset_filter bitset_filter; - filter_type bloom; - rmm::device_uvector bloom_payload; - cuvs::neighbors::filtering::bloom_filter bloom_filter; - float filtering_rate{0.0f}; -}; - -filter_assets make_filters(raft::device_resources const& res, - int64_t n_rows, - int valid_pct, - rmm::cuda_stream_view stream) -{ - std::vector valid_ids_host; - std::vector removed_ids_host; - valid_ids_host.reserve(static_cast(n_rows)); - removed_ids_host.reserve(static_cast(n_rows)); - - for (int64_t i = 0; i < n_rows; ++i) { - auto const row = static_cast(i); - if (is_valid_row(row, valid_pct)) { - valid_ids_host.push_back(row); - } else { - removed_ids_host.push_back(i); - } - } - - auto removed_ids = - raft::make_device_vector(res, static_cast(removed_ids_host.size())); - if (!removed_ids_host.empty()) { - raft::copy(removed_ids.data_handle(), removed_ids_host.data(), removed_ids_host.size(), stream); - } - - auto removed_bitset = cuvs::core::bitset(res, removed_ids.view(), n_rows); - auto bitset_filter = - cuvs::neighbors::filtering::bitset_filter(removed_bitset.view()); - auto bloom = filter_type{bloom_num_blocks(valid_ids_host.size()), {}, {}, {}, stream}; - auto payload_device = rmm::device_uvector{1, stream}; - float const filtering_rate = static_cast(100 - valid_pct) / 100.0f; - - if (!valid_ids_host.empty()) { - rmm::device_uvector valid_ids_device(valid_ids_host.size(), stream); - raft::copy(valid_ids_device.data(), valid_ids_host.data(), valid_ids_host.size(), stream); - bloom.add_async( - valid_ids_device.data(), valid_ids_device.data() + valid_ids_device.size(), stream); - } - - bloom_payload host_payload{bloom.ref()}; - raft::copy(payload_device.data(), &host_payload, 1, stream); - auto bloom_filter_obj = - cuvs::neighbors::filtering::bloom_filter(payload_device.data(), filtering_rate); - - raft::resource::sync_stream(res); - return filter_assets{std::move(removed_bitset), - std::move(bitset_filter), - std::move(bloom), - std::move(payload_device), - std::move(bloom_filter_obj), - filtering_rate}; -} - -struct benchmark_case { - int64_t build_n_rows; - int64_t build_n_cols; - int64_t search_n_rows; - int64_t search_n_cols; - int valid_pct; - int64_t k; -}; - -struct csv_row { - benchmark_case config; - std::string filter_name; - double build_time_ms; - double avg_search_latency_ms; - double avg_latency_per_query_ms; - double recall; -}; - -void write_csv_header(std::ostream& os) -{ - os << "build_n_rows,build_n_cols,search_n_rows,search_n_cols,valid_pct,filter_type," - "build_time_ms,avg_search_latency_ms,avg_latency_per_query_ms,recall,k,warmup_runs," - "timed_runs\n"; -} - -void write_csv_row(std::ostream& os, csv_row const& row) -{ - os << row.config.build_n_rows << ',' << row.config.build_n_cols << ',' << row.config.search_n_rows - << ',' << row.config.search_n_cols << ',' << row.config.valid_pct << ',' << row.filter_name - << ',' << row.build_time_ms << ',' << row.avg_search_latency_ms << ',' - << row.avg_latency_per_query_ms << ',' << row.recall << ',' << row.config.k << ',' - << k_warmup_runs << ',' << k_timed_runs << '\n'; -} - -template -double time_cuda_ms(raft::device_resources const& res, int runs, Fn&& fn) -{ - auto stream = raft::resource::get_cuda_stream(res); - cudaEvent_t start{}; - cudaEvent_t stop{}; - RAFT_CUDA_TRY(cudaEventCreate(&start)); - RAFT_CUDA_TRY(cudaEventCreate(&stop)); - - RAFT_CUDA_TRY(cudaEventRecord(start, stream)); - for (int i = 0; i < runs; ++i) { - fn(); - } - RAFT_CUDA_TRY(cudaEventRecord(stop, stream)); - RAFT_CUDA_TRY(cudaEventSynchronize(stop)); - - float elapsed_ms = 0.0f; - RAFT_CUDA_TRY(cudaEventElapsedTime(&elapsed_ms, start, stop)); - RAFT_CUDA_TRY(cudaEventDestroy(start)); - RAFT_CUDA_TRY(cudaEventDestroy(stop)); - return static_cast(elapsed_ms) / static_cast(runs); -} - -void append_cases(std::vector& cases, - std::vector const& build_rows, - std::vector const& build_cols, - std::vector const& search_rows, - std::vector const& search_cols, - std::vector const& valid_pcts, - std::vector const& k_sweep) -{ - for (auto build_n_rows : build_rows) { - for (auto build_n_cols : build_cols) { - for (auto search_n_rows : search_rows) { - for (auto search_n_cols : search_cols) { - if (search_n_cols != build_n_cols) { continue; } - for (auto valid_pct : valid_pcts) { - for (auto k : k_sweep) { - cases.push_back(benchmark_case{ - build_n_rows, build_n_cols, search_n_rows, search_n_cols, valid_pct, k}); - } - } - } - } - } - } -} - -std::vector make_cases(bool quick) -{ - std::vector cases; - if (quick) { - append_cases(cases, {100'000}, {128}, {10'000}, {128}, {1, 50}, {64}); - } else { - append_cases(cases, - {k_build_rows.begin(), k_build_rows.end()}, - {k_build_cols.begin(), k_build_cols.end()}, - {k_search_rows.begin(), k_search_rows.end()}, - {k_search_cols.begin(), k_search_cols.end()}, - {k_valid_pcts.begin(), k_valid_pcts.end()}, - {k_values.begin(), k_values.end()}); - } - return cases; -} - -constexpr std::size_t k_max_bf_bytes = 20ULL << 30; // skip BF recall above this estimate -constexpr std::size_t k_max_bf_chunk_bytes = 2ULL << 30; // cap each BF chunk when computing recall - -std::size_t estimate_bf_distance_matrix_bytes(int64_t n_queries, int64_t n_dataset) -{ - return static_cast(n_queries) * static_cast(n_dataset) * sizeof(float); -} - -bool should_compute_bf_recall(int64_t n_queries, int64_t n_dataset) -{ - return estimate_bf_distance_matrix_bytes(n_queries, n_dataset) <= k_max_bf_bytes; -} - -int64_t choose_gt_chunk_queries(int64_t n_dataset) -{ - int64_t chunk = 256; - while (chunk > 1 && estimate_bf_distance_matrix_bytes(chunk, n_dataset) > k_max_bf_chunk_bytes) { - chunk /= 2; - } - return chunk; -} - -std::vector brute_force_ground_truth( - raft::device_resources const& res, - cuvs::neighbors::brute_force::index& bf_index, - cuvs::neighbors::brute_force::search_params const& bf_search_params, - raft::device_matrix_view queries, - cuvs::neighbors::filtering::bitset_filter const& bitset_filter, - int64_t k, - int64_t gt_chunk_queries) -{ - int64_t const n_queries = queries.extent(0); - std::vector gt_host(static_cast(n_queries * k)); - auto stream = raft::resource::get_cuda_stream(res); - - for (int64_t query_offset = 0; query_offset < n_queries; query_offset += gt_chunk_queries) { - int64_t const chunk_queries = std::min(gt_chunk_queries, n_queries - query_offset); - auto query_chunk = raft::make_device_matrix_view( - queries.data_handle() + query_offset * queries.extent(1), chunk_queries, queries.extent(1)); - auto gt_neighbors = raft::make_device_matrix(res, chunk_queries, k); - auto gt_distances = raft::make_device_matrix(res, chunk_queries, k); - - cuvs::neighbors::brute_force::search(res, - bf_search_params, - bf_index, - raft::make_const_mdspan(query_chunk), - gt_neighbors.view(), - gt_distances.view(), - bitset_filter); - raft::resource::sync_stream(res); - - std::vector chunk_host(static_cast(chunk_queries * k)); - raft::copy(chunk_host.data(), gt_neighbors.data_handle(), chunk_host.size(), stream); - raft::resource::sync_stream(res); - - for (int64_t q = 0; q < chunk_queries; ++q) { - for (int64_t ki = 0; ki < k; ++ki) { - gt_host[static_cast((query_offset + q) * k + ki)] = - static_cast(chunk_host[static_cast(q * k + ki)]); - } - } - } - - return gt_host; -} - -} // namespace - -int main(int argc, char** argv) -{ - std::string output_path = "cagra_filter_benchmark_results.csv"; - bool quick = false; - bool compute_ground_truth = false; - for (int i = 1; i < argc; ++i) { - std::string arg = argv[i]; - if (arg == "--quick") { - quick = true; - } else if (arg == "--ground-truth") { - compute_ground_truth = true; - } else if (arg == "--skip-ground-truth") { - compute_ground_truth = false; - } else if (arg == "--output" && i + 1 < argc) { - output_path = argv[++i]; - } else if (arg == "--help" || arg == "-h") { - std::cout << "Usage: " << argv[0] - << " [--quick] [--ground-truth] [--skip-ground-truth] [--output path.csv]\n" - << "\n" - << "Brute-force recall is skipped by default. Pass --ground-truth to compute it.\n"; - return 0; - } else { - output_path = arg; - } - } - - auto cases = make_cases(quick); - std::cout << "Running " << cases.size() << " benchmark configurations" - << (quick ? " (quick mode)" : "") - << (compute_ground_truth ? " (with ground-truth recall)" : " (ground-truth skipped)") - << std::endl; - - std::ofstream csv(output_path); - if (!csv) { - std::cerr << "Failed to open output file: " << output_path << std::endl; - return 1; - } - write_csv_header(csv); - - raft::device_resources res; - auto stream = raft::resource::get_cuda_stream(res); - - // Large enough for the biggest benchmark configuration (1M x 1024 dataset + index overhead). - rmm::mr::pool_memory_resource pool_mr(rmm::mr::get_current_device_resource_ref(), 16ULL << 30); - rmm::mr::set_current_device_resource(pool_mr); - - int64_t prev_build_rows = -1; - int64_t prev_build_cols = -1; - double last_build_time_ms = 0.0; - - std::optional> index; - std::optional> dataset; - std::optional> bf_index; - - cuvs::neighbors::cagra::index_params index_params; - index_params.metric = cuvs::distance::DistanceType::L2Expanded; - index_params.graph_degree = 32; - index_params.intermediate_graph_degree = 64; - index_params.graph_build_params = cuvs::neighbors::cagra::graph_build_params::nn_descent_params( - index_params.intermediate_graph_degree); - - cuvs::neighbors::cagra::search_params search_params; - search_params.algo = cuvs::neighbors::cagra::search_algo::MULTI_CTA; - search_params.itopk_size = 128; - search_params.thread_block_size = 256; - - cuvs::neighbors::brute_force::index_params bf_index_params; - cuvs::neighbors::brute_force::search_params bf_search_params; - - for (std::size_t case_idx = 0; case_idx < cases.size(); ++case_idx) { - auto const& cfg = cases[case_idx]; - - if (cfg.build_n_rows != prev_build_rows || cfg.build_n_cols != prev_build_cols) { - std::cout << "Building CAGRA index: n_rows=" << cfg.build_n_rows - << " n_cols=" << cfg.build_n_cols << std::endl; - - dataset.emplace( - raft::make_device_matrix(res, cfg.build_n_rows, cfg.build_n_cols)); - raft::random::RngState rng( - static_cast(cfg.build_n_rows * 17 + cfg.build_n_cols)); - raft::random::uniform(res, rng, dataset->data_handle(), dataset->size(), -1.0f, 1.0f); - - auto build_start = std::chrono::steady_clock::now(); - index.emplace( - cuvs::neighbors::cagra::build(res, index_params, raft::make_const_mdspan(dataset->view()))); - bf_index.reset(); - raft::resource::sync_stream(res); - auto build_end = std::chrono::steady_clock::now(); - last_build_time_ms = - std::chrono::duration(build_end - build_start).count(); - - prev_build_rows = cfg.build_n_rows; - prev_build_cols = cfg.build_n_cols; - } - - int64_t const total_queries = cfg.search_n_rows; - search_params.max_queries = total_queries; - search_params.itopk_size = static_cast(cfg.k); - - std::cout << "Case " << (case_idx + 1) << '/' << cases.size() - << ": build_n_rows=" << cfg.build_n_rows << " search_n_rows=" << total_queries - << " k=" << cfg.k << " valid_pct=" << cfg.valid_pct << '%' << std::endl; - - try { - auto queries = - raft::make_device_matrix(res, total_queries, cfg.search_n_cols); - raft::random::RngState query_rng(static_cast( - cfg.build_n_rows * 31 + cfg.search_n_rows * 17 + cfg.search_n_cols + cfg.valid_pct)); - raft::random::uniform(res, query_rng, queries.data_handle(), queries.size(), -1.0f, 1.0f); - - auto neighbors = raft::make_device_matrix(res, total_queries, cfg.k); - auto distances = raft::make_device_matrix(res, total_queries, cfg.k); - - auto filters = make_filters(res, cfg.build_n_rows, cfg.valid_pct, stream); - - bool const run_bf_recall = - compute_ground_truth && should_compute_bf_recall(total_queries, cfg.build_n_rows); - std::optional> gt_host; - if (run_bf_recall) { - if (!bf_index.has_value()) { - bf_index.emplace(cuvs::neighbors::brute_force::build( - res, bf_index_params, raft::make_const_mdspan(dataset->view()))); - } - gt_host.emplace(brute_force_ground_truth(res, - *bf_index, - bf_search_params, - queries.view(), - filters.bitset_filter, - cfg.k, - choose_gt_chunk_queries(cfg.build_n_rows))); - } else if (compute_ground_truth) { - auto const est_gib = - static_cast(estimate_bf_distance_matrix_bytes(total_queries, cfg.build_n_rows)) / - static_cast(1ULL << 30); - std::cout << " skipping brute-force recall (estimated " << est_gib - << " GiB distance matrix > 20 GiB limit)" << std::endl; - } - - auto run_cagra_search = [&](cuvs::neighbors::filtering::base_filter const& filter) { - cuvs::neighbors::cagra::search(res, - search_params, - *index, - raft::make_const_mdspan(queries.view()), - neighbors.view(), - distances.view(), - filter); - raft::resource::sync_stream(res); - }; - - struct filter_run { - std::string name; - cuvs::neighbors::filtering::base_filter const* filter; - }; - std::vector filter_runs{ - {"bitset", &filters.bitset_filter}, - {"bloom_filter", &filters.bloom_filter}, - }; - - for (auto const& fr : filter_runs) { - for (int w = 0; w < k_warmup_runs; ++w) { - run_cagra_search(*fr.filter); - } - - double const avg_search_ms = - time_cuda_ms(res, k_timed_runs, [&] { run_cagra_search(*fr.filter); }); - double const avg_per_query_ms = avg_search_ms / static_cast(total_queries); - - double const recall = [&]() { - if (!gt_host.has_value()) { return std::numeric_limits::quiet_NaN(); } - auto result_host = copy_neighbors_to_host(res, neighbors.view()); - return compute_recall(*gt_host, result_host, total_queries, cfg.k, cfg.build_n_rows); - }(); - - write_csv_row( - csv, csv_row{cfg, fr.name, last_build_time_ms, avg_search_ms, avg_per_query_ms, recall}); - csv.flush(); - - std::cout << " " << fr.name << ": search_ms=" << avg_search_ms - << " per_query_ms=" << avg_per_query_ms << " recall="; - if (gt_host.has_value()) { - std::cout << recall; - } else { - std::cout << "n/a"; - } - std::cout << std::endl; - } - } catch (std::exception const& ex) { - std::cerr << " case failed: " << ex.what() << std::endl; - for (auto const* filter_name : {"bitset", "bloom_filter"}) { - write_csv_row(csv, - csv_row{cfg, - filter_name, - last_build_time_ms, - std::numeric_limits::quiet_NaN(), - std::numeric_limits::quiet_NaN(), - std::numeric_limits::quiet_NaN()}); - csv.flush(); - } - } - } - - std::cout << "Wrote results to " << output_path << std::endl; - return 0; -} From 56d959a18838315a6e79f9ffd932891cd38cf07e Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Wed, 8 Jul 2026 21:47:41 +0000 Subject: [PATCH 04/11] copyright --- cpp/include/cuvs/core/bloom_filter.hpp | 2 +- cpp/include/cuvs/detail/jit_lto/common_fragments.hpp | 2 +- cpp/include/cuvs/neighbors/common.hpp | 2 +- cpp/src/core/bloom_filter.cu | 2 +- cpp/src/neighbors/cagra.cuh | 2 +- cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp | 2 +- cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in | 2 +- cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in | 2 +- .../detail/cagra/search_single_cta_kernel_launcher_jit.cuh | 2 +- cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp | 2 +- examples/cpp/src/cagra_bloom_filter_example.cu | 2 +- 11 files changed, 11 insertions(+), 11 deletions(-) diff --git a/cpp/include/cuvs/core/bloom_filter.hpp b/cpp/include/cuvs/core/bloom_filter.hpp index e9d7de56df..dee370f5da 100644 --- a/cpp/include/cuvs/core/bloom_filter.hpp +++ b/cpp/include/cuvs/core/bloom_filter.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp index 3ae94a4b13..c1a73687c2 100644 --- a/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp +++ b/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 403091f3fb..0ad32b853b 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/core/bloom_filter.cu b/cpp/src/core/bloom_filter.cu index cb493c7fd1..00972ac4f4 100644 --- a/cpp/src/core/bloom_filter.cu +++ b/cpp/src/core/bloom_filter.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index af8331a9f1..b8900cfe11 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -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 */ diff --git a/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp b/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp index aa927215ef..39c59cb6b7 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp +++ b/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * reserved. SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in b/cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in index e60af91046..1f097b5608 100644 --- a/cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in +++ b/cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in @@ -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 */ diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in b/cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in index c9eac33d44..6f78e61976 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in @@ -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 */ diff --git a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel_launcher_jit.cuh b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel_launcher_jit.cuh index 28d660e2e1..cb2149af60 100644 --- a/cpp/src/neighbors/detail/cagra/search_single_cta_kernel_launcher_jit.cuh +++ b/cpp/src/neighbors/detail/cagra/search_single_cta_kernel_launcher_jit.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp b/cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp index dd118757e4..b01aa6125d 100644 --- a/cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp +++ b/cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ diff --git a/examples/cpp/src/cagra_bloom_filter_example.cu b/examples/cpp/src/cagra_bloom_filter_example.cu index 2d2f6ce9e9..eff960f930 100644 --- a/examples/cpp/src/cagra_bloom_filter_example.cu +++ b/examples/cpp/src/cagra_bloom_filter_example.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ From 93db3002ded7c83a046fdf9a4dc70143a7f54477 Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Thu, 9 Jul 2026 00:07:19 +0000 Subject: [PATCH 05/11] fix build --- cpp/tests/neighbors/ann_cagra.cuh | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index 09dcb3068c..bf22006c64 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -998,7 +998,7 @@ class AnnCagraFilterTest : public ::testing::TestWithParam { bloom_distances_host.data(), distances_dev.data(), queries_size, stream_); raft::resource::sync_stream(handle_); - auto [bloom_recall, bloom_match_count, bloom_total_count] = + auto [bloom_recall, bloom_index_recall, bloom_match_count, bloom_total_count] = calc_recall(indices_naive, bloom_indices_host, distances_naive, @@ -1011,6 +1011,7 @@ class AnnCagraFilterTest : public ::testing::TestWithParam { bloom_match_count, bloom_total_count, target_false_positive_rate); + RAFT_LOG_INFO("Bloom filter index recall = %f", bloom_index_recall); bloom_recalls.push_back(bloom_recall); } if (bloom_recalls.size() > 1) { From 4e6aa9c40990a5b751f6ace0c92722b3d89f38f0 Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Thu, 9 Jul 2026 00:30:37 +0000 Subject: [PATCH 06/11] naming --- cpp/include/cuvs/core/bloom_filter.hpp | 8 ++++---- cpp/src/core/bloom_filter.cu | 25 ++++++++++++------------- 2 files changed, 16 insertions(+), 17 deletions(-) diff --git a/cpp/include/cuvs/core/bloom_filter.hpp b/cpp/include/cuvs/core/bloom_filter.hpp index dee370f5da..ca5827eca4 100644 --- a/cpp/include/cuvs/core/bloom_filter.hpp +++ b/cpp/include/cuvs/core/bloom_filter.hpp @@ -34,11 +34,11 @@ class CUVS_EXPORT bloom_filter { * may grow above that floor to satisfy the requested false-positive rate. * * The primary tuning knobs are: - * - @p expected_valid_rate: expected fraction of dataset rows that will be inserted as valid ids. + * - @p filtering_rate: expected fraction of dataset rows that will be inserted as valid ids. * - @p target_false_positive_rate: desired Bloom filter false-positive probability. * * Sizing math used internally: - * - `expected_insertions = ceil(dataset_rows * expected_valid_rate)` + * - `expected_insertions = ceil(dataset_rows * filtering_rate)` * - `required_bits = -expected_insertions * ln(target_false_positive_rate) / (ln(2)^2)` * - `required_blocks = ceil(required_bits / 256)` (default cuco policy uses 256-bit blocks) * - `final_blocks = max(num_blocks, required_blocks)` @@ -46,11 +46,11 @@ class CUVS_EXPORT bloom_filter { * Practical knob behavior: * - Lower @p target_false_positive_rate -> larger filter, fewer false positives, typically higher * filtered-search recall. - * - Higher @p expected_valid_rate -> larger filter for the same target false-positive rate. + * - Higher @p filtering_rate -> larger filter for the same target false-positive rate. * - @p num_blocks is an expert floor; keep default unless you need a hard minimum memory budget. */ bloom_filter(raft::resources const& res, - float expected_valid_rate = 1.0f, + float filtering_rate = 1.0f, float target_false_positive_rate = 0.01f, std::size_t num_blocks = 256); ~bloom_filter(); diff --git a/cpp/src/core/bloom_filter.cu b/cpp/src/core/bloom_filter.cu index 00972ac4f4..3554578d7e 100644 --- a/cpp/src/core/bloom_filter.cu +++ b/cpp/src/core/bloom_filter.cu @@ -22,13 +22,13 @@ namespace cuvs::core { std::size_t compute_num_blocks_from_rates(std::size_t dataset_rows, - float expected_valid_rate, + float filtering_rate, float target_false_positive_rate) { RAFT_EXPECTS(dataset_rows > 0, "dataset_rows must be greater than zero when deriving bloom size."); - RAFT_EXPECTS(expected_valid_rate > 0.0f && expected_valid_rate <= 1.0f, - "expected_valid_rate must be in (0, 1]."); + RAFT_EXPECTS(filtering_rate > 0.0f && filtering_rate <= 1.0f, + "filtering_rate must be in (0, 1]."); RAFT_EXPECTS(target_false_positive_rate > 0.0f && target_false_positive_rate < 1.0f, "target_false_positive_rate must be in (0, 1)."); @@ -38,8 +38,7 @@ std::size_t compute_num_blocks_from_rates(std::size_t dataset_rows, constexpr double kLn2Sq = kLn2 * kLn2; auto expected_insertions = std::max( - 1, - static_cast(std::ceil(static_cast(dataset_rows) * expected_valid_rate))); + 1, static_cast(std::ceil(static_cast(dataset_rows) * filtering_rate))); auto required_bits = -static_cast(expected_insertions) * std::log(static_cast(target_false_positive_rate)) / kLn2Sq; return std::max(1, @@ -54,16 +53,16 @@ struct bloom_filter::impl { cuco_filter_type filter; rmm::device_uvector payload; std::optional configured_dataset_rows; - float expected_valid_rate; + float filtering_rate; float target_false_positive_rate; impl(raft::resources const& res, std::size_t num_blocks, - float expected_valid_rate_, + float filtering_rate_, float target_false_positive_rate_) : filter(num_blocks, {}, {}, {}, raft::resource::get_cuda_stream(res)), payload(1, raft::resource::get_cuda_stream(res)), - expected_valid_rate(expected_valid_rate_), + filtering_rate(filtering_rate_), target_false_positive_rate(target_false_positive_rate_) { auto stream = raft::resource::get_cuda_stream(res); @@ -88,7 +87,7 @@ struct bloom_filter::impl { if (!configured_dataset_rows.has_value()) { configured_dataset_rows = inferred_dataset_rows; auto required_num_blocks = compute_num_blocks_from_rates( - inferred_dataset_rows, expected_valid_rate, target_false_positive_rate); + inferred_dataset_rows, filtering_rate, target_false_positive_rate); auto target_num_blocks = std::max(filter.block_extent(), required_num_blocks); if (target_num_blocks != filter.block_extent()) { rebuild_filter(target_num_blocks, stream); } return; @@ -101,14 +100,14 @@ struct bloom_filter::impl { }; bloom_filter::bloom_filter(raft::resources const& res, - float expected_valid_rate, + float filtering_rate, float target_false_positive_rate, std::size_t num_blocks) : impl_(std::make_unique( - res, std::max(1, num_blocks), expected_valid_rate, target_false_positive_rate)) + res, std::max(1, num_blocks), filtering_rate, target_false_positive_rate)) { - RAFT_EXPECTS(expected_valid_rate > 0.0f && expected_valid_rate <= 1.0f, - "expected_valid_rate must be in (0, 1]."); + RAFT_EXPECTS(filtering_rate > 0.0f && filtering_rate <= 1.0f, + "filtering_rate must be in (0, 1]."); RAFT_EXPECTS(target_false_positive_rate > 0.0f && target_false_positive_rate < 1.0f, "target_false_positive_rate must be in (0, 1)."); } From cf40ea16bed2fc468fd6d8a1a6cc55b5bc190047 Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Thu, 9 Jul 2026 00:37:18 +0000 Subject: [PATCH 07/11] payload management --- cpp/include/cuvs/core/bloom_filter.hpp | 8 +++---- cpp/include/cuvs/neighbors/common.hpp | 17 ++++++++++---- cpp/src/core/bloom_filter.cu | 22 +++++++++---------- .../detail/cagra/cagra_filter_payload.hpp | 22 ++++++++++++++++++- cpp/tests/neighbors/ann_cagra.cuh | 2 +- .../cpp/src/cagra_bloom_filter_example.cu | 2 +- 6 files changed, 50 insertions(+), 23 deletions(-) diff --git a/cpp/include/cuvs/core/bloom_filter.hpp b/cpp/include/cuvs/core/bloom_filter.hpp index ca5827eca4..fdfbc03a5c 100644 --- a/cpp/include/cuvs/core/bloom_filter.hpp +++ b/cpp/include/cuvs/core/bloom_filter.hpp @@ -77,12 +77,12 @@ class CUVS_EXPORT bloom_filter { [[nodiscard]] std::size_t num_blocks() const noexcept; /** - * @brief Device pointer to the CAGRA JIT sample-filter payload. + * @brief Export host payload used by CAGRA JIT filter internals. * - * The pointed object is device memory owned by this wrapper and remains valid while this object - * is alive. + * This API is intended for cuVS internal use. The caller must provide storage matching the + * payload layout expected by CAGRA. */ - [[nodiscard]] void* filter_data() const noexcept; + void export_payload(void* payload_out, std::size_t payload_bytes) const; private: struct impl; diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 0ad32b853b..67585aabf8 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -33,6 +33,9 @@ #endif namespace CUVS_EXPORT cuvs { +namespace core { +class bloom_filter; +} namespace neighbors { /** * @addtogroup cagra_cpp_index_params @@ -620,10 +623,10 @@ struct bitset_filter : public base_filter { /** * @brief Filter CAGRA candidates with a global @c cuvs::core::bloom_filter over the index. * - * Build the filter once on the host with bulk @c add() over the allowed dataset row ids, obtain a - * @c ref() from the owning @c cuvs::core::bloom_filter, copy that ref to device memory, and pass - * the device pointer as @c filter_data. The linked JIT-LTO fragment probes the same filter for - * every query and candidate, similar to @ref bitset_filter but with probabilistic membership tests. + * Build the filter once on the host with bulk @c add() over the allowed dataset row ids and pass + * the owning @c cuvs::core::bloom_filter to this wrapper. CAGRA internals build/cache the device + * payload, similar to @ref bitset_filter, and the linked JIT-LTO fragment probes the same filter + * for every query and candidate with probabilistic membership tests. * * Bloom filters have no false negatives: if a row was inserted, @c contains returns @c true. False * positives are possible, so highly selective predicates may still need a bitset or UDF for exact @@ -631,6 +634,7 @@ struct bitset_filter : public base_filter { */ struct bloom_filter : public base_filter { void* filter_data{nullptr}; + const cuvs::core::bloom_filter* bloom_filter_ptr{nullptr}; float filtering_rate{-1.0f}; bloom_filter() = default; @@ -640,6 +644,11 @@ struct bloom_filter : public base_filter { { } + explicit bloom_filter(const cuvs::core::bloom_filter& bloom_filter, float filtering_rate = -1.0f) + : bloom_filter_ptr(&bloom_filter), filtering_rate(filtering_rate) + { + } + FilterType get_filter_type() const override { return FilterType::Bloom; } }; diff --git a/cpp/src/core/bloom_filter.cu b/cpp/src/core/bloom_filter.cu index 3554578d7e..86c1d6906b 100644 --- a/cpp/src/core/bloom_filter.cu +++ b/cpp/src/core/bloom_filter.cu @@ -8,12 +8,9 @@ #include -#include #include #include -#include - #include #include #include @@ -51,7 +48,6 @@ struct bloom_filter::impl { using sample_filter_payload = cuvs::neighbors::detail::bloom_filter_data_t; cuco_filter_type filter; - rmm::device_uvector payload; std::optional configured_dataset_rows; float filtering_rate; float target_false_positive_rate; @@ -61,21 +57,14 @@ struct bloom_filter::impl { float filtering_rate_, float target_false_positive_rate_) : filter(num_blocks, {}, {}, {}, raft::resource::get_cuda_stream(res)), - payload(1, raft::resource::get_cuda_stream(res)), filtering_rate(filtering_rate_), target_false_positive_rate(target_false_positive_rate_) { - auto stream = raft::resource::get_cuda_stream(res); - sample_filter_payload host_payload{filter.ref()}; - raft::copy(payload.data(), &host_payload, 1, stream); - raft::resource::sync_stream(res); } void rebuild_filter(std::size_t num_blocks, cudaStream_t stream) { filter = cuco_filter_type(num_blocks, {}, {}, {}, stream); - sample_filter_payload host_payload{filter.ref()}; - raft::copy(payload.data(), &host_payload, 1, stream); } void configure_or_validate_dataset_rows(raft::device_vector_view keys, @@ -164,6 +153,15 @@ void bloom_filter::contains_async(raft::resources const& res, std::size_t bloom_filter::num_blocks() const noexcept { return impl_->filter.block_extent(); } -void* bloom_filter::filter_data() const noexcept { return impl_->payload.data(); } +void bloom_filter::export_payload(void* payload_out, std::size_t payload_bytes) const +{ + using sample_filter_payload = impl::sample_filter_payload; + RAFT_EXPECTS(payload_out != nullptr, "payload_out must not be null."); + RAFT_EXPECTS(payload_bytes == sizeof(sample_filter_payload), + "payload_bytes must match bloom filter payload size."); + + auto* typed_payload = static_cast(payload_out); + *typed_payload = sample_filter_payload{impl_->filter.ref()}; +} } // namespace cuvs::core diff --git a/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp b/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp index 39c59cb6b7..5de6969f31 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp +++ b/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp @@ -8,6 +8,7 @@ #include "../sample_filter_data.cuh" #include "jit_lto_kernels/cagra_filter_payload.cuh" +#include #include #include @@ -162,6 +163,18 @@ ::cuvs::neighbors::detail::bitset_filter_data_t make_cagra_bitset_ static_cast(bitset_view.get_original_nbits())}; } +template +::cuvs::neighbors::detail::bloom_filter_data_t make_cagra_bloom_filter_storage( + const FilterT& filter) +{ + using payload_t = ::cuvs::neighbors::detail::bloom_filter_data_t; + payload_t payload{}; + RAFT_EXPECTS(filter.bloom_filter_ptr != nullptr, + "bloom_filter requires a cuvs::core::bloom_filter object or prebuilt filter_data."); + filter.bloom_filter_ptr->export_payload(&payload, sizeof(payload)); + return payload; +} + template void* get_cagra_device_payload(PayloadT payload, cudaStream_t stream) { @@ -184,7 +197,11 @@ void fill_cagra_sample_filter(cagra_sample_filter& out, if constexpr (is_bitset_filter::value) { out.filter_data = make_cagra_bitset_filter_payload(filter, stream); } else if constexpr (is_bloom_filter::value) { - out.filter_data = filter.filter_data; + if (filter.filter_data != nullptr) { + out.filter_data = filter.filter_data; + } else { + out.filter_data = get_cagra_device_payload(make_cagra_bloom_filter_storage(filter), stream); + } } else if constexpr (is_udf_filter::value) { out.filter_data = filter.filter_data; } @@ -196,6 +213,9 @@ std::uint64_t cagra_filter_payload_hash(const FilterT& filter) using DecayedFilter = std::decay_t; if constexpr (is_bitset_filter::value) { return cagra_payload_hash(make_cagra_bitset_filter_storage(filter)); + } else if constexpr (is_bloom_filter::value) { + if (filter.filter_data != nullptr) { return 0; } + return cagra_payload_hash(make_cagra_bloom_filter_storage(filter)); } else if constexpr (requires { filter.filter; }) { return cagra_filter_payload_hash(filter.filter); } else { diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index bf22006c64..be02524090 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -980,7 +980,7 @@ class AnnCagraFilterTest : public ::testing::TestWithParam { raft::resource::sync_stream(handle_); auto bloom_filter_obj = cuvs::neighbors::filtering::bloom_filter( - global_bloom_filter.filter_data(), + global_bloom_filter, static_cast(test_cagra_sample_filter::offset) / static_cast(ps.n_rows)); cagra::search(handle_, diff --git a/examples/cpp/src/cagra_bloom_filter_example.cu b/examples/cpp/src/cagra_bloom_filter_example.cu index eff960f930..718168d42c 100644 --- a/examples/cpp/src/cagra_bloom_filter_example.cu +++ b/examples/cpp/src/cagra_bloom_filter_example.cu @@ -104,7 +104,7 @@ int main() search_params.thread_block_size = 256; // ~50% of rows are rejected by the global even-id predicate. - auto filter = cuvs::neighbors::filtering::bloom_filter(allowed_rows.filter_data(), 0.5f); + auto filter = cuvs::neighbors::filtering::bloom_filter(allowed_rows, 0.5f); cuvs::neighbors::cagra::search(res, search_params, From c57571710f5fbd668e5d5e7396ced52041dc9b1c Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Thu, 9 Jul 2026 00:53:29 +0000 Subject: [PATCH 08/11] cleanup payload mgmt, estimate filtering rate --- cpp/include/cuvs/core/bloom_filter.hpp | 2 ++ cpp/include/cuvs/neighbors/common.hpp | 11 ++------ cpp/src/core/bloom_filter.cu | 25 +++++++++++++++++++ cpp/src/neighbors/cagra.cuh | 13 ++++++---- .../detail/cagra/cagra_filter_payload.hpp | 14 ++++------- cpp/tests/neighbors/ann_cagra.cuh | 4 +-- .../cpp/src/cagra_bloom_filter_example.cu | 2 +- 7 files changed, 44 insertions(+), 27 deletions(-) diff --git a/cpp/include/cuvs/core/bloom_filter.hpp b/cpp/include/cuvs/core/bloom_filter.hpp index fdfbc03a5c..e3abc10616 100644 --- a/cpp/include/cuvs/core/bloom_filter.hpp +++ b/cpp/include/cuvs/core/bloom_filter.hpp @@ -75,6 +75,8 @@ class CUVS_EXPORT bloom_filter { raft::device_vector_view output) const; [[nodiscard]] std::size_t num_blocks() const noexcept; + [[nodiscard]] float estimate_filtering_rate(raft::resources const& res, + std::size_t dataset_rows) const; /** * @brief Export host payload used by CAGRA JIT filter internals. diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 67585aabf8..6330d796bd 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -634,18 +634,11 @@ struct bitset_filter : public base_filter { */ struct bloom_filter : public base_filter { void* filter_data{nullptr}; - const cuvs::core::bloom_filter* bloom_filter_ptr{nullptr}; - float filtering_rate{-1.0f}; bloom_filter() = default; - explicit bloom_filter(void* filter_data, float filtering_rate = -1.0f) - : filter_data(filter_data), filtering_rate(filtering_rate) - { - } - - explicit bloom_filter(const cuvs::core::bloom_filter& bloom_filter, float filtering_rate = -1.0f) - : bloom_filter_ptr(&bloom_filter), filtering_rate(filtering_rate) + explicit bloom_filter(const cuvs::core::bloom_filter& bloom_filter) + : filter_data(const_cast(&bloom_filter)) { } diff --git a/cpp/src/core/bloom_filter.cu b/cpp/src/core/bloom_filter.cu index 86c1d6906b..c7e7bfd6a8 100644 --- a/cpp/src/core/bloom_filter.cu +++ b/cpp/src/core/bloom_filter.cu @@ -10,6 +10,11 @@ #include #include +#include + +#include +#include +#include #include #include @@ -153,6 +158,26 @@ void bloom_filter::contains_async(raft::resources const& res, std::size_t bloom_filter::num_blocks() const noexcept { return impl_->filter.block_extent(); } +float bloom_filter::estimate_filtering_rate(raft::resources const& res, + std::size_t dataset_rows) const +{ + if (dataset_rows == 0) { return 0.0f; } + auto stream = raft::resource::get_cuda_stream(res); + auto policy = raft::resource::get_thrust_policy(res); + + rmm::device_uvector hits(dataset_rows, stream); + + auto first_id = thrust::counting_iterator(0); + impl_->filter.contains_async(first_id, first_id + dataset_rows, hits.data(), stream); + + auto positives = thrust::count_if( + policy, hits.begin(), hits.end(), [] __device__(std::uint8_t v) { return v != 0; }); + raft::resource::sync_stream(res); + auto filtering_rate = static_cast(dataset_rows - static_cast(positives)) / + static_cast(dataset_rows); + return std::clamp(filtering_rate, 0.0f, 0.999f); +} + void bloom_filter::export_payload(void* payload_out, std::size_t payload_bytes) const { using sample_filter_payload = impl::sample_filter_payload; diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index b8900cfe11..3a41380d90 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -389,14 +389,17 @@ void search(raft::resources const& res, auto& sample_filter = dynamic_cast(sample_filter_ref); search_params params_copy = params; - if (params.filtering_rate < 0.0) { + if (params.filtering_rate <= 0.0f) { const float min_filtering_rate = 0.0f; const float max_filtering_rate = 0.999f; + auto const* bloom_filter_obj = + static_cast(sample_filter.filter_data); + RAFT_EXPECTS(bloom_filter_obj != nullptr, + "bloom_filter must carry a valid cuvs::core::bloom_filter handle."); + params_copy.filtering_rate = bloom_filter_obj->estimate_filtering_rate( + res, static_cast(idx.data().n_rows())); params_copy.filtering_rate = - sample_filter.filtering_rate < 0.0f - ? 0.0f - : std::min(std::max(sample_filter.filtering_rate, min_filtering_rate), - max_filtering_rate); + std::min(std::max(params_copy.filtering_rate, min_filtering_rate), max_filtering_rate); } auto sample_filter_copy = sample_filter; return search_with_filtering( diff --git a/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp b/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp index 5de6969f31..933fb8fa13 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp +++ b/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp @@ -169,9 +169,10 @@ ::cuvs::neighbors::detail::bloom_filter_data_t make_cagra_bloom_f { using payload_t = ::cuvs::neighbors::detail::bloom_filter_data_t; payload_t payload{}; - RAFT_EXPECTS(filter.bloom_filter_ptr != nullptr, - "bloom_filter requires a cuvs::core::bloom_filter object or prebuilt filter_data."); - filter.bloom_filter_ptr->export_payload(&payload, sizeof(payload)); + RAFT_EXPECTS(filter.filter_data != nullptr, + "bloom_filter requires a cuvs::core::bloom_filter object."); + auto const* bloom_filter_obj = static_cast(filter.filter_data); + bloom_filter_obj->export_payload(&payload, sizeof(payload)); return payload; } @@ -197,11 +198,7 @@ void fill_cagra_sample_filter(cagra_sample_filter& out, if constexpr (is_bitset_filter::value) { out.filter_data = make_cagra_bitset_filter_payload(filter, stream); } else if constexpr (is_bloom_filter::value) { - if (filter.filter_data != nullptr) { - out.filter_data = filter.filter_data; - } else { - out.filter_data = get_cagra_device_payload(make_cagra_bloom_filter_storage(filter), stream); - } + out.filter_data = get_cagra_device_payload(make_cagra_bloom_filter_storage(filter), stream); } else if constexpr (is_udf_filter::value) { out.filter_data = filter.filter_data; } @@ -214,7 +211,6 @@ std::uint64_t cagra_filter_payload_hash(const FilterT& filter) if constexpr (is_bitset_filter::value) { return cagra_payload_hash(make_cagra_bitset_filter_storage(filter)); } else if constexpr (is_bloom_filter::value) { - if (filter.filter_data != nullptr) { return 0; } return cagra_payload_hash(make_cagra_bloom_filter_storage(filter)); } else if constexpr (requires { filter.filter; }) { return cagra_filter_payload_hash(filter.filter); diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index be02524090..4047605469 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -979,9 +979,7 @@ class AnnCagraFilterTest : public ::testing::TestWithParam { global_bloom_filter.add_async(handle_, valid_ids_view); raft::resource::sync_stream(handle_); - auto bloom_filter_obj = cuvs::neighbors::filtering::bloom_filter( - global_bloom_filter, - static_cast(test_cagra_sample_filter::offset) / static_cast(ps.n_rows)); + auto bloom_filter_obj = cuvs::neighbors::filtering::bloom_filter(global_bloom_filter); cagra::search(handle_, search_params, diff --git a/examples/cpp/src/cagra_bloom_filter_example.cu b/examples/cpp/src/cagra_bloom_filter_example.cu index 718168d42c..34f5f74393 100644 --- a/examples/cpp/src/cagra_bloom_filter_example.cu +++ b/examples/cpp/src/cagra_bloom_filter_example.cu @@ -104,7 +104,7 @@ int main() search_params.thread_block_size = 256; // ~50% of rows are rejected by the global even-id predicate. - auto filter = cuvs::neighbors::filtering::bloom_filter(allowed_rows, 0.5f); + auto filter = cuvs::neighbors::filtering::bloom_filter(allowed_rows); cuvs::neighbors::cagra::search(res, search_params, From 2bebdb9fc3dabf08ef6c9dcf558e63896772bd34 Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Fri, 10 Jul 2026 01:40:40 +0000 Subject: [PATCH 09/11] fix build --- cpp/src/core/bloom_filter.cu | 4 ++-- .../neighbors/detail/cagra/cagra_filter_payload.hpp | 10 +++++++--- cpp/src/neighbors/detail/sample_filter_data.cuh | 4 +++- 3 files changed, 12 insertions(+), 6 deletions(-) diff --git a/cpp/src/core/bloom_filter.cu b/cpp/src/core/bloom_filter.cu index c7e7bfd6a8..3e8710f5eb 100644 --- a/cpp/src/core/bloom_filter.cu +++ b/cpp/src/core/bloom_filter.cu @@ -18,6 +18,7 @@ #include #include +#include #include #include @@ -185,8 +186,7 @@ void bloom_filter::export_payload(void* payload_out, std::size_t payload_bytes) RAFT_EXPECTS(payload_bytes == sizeof(sample_filter_payload), "payload_bytes must match bloom filter payload size."); - auto* typed_payload = static_cast(payload_out); - *typed_payload = sample_filter_payload{impl_->filter.ref()}; + ::new (payload_out) sample_filter_payload{impl_->filter.ref()}; } } // namespace cuvs::core diff --git a/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp b/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp index 933fb8fa13..b7b1dc3f20 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp +++ b/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp @@ -18,6 +18,7 @@ #include #include #include +#include #include #include @@ -41,7 +42,7 @@ std::uint64_t cagra_payload_hash(PayloadT const& payload) template struct cagra_device_payload_owner { struct state { - PayloadT host_payload{}; + PayloadT host_payload; PayloadT* device_payload{nullptr}; cudaStream_t stream{}; cudaEvent_t ready_event{}; @@ -168,11 +169,14 @@ ::cuvs::neighbors::detail::bloom_filter_data_t make_cagra_bloom_f const FilterT& filter) { using payload_t = ::cuvs::neighbors::detail::bloom_filter_data_t; - payload_t payload{}; RAFT_EXPECTS(filter.filter_data != nullptr, "bloom_filter requires a cuvs::core::bloom_filter object."); auto const* bloom_filter_obj = static_cast(filter.filter_data); - bloom_filter_obj->export_payload(&payload, sizeof(payload)); + std::aligned_storage_t storage; + bloom_filter_obj->export_payload(&storage, sizeof(payload_t)); + auto* payload_ptr = std::launder(reinterpret_cast(&storage)); + payload_t payload = *payload_ptr; + payload_ptr->~payload_t(); return payload; } diff --git a/cpp/src/neighbors/detail/sample_filter_data.cuh b/cpp/src/neighbors/detail/sample_filter_data.cuh index 3f2e412681..e5d04d5734 100644 --- a/cpp/src/neighbors/detail/sample_filter_data.cuh +++ b/cpp/src/neighbors/detail/sample_filter_data.cuh @@ -28,7 +28,9 @@ template struct bloom_filter_data_t { using ref_type = typename cuco::bloom_filter::ref_type<>; - ref_type filter{}; + explicit bloom_filter_data_t(ref_type filter) : filter(filter) {} + + ref_type filter; }; } // namespace cuvs::neighbors::detail From 863e9fbbb2858d0d21f12ebc8b377d7434f55436 Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Fri, 10 Jul 2026 19:52:29 +0000 Subject: [PATCH 10/11] fix examples cmake --- examples/cpp/CMakeLists.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/cpp/CMakeLists.txt b/examples/cpp/CMakeLists.txt index 334602764c..50b4572dea 100644 --- a/examples/cpp/CMakeLists.txt +++ b/examples/cpp/CMakeLists.txt @@ -61,7 +61,7 @@ target_link_libraries( target_link_libraries( CAGRA_BLOOM_FILTER_EXAMPLE PRIVATE cuvs::cuvs cuco::cuco $ ) -target_include_directories( +target_link_libraries( CAGRA_HNSW_ACE_BUILD_EXAMPLE PRIVATE cuvs::cuvs $ ) target_link_libraries(CAGRA_HNSW_ACE_EXAMPLE PRIVATE cuvs::cuvs $) From 0569eb7a58ed9a60b68ec25019d4103467fee499 Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Fri, 10 Jul 2026 20:21:09 +0000 Subject: [PATCH 11/11] simplify payload export --- cpp/include/cuvs/core/bloom_filter.hpp | 13 ++++--------- cpp/src/core/bloom_filter.cu | 19 ++++++++++--------- .../detail/cagra/cagra_filter_payload.hpp | 9 +-------- .../neighbors/detail/sample_filter_data.cuh | 8 ++++++++ 4 files changed, 23 insertions(+), 26 deletions(-) diff --git a/cpp/include/cuvs/core/bloom_filter.hpp b/cpp/include/cuvs/core/bloom_filter.hpp index e3abc10616..706b6e0af1 100644 --- a/cpp/include/cuvs/core/bloom_filter.hpp +++ b/cpp/include/cuvs/core/bloom_filter.hpp @@ -23,6 +23,9 @@ namespace core { * The wrapper supports the expected bulk host APIs used by ANN workflows. */ class CUVS_EXPORT bloom_filter { + private: + struct impl; + public: using key_type = std::uint32_t; @@ -77,17 +80,9 @@ class CUVS_EXPORT bloom_filter { [[nodiscard]] std::size_t num_blocks() const noexcept; [[nodiscard]] float estimate_filtering_rate(raft::resources const& res, std::size_t dataset_rows) const; - - /** - * @brief Export host payload used by CAGRA JIT filter internals. - * - * This API is intended for cuVS internal use. The caller must provide storage matching the - * payload layout expected by CAGRA. - */ - void export_payload(void* payload_out, std::size_t payload_bytes) const; + [[nodiscard]] impl const& get_impl() const noexcept; private: - struct impl; std::unique_ptr impl_; }; diff --git a/cpp/src/core/bloom_filter.cu b/cpp/src/core/bloom_filter.cu index 3e8710f5eb..94ba3797ad 100644 --- a/cpp/src/core/bloom_filter.cu +++ b/cpp/src/core/bloom_filter.cu @@ -18,7 +18,6 @@ #include #include -#include #include #include @@ -159,6 +158,8 @@ void bloom_filter::contains_async(raft::resources const& res, std::size_t bloom_filter::num_blocks() const noexcept { return impl_->filter.block_extent(); } +auto bloom_filter::get_impl() const noexcept -> impl const& { return *impl_; } + float bloom_filter::estimate_filtering_rate(raft::resources const& res, std::size_t dataset_rows) const { @@ -179,14 +180,14 @@ float bloom_filter::estimate_filtering_rate(raft::resources const& res, return std::clamp(filtering_rate, 0.0f, 0.999f); } -void bloom_filter::export_payload(void* payload_out, std::size_t payload_bytes) const -{ - using sample_filter_payload = impl::sample_filter_payload; - RAFT_EXPECTS(payload_out != nullptr, "payload_out must not be null."); - RAFT_EXPECTS(payload_bytes == sizeof(sample_filter_payload), - "payload_bytes must match bloom filter payload size."); +} // namespace cuvs::core - ::new (payload_out) sample_filter_payload{impl_->filter.ref()}; +namespace cuvs::neighbors::detail { + +bloom_filter_data_t bloom_filter_factory::make( + cuvs::core::bloom_filter const& filter) +{ + return bloom_filter_data_t{filter.get_impl().filter.ref()}; } -} // namespace cuvs::core +} // namespace cuvs::neighbors::detail diff --git a/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp b/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp index b7b1dc3f20..5c5d9054b4 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp +++ b/cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp @@ -18,7 +18,6 @@ #include #include #include -#include #include #include @@ -168,16 +167,10 @@ template ::cuvs::neighbors::detail::bloom_filter_data_t make_cagra_bloom_filter_storage( const FilterT& filter) { - using payload_t = ::cuvs::neighbors::detail::bloom_filter_data_t; RAFT_EXPECTS(filter.filter_data != nullptr, "bloom_filter requires a cuvs::core::bloom_filter object."); auto const* bloom_filter_obj = static_cast(filter.filter_data); - std::aligned_storage_t storage; - bloom_filter_obj->export_payload(&storage, sizeof(payload_t)); - auto* payload_ptr = std::launder(reinterpret_cast(&storage)); - payload_t payload = *payload_ptr; - payload_ptr->~payload_t(); - return payload; + return ::cuvs::neighbors::detail::bloom_filter_factory::make(*bloom_filter_obj); } template diff --git a/cpp/src/neighbors/detail/sample_filter_data.cuh b/cpp/src/neighbors/detail/sample_filter_data.cuh index e5d04d5734..9ff935409c 100644 --- a/cpp/src/neighbors/detail/sample_filter_data.cuh +++ b/cpp/src/neighbors/detail/sample_filter_data.cuh @@ -10,6 +10,10 @@ #include #include +namespace cuvs::core { +class bloom_filter; +} + namespace cuvs::neighbors::detail { /// Bitset (and length metadata) for linked @c sample_filter in JIT LTO; passed by value to @@ -33,4 +37,8 @@ struct bloom_filter_data_t { ref_type filter; }; +struct bloom_filter_factory { + static bloom_filter_data_t make(cuvs::core::bloom_filter const& filter); +}; + } // namespace cuvs::neighbors::detail