diff --git a/cpp/include/cuvs/neighbors/brute_force.hpp b/cpp/include/cuvs/neighbors/brute_force.hpp index d30ea21320..b4addb22dd 100644 --- a/cpp/include/cuvs/neighbors/brute_force.hpp +++ b/cpp/include/cuvs/neighbors/brute_force.hpp @@ -21,7 +21,22 @@ namespace brute_force { struct index_params : cuvs::neighbors::index_params {}; -struct search_params : cuvs::neighbors::search_params {}; +/** @brief Parameters for brute-force search. */ +struct search_params : cuvs::neighbors::search_params { + /** + * A hint for the fraction of entries the filter removes, i.e. + * `(n_queries * n_dataset - n_set_bits) / (n_queries * n_dataset)`; for a `bitset_filter` this is + * `(n_dataset - n_set_bits) / n_dataset`. + * + * - Negative (default): the rate is computed from the filter, which runs a popcount kernel and + * synchronizes the stream on every filtered search. + * - In `[0.0, 1.0)`: the value is trusted and the popcount is skipped when the dense path is + * selected. The sparse and gather paths still count the filter, because they need the exact + * number of passing entries. + * - Any other value (including NaN) is rejected. + */ + float filtering_rate = -1.0; +}; /** * @defgroup bruteforce_cpp_index Bruteforce index diff --git a/cpp/src/neighbors/brute_force.cu b/cpp/src/neighbors/brute_force.cu index e4ee656dff..f664968243 100644 --- a/cpp/src/neighbors/brute_force.cu +++ b/cpp/src/neighbors/brute_force.cu @@ -215,7 +215,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, \ @@ -225,7 +225,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, \ @@ -236,7 +236,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, \ @@ -246,7 +246,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 cac660201f..016a2ed9ab 100644 --- a/cpp/src/neighbors/detail/knn_brute_force.cuh +++ b/cpp/src/neighbors/detail/knn_brute_force.cuh @@ -736,6 +736,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, @@ -759,6 +760,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); @@ -771,7 +781,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; @@ -782,22 +795,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); @@ -906,6 +948,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, @@ -926,7 +969,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&) { } @@ -935,7 +978,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 dfe27e3978..f02b0d9e5c 100644 --- a/cpp/tests/neighbors/brute_force_prefiltered.cu +++ b/cpp/tests/neighbors/brute_force_prefiltered.cu @@ -567,6 +567,88 @@ class PrefilteredBruteForceOnBitmapTest true)); } + // Same as Run(), but passes the filter's true filtering_rate as a hint to the params-taking + // search overload. Confirms results match auto-detection. The rate comes from the measured + // `nnz`, not `params.sparsity`: RMAT can draw the same edge twice, so fewer bits may be set. + 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 = + static_cast(1.0 - static_cast(nnz) / (static_cast(params.n_queries) * + static_cast(params.n_dataset))); + + 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, std::numeric_limits::quiet_NaN()}) { + 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; @@ -948,6 +1030,53 @@ class PrefilteredBruteForceOnBitsetTest true)); } + // Same as Run(), but passes the filter's true filtering_rate as a hint to the params-taking + // search overload. Confirms results match auto-detection. The rate comes from the measured + // `nnz`, not `params.sparsity`: RMAT can draw the same edge twice, so fewer bits may be set. + 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 = + static_cast(1.0 - static_cast(nnz) / (static_cast(params.n_queries) * + static_cast(params.n_dataset))); + + 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; @@ -970,18 +1099,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 = { diff --git a/fern/pages/cpp_api/cpp-api-neighbors-brute-force.md b/fern/pages/cpp_api/cpp-api-neighbors-brute-force.md index b5bdd6cc78..fdce39b85b 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-brute-force.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-brute-force.md @@ -6,6 +6,25 @@ slug: api-reference/cpp-api-neighbors-brute-force _Source header: `cuvs/neighbors/brute_force.hpp`_ +## Types + + +### neighbors::brute_force::search_params + +Parameters for brute-force search. + +```cpp +struct search_params : cuvs::neighbors::search_params { + float filtering_rate; +}; +``` + +**Fields** + +| Name | Type | Description | +| --- | --- | --- | +| `filtering_rate` | `float` | A hint for the fraction of entries the filter removes, i.e. `(n_queries * n_dataset - n_set_bits) / (n_queries * n_dataset)`; for a `bitset_filter` this is `(n_dataset - n_set_bits) / n_dataset`.

- Negative (default): the rate is computed from the filter, which runs a popcount kernel and synchronizes the stream on every filtered search.
- In `[0.0, 1.0)`: the value is trusted and the popcount is skipped when the dense path is selected. The sparse and gather paths still count the filter, because they need the exact number of passing entries.
- Any other value (including NaN) is rejected. | + ## Bruteforce index diff --git a/fern/pages/neighbors/bruteforce.md b/fern/pages/neighbors/bruteforce.md index 1219288576..f76c96ff2e 100644 --- a/fern/pages/neighbors/bruteforce.md +++ b/fern/pages/neighbors/bruteforce.md @@ -542,7 +542,9 @@ SearchResults results = index.search(query); ### Search parameters -The C++ `brute_force::search_params` struct currently has no tunable fields. +| Name | Default | Description | +| --- | --- | --- | +| `filtering_rate` | `-1.0` | C++ only. Expected fraction of query-vector pairs removed by the filter. Negative values make filtered search count the filter on every call, which launches a GPU reduction and synchronizes the stream. A value in `[0.0, 1.0)` skips that count when the dense path is selected. Other values are rejected. | Filters are passed as search function arguments in bindings that expose filtered search, not as fields in `brute_force::search_params`.