Commit e626144
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: #21281 parent 8ecfa65 commit e626144
1 file changed
Lines changed: 8 additions & 8 deletions
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
1 | 1 | | |
2 | | - | |
| 2 | + | |
3 | 3 | | |
4 | 4 | | |
5 | 5 | | |
| |||
671 | 671 | | |
672 | 672 | | |
673 | 673 | | |
674 | | - | |
675 | | - | |
676 | | - | |
677 | | - | |
678 | | - | |
679 | | - | |
680 | | - | |
| 674 | + | |
681 | 675 | | |
682 | 676 | | |
683 | 677 | | |
| |||
714 | 708 | | |
715 | 709 | | |
716 | 710 | | |
| 711 | + | |
| 712 | + | |
| 713 | + | |
| 714 | + | |
| 715 | + | |
| 716 | + | |
717 | 717 | | |
718 | 718 | | |
719 | 719 | | |
| |||
0 commit comments