diff --git a/cpp/src/neighbors/detail/knn_brute_force.cuh b/cpp/src/neighbors/detail/knn_brute_force.cuh index 34ef8dd937..0d97f22296 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,151 @@ void brute_force_search( query_norms ? query_norms->data_handle() : nullptr); } +/** + * @brief The three strategies available for a filtered brute-force search. + * + * - `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 are empirical. + * + * @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, 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. + 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; + + 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 (!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. + const double amortized = 1.0 - kGatherMinDataset / n_dataset_d; + const double gather_max_selectivity = + amortized * + std::min(kGatherSelectivityScale / std::sqrt(static_cast(dim)), kGatherMaxSelectivity); + + // 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; +} + +/** + * @brief Filtered search over the rows a bitset selects, by compacting them first. + * + * 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( + 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 +770,10 @@ void brute_force_search_filtered( const cuvs::core::bitset_view>> filter_view; - IdxT nnz_h = 0; - float sparsity = 0.0f; + IdxT nnz_h = 0; + // 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; @@ -628,21 +781,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); } 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); } 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 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); + } 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},