diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 96c4e70cfd..e661fb4b6d 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) @@ -1328,6 +1330,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..706b6e0af1 --- /dev/null +++ b/cpp/include/cuvs/core/bloom_filter.hpp @@ -0,0 +1,90 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * 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 { + private: + struct impl; + + 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 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 * 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)` + * + * Practical knob behavior: + * - Lower @p target_false_positive_rate -> larger filter, fewer false positives, typically higher + * filtered-search recall. + * - 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 filtering_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; + [[nodiscard]] float estimate_filtering_rate(raft::resources const& res, + std::size_t dataset_rows) const; + [[nodiscard]] impl const& get_impl() const noexcept; + + private: + std::unique_ptr impl_; +}; + +} // namespace core +} // namespace CUVS_EXPORT cuvs diff --git a/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp b/cpp/include/cuvs/detail/jit_lto/common_fragments.hpp index ef2a8e6002..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 */ @@ -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..6330d796bd 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 */ @@ -33,6 +33,9 @@ #endif namespace CUVS_EXPORT cuvs { +namespace core { +class bloom_filter; +} namespace neighbors { /** * @addtogroup cagra_cpp_index_params @@ -497,7 +500,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 +620,31 @@ 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 cuvs::core::bloom_filter over the index. + * + * 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 + * filtering. + */ +struct bloom_filter : public base_filter { + void* filter_data{nullptr}; + + bloom_filter() = default; + + explicit bloom_filter(const cuvs::core::bloom_filter& bloom_filter) + : filter_data(const_cast(&bloom_filter)) + { + } + + FilterType get_filter_type() const override { return FilterType::Bloom; } +}; + /** * @brief JIT-LTO user-defined filter predicate. * diff --git a/cpp/src/core/bloom_filter.cu b/cpp/src/core/bloom_filter.cu new file mode 100644 index 0000000000..94ba3797ad --- /dev/null +++ b/cpp/src/core/bloom_filter.cu @@ -0,0 +1,193 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../neighbors/detail/sample_filter_data.cuh" +#include + +#include + +#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 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(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)."); + + // 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) * filtering_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; + std::optional configured_dataset_rows; + float filtering_rate; + float target_false_positive_rate; + + impl(raft::resources const& res, + std::size_t num_blocks, + float filtering_rate_, + float target_false_positive_rate_) + : filter(num_blocks, {}, {}, {}, raft::resource::get_cuda_stream(res)), + filtering_rate(filtering_rate_), + target_false_positive_rate(target_false_positive_rate_) + { + } + + void rebuild_filter(std::size_t num_blocks, cudaStream_t stream) + { + filter = cuco_filter_type(num_blocks, {}, {}, {}, 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, 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; + } + + 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 filtering_rate, + float target_false_positive_rate, + std::size_t num_blocks) + : impl_(std::make_unique( + res, std::max(1, num_blocks), filtering_rate, target_false_positive_rate)) +{ + 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)."); +} + +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(); } + +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 +{ + 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); +} + +} // namespace cuvs::core + +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::neighbors::detail diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index ee87c2c0ab..3a41380d90 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 */ @@ -385,6 +385,28 @@ 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.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 = + std::min(std::max(params_copy.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..5c5d9054b4 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 @@ -8,6 +8,7 @@ #include "../sample_filter_data.cuh" #include "jit_lto_kernels/cagra_filter_payload.cuh" +#include #include #include @@ -40,7 +41,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{}; @@ -139,6 +140,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 {}; @@ -156,6 +163,16 @@ ::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) +{ + 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); + return ::cuvs::neighbors::detail::bloom_filter_factory::make(*bloom_filter_obj); +} + template void* get_cagra_device_payload(PayloadT payload, cudaStream_t stream) { @@ -177,6 +194,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 = get_cagra_device_payload(make_cagra_bloom_filter_storage(filter), stream); } else if constexpr (is_udf_filter::value) { out.filter_data = filter.filter_data; } @@ -188,6 +207,8 @@ 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) { + return cagra_payload_hash(make_cagra_bloom_filter_storage(filter)); } else if constexpr (requires { filter.filter; }) { return cagra_filter_payload_hash(filter.filter); } else { @@ -199,7 +220,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..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 */ @@ -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..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 */ @@ -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..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 */ @@ -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..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 */ @@ -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..9ff935409c 100644 --- a/cpp/src/neighbors/detail/sample_filter_data.cuh +++ b/cpp/src/neighbors/detail/sample_filter_data.cuh @@ -5,9 +5,15 @@ #pragma once +#include + #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 @@ -21,4 +27,18 @@ 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<>; + + explicit bloom_filter_data_t(ref_type filter) : filter(filter) {} + + ref_type filter; +}; + +struct bloom_filter_factory { + static bloom_filter_data_t make(cuvs::core::bloom_filter const& filter); +}; + } // namespace cuvs::neighbors::detail diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index 7b86cc70ad..4047605469 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -11,6 +11,7 @@ #include "naive_knn.cuh" +#include #include #include #include @@ -35,6 +36,7 @@ #include +#include #include #include #include @@ -946,6 +948,77 @@ 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); + + 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_index_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); + RAFT_LOG_INFO("Bloom filter index recall = %f", bloom_index_recall); + 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 54e4ff97b9..50b4572dea 100644 --- a/examples/cpp/CMakeLists.txt +++ b/examples/cpp/CMakeLists.txt @@ -1,6 +1,6 @@ # ============================================================================= # cmake-format: off -# 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 # cmake-format: on @@ -32,11 +32,14 @@ set(CMAKE_CUDA_STANDARD_REQUIRED ON) 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_BUILD_EXAMPLE src/cagra_hnsw_ace_build.cu) add_executable(CAGRA_HNSW_ACE_EXAMPLE src/cagra_hnsw_ace_example.cu) add_executable(CAGRA_PERSISTENT_EXAMPLE src/cagra_persistent_example.cu) @@ -55,6 +58,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_BUILD_EXAMPLE PRIVATE cuvs::cuvs $ ) 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..34f5f74393 --- /dev/null +++ b/examples/cpp/src/cagra_bloom_filter_example.cu @@ -0,0 +1,160 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * 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 std::size_t sub_filters = 256; + +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; } + +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); + + 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; + + 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(allowed_rows); + + 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); + 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()); + 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; +}