diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 175669a620..31d27815e7 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -365,7 +365,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) @@ -1326,6 +1328,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..c622532f76 --- /dev/null +++ b/cpp/include/cuvs/core/bloom_filter.hpp @@ -0,0 +1,89 @@ +/* + * 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. + * + * @p dataset_rows is the number of rows in the indexed dataset. The filter uses it with + * @p filtering_rate to estimate the number of inserted valid ids and compute a target filter + * size that satisfies 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) + * + * 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. + */ + bloom_filter(raft::resources const& res, + std::size_t dataset_rows, + float filtering_rate = 1.0f, + float target_false_positive_rate = 0.01f); + ~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; + + private: + friend impl const& get_bloom_filter_impl(bloom_filter const& filter) noexcept; + + 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..8e63b6c304 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 : int { None = 0, Bitmap = 1, Bitset = 2, Bloom = 3, UDF = 100 }; struct base_filter { ~base_filter() = default; @@ -617,6 +620,34 @@ 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. + * + * This adapter is non-owning. The referenced @c cuvs::core::bloom_filter must outlive the adapter + * and any searches that use it, and must not be moved or mutated concurrently with a search. + */ +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..751ef22b32 --- /dev/null +++ b/cpp/src/core/bloom_filter.cu @@ -0,0 +1,184 @@ +/* + * 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 + +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::size_t dataset_rows; + float filtering_rate; + float target_false_positive_rate; + + impl(raft::resources const& res, + std::size_t num_blocks, + std::size_t dataset_rows_, + float filtering_rate_, + float target_false_positive_rate_) + : filter(num_blocks, {}, {}, {}, raft::resource::get_cuda_stream(res)), + dataset_rows(dataset_rows_), + filtering_rate(filtering_rate_), + target_false_positive_rate(target_false_positive_rate_) + { + } + + void validate_keys(raft::device_vector_view keys) const + { + if (keys.extent(0) == 0) { return; } + auto num_insertions = static_cast(keys.extent(0)); + RAFT_EXPECTS(num_insertions <= dataset_rows, + "Number of keys to add must not exceed dataset_rows."); + } +}; + +bloom_filter::bloom_filter(raft::resources const& res, + std::size_t dataset_rows, + float filtering_rate, + float target_false_positive_rate) + : impl_(std::make_unique( + res, + compute_num_blocks_from_rates(dataset_rows, filtering_rate, target_false_positive_rate), + dataset_rows, + filtering_rate, + target_false_positive_rate)) +{ +} + +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_->validate_keys(keys); + 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_->validate_keys(keys); + 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 +{ + RAFT_EXPECTS(output.extent(0) >= keys.extent(0), + "Bloom filter contains output size must be at least keys.size()."); + 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 +{ + RAFT_EXPECTS(output.extent(0) >= keys.extent(0), + "Bloom filter contains output size must be at least keys.size()."); + 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(); } + +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); +} + +auto get_bloom_filter_impl(bloom_filter const& filter) noexcept -> bloom_filter::impl const& +{ + return *filter.impl_; +} + +} // 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{get_bloom_filter_impl(filter).filter.ref()}; +} + +} // namespace cuvs::neighbors::detail diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index ee87c2c0ab..2e7e9357cb 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 984ff9ff4d..e0406f5d1a 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 @@ -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/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index 1dd4cbe075..e1d19d936d 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.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 */ #pragma once @@ -44,6 +44,8 @@ index merge(raft::resources const& handle, RAFT_EXPECTS(row_filter.get_filter_type() != cuvs::neighbors::filtering::FilterType::Bitmap, "Bitmap filter isn't supported inside cagra::merge"); + RAFT_EXPECTS(row_filter.get_filter_type() != cuvs::neighbors::filtering::FilterType::Bloom, + "Bloom filter isn't supported inside cagra::merge"); for (cagra_index_t* index : indices) { RAFT_EXPECTS(index != nullptr, 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 79bf339827..a5b2545aa1 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,10 +36,13 @@ #include +#include #include +#include #include #include #include +#include #include namespace cuvs::neighbors::cagra { @@ -946,6 +950,113 @@ 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 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}; + } + + auto bloom_valid_fraction = + static_cast(valid_ids_host.size()) / static_cast(ps.n_rows); + for (auto target_false_positive_rate : candidate_fprs) { + cuvs::core::bloom_filter global_bloom_filter(handle_, + static_cast(ps.n_rows), + bloom_valid_fraction, + target_false_positive_rate); + 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_); + + std::vector bloom_indices_as_keys_host(queries_size); + bool bloom_out_of_domain = false; + for (size_t i = 0; i < queries_size; ++i) { + const auto id = bloom_indices_host[i]; + bool negative_id = false; + if constexpr (std::is_signed_v) { negative_id = id < IdxT{0}; } + if (negative_id || + static_cast(id) >= static_cast(ps.n_rows)) { + bloom_out_of_domain = true; + bloom_indices_as_keys_host[i] = 0; + } else { + bloom_indices_as_keys_host[i] = static_cast(id); + } + } + EXPECT_FALSE(bloom_out_of_domain); + + rmm::device_uvector bloom_indices_as_keys_device(queries_size, stream_); + rmm::device_uvector bloom_accepts_device(queries_size, stream_); + std::vector bloom_accepts_host(queries_size); + raft::copy(bloom_indices_as_keys_device.data(), + bloom_indices_as_keys_host.data(), + queries_size, + stream_); + auto bloom_indices_as_keys_view = + raft::make_device_vector_view( + bloom_indices_as_keys_device.data(), static_cast(queries_size)); + auto bloom_accepts_view = raft::make_device_vector_view( + bloom_accepts_device.data(), static_cast(queries_size)); + global_bloom_filter.contains(handle_, bloom_indices_as_keys_view, bloom_accepts_view); + raft::update_host( + bloom_accepts_host.data(), bloom_accepts_device.data(), queries_size, stream_); + raft::resource::sync_stream(handle_); + EXPECT_TRUE(std::all_of(bloom_accepts_host.begin(), + bloom_accepts_host.end(), + [](std::uint8_t accepted) { return accepted != 0; })); + + auto bloom_recall_result = calc_recall(indices_naive, + bloom_indices_host, + distances_naive, + bloom_distances_host, + ps.n_queries, + ps.k, + 0.003); + auto bloom_recall = std::get<0>(bloom_recall_result); + auto bloom_match_count = std::get<2>(bloom_recall_result); + auto bloom_total_count = std::get<3>(bloom_recall_result); + 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); + auto bloom_recall_slack = std::max(1.0, 0.01 * static_cast(bloom_total_count)) / + static_cast(bloom_total_count); + auto bloom_min_recall = std::max( + 0.0, + ps.min_recall - static_cast(target_false_positive_rate) - bloom_recall_slack); + EXPECT_GE(bloom_recall, bloom_min_recall); + } } // 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..39426ad2fd --- /dev/null +++ b/examples/cpp/src/cagra_bloom_filter_example.cu @@ -0,0 +1,154 @@ +/* + * 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; + +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, n_rows, 0.5f, 0.01f); + 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 the Bloom filter'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; + } + } + + std::cout << "CAGRA bloom filter example produced bloom-accepted neighbors." << std::endl; + return 0; +}