From 233d5b5cabeb8a78ef1e74e26c07edae9408b3e2 Mon Sep 17 00:00:00 2001 From: Max Buckley Date: Tue, 14 Jul 2026 21:05:13 +0200 Subject: [PATCH 1/4] perf(brute_force): add scatter-gather path and fix the filtered-search dispatch thresholds Filtered brute-force search picks between a sparse (SDDMM) path and a dense post-filter path on `sparsity < 0.9`, i.e. it runs SDDMM whenever <=10% of the dataset passes the filter. Two things are wrong with that. First, the threshold is in the wrong place. SDDMM's cost grows with the number of passing rows while the dense path's is flat, so across a sweep of selectivity x n_dataset x dim (fp16, k=64, 100 queries, RTX 5090), SDDMM is already 2.2-4.4x slower than dense by the time it hands off at 10%. The real SDDMM/dense crossover is ~2.5-4%. Second, and more importantly, the threshold is the wrong *kind*. SDDMM stops paying off at an absolute number of passing rows -- roughly constant in n_dataset -- not at a fraction of the dataset, because it competes against a fixed setup cost rather than against work that scales with n_dataset. Expressed as a fraction, the correct handoff therefore falls as n_dataset grows: at n_dataset=10M it is ~0.11%, where the current rule waits until 10%. The error grows with the dataset. This adds a third path for the band the other two leave uncovered. For a bitset filter (which selects the same rows for every query) the passing rows are enumerated, compacted into a dense [n_pass x dim] matrix along with their precomputed norms, searched with the ordinary unfiltered kernel, and the resulting indices mapped back. Its cost scales with n_pass rather than n_dataset, so it wins over most of the 0.1-10% band that SDDMM currently owns. Measured against the best of the three paths, the current dispatch is up to 12.6x slower than it needs to be (n_dataset=10M, dim=128, 10% selectivity). A bitmap filter selects different rows per query, so there is no single set of rows to compact; bitmaps keep a two-way dispatch, but with the corrected SDDMM/dense threshold rather than the flat 10%. Co-Authored-By: Claude Opus 4.8 --- cpp/src/neighbors/detail/knn_brute_force.cuh | 221 +++++++++++++++++- .../neighbors/brute_force_prefiltered.cu | 10 + 2 files changed, 224 insertions(+), 7 deletions(-) diff --git a/cpp/src/neighbors/detail/knn_brute_force.cuh b/cpp/src/neighbors/detail/knn_brute_force.cuh index 34ef8dd937..ed062a43a7 100644 --- a/cpp/src/neighbors/detail/knn_brute_force.cuh +++ b/cpp/src/neighbors/detail/knn_brute_force.cuh @@ -28,6 +28,7 @@ #include #include #include +#include #include #include #include @@ -41,8 +42,13 @@ #include #include #include +#include #include +#include +#include +#include +#include #include #include #include @@ -581,6 +587,196 @@ void brute_force_search( query_norms ? query_norms->data_handle() : nullptr); } +/** + * @brief The three strategies available for a filtered brute-force search. + * + * They differ in what work they do per row that passes the filter: + * + * - `sddmm`: build a CSR of the passing (query, row) pairs and evaluate only those + * distances via a masked matmul. Cost grows with the number of passing + * pairs, so it is the cheapest option only while very few rows pass. + * - `gather`: compact the passing rows into a dense [n_pass x dim] matrix and run an + * ordinary unfiltered search over it. Cost grows with n_pass, on top of a + * fixed setup cost (enumerating the rows, allocating, launching). + * - `dense`: GEMM against the whole dataset and mask the filtered rows afterwards. + * Cost is independent of how many rows pass. + */ +enum class filtered_search_path { sddmm, gather, dense }; + +/** + * @brief Pick the cheapest search path for a filtered brute-force query. + * + * The thresholds below are empirical. They were measured on an RTX 5090 (sm_120) with + * fp16 data, k=64 and 100 queries, sweeping selectivity over 0.1-60% for n_dataset in + * [2e5, 1e7] and dim in [128, 1024]. Both the shapes of the fits and the constants come + * from that sweep; see cpp/bench/prims/core/results/bf_sweep/README.md. + * + * Two facts drive the choice: + * + * 1. SDDMM stops paying off at an *absolute* number of passing rows, not at a fraction + * of the dataset. Its cost is proportional to the passing pairs it evaluates, while + * the gather path pays a fixed setup cost before doing any work; SDDMM wins exactly + * while its work still fits under that fixed cost. The crossover is therefore roughly + * constant in n_dataset, and falls off as dim^-0.4 because wider rows make each + * passing pair more expensive. + * + * 2. The gather path reads n_pass rows scattered across memory (traffic proportional to + * dim) to avoid a GEMM proportional to n_dataset. So its crossover with the dense + * path falls off as dim^-0.6, and rises with n_dataset as the fixed setup cost is + * amortized over more saved work. Below roughly `kGatherSetupEquivRows` rows the + * setup alone costs more than the entire dense search, and gathering never wins. + * + * @param[in] n_dataset rows in the dataset + * @param[in] dim columns in the dataset + * @param[in] n_pass rows passing the filter (for a bitmap, averaged over queries) + * @param[in] k neighbors requested + * @param[in] gather_possible whether the gather path is applicable at all (see below) + */ +inline filtered_search_path select_filtered_search_path( + int64_t n_dataset, int64_t dim, int64_t n_pass, int64_t k, bool gather_possible) +{ + // SDDMM vs gather: an absolute passing-row count, ~constant in n_dataset (fact 1). + constexpr double kSddmmRowsAtDim128 = 11000.0; + constexpr double kSddmmDimExponent = 0.4; + // Gather vs dense: a selectivity, falling with dim and rising with n_dataset (fact 2). + constexpr double kGatherSelNumerator = 9.3; + constexpr double kGatherDimExponent = 0.6; + constexpr double kGatherSetupEquivRows = 212500.0; + // The fit above was calibrated over dim in [128, 1024]. Extrapolated below that range it + // exceeds 1, i.e. it would claim gathering wins at every selectivity -- which cannot be + // true, since gathering every row and then running the same GEMM over it is strictly more + // work than the dense path. Cap at the widest crossover actually observed. + constexpr double kGatherMaxSelectivity = 0.60; + // SDDMM vs dense, used when the gather path is unavailable. Both costs scale with + // n_dataset, so unlike the above this crossover is a selectivity rather than a count. + constexpr double kSddmmDenseSelNumerator = 0.115; + constexpr double kSddmmDenseDimExponent = 0.22; + + const auto n_pass_d = static_cast(n_pass); + const auto n_dataset_d = static_cast(n_dataset); + const double selectivity = n_pass_d / n_dataset_d; + + // Evaluate the fits only over the range they were calibrated on. Outside it the power + // laws are unconstrained, so hold them at the nearest measured boundary rather than + // extrapolating -- brute force is routinely run at dim well below 128. + const double dim_d = std::clamp(static_cast(dim), 128.0, 1024.0); + + const double sddmm_vs_dense_sel = + kSddmmDenseSelNumerator / std::pow(dim_d, kSddmmDenseDimExponent); + + if (!gather_possible) { + return selectivity < sddmm_vs_dense_sel ? filtered_search_path::sddmm + : filtered_search_path::dense; + } + + // SDDMM has to beat both of the others to be worth running. + const double sddmm_max_rows = kSddmmRowsAtDim128 * std::pow(128.0 / dim_d, kSddmmDimExponent); + if (n_pass_d < sddmm_max_rows && selectivity < sddmm_vs_dense_sel) { + return filtered_search_path::sddmm; + } + + // The fraction of the dense search that survives paying the gather setup cost. At + // small n_dataset this goes to zero (or negative) and gathering can never win. + const double amortized = 1.0 - kGatherSetupEquivRows / n_dataset_d; + const double gather_vs_dense_sel = + amortized * + std::min(kGatherSelNumerator / std::pow(dim_d, kGatherDimExponent), kGatherMaxSelectivity); + + // n_pass > k keeps the compacted dataset big enough to actually contain k neighbors. + if (amortized > 0.0 && selectivity < gather_vs_dense_sel && n_pass > k) { + return filtered_search_path::gather; + } + return filtered_search_path::dense; +} + +/** + * @brief Filtered search over the rows a bitset selects, by compacting them first. + * + * Enumerates the passing rows, gathers them (and their precomputed norms) into a dense + * matrix, runs an ordinary unfiltered brute-force search over that matrix, and maps the + * resulting indices back into the original dataset's numbering. + * + * Only meaningful for a bitset filter, which selects the same rows for every query. A + * bitmap filter selects different rows per query, so there is no single set of rows to + * compact. + */ +template +void brute_force_search_gathered( + raft::resources const& res, + const cuvs::neighbors::brute_force::index& idx, + raft::device_matrix_view queries, + const cuvs::core::bitset_view& filter_view, + IdxT n_pass, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + std::optional> query_norms = std::nullopt) +{ + auto stream = raft::resource::get_cuda_stream(res); + auto policy = raft::resource::get_thrust_policy(res); + IdxT n_queries = queries.extent(0); + IdxT n_dataset = idx.dataset().extent(0); + IdxT dim = idx.dataset().extent(1); + IdxT k = neighbors.extent(1); + + // 1. enumerate the rows the filter keeps + rmm::device_uvector passing(n_pass, stream); + thrust::copy_if(policy, + thrust::make_counting_iterator(0), + thrust::make_counting_iterator(n_dataset), + passing.data(), + [filter_view] __device__(IdxT i) { return filter_view.test(i); }); + + // 2. compact those rows, and their norms, into a dense dataset + rmm::device_uvector gathered(static_cast(n_pass) * dim, stream); + raft::matrix::gather( + res, + raft::make_device_matrix_view( + idx.dataset().data_handle(), n_dataset, dim), + raft::make_device_vector_view(passing.data(), n_pass), + raft::make_device_matrix_view(gathered.data(), n_pass, dim)); + + rmm::device_uvector gathered_norms(idx.has_norms() ? n_pass : 0, stream); + if (idx.has_norms()) { + thrust::gather(policy, + passing.data(), + passing.data() + n_pass, + idx.norms().data_handle(), + gathered_norms.data()); + } + + // 3. ordinary unfiltered search over the compacted dataset + rmm::device_uvector compact_neighbors(static_cast(n_queries) * k, stream); + std::vector dataset = {gathered.data()}; + std::vector sizes = {static_cast(n_pass)}; + std::vector norms; + if (idx.has_norms()) { norms.push_back(gathered_norms.data()); } + + brute_force_knn_impl( + res, + dataset, + sizes, + dim, + const_cast(queries.data_handle()), + n_queries, + compact_neighbors.data(), + distances.data_handle(), + k, + true, + true, + nullptr, + idx.metric(), + idx.metric_arg(), + norms.size() ? &norms : nullptr, + query_norms ? query_norms->data_handle() : nullptr); + + // 4. translate compacted row numbers back into the original dataset's numbering + thrust::gather(policy, + compact_neighbors.data(), + compact_neighbors.data() + static_cast(n_queries) * k, + passing.data(), + neighbors.data_handle()); +} + template void brute_force_search_filtered( raft::resources const& res, @@ -619,8 +815,12 @@ void brute_force_search_filtered( const cuvs::core::bitset_view>> filter_view; - IdxT nnz_h = 0; - float sparsity = 0.0f; + IdxT nnz_h = 0; + // Rows passing the filter. A bitset selects the same rows for every query; a bitmap + // selects different rows per query, so there this is the per-query average. + IdxT n_pass = 0; + // A bitmap's rows differ per query, so there is no single set of rows to compact. + bool gather_possible = false; const BitsT* filter_data = nullptr; @@ -628,21 +828,28 @@ void brute_force_search_filtered( auto actual_filter = dynamic_cast*>(filter); filter_view.emplace(actual_filter->view()); - nnz_h = actual_filter->view().count(res); - sparsity = 1.0 - nnz_h / (1.0 * n_queries * n_dataset); + nnz_h = actual_filter->view().count(res); + n_pass = nnz_h / n_queries; } else if (filter_type == cuvs::neighbors::filtering::FilterType::Bitset) { auto actual_filter = dynamic_cast*>(filter); filter_view.emplace(actual_filter->view()); - nnz_h = n_queries * actual_filter->view().count(res); - sparsity = 1.0 - nnz_h / (1.0 * n_queries * n_dataset); + n_pass = actual_filter->view().count(res); + nnz_h = n_queries * n_pass; + gather_possible = true; } else { RAFT_FAIL("Unsupported sample filter type"); } std::visit([&](const auto& actual_view) { filter_data = actual_view.data(); }, *filter_view); - if (sparsity < 0.9f) { + const auto path = select_filtered_search_path(n_dataset, dim, n_pass, k, gather_possible); + + if (path == filtered_search_path::gather) { + auto bitset_view = std::get>(*filter_view); + brute_force_search_gathered( + res, idx, queries, bitset_view, n_pass, neighbors, distances, query_norms); + } else if (path == filtered_search_path::dense) { raft::resources stream_pool_handle(res); raft::resource::set_cuda_stream(stream_pool_handle, stream); auto idx_norm = idx.has_norms() ? const_cast(idx.norms().data_handle()) : nullptr; diff --git a/cpp/tests/neighbors/brute_force_prefiltered.cu b/cpp/tests/neighbors/brute_force_prefiltered.cu index c15f08363f..c090dbc40d 100644 --- a/cpp/tests/neighbors/brute_force_prefiltered.cu +++ b/cpp/tests/neighbors/brute_force_prefiltered.cu @@ -978,6 +978,16 @@ TEST_P(PrefilteredBruteForceTestOnBitset_half_int64, Result) { Run(); } template const std::vector> selectk_inputs = { + // Large enough (n_dataset, dim) to reach each of the three dispatch paths in + // select_filtered_search_path. With a bitset filter these land on, respectively: + // sddmm (300 rows pass), gather (24000 pass), and dense (90000 pass). + {4, 300000, 128, 16, 0.001, cuvs::distance::DistanceType::L2Expanded}, + {4, 300000, 128, 16, 0.08, cuvs::distance::DistanceType::L2Expanded}, + {4, 300000, 128, 16, 0.08, cuvs::distance::DistanceType::InnerProduct}, + {4, 300000, 128, 16, 0.08, cuvs::distance::DistanceType::L2SqrtExpanded}, + {4, 300000, 128, 16, 0.08, cuvs::distance::DistanceType::CosineExpanded}, + {4, 300000, 128, 16, 0.30, cuvs::distance::DistanceType::L2Expanded}, + {8, 131072, 255, 255, 0.01, cuvs::distance::DistanceType::L2Expanded}, {8, 131072, 255, 255, 0.01, cuvs::distance::DistanceType::InnerProduct}, {8, 131072, 255, 255, 0.01, cuvs::distance::DistanceType::L2SqrtExpanded}, From 20c5d5266f15fac95680a1e3584cd4b1b2e7a2fa Mon Sep 17 00:00:00 2001 From: Max Buckley Date: Tue, 14 Jul 2026 21:16:48 +0200 Subject: [PATCH 2/4] perf(brute_force): simplify dispatch thresholds to round, hardware-tolerant constants The previous constants were power-law fits to a single GPU (dim^0.4, dim^0.6, dim^0.22, 9.3, 212500, 11000). That is false precision: the exponents are not derivable from a cost model -- solving one against the measured crossovers yields a negative gather-traffic coefficient -- so they are curve-fits, and three significant figures of curve-fit do not transfer to other hardware. Replace them with four round constants. Over the measured sweep this is not a regression: mean penalty against a perfect oracle is 1.01x either way, and worst case improves slightly (1.52x -> 1.49x). All the extra precision bought nothing. The simpler rule also degrades gracefully. Perturbing any constant by 2x, which is roughly the spread expected across GPUs, keeps the mean penalty at 1.01-1.05x and the worst case at or below 2.7x -- still far better than the 1.65x mean and 12.6x worst case of the `sparsity < 0.9` rule it replaces. Also drops the commentary explaining the fits, which no longer exist. Co-Authored-By: Claude Opus 4.8 --- cpp/src/neighbors/detail/knn_brute_force.cuh | 121 ++++++------------- 1 file changed, 38 insertions(+), 83 deletions(-) diff --git a/cpp/src/neighbors/detail/knn_brute_force.cuh b/cpp/src/neighbors/detail/knn_brute_force.cuh index ed062a43a7..cc9b4b248c 100644 --- a/cpp/src/neighbors/detail/knn_brute_force.cuh +++ b/cpp/src/neighbors/detail/knn_brute_force.cuh @@ -590,100 +590,62 @@ void brute_force_search( /** * @brief The three strategies available for a filtered brute-force search. * - * They differ in what work they do per row that passes the filter: - * - * - `sddmm`: build a CSR of the passing (query, row) pairs and evaluate only those - * distances via a masked matmul. Cost grows with the number of passing - * pairs, so it is the cheapest option only while very few rows pass. - * - `gather`: compact the passing rows into a dense [n_pass x dim] matrix and run an - * ordinary unfiltered search over it. Cost grows with n_pass, on top of a - * fixed setup cost (enumerating the rows, allocating, launching). - * - `dense`: GEMM against the whole dataset and mask the filtered rows afterwards. - * Cost is independent of how many rows pass. + * - `sddmm`: evaluate only the passing (query, row) pairs via a masked matmul. Cost + * grows with the number of passing rows. + * - `gather`: compact the passing rows into a dense [n_pass x dim] matrix and search + * that. Cost grows with the passing fraction, over a fixed setup cost. + * - `dense`: GEMM against the whole dataset, then mask. Cost is independent of the + * filter. */ enum class filtered_search_path { sddmm, gather, dense }; /** * @brief Pick the cheapest search path for a filtered brute-force query. * - * The thresholds below are empirical. They were measured on an RTX 5090 (sm_120) with - * fp16 data, k=64 and 100 queries, sweeping selectivity over 0.1-60% for n_dataset in - * [2e5, 1e7] and dim in [128, 1024]. Both the shapes of the fits and the constants come - * from that sweep; see cpp/bench/prims/core/results/bf_sweep/README.md. - * - * Two facts drive the choice: - * - * 1. SDDMM stops paying off at an *absolute* number of passing rows, not at a fraction - * of the dataset. Its cost is proportional to the passing pairs it evaluates, while - * the gather path pays a fixed setup cost before doing any work; SDDMM wins exactly - * while its work still fits under that fixed cost. The crossover is therefore roughly - * constant in n_dataset, and falls off as dim^-0.4 because wider rows make each - * passing pair more expensive. - * - * 2. The gather path reads n_pass rows scattered across memory (traffic proportional to - * dim) to avoid a GEMM proportional to n_dataset. So its crossover with the dense - * path falls off as dim^-0.6, and rises with n_dataset as the fixed setup cost is - * amortized over more saved work. Below roughly `kGatherSetupEquivRows` rows the - * setup alone costs more than the entire dense search, and gathering never wins. + * The constants are empirical (fp16, RTX 5090; see + * cpp/bench/prims/core/results/bf_sweep/README.md) and are deliberately round: over the + * measured sweep, perturbing any of them by 2x -- roughly the spread expected across + * GPUs -- leaves the mean penalty against a perfect oracle at 1.01-1.05x, versus 1.65x + * mean and 12.6x worst case for the old `sparsity < 0.9` rule. * * @param[in] n_dataset rows in the dataset * @param[in] dim columns in the dataset * @param[in] n_pass rows passing the filter (for a bitmap, averaged over queries) * @param[in] k neighbors requested - * @param[in] gather_possible whether the gather path is applicable at all (see below) + * @param[in] gather_possible whether the gather path applies (bitset filters only) */ inline filtered_search_path select_filtered_search_path( int64_t n_dataset, int64_t dim, int64_t n_pass, int64_t k, bool gather_possible) { - // SDDMM vs gather: an absolute passing-row count, ~constant in n_dataset (fact 1). - constexpr double kSddmmRowsAtDim128 = 11000.0; - constexpr double kSddmmDimExponent = 0.4; - // Gather vs dense: a selectivity, falling with dim and rising with n_dataset (fact 2). - constexpr double kGatherSelNumerator = 9.3; - constexpr double kGatherDimExponent = 0.6; - constexpr double kGatherSetupEquivRows = 212500.0; - // The fit above was calibrated over dim in [128, 1024]. Extrapolated below that range it - // exceeds 1, i.e. it would claim gathering wins at every selectivity -- which cannot be - // true, since gathering every row and then running the same GEMM over it is strictly more - // work than the dense path. Cap at the widest crossover actually observed. - constexpr double kGatherMaxSelectivity = 0.60; - // SDDMM vs dense, used when the gather path is unavailable. Both costs scale with - // n_dataset, so unlike the above this crossover is a selectivity rather than a count. - constexpr double kSddmmDenseSelNumerator = 0.115; - constexpr double kSddmmDenseDimExponent = 0.22; - - const auto n_pass_d = static_cast(n_pass); - const auto n_dataset_d = static_cast(n_dataset); - const double selectivity = n_pass_d / n_dataset_d; + // SDDMM competes against a fixed setup cost, so it stops paying off at an absolute row + // count rather than at a fraction of the dataset. + constexpr int64_t kSddmmMaxPassingRows = 10000; + constexpr double kSddmmMaxSelectivity = 0.03; + // Gathering trades memory traffic (proportional to dim) against a GEMM (proportional to + // n_dataset), and cannot beat simply GEMMing everything once most rows pass. + constexpr double kGatherSelectivityScale = 5.5; + constexpr double kGatherMaxSelectivity = 0.5; + // Below this, the gather setup costs more than the entire dense search. + constexpr double kGatherMinDataset = 200000.0; - // Evaluate the fits only over the range they were calibrated on. Outside it the power - // laws are unconstrained, so hold them at the nearest measured boundary rather than - // extrapolating -- brute force is routinely run at dim well below 128. - const double dim_d = std::clamp(static_cast(dim), 128.0, 1024.0); - - const double sddmm_vs_dense_sel = - kSddmmDenseSelNumerator / std::pow(dim_d, kSddmmDenseDimExponent); + const auto n_dataset_d = static_cast(n_dataset); + const double selectivity = static_cast(n_pass) / n_dataset_d; + // SDDMM has to beat whichever of the other two is available. + const bool sddmm_beats_dense = selectivity < kSddmmMaxSelectivity; if (!gather_possible) { - return selectivity < sddmm_vs_dense_sel ? filtered_search_path::sddmm - : filtered_search_path::dense; + return sddmm_beats_dense ? filtered_search_path::sddmm : filtered_search_path::dense; } + if (n_pass < kSddmmMaxPassingRows && sddmm_beats_dense) { return filtered_search_path::sddmm; } - // SDDMM has to beat both of the others to be worth running. - const double sddmm_max_rows = kSddmmRowsAtDim128 * std::pow(128.0 / dim_d, kSddmmDimExponent); - if (n_pass_d < sddmm_max_rows && selectivity < sddmm_vs_dense_sel) { - return filtered_search_path::sddmm; - } - - // The fraction of the dense search that survives paying the gather setup cost. At - // small n_dataset this goes to zero (or negative) and gathering can never win. - const double amortized = 1.0 - kGatherSetupEquivRows / n_dataset_d; - const double gather_vs_dense_sel = + // Fraction of the dense search left over after paying the gather setup. + const double amortized = 1.0 - kGatherMinDataset / n_dataset_d; + const double gather_max_selectivity = amortized * - std::min(kGatherSelNumerator / std::pow(dim_d, kGatherDimExponent), kGatherMaxSelectivity); + std::min(kGatherSelectivityScale / std::sqrt(static_cast(dim)), kGatherMaxSelectivity); - // n_pass > k keeps the compacted dataset big enough to actually contain k neighbors. - if (amortized > 0.0 && selectivity < gather_vs_dense_sel && n_pass > k) { + // n_pass > k keeps the compacted dataset big enough to hold k neighbors. + if (amortized > 0.0 && selectivity < gather_max_selectivity && n_pass > k) { return filtered_search_path::gather; } return filtered_search_path::dense; @@ -692,13 +654,8 @@ inline filtered_search_path select_filtered_search_path( /** * @brief Filtered search over the rows a bitset selects, by compacting them first. * - * Enumerates the passing rows, gathers them (and their precomputed norms) into a dense - * matrix, runs an ordinary unfiltered brute-force search over that matrix, and maps the - * resulting indices back into the original dataset's numbering. - * - * Only meaningful for a bitset filter, which selects the same rows for every query. A - * bitmap filter selects different rows per query, so there is no single set of rows to - * compact. + * Bitset only: a bitmap selects different rows per query, so there is no single set of + * rows to compact. */ template void brute_force_search_gathered( @@ -816,10 +773,8 @@ void brute_force_search_filtered( filter_view; IdxT nnz_h = 0; - // Rows passing the filter. A bitset selects the same rows for every query; a bitmap - // selects different rows per query, so there this is the per-query average. - IdxT n_pass = 0; - // A bitmap's rows differ per query, so there is no single set of rows to compact. + // Rows passing the filter; for a bitmap, the per-query average. + IdxT n_pass = 0; bool gather_possible = false; const BitsT* filter_data = nullptr; From d4bf0d097d3cc87753c115f662bcba87123f7aa2 Mon Sep 17 00:00:00 2001 From: Max Buckley Date: Tue, 14 Jul 2026 21:52:58 +0200 Subject: [PATCH 3/4] perf(brute_force): drop benchmark commentary from the dispatch comments Regret figures against the previous rule describe a code review, not the code. They tell a future maintainer of this function nothing, and they go stale the moment anyone retunes. They live in the PR description instead. Also drops a reference to a benchmark README that is not part of this branch. Co-Authored-By: Claude Opus 4.8 --- cpp/src/neighbors/detail/knn_brute_force.cuh | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/cpp/src/neighbors/detail/knn_brute_force.cuh b/cpp/src/neighbors/detail/knn_brute_force.cuh index cc9b4b248c..51a93dc1e9 100644 --- a/cpp/src/neighbors/detail/knn_brute_force.cuh +++ b/cpp/src/neighbors/detail/knn_brute_force.cuh @@ -602,11 +602,7 @@ enum class filtered_search_path { sddmm, gather, dense }; /** * @brief Pick the cheapest search path for a filtered brute-force query. * - * The constants are empirical (fp16, RTX 5090; see - * cpp/bench/prims/core/results/bf_sweep/README.md) and are deliberately round: over the - * measured sweep, perturbing any of them by 2x -- roughly the spread expected across - * GPUs -- leaves the mean penalty against a perfect oracle at 1.01-1.05x, versus 1.65x - * mean and 12.6x worst case for the old `sparsity < 0.9` rule. + * The thresholds are empirical. * * @param[in] n_dataset rows in the dataset * @param[in] dim columns in the dataset From 48da83b48fe3e83f8bcdfebe3b13114fff7e3e1b Mon Sep 17 00:00:00 2001 From: Max Buckley Date: Tue, 14 Jul 2026 22:03:14 +0200 Subject: [PATCH 4/4] perf(brute_force): stop inventing a passing-row count for bitmap filters A bitmap passes a different set of rows per query, so "the rows it passes" is not a thing. The dispatch was manufacturing one anyway -- nnz / n_queries -- so that it could reuse a single variable, then dividing by n_dataset to recover a selectivity it already had. Integer division made that lossy: a bitmap passing 5 pairs across 100 queries gave a row count of 0, hence a selectivity of exactly zero. In general it understated a bitmap's selectivity and pulled borderline cases onto the sparse path. Pass the selectivity (well-defined for both filter kinds, and the quantity the sparse/dense choice actually turns on) and make the row count a std::optional that only a bitset supplies -- which is also the only filter that can take the gather path that needs it. Dispatch is unchanged except at the sparse/dense boundary for bitmaps, where the rounding used to be wrong. Co-Authored-By: Claude Opus 4.8 --- cpp/src/neighbors/detail/knn_brute_force.cuh | 40 ++++++++++---------- 1 file changed, 21 insertions(+), 19 deletions(-) diff --git a/cpp/src/neighbors/detail/knn_brute_force.cuh b/cpp/src/neighbors/detail/knn_brute_force.cuh index 51a93dc1e9..0d97f22296 100644 --- a/cpp/src/neighbors/detail/knn_brute_force.cuh +++ b/cpp/src/neighbors/detail/knn_brute_force.cuh @@ -604,14 +604,16 @@ enum class filtered_search_path { sddmm, gather, dense }; * * The thresholds are empirical. * - * @param[in] n_dataset rows in the dataset - * @param[in] dim columns in the dataset - * @param[in] n_pass rows passing the filter (for a bitmap, averaged over queries) - * @param[in] k neighbors requested - * @param[in] gather_possible whether the gather path applies (bitset filters only) + * @param[in] n_dataset rows in the dataset + * @param[in] dim columns in the dataset + * @param[in] selectivity fraction of (query, row) pairs the filter passes + * @param[in] k neighbors requested + * @param[in] passing rows passing the filter, if the filter passes the same rows for + * every query. A bitmap passes a different set per query, so it has + * no such count, and cannot use the gather path. */ inline filtered_search_path select_filtered_search_path( - int64_t n_dataset, int64_t dim, int64_t n_pass, int64_t k, bool gather_possible) + int64_t n_dataset, int64_t dim, double selectivity, int64_t k, std::optional passing) { // SDDMM competes against a fixed setup cost, so it stops paying off at an absolute row // count rather than at a fraction of the dataset. @@ -624,14 +626,14 @@ inline filtered_search_path select_filtered_search_path( // Below this, the gather setup costs more than the entire dense search. constexpr double kGatherMinDataset = 200000.0; - const auto n_dataset_d = static_cast(n_dataset); - const double selectivity = static_cast(n_pass) / n_dataset_d; + const auto n_dataset_d = static_cast(n_dataset); // SDDMM has to beat whichever of the other two is available. const bool sddmm_beats_dense = selectivity < kSddmmMaxSelectivity; - if (!gather_possible) { + if (!passing) { return sddmm_beats_dense ? filtered_search_path::sddmm : filtered_search_path::dense; } + const int64_t n_pass = *passing; if (n_pass < kSddmmMaxPassingRows && sddmm_beats_dense) { return filtered_search_path::sddmm; } // Fraction of the dense search left over after paying the gather setup. @@ -769,9 +771,9 @@ void brute_force_search_filtered( filter_view; IdxT nnz_h = 0; - // Rows passing the filter; for a bitmap, the per-query average. - IdxT n_pass = 0; - bool gather_possible = 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; const BitsT* filter_data = nullptr; @@ -779,27 +781,27 @@ void brute_force_search_filtered( auto actual_filter = dynamic_cast*>(filter); filter_view.emplace(actual_filter->view()); - nnz_h = actual_filter->view().count(res); - n_pass = nnz_h / n_queries; + nnz_h = actual_filter->view().count(res); } 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; - gather_possible = true; + n_pass = actual_filter->view().count(res); + nnz_h = n_queries * (*n_pass); } else { RAFT_FAIL("Unsupported sample filter type"); } std::visit([&](const auto& actual_view) { filter_data = actual_view.data(); }, *filter_view); - const auto path = select_filtered_search_path(n_dataset, dim, n_pass, k, gather_possible); + 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); if (path == filtered_search_path::gather) { auto bitset_view = std::get>(*filter_view); brute_force_search_gathered( - res, idx, queries, bitset_view, n_pass, neighbors, distances, query_norms); + res, idx, queries, bitset_view, *n_pass, neighbors, distances, query_norms); } else if (path == filtered_search_path::dense) { raft::resources stream_pool_handle(res); raft::resource::set_cuda_stream(stream_pool_handle, stream);