diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 43ae7a6235..2f8b9b0720 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -400,15 +400,11 @@ struct search_params : cuvs::neighbors::search_params { */ float persistent_device_usage = 1.0; - /** - * A parameter indicating the rate of nodes to be filtered-out, when filtering is used. - * The value must be equal to or greater than 0.0 and less than 1.0. Default value is - * negative, in which case the filtering rate is automatically calculated when possible. - * For `filtering::udf_filter`, CAGRA uses `udf_filter::filtering_rate` when this value is - * negative. If both values are negative, CAGRA assumes 0.0 because a UDF's selectivity cannot be - * inferred from the source string. - */ - float filtering_rate = -1.0; + // `filtering_rate` is inherited from `cuvs::neighbors::search_params`. CAGRA uses it to size + // `itopk_size`; supplying a non-negative value avoids a per-search popcount + host sync on the + // `bitset_filter` path. For `filtering::udf_filter`, CAGRA falls back to + // `udf_filter::filtering_rate` when this value is negative. If both are negative, CAGRA assumes + // 0.0 because a UDF's selectivity cannot be inferred from the source string. /** Data type of the query vector and codebook table on shared memory. Currently, only VPQ * supports FP8. **/ diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 935938c9b0..8a9800f4e7 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -129,7 +129,23 @@ struct index_params { float metric_arg = 2.0f; }; -struct search_params {}; +struct search_params { + /** + * A hint indicating the rate at which the sample filter is expected to filter out items + * (i.e. `(n_dataset - n_set_bits) / n_dataset` for a `bitset_filter`). + * + * Algorithms that benefit from knowing the filter's selectivity may use this hint to tune + * internal parameters or skip a hidden popcount kernel + host sync that would otherwise be + * required to derive it from the filter on every search call. + * + * - Negative (default): the algorithm auto-detects the rate from the filter when needed. + * This launches a GPU popcount reduction and synchronizes the stream per search call. + * - In `[0.0, 1.0)`: the algorithm trusts the supplied value and skips the auto-detection. + * + * Algorithms that do not use this hint (e.g. `ivf_flat`, `ivf_pq`) ignore it. + */ + float filtering_rate = -1.0; +}; /** * @brief Strategy for merging indices. diff --git a/cpp/src/neighbors/brute_force.cu b/cpp/src/neighbors/brute_force.cu index 2f9000acf7..9712a360e7 100644 --- a/cpp/src/neighbors/brute_force.cu +++ b/cpp/src/neighbors/brute_force.cu @@ -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 */ @@ -199,7 +199,7 @@ void index::update_dataset( const cuvs::neighbors::filtering::base_filter& sample_filter) \ { \ detail::search( \ - res, idx, queries, neighbors, distances, sample_filter); \ + res, params, idx, queries, neighbors, distances, sample_filter); \ } \ void search(raft::resources const& res, \ const cuvs::neighbors::brute_force::index& idx, \ @@ -209,7 +209,7 @@ void index::update_dataset( const cuvs::neighbors::filtering::base_filter& sample_filter) \ { \ detail::search( \ - res, idx, queries, neighbors, distances, sample_filter); \ + res, {}, idx, queries, neighbors, distances, sample_filter); \ } \ void search(raft::resources const& res, \ const cuvs::neighbors::brute_force::search_params& params, \ @@ -220,7 +220,7 @@ void index::update_dataset( const cuvs::neighbors::filtering::base_filter& sample_filter) \ { \ detail::search( \ - res, idx, queries, neighbors, distances, sample_filter); \ + res, params, idx, queries, neighbors, distances, sample_filter); \ } \ void search(raft::resources const& res, \ const cuvs::neighbors::brute_force::index& idx, \ @@ -230,7 +230,7 @@ void index::update_dataset( const cuvs::neighbors::filtering::base_filter& sample_filter) \ { \ detail::search( \ - res, idx, queries, neighbors, distances, sample_filter); \ + res, {}, idx, queries, neighbors, distances, sample_filter); \ } CUVS_INST_BFKNN(float, float); diff --git a/cpp/src/neighbors/detail/knn_brute_force.cuh b/cpp/src/neighbors/detail/knn_brute_force.cuh index 76ed71fe57..210150d9bd 100644 --- a/cpp/src/neighbors/detail/knn_brute_force.cuh +++ b/cpp/src/neighbors/detail/knn_brute_force.cuh @@ -734,6 +734,7 @@ void brute_force_search_gathered( template void brute_force_search_filtered( raft::resources const& res, + const cuvs::neighbors::brute_force::search_params& params, const cuvs::neighbors::brute_force::index& idx, raft::device_matrix_view queries, const cuvs::neighbors::filtering::base_filter* filter, @@ -757,6 +758,15 @@ void brute_force_search_filtered( metric == cuvs::distance::DistanceType::CosineExpanded), "Index must has norms when using Euclidean, IP, and Cosine!"); + // Negative means "auto-detect from the filter"; [0.0, 1.0) is trusted as-is. Anything else + // (including NaN, which compares false against both bounds) is out of contract. + const bool auto_filtering_rate = params.filtering_rate < 0.0f; + const bool use_hint = params.filtering_rate >= 0.0f && params.filtering_rate < 1.0f; + RAFT_EXPECTS(auto_filtering_rate || use_hint, + "search_params::filtering_rate must be negative (auto-detect) or in [0.0, 1.0), " + "got %f", + params.filtering_rate); + IdxT n_queries = queries.extent(0); IdxT n_dataset = idx.dataset().extent(0); IdxT dim = idx.dataset().extent(1); @@ -769,7 +779,10 @@ void brute_force_search_filtered( const cuvs::core::bitset_view>> filter_view; - IdxT nnz_h = 0; + IdxT nnz_h = 0; + double selectivity = 0.0; + // Whether nnz_h / n_pass came from an actual popcount rather than the filtering_rate hint. + bool counted = false; // A bitset passes the same rows for every query, so it has a row count; a bitmap passes a // different set per query and has none. std::optional n_pass; @@ -780,22 +793,51 @@ void brute_force_search_filtered( auto actual_filter = dynamic_cast*>(filter); filter_view.emplace(actual_filter->view()); - nnz_h = actual_filter->view().count(res); + if (use_hint) { + selectivity = 1.0 - params.filtering_rate; + } else { + nnz_h = actual_filter->view().count(res); + selectivity = static_cast(nnz_h) / + (static_cast(n_queries) * static_cast(n_dataset)); + counted = true; + } } else if (filter_type == cuvs::neighbors::filtering::FilterType::Bitset) { auto actual_filter = dynamic_cast*>(filter); filter_view.emplace(actual_filter->view()); - n_pass = actual_filter->view().count(res); - nnz_h = n_queries * (*n_pass); + if (use_hint) { + selectivity = 1.0 - params.filtering_rate; + // Path selection also needs a row count; the exact one is only fetched below, and only + // if the path it picks turns out to need it. + n_pass = static_cast(std::llround(selectivity * static_cast(n_dataset))); + } else { + n_pass = actual_filter->view().count(res); + nnz_h = n_queries * (*n_pass); + selectivity = static_cast(nnz_h) / + (static_cast(n_queries) * static_cast(n_dataset)); + counted = true; + } } else { RAFT_FAIL("Unsupported sample filter type"); } std::visit([&](const auto& actual_view) { filter_data = actual_view.data(); }, *filter_view); - const double selectivity = - static_cast(nnz_h) / (static_cast(n_queries) * static_cast(n_dataset)); - const auto path = select_filtered_search_path(n_dataset, dim, selectivity, k, n_pass); + auto path = select_filtered_search_path(n_dataset, dim, selectivity, k, n_pass); + + // Only the dense path can run on a hinted selectivity alone: sddmm sizes its CSR from the + // exact nnz, and gather sizes the compacted dataset from the exact row count. If we have to + // popcount after all, spend the exact numbers on the decision too. + if (!counted && path != filtered_search_path::dense) { + std::visit([&](const auto& actual_view) { nnz_h = actual_view.count(res); }, *filter_view); + if (n_pass) { + n_pass = nnz_h; + nnz_h = n_queries * (*n_pass); + } + selectivity = static_cast(nnz_h) / + (static_cast(n_queries) * static_cast(n_dataset)); + path = select_filtered_search_path(n_dataset, dim, selectivity, k, n_pass); + } if (path == filtered_search_path::gather) { auto bitset_view = std::get>(*filter_view); @@ -898,6 +940,7 @@ void brute_force_search_filtered( template void search(raft::resources const& res, + const cuvs::neighbors::brute_force::search_params& params, const cuvs::neighbors::brute_force::index& idx, raft::device_matrix_view queries, raft::device_matrix_view neighbors, @@ -918,7 +961,7 @@ void search(raft::resources const& res, dynamic_cast&>( sample_filter_ref); return brute_force_search_filtered( - res, idx, queries, &sample_filter, neighbors, distances); + res, params, idx, queries, &sample_filter, neighbors, distances); } catch (const std::bad_cast&) { } @@ -927,7 +970,7 @@ void search(raft::resources const& res, dynamic_cast&>( sample_filter_ref); return brute_force_search_filtered( - res, idx, queries, &sample_filter, neighbors, distances); + res, params, idx, queries, &sample_filter, neighbors, distances); } catch (const std::bad_cast&) { RAFT_FAIL("Unsupported sample filter type"); } diff --git a/cpp/tests/neighbors/brute_force_prefiltered.cu b/cpp/tests/neighbors/brute_force_prefiltered.cu index 1c9f9966a7..7e00f2e54b 100644 --- a/cpp/tests/neighbors/brute_force_prefiltered.cu +++ b/cpp/tests/neighbors/brute_force_prefiltered.cu @@ -520,6 +520,86 @@ class PrefilteredBruteForceOnBitmapTest true)); } + // Same as Run(), but passes the true sparsity as a filtering_rate hint to the + // params-taking search overload. Confirms results match auto-detection. + // `params.sparsity` is the density of kept entries, so the hint is its complement. + void RunWithFilteringRateHint() + { + auto dataset_raw = raft::make_device_matrix_view( + (const value_t*)dataset_d.data(), params.n_dataset, params.dim); + + auto queries = raft::make_device_matrix_view( + (const value_t*)queries_d.data(), params.n_queries, params.dim); + + auto dataset = brute_force::build(handle, dataset_raw, params.metric); + + auto filter = cuvs::core::bitmap_view( + (bitmap_t*)filter_d.data(), params.n_queries, params.n_dataset); + + auto out_val = raft::make_device_matrix_view( + out_val_d.data(), params.n_queries, params.top_k); + auto out_idx = raft::make_device_matrix_view( + out_idx_d.data(), params.n_queries, params.top_k); + + cuvs::neighbors::brute_force::search_params search_params; + search_params.filtering_rate = 1.0f - params.sparsity; + + brute_force::search(handle, + search_params, + dataset, + queries, + out_idx, + out_val, + cuvs::neighbors::filtering::bitmap_filter(filter)); + + raft::resource::sync_stream(handle); + + ASSERT_TRUE(cuvs::neighbors::devArrMatchKnnPair(out_idx_expected_d.data(), + out_idx.data_handle(), + out_val_expected_d.data(), + out_val.data_handle(), + params.n_queries, + params.top_k, + 0.001f, + stream, + true)); + } + + // A filtering_rate that is neither the negative auto-detect sentinel nor in [0.0, 1.0) is + // out of contract and must be rejected rather than silently steering path selection. + void RunWithInvalidFilteringRate() + { + auto dataset_raw = raft::make_device_matrix_view( + (const value_t*)dataset_d.data(), params.n_dataset, params.dim); + + auto queries = raft::make_device_matrix_view( + (const value_t*)queries_d.data(), params.n_queries, params.dim); + + auto dataset = brute_force::build(handle, dataset_raw, params.metric); + + auto filter = cuvs::core::bitmap_view( + (bitmap_t*)filter_d.data(), params.n_queries, params.n_dataset); + + auto out_val = raft::make_device_matrix_view( + out_val_d.data(), params.n_queries, params.top_k); + auto out_idx = raft::make_device_matrix_view( + out_idx_d.data(), params.n_queries, params.top_k); + + for (float invalid_rate : {1.0f, 1.5f, 50.0f}) { + cuvs::neighbors::brute_force::search_params search_params; + search_params.filtering_rate = invalid_rate; + + EXPECT_THROW(brute_force::search(handle, + search_params, + dataset, + queries, + out_idx, + out_val, + cuvs::neighbors::filtering::bitmap_filter(filter)), + raft::logic_error); + } + } + protected: raft::resources handle; cudaStream_t stream; @@ -941,6 +1021,51 @@ class PrefilteredBruteForceOnBitsetTest true)); } + // Same as Run(), but passes the true sparsity as a filtering_rate hint to the + // params-taking search overload. Confirms results match auto-detection. + // `params.sparsity` is the density of kept entries, so the hint is its complement. + void RunWithFilteringRateHint() + { + auto dataset_raw = raft::make_device_matrix_view( + (const value_t*)dataset_d.data(), params.n_dataset, params.dim); + + auto queries = raft::make_device_matrix_view( + (const value_t*)queries_d.data(), params.n_queries, params.dim); + + auto dataset = brute_force::build(handle, dataset_raw, params.metric); + + auto filter = + cuvs::core::bitset_view((bitset_t*)filter_d.data(), params.n_dataset); + + auto out_val = raft::make_device_matrix_view( + out_val_d.data(), params.n_queries, params.top_k); + auto out_idx = raft::make_device_matrix_view( + out_idx_d.data(), params.n_queries, params.top_k); + + cuvs::neighbors::brute_force::search_params search_params; + search_params.filtering_rate = 1.0f - params.sparsity; + + brute_force::search(handle, + search_params, + dataset, + queries, + out_idx, + out_val, + cuvs::neighbors::filtering::bitset_filter(filter)); + + raft::resource::sync_stream(handle); + + ASSERT_TRUE(cuvs::neighbors::devArrMatchKnnPair(out_idx_expected_d.data(), + out_idx.data_handle(), + out_val_expected_d.data(), + out_val.data_handle(), + params.n_queries, + params.top_k, + 0.001f, + stream, + true)); + } + protected: raft::resources handle; cudaStream_t stream; @@ -963,18 +1088,38 @@ class PrefilteredBruteForceOnBitsetTest using PrefilteredBruteForceTestOnBitmap_float_int64 = PrefilteredBruteForceOnBitmapTest; TEST_P(PrefilteredBruteForceTestOnBitmap_float_int64, Result) { Run(); } +TEST_P(PrefilteredBruteForceTestOnBitmap_float_int64, ResultWithFilteringRateHint) +{ + RunWithFilteringRateHint(); +} +TEST_P(PrefilteredBruteForceTestOnBitmap_float_int64, InvalidFilteringRateThrows) +{ + RunWithInvalidFilteringRate(); +} using PrefilteredBruteForceTestOnBitmap_half_int64 = PrefilteredBruteForceOnBitmapTest; TEST_P(PrefilteredBruteForceTestOnBitmap_half_int64, Result) { Run(); } +TEST_P(PrefilteredBruteForceTestOnBitmap_half_int64, ResultWithFilteringRateHint) +{ + RunWithFilteringRateHint(); +} using PrefilteredBruteForceTestOnBitset_float_int64 = PrefilteredBruteForceOnBitsetTest; TEST_P(PrefilteredBruteForceTestOnBitset_float_int64, Result) { Run(); } +TEST_P(PrefilteredBruteForceTestOnBitset_float_int64, ResultWithFilteringRateHint) +{ + RunWithFilteringRateHint(); +} using PrefilteredBruteForceTestOnBitset_half_int64 = PrefilteredBruteForceOnBitsetTest; TEST_P(PrefilteredBruteForceTestOnBitset_half_int64, Result) { Run(); } +TEST_P(PrefilteredBruteForceTestOnBitset_half_int64, ResultWithFilteringRateHint) +{ + RunWithFilteringRateHint(); +} template const std::vector> selectk_inputs = {