perf(brute_force): drop the duplicate bitset/bitmap→CSR conversion on the SDDMM path - #2505
Conversation
…n the SDDMM path brute_force_search_filtered builds the CSR structure from the filter via to_csr, then hands the same CSR to raft::sparse::linalg::masked_matmul, which converts the mask into that structure again. masked_matmul's temporary is sparsity-preserving, so it shares indptr/indices with the CSR cuVS just filled and writes identical values over it. Every SDDMM-path search ran the conversion twice. nsys over 21 SDDMM searches (1M x 128, 0.5% selectivity, RTX 5090), before -> after: calc_nnz_by_rows 42 -> 21, repeat_csr 42 -> 21, exclusive-scan 42 -> 21, value fill 42 -> 21, and fill_indices_by_rows<check_nnz=false> -- the sparsity-owning instantiation only cuVS's to_csr could produce -- 21 -> 0. sddmm_ker and csr_to_coo stay at 21, so no real work is removed. Dropping cuVS's call also removes a host sync per search: to_csr on a sparsity-owning matrix reads the computed nnz back to the host and syncs the stream, while masked_matmul's conversion runs on a sparsity-preserving matrix and takes neither branch. 8-byte device-to-host copies per 29 filtered searches fall from 50 to 29 -- exactly the 21 removed readbacks. Paired A/B on a rebuilt libcuvs.so, 10 alternating rounds x 200 reps, dedicated idle RTX 5090, fp16/L2Expanded/k=64/100 queries. The saving is a roughly fixed ~150 us per search, so it is largest where the search is cheapest: N dim sel baseline patched speedup saving (95% CI) 1M 128 0.1% 0.383 ms 0.238 ms 1.61x +136 us [ 93, 180] 1M 1024 0.1% 0.468 ms 0.323 ms 1.45x +146 us [109, 183] 1M 128 1.0% 0.708 ms 0.579 ms 1.22x +133 us [115, 151] 10M 128 0.1% 0.750 ms 0.589 ms 1.27x +158 us [134, 182] 10M 128 3.0% 14.084 ms 13.760 ms 1.02x +335 us [252, 419] All 12 SDDMM-path configs improve with the interval excluding zero (t = 6.7 to 26.5). Three dense-path configs act as a null control and are flat, as they must be: +1.8 us, -57.6 us, +96.7 us, every interval spanning zero. Top-k ids and distances are bit-identical to the previous build across 14 configs spanning both dispatch paths, and the prefiltered brute-force tests pass. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
@lowener would you mind taking a look? This is the same corner of It sits inside the The one thing worth a reviewer's eye is the assumption the change rests on: that (No permission to add reviewers from a fork, hence the mention.) |
1 participant
The problem
On the sparse (SDDMM) branch of
brute_force_search_filtered, the filter is converted to CSR twice per search:masked_matmulbuilds a sparsity-preserving CSR over the caller's structure view and runsbitset_to_csron it, which writes through to the sameindptr/indicescuVS just filled:So (2) recomputes (1), identically.
The change
Drop cuVS's
to_csrcall and letmasked_matmuldo the conversion once.Both conversions really were running
nsysover 21 SDDMM-path searches (1M x 128, 0.5% selectivity), launch counts before → after:calc_nnz_by_rows_kernelfill_indices_by_rows_kernel<check_nnz=false>fill_indices_by_rows_kernel<check_nnz=true>repeat_csr_kernelDeviceScanKernel)cusparse::sddmm_kercsr_to_coo_kernelEvery conversion kernel halves while
sddmm_kerandcsr_to_cooare unchanged, i.e. no actual work is removed.check_nnz=falseis the sparsity-owning instantiation, which only cuVS'sto_csrcould produce; it disappears entirely.A host sync goes with it:
to_csron a sparsity-owning matrix readsnnzback to the host to size the matrix, which the sparsity-preserving path insidemasked_matmulnever does. 8-byte device-to-host copies over 29 searches drop from 50 to 29.Why relying on
masked_matmulis safeSame problem run twice, once with C's structure pre-filled correctly and once with a valid but deliberately wrong pattern (same row counts, every column index 0):
Bit-identical on both dispatch branches, so
masked_matmulunconditionally derives C's sparsity from the mask. This is deliberate in RAFT —bitset_to_csrcarries aconstexpr bool check_nnz = is_device_csr_sparsity_preserving_v<csr_matrix_t>;branch for exactly this case — but it is not documented:masked_matmulmarksC[inout]only with respect to its values, and the sparsity-preserving CSR it constructs is documented as "allowing the values in the sparsity to change but not the sparsity itself". Happy to hold this PR if reviewers would rather a RAFT doc fix landed first.(An earlier revision of this description said RAFT's
masked_matmultest depends on the overwrite. It does not — the test builds C's structure from the same mask it passes in, so it cannot discriminate. Hence the experiment above.)Measurements
Paired A/B, arms alternating within every round and every (N, dim) block so drift hits both equally: 10 rounds x 200 reps per config, separately built
libcuvs.soswapped viaLD_LIBRARY_PATH, dedicated idle RTX 5090 (sm_120, CUDA 12.9, WSL2), fp16 / L2Expanded / k=64 / 100 queries / shared bitset. CIs are Student-t over the 10 paired round-medians.All 12 SDDMM-path configs improve with the interval excluding zero. The saving is a roughly fixed ~150 us per search rather than a ratio, so it is largest where the search is cheapest — which post-#2321 is exactly where the sparse path lives, below ~2000 passing rows.
Null control. Dense-path configs cannot be affected and are flat: +1.8 us [-11.3, +15.0] at 1M/128/25%, -57.6 us [-134.0, +18.9] at 1M/1024/25%, +96.7 us [-277.8, +471.1] at 10M/128/25%. All three span zero. 10M x 1024 is drift-dominated on this host (±2 ms swings in both arms, including the control) and is excluded rather than reported.
The saving is not the removed sync. Removing the other per-search sync (
bitset_view::count, fed instead from an environment variable, results bit-identical) is worth median +0.0 us, 0 of 15 configs significant — a sync on an already-idle stream costs nothing. The win is the removed conversion work plus the pipeline bubble from draining the stream immediately before a block of host-side cuSPARSE setup.Correctness
NEIGHBORS_TEST --gtest_filter='*Prefiltered*': 156/156 pass. That run is on a pre-perf(brute_force): add scatter-gather path and fix filtered-search dispatch thresholds #2321 checkout, so 39 parameterisations against 45 onmain; the 6 that perf(brute_force): add scatter-gather path and fix filtered-search dispatch thresholds #2321 added cover the gather and dense paths, which this change does not touch.mainwith the build's own-Werror=all-warnings.RAFT-side follow-ups, not in this PR
masked_matmulallocatesnnz * sizeof(output_t)forC_matrix's element buffer that nothing ever reads — bothsddmmandfaster_dot_on_csrwrite throughC.get_elements(). Memory-only, no measurable latency change.masked_matmulshould document that it overwritesC's sparsity pattern from the mask, and its test should pass a structure that differs from the expected one so it actually covers the behaviour relied on here.