CAGRA Bloom Filter by divyegala · Pull Request #2236 · NVIDIA/cuvs

divyegala · 2026-06-10T23:34:32Z

No description provided.

coderabbitai · 2026-06-10T23:40:14Z

📝 Walkthrough

Summary by CodeRabbit

New Features
- Added bloom-filter-based filtering support to CAGRA nearest-neighbor search
- Added a new bloom filter example for CAGRA
- Added bloom-vs-bitset filtering benchmark and a plotting script to visualize results
Bug Fixes
- Fixed filter type selection so UDF sample filters now return the correct type id in JIT selection
Chores
- Added cuco library dependency to enable bloom filter functionality

Walkthrough

This PR integrates cuco bloom filters into CAGRA search by adding a new bloom_filter struct, compile-time type detection, device-side membership kernels, JIT fragment routing, explicit kernel instantiation, runtime dispatch via dynamic_cast, CMake dependency linking, and complete example and benchmark implementations demonstrating construction, search, and performance comparison.

Changes

Bloom Filter Support for CAGRA Search

Layer / File(s)	Summary
Public API: Bloom Filter Type Contract `cpp/include/cuvs/neighbors/common.hpp`, `cpp/include/cuvs/detail/jit_lto/common_fragments.hpp`	Adds `cuvs::neighbors::filtering::bloom_filter` struct storing a device void pointer and optional filtering rate, with `get_filter_type()` returning `FilterType::Bloom`. Introduces `tag_filter_bloom_filter` JIT fragment tag.
Compile-time Bloom Detection and JIT Fragment Routing `cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp`, `cpp/src/neighbors/detail/cagra/shared_launcher_jit.hpp`	Adds `is_bloom_filter` type trait and updates `sample_filter_jit_tag` to recognize bloom filters and map them to the bloom filter JIT fragment for both direct and nested filter types.
Device-side Bloom Filter Data and Kernel Implementation `cpp/src/neighbors/detail/sample_filter_data.cuh`, `cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_impl.cuh`	Defines `bloom_filter_data_t` template holding a cuco bloom filter reference, and implements the `sample_filter_bloom_filter_impl` device kernel that checks bloom filter membership with null-data passthrough.
CAGRA Payload and Filter Pointer Wiring `cpp/src/neighbors/detail/cagra/cagra_filter_payload.hpp`	Updates `fill_cagra_sample_filter` and `cagra_filter_data_ptr` to handle bloom filters by extracting the filter_data pointer and treating bloom filters equivalently to UDF-style filters.
Kernel Instantiation and JIT Configuration `cpp/src/neighbors/detail/cagra/jit_lto_kernels/sample_filter_matrix.json`, `cpp/src/neighbors/detail/cagra/search_single_cta_inst.cu.in`, `cpp/src/neighbors/detail/cagra/search_multi_cta_inst.cu.in`, `cpp/src/neighbors/detail/cagra/search_single_cta_kernel_launcher_jit.cuh`	Adds `bloom_filter` to the JSON filter matrix and defines `bloom_filter_t` type alias in both single and multi-CTA instantiation templates. Fixes UDF filter type-id assignment to explicitly return 3.
Runtime Search Dispatch for Bloom Filters `cpp/src/neighbors/cagra.cuh`	Extends `cagra::search` with dynamic_cast-based branch detecting `bloom_filter`, clamping `filtering_rate` to `[0.0f, 0.999f]`, and dispatching to `search_with_filtering`.
Build System: CMake cuco Dependency `cpp/CMakeLists.txt`, `examples/cpp/CMakeLists.txt`	Adds cuco as a CPM-managed dependency and links `cuco::cuco` to both the core `jit_lto_kernel_usage_requirements` target and example/benchmark targets.
Example: CAGRA Bloom Filter Search `examples/cpp/src/cagra_bloom_filter_example.cu`	Complete CUDA example constructing a CAGRA index, building a cuco bloom filter for even-ID filtering, executing multi-query search with bloom filtering enabled, and validating returned neighbors using cuco's `contains_async` API.
Benchmark Program: Bitset vs Bloom Filter Comparison `examples/cpp/src/cagra_filter_benchmark.cu`	CUDA benchmark program evaluating CAGRA search latency and recall across bitset and bloom filter strategies with configurable dataset sizes, validity percentages, optional ground-truth computation, and CSV result output.
Benchmark Visualization: Python Plotting Tool `examples/cpp/plot_filter_benchmark.py`	Python script generating multi-row latency and recall grids plus overview sweeps from benchmark CSV, with grouped bars, legends, and optional recall@k visualization via `--with-recall` flag.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

Possibly related PRs

