CAGRA Bloom Filter

cpp/CMakeLists.txt

@@ -367,7 +367,9 @@ if(NOT BUILD_CPU_ONLY)
                                                 "$<$<COMPILE_LANGUAGE:CUDA>:${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

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 <cuvs/core/export.hpp>
+#include <raft/core/device_mdarray.hpp>
+#include <raft/core/resources.hpp>
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+
+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<const key_type, int64_t> keys);
+  void add_async(raft::resources const& res,
+                 raft::device_vector_view<const key_type, int64_t> keys);
+
+  void contains(raft::resources const& res,
+                raft::device_vector_view<const key_type, int64_t> keys,
+                raft::device_vector_view<std::uint8_t, int64_t> output) const;
+  void contains_async(raft::resources const& res,
+                      raft::device_vector_view<const key_type, int64_t> keys,
+                      raft::device_vector_view<std::uint8_t, int64_t> 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> impl_;
+};
+
+}  // namespace core
+}  // namespace CUVS_EXPORT cuvs

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 {};

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<cuvs::core::bloom_filter*>(&bloom_filter))
+  {
+  }
+
+  FilterType get_filter_type() const override { return FilterType::Bloom; }
+};
+
 /**
  * @brief JIT-LTO user-defined filter predicate.
  *

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 <cuvs/core/bloom_filter.hpp>
+
+#include <cuco/bloom_filter.cuh>
+
+#include <raft/core/error.hpp>
+#include <raft/core/resource/cuda_stream.hpp>
+#include <raft/core/resource/thrust_policy.hpp>
+
+#include <rmm/device_uvector.hpp>
+#include <thrust/count.h>
+#include <thrust/iterator/counting_iterator.h>
+
+#include <algorithm>
+#include <cmath>
+#include <optional>
+#include <utility>
+
+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<std::size_t>(
+    1, static_cast<std::size_t>(std::ceil(static_cast<double>(dataset_rows) * filtering_rate)));
+  auto required_bits = -static_cast<double>(expected_insertions) *
+                       std::log(static_cast<double>(target_false_positive_rate)) / kLn2Sq;
+  return std::max<std::size_t>(1,
+                               static_cast<std::size_t>(std::ceil(required_bits / kBitsPerBlock)));
+}
+
+struct bloom_filter::impl {
+  using key_type              = bloom_filter::key_type;
+  using cuco_filter_type      = cuco::bloom_filter<key_type>;
+  using sample_filter_payload = cuvs::neighbors::detail::bloom_filter_data_t<key_type>;
+
+  cuco_filter_type filter;
+  std::optional<std::size_t> 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<const key_type, int64_t> keys,
+                                          cudaStream_t stream)
+  {
+    if (keys.extent(0) == 0) { return; }
+
+    auto inferred_dataset_rows = static_cast<std::size_t>(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<std::size_t>(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<impl>(
+      res, std::max<std::size_t>(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<const key_type, int64_t> 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<const key_type, int64_t> 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<const key_type, int64_t> keys,
+                            raft::device_vector_view<std::uint8_t, int64_t> 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<const key_type, int64_t> keys,
+                                  raft::device_vector_view<std::uint8_t, int64_t> 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<std::uint8_t> hits(dataset_rows, stream);
+
+  auto first_id = thrust::counting_iterator<key_type>(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<float>(dataset_rows - static_cast<std::size_t>(positives)) /
+                        static_cast<float>(dataset_rows);
+  return std::clamp(filtering_rate, 0.0f, 0.999f);
+}
+
+}  // namespace cuvs::core
+
+namespace cuvs::neighbors::detail {
+
+bloom_filter_data_t<std::uint32_t> bloom_filter_factory::make(
+  cuvs::core::bloom_filter const& filter)
+{
+  return bloom_filter_data_t<std::uint32_t>{filter.get_impl().filter.ref()};
+}
+
+}  // namespace cuvs::neighbors::detail