Skip to content

Commit e626144

Browse files
authored
perf(brute_force): skip csr_to_coo on inner-product filtered search (#2128)
## Summary Small perf cleanup for the filtered brute-force CSR path in `knn_brute_force.cuh`. - When sparsity ≥ 0.9, filtered brute-force takes the sparse CSR / `masked_matmul` path. As part of setup it builds a per-nonzero `rows` array via `raft::sparse::convert::csr_to_coo`. - That `rows` array is only consumed by `cuvs::neighbors::detail::epilogue_on_csr`, which only runs for `L2Expanded` / `L2SqrtExpanded` / `CosineExpanded` (it combines masked inner products with precomputed norms). - For `InnerProduct`, the epilogue is skipped — so the allocation and the `csr_to_coo` kernel launch were dead work on every IP-metric filtered search hitting the CSR path. This PR hoists the `rmm::device_uvector<IdxT> rows(...)` allocation and the `csr_to_coo` call inside the L2/Cosine branch, next to their only consumer. ## What changes - No behavior change for `L2Expanded`, `L2SqrtExpanded`, `CosineExpanded` — same kernels, same order. - For `InnerProduct` filtered searches that hit the sparse CSR branch: saves one device allocation of `nnz * sizeof(IdxT)` and one kernel launch + write pass per call. ## Test plan - [ ] Existing filtered brute-force tests pass (`BRUTE_FORCE_TEST`, prefiltered variants) for `InnerProduct`, `L2Expanded`, `L2SqrtExpanded`, `CosineExpanded`. - [x] No diff in output for L2 / Cosine vs. base. 🤖 Generated with [Claude Code](https://claude.com/claude-code) Authors: - Max Buckley (https://github.com/maxwbuckley) - Micka (https://github.com/lowener) - Anupam (https://github.com/aamijar) Approvers: - Micka (https://github.com/lowener) URL: #2128
1 parent 8ecfa65 commit e626144

1 file changed

Lines changed: 8 additions & 8 deletions

File tree

cpp/src/neighbors/detail/knn_brute_force.cuh

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
2+
* SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
33
* SPDX-License-Identifier: Apache-2.0
44
*/
55

@@ -671,13 +671,7 @@ void brute_force_search_filtered(
671671

672672
// create filter csr view
673673
auto compressed_csr_view = csr.structure_view();
674-
rmm::device_uvector<IdxT> rows(compressed_csr_view.get_nnz(), stream);
675-
raft::sparse::convert::csr_to_coo(compressed_csr_view.get_indptr().data(),
676-
compressed_csr_view.get_n_rows(),
677-
rows.data(),
678-
compressed_csr_view.get_nnz(),
679-
stream);
680-
auto dataset_view = raft::make_device_matrix_view<const T, IdxT, raft::row_major>(
674+
auto dataset_view = raft::make_device_matrix_view<const T, IdxT, raft::row_major>(
681675
idx.dataset().data_handle(), n_dataset, dim);
682676

683677
auto csr_view = raft::make_device_csr_matrix_view<DistanceT, IdxT, IdxT, IdxT>(
@@ -714,6 +708,12 @@ void brute_force_search_filtered(
714708
query_norms_->view());
715709
}
716710
}
711+
rmm::device_uvector<IdxT> rows(compressed_csr_view.get_nnz(), stream);
712+
raft::sparse::convert::csr_to_coo(compressed_csr_view.get_indptr().data(),
713+
compressed_csr_view.get_n_rows(),
714+
rows.data(),
715+
compressed_csr_view.get_nnz(),
716+
stream);
717717
cuvs::neighbors::detail::epilogue_on_csr(
718718
res,
719719
csr.get_elements().data(),

0 commit comments

Comments
 (0)