rapidsai/cuvs#2132: Establishes the UDF filter pattern that bloom filter integration reuses for payload routing in shared_launcher_jit.hpp and cagra_filter_payload.hpp.
rapidsai/cuvs#1807: Introduces the JIT-LTO CAGRA search infrastructure that this PR extends with bloom filter support in the fragment selection and routing pipeline.

Suggested labels

C++, benchmarking

Suggested reviewers

KyleFromNVIDIA
cjnolet

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 22.22% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Description check	❓ Inconclusive	No description was provided by the author; the pull request lacks any accompanying explanation of the changes.	Add a pull request description explaining the motivation, implementation approach, and key changes for bloom filter support in CAGRA.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title 'CAGRA Bloom Filter' directly summarizes the main change—adding bloom filter support to CAGRA search functionality across multiple files.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

Create PR with unit tests

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

coderabbitai

Actionable comments posted: 1

🧹 Nitpick comments (1)

examples/cpp/plot_filter_benchmark.py (1)

120-153: 💤 Low value

plot_panel_lines appears unused.

This function is defined but never called in the script. Consider removing it if not needed, or adding a comment indicating its intended future use.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@examples/cpp/plot_filter_benchmark.py` around lines 120 - 153, The function
plot_panel_lines is defined but is never called anywhere in the script. Either
remove this function entirely if it is not needed, or if it is intended for
future use, add a docstring or comment at the top of the function that clearly
explains its intended use case and why it is being retained. This will help
clarify whether the function is dead code or work-in-progress that should remain
in the codebase.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@examples/cpp/src/cagra_filter_benchmark.cu`:
- Around line 146-151: The `valid_ids_device` variable is destroyed when the if
block ends at line 151, but the `bloom.add_async()` call is asynchronous and may
still access this memory after it's freed. To fix this, either move the stream
synchronization call (currently at line 158) inside the if block before the
closing brace to ensure the async operation completes before valid_ids_device is
destroyed, or replace the `bloom.add_async()` call with the synchronous
`bloom.add()` method to avoid the race condition entirely.

---

Nitpick comments:
In `@examples/cpp/plot_filter_benchmark.py`:
- Around line 120-153: The function plot_panel_lines is defined but is never
called anywhere in the script. Either remove this function entirely if it is not
needed, or if it is intended for future use, add a docstring or comment at the
top of the function that clearly explains its intended use case and why it is
being retained. This will help clarify whether the function is dead code or
work-in-progress that should remain in the codebase.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

Push a commit to this branch (recommended)
Create a new PR with the fixes

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 232cce46-f7f4-454f-98de-3f4904ed45c9

📥 Commits

Reviewing files that changed from the base of the PR and between 21e5f20 and 16946a4.

📒 Files selected for processing (3)

examples/cpp/CMakeLists.txt
examples/cpp/plot_filter_benchmark.py
examples/cpp/src/cagra_filter_benchmark.cu

🚧 Files skipped from review as they are similar to previous changes (1)

examples/cpp/CMakeLists.txt

dantegd

One thing I noticed: cagra::merge() accepts a base_filter, but its implementation only handles Bitset. A Bloom filter would fall through to the unfiltered branch and keep every row. Should we reject Bloom there for now, like Bitmap, unless probabilistic merge filtering is meant to be supported here?

dantegd · 2026-07-13T16:50:12Z

+  {
+    if (keys.extent(0) == 0) { return; }
+
+    auto inferred_dataset_rows = static_cast<std::size_t>(keys.extent(0));


Should we separate index cardinality from insertion-batch size here? For example, if a 1,000-row index has 500 valid IDs and the caller uses the documented filtering_rate = 0.5, this treats the dataset size as 500 and allocates for only 250 insertions, even though all 500 keys are added. That invalidates the requested FPR and can reduce CAGRA recall, right? What do you think about accepting the index row count or expected insertion count explicitly?

Hmm, good catch. I think accepting the index row count is a far cleaner design and tightens the contract by not having to infer the size of the bloom filter.

dantegd · 2026-07-13T16:50:39Z

+    }
+
+    RAFT_EXPECTS(
+      inferred_dataset_rows == *configured_dataset_rows,


Is requiring every add() batch to have exactly the same length intentional? This prevents normal chunked insertion and also prevents clearing and repopulating with a different cardinality. Could sizing capacity be fixed independently while allowing arbitrary add-batch sizes?

Fixed as a consequence of above comment

dantegd · 2026-07-13T16:52:10Z

+  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);


The default CAGRA path calls this estimate on every search, so each search allocates an N-byte buffer, probes every index row, and synchronizes before ANN search begins, that could be costly, no? Could we cache an estimate invalidated by add/clear, retain insertion metadata, or use bounded sampling instead?

Also fixed by bounding to size ahead of time in the constructor. Very, very good catch.

