Merge branch 'main' of https://github.com/rapidsai/cuvs into cuvs-bench-fixup

Signed-off-by: James Bourbeau <jbourbeau@nvidia.com>

Author: jrbourbeau

c/include/cuvs/neighbors/cagra.h

@@ -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
  */
 
@@ -299,6 +299,27 @@ CUVS_EXPORT cuvsError_t cuvsCagraIndexParamsFromHnswParams(cuvsCagraIndexParams_
                                                 enum cuvsCagraHnswHeuristicType heuristic,
                                                 cuvsDistanceType metric);
 
+/**
+ * @brief Create CAGRA index parameters heuristically tuned for a dataset
+ *
+ * This factory function selects the graph build algorithm and its parameters based on the shape of
+ * the dataset.
+ *
+ * @param[out] params The CAGRA index params to populate
+ * @param[in] n_rows Number of rows in the dataset
+ * @param[in] dim Number of dimensions in the dataset
+ * @param[in] graph_degree Degree of the output graph
+ * @param[in] metric Distance metric to use
+ * @param[in] build_quality Higher values increase build quality (and cost) up to a point
+ * @return cuvsError_t
+ */
+CUVS_EXPORT cuvsError_t cuvsCagraIndexParamsFromDataset(cuvsCagraIndexParams_t params,
+                                                int64_t n_rows,
+                                                int64_t dim,
+                                                size_t graph_degree,
+                                                cuvsDistanceType metric,
+                                                size_t build_quality);
+
 /**
  * @}
  */

c/src/neighbors/cagra.cpp

@@ -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
  */
 
@@ -834,6 +834,22 @@ extern "C" cuvsError_t cuvsCagraIndexParamsFromHnswParams(cuvsCagraIndexParams_t
   });
 }
 
+extern "C" cuvsError_t cuvsCagraIndexParamsFromDataset(cuvsCagraIndexParams_t params,
+                                                       int64_t n_rows,
+                                                       int64_t dim,
+                                                       size_t graph_degree,
+                                                       cuvsDistanceType metric,
+                                                       size_t build_quality)
+{
+  return cuvs::core::translate_exceptions([=] {
+    auto cpp_metric = static_cast<cuvs::distance::DistanceType>((int)metric);
+    auto cpp_params = cuvs::neighbors::cagra::index_params::from_dataset(
+      raft::matrix_extent<int64_t>(n_rows, dim), graph_degree, cpp_metric, build_quality);
+
+    _populate_cagra_index_params_from_cpp(params, cpp_params);
+  });
+}
+
 extern "C" cuvsError_t cuvsCagraExtendParamsCreate(cuvsCagraExtendParams_t* params)
 {
   return cuvs::core::translate_exceptions(

cpp/CMakeLists.txt

@@ -365,7 +365,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)
@@ -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

cpp/include/cuvs/core/bloom_filter.hpp

@@ -0,0 +1,99 @@
+/*
+ * 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.
+   *
+   * @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)`
+   * - The default policy uses 256-bit blocks split into eight 32-bit words and sets one bit in each
+   *   word per inserted key.
+   * - For each candidate block count, the expected false-positive rate accounts for the binomial
+   *   distribution of inserted keys across blocks and the fixed eight-bit fingerprint.
+   * - The smallest block count whose expected false-positive rate meets
+   *   @p target_false_positive_rate is selected.
+   *
+   * 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<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;
+
+  /**
+   * @brief Return the estimated fraction of dataset rows rejected by this filter.
+   *
+   * The estimate is derived at construction from the configured valid-row fraction and the
+   * expected false-positive rate of the selected filter geometry. It performs no device work.
+   */
+  [[nodiscard]] float estimate_filtering_rate() const noexcept;
+
+ private:
+  friend impl const& get_bloom_filter_impl(bloom_filter const& filter) noexcept;
+
+  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/cagra.hpp

@@ -112,6 +112,12 @@ namespace graph_build_params = cuvs::neighbors::graph_build_params;
  * @{
  */
 
+using graph_build_params_t = std::variant<std::monostate,
+                                          graph_build_params::ivf_pq_params,
+                                          graph_build_params::nn_descent_params,
+                                          graph_build_params::ace_params,
+                                          graph_build_params::iterative_search_params>;
+
 /**
  * @brief A strategy for selecting the graph build parameters based on similar HNSW index
  * parameters.
@@ -181,12 +187,7 @@ struct index_params : cuvs::neighbors::index_params {
    * cagra::graph_build_params::iterative_search_params();
    * @endcode
    */
-  std::variant<std::monostate,
-               graph_build_params::ivf_pq_params,
-               graph_build_params::nn_descent_params,
-               graph_build_params::ace_params,
-               graph_build_params::iterative_search_params>
-    graph_build_params;
+  graph_build_params_t graph_build_params;
   /**
    * Whether to use MST optimization to guarantee graph connectivity.
    */
@@ -223,6 +224,55 @@ struct index_params : cuvs::neighbors::index_params {
    */
   bool attach_dataset_on_build = true;
 
+  /**
+   * @brief Select the graph build algorithm and its parameters for a dataset.
+   *
+   * This is the main CAGRA build heuristic: it chooses between NN-descent and IVF-PQ based on the
+   * dataset size and tunes their parameters based on the target intermediate graph degree and the
+   * requested build quality. It returns the `graph_build_params` variant only; the caller is
+   * responsible for setting `graph_degree` / `intermediate_graph_degree`.
+   *
+   * @param dataset The shape of the input dataset
+   * @param intermediate_graph_degree The intermediate (kNN) graph degree the build should target.
+   *  Note: the intermediate graph degree must be not smaller than the output graph degree; a good
+   *        practice is to have it 1.5x to 2x of the desired graph_degree and a multiple of 32.
+   * @param metric The distance metric to search
+   * @param build_quality Higher values increase the build quality (and cost) up to a point.
+   *        Any value is valid, but values below 20 are the most practical (default = 7).
+   */
+  static graph_build_params_t graph_build_heuristic(
+    raft::matrix_extent<int64_t> dataset,
+    size_t intermediate_graph_degree,
+    cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded,
+    size_t build_quality                = 7);
+
+  /**
+   * @brief Create CAGRA index parameters heuristically tuned for a dataset.
+   *
+   * Returns default CAGRA `index_params` with `graph_build_params` selected by
+   * `graph_build_heuristic` for the given dataset.
+   *
+   * @param dataset The shape of the input dataset
+   * @param graph_degree Degree of the output graph.
+   * @param metric The distance metric to search
+   * @param build_quality Higher values increase the build quality (and cost) up to a point.
+   *        Any value is valid, but values below 20 are the most practical (default = 7).
+   *
+   * Usage example:
+   * @code{.cpp}
+   *   using namespace cuvs::neighbors;
+   *   raft::resources res;
+   *   auto dataset = raft::make_device_matrix<float, int64_t>(res, N, D);
+   *   auto cagra_params = cagra::index_params::from_dataset(dataset.extents());
+   *   auto cagra_index = cagra::build(res, cagra_params, dataset);
+   * @endcode
+   */
+  static cagra::index_params from_dataset(
+    raft::matrix_extent<int64_t> dataset,
+    size_t graph_degree                 = 64,
+    cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded,
+    size_t build_quality                = 7);
+
   /**
    * @brief Create a CAGRA index parameters compatible with HNSW index
    *

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