dantegd · 2026-07-13T16:53:07Z

+          RAFT_LOG_INFO("Bloom filter index recall = %f", bloom_index_recall);
+          bloom_recalls.push_back(bloom_recall);
+        }
+        if (bloom_recalls.size() > 1) {


Should we make every Bloom run assert that each returned ID is accepted by global_bloom_filter.contains, and enforce an absolute recall floor against the exact filtered ground truth?

Currently the normal one-FPR path makes no assertions, and the 4096-block floor means the FPR sweep does not exercise different filter sizes.

This brings me to another question: should we just remove num_blocks from the constructor? I think so. We can always just infer it for the user and make the API less confusing to use.

dantegd · 2026-07-13T17:00:51Z

 */

-enum class FilterType { None, Bitmap, Bitset, UDF };
+enum class FilterType { None, Bitmap, Bitset, Bloom, UDF };


Would it make sense to append Bloom after UDF, or assign explicit values, so the existing UDF ordinal does not change? The C++ API is not the stable ABI boundary, so this is not necessarily blocking, but preserving the value is inexpensive.

Ah, let me bound this by giving them actual values also then to tighten this for the future.

dantegd · 2026-07-13T17:03:10Z

+  [[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;


Could the internal factory be made a friend instead? Exposing impl through the public API weakens the intended opacity and couples the core wrapper to neighbors internals.

I went down the friend class route, and did not like the idea of forward declaring an internal type in a public header to then mark it as friend. Good idea to just declare a friend function instead, I did not think of that.

dantegd · 2026-07-13T17:04:07Z

+
+  bloom_filter() = default;
+
+  explicit bloom_filter(const cuvs::core::bloom_filter& bloom_filter)


The referenced core filter must outlive the adapter and all searches, and must not be moved or mutated concurrently, right? The stored pointer otherwise becomes invalid or refers to a moved-from object without a diagnostic, so we should document this if I'm right.

Yes, I'll add the note but it is a problem with all our filters as we separate the implementation from the interface.

divyegala · 2026-07-13T23:21:49Z

Should we reject Bloom there for now, like Bitmap, unless probabilistic merge filtering is meant to be supported here?

@dantegd good question! Perhaps it is for yourself and Ben (I can't find his @) to answer? I'll explicitly reject it for now.

bloom filter

21e5f20

divyegala self-assigned this Jun 10, 2026

divyegala added the feature request New feature or request label Jun 10, 2026

divyegala requested review from a team as code owners June 10, 2026 23:34

divyegala added the non-breaking Introduces a non-breaking change label Jun 10, 2026

github-project-automation Bot added this to Unstructured Data Processing Jun 10, 2026

benchmarks

16946a4

coderabbitai Bot reviewed Jun 15, 2026

View reviewed changes

Comment thread examples/cpp/src/cagra_filter_benchmark.cu Outdated

cjnolet moved this to Todo in Unstructured Data Processing Jun 23, 2026

cjnolet moved this from Todo to In Progress in Unstructured Data Processing Jun 23, 2026

josephine-wolf-oberholtzer added this to Unstructured Data Processing Jul 1, 2026

josephine-wolf-oberholtzer moved this to In Progress in Unstructured Data Processing Jul 1, 2026

cjnolet reviewed Jul 1, 2026

View reviewed changes

Comment thread examples/cpp/plot_filter_benchmark.py Outdated

divyegala added 3 commits July 8, 2026 21:43

add api wrapper

a3d7b96

merge upstream

80027b5

copyright

56d959a

divyegala requested review from a team as code owners July 8, 2026 21:47

divyegala added 8 commits July 9, 2026 00:07

fix build

93db300

naming

4e6aa9c

payload management

cf40ea1

cleanup payload mgmt, estimate filtering rate

c575717

fix build

2bebdb9

Merge remote-tracking branch 'origin/main' into cagra-bloom-filter

4b33230

fix examples cmake

863e9fb

simplify payload export

0569eb7

dantegd requested changes Jul 13, 2026

View reviewed changes

divyegala added 2 commits July 14, 2026 00:21

address review

241ca5d

merge upstream

dfe8f9f


		bloom_filter() = default;

		explicit bloom_filter(const cuvs::core::bloom_filter& bloom_filter)

Uh oh!

Conversation

divyegala commented Jun 10, 2026

Uh oh!

coderabbitai Bot commented Jun 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary by CodeRabbit

Walkthrough

Changes

Estimated code review effort

Possibly related PRs

Suggested labels

Suggested reviewers

❌ Failed checks (1 warning, 1 inconclusive)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

dantegd left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

divyegala commented Jul 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

coderabbitai Bot commented Jun 10, 2026 •

edited

Loading