Warp-aggregate the attractive-force atomics in t-SNE

[maxwbuckley]

Problem

Both t-SNE attractive-force kernels apply one atomicAdd per non-zero to
attr_forces[row]:

• BH::attractive_kernel_bh (Barnes-Hut)
• FFT::compute_Pij_x_Qij_kernel (FFT, unseeded path)

raft::sparse::linalg::from_knn_symmetrize_matrix writes the symmetrized COO
into per-row spans, so it comes out grouped by row. A warp therefore normally
spans one or two rows, all 32 lanes target the same address, and the adds
serialize in L2. BH::RepulsionKernel has the same problem in a sharper form:
one atomicAdd per body onto the single Z_norm scalar, i.e. ~n serialized
adds to one address per launch.

Profiling TSNE_fit on an RTX 5090 (n=50,000, 1000 iterations, the parameters
cuml.TSNE configures by default) put compute_Pij_x_Qij_kernel at 61% of
all GPU time on the FFT path and attractive_kernel_bh at 26% on the
Barnes-Hut path.

Change

Sum each run of equal row indices inside the warp first and let the run's last
lane issue one atomic per row; fold Z_norm into one atomic per warp.

Two details worth reviewing:

• The reduction is a segmented Hillis–Steele scan driven by lane position
  (run_start, derived from a ballot of run heads), not by comparing row
  indices pairwise. A naive other == i test would double-count a row that
  happens to appear in two separate runs within one warp; this does not.
• Lanes past the end of the COO no longer return early. They stay in the warp
  carrying a row index no real edge can match, so every warp-wide primitive
  sees a full mask.

Results

RTX 5090 (sm_120), CUDA 13.2, driver 580.126.09, -O3, CMAKE_CUDA_ARCHITECTURES=120-real.
Baseline and candidate runs interleaved in the same process sequence, median of
3 rounds of 5 timed iterations; run-to-run spread was under 1%. Workload is 20
Gaussian clusters, random_state unset, with the n_neighbors /
learning-rate schedule cuml.TSNE derives from n (its default
learning_rate_method="adaptive"), not the raw TSNEParams defaults.

algorithm	n	before	after	speedup
Barnes-Hut	20,000	604 ms	444 ms	1.36x
Barnes-Hut	50,000	760 ms	601 ms	1.27x
Barnes-Hut	200,000	2105 ms	1551 ms	1.36x
FFT	20,000	392 ms	193 ms	2.03x
FFT	50,000	424 ms	268 ms	1.58x
FFT	200,000	1699 ms	1185 ms	1.43x

At the kernel level, attractive_kernel_bh goes from 225 µs to 38 µs per
launch (5.9x).

The mechanism is not architecture-specific — same-address atomicAdd is
serialized on every supported architecture, and the intrinsics used
(__shfl_up_sync, __ballot_sync, __clz) are available on all of them — but
I only have sm_120 hardware, so the numbers above are the only measured ones.
The magnitude will differ elsewhere: this GPU's 96 MB L2 holds most of the COO
working set, so the post-change kernel runs largely out of L2 here.

Correctness

• Quality parity. Trustworthiness across n ∈ {2,000, 20,000, 50,000} ×
  {Barnes-Hut, FFT} × 2 seeds: maximum delta 0.008, inside the baseline's own
  seed-to-seed spread. No non-finite outputs.
• No contract change. Both modified kernels were already non-deterministic
  (atomicAdd ordering; Barnes-Hut also builds its tree with atomicCAS), so
  random_state never made these paths reproducible. The seeded FFT path uses
  compute_Pij_x_Qij_deterministic_rows, which walks rows without atomics — it
  is untouched and its output is bit-identical before and after. The warp
  reduction is a tree sum, so it is if anything slightly better conditioned
  than the previous serialized adds.
• compute-sanitizer: memcheck and synccheck report 0 errors for both
  algorithms. racecheck reports exactly the same hazard counts as the
  baseline build (Barnes-Hut 9, FFT 16 — all pre-existing warnings, 0 errors),
  so no new hazards.
• Edge cases: n ∈ {50, 101, 1000, 3001} with odd p, across Barnes-Hut,
  FFT and exact — finite KL divergence, no failures. Partial final warps
  (NNZ not a multiple of 32) are exercised by every one of these.
• pre-commit run is clean on both files.

Not addressed here

FFT::compute_interpolated_indices (~9% of unseeded FFT GPU time) has the same
contention pattern, but its colliding keys are strided rather than contiguous,
so it needs __match_any_sync rather than this run-based reduction. Left for a
follow-up.

The ~18-line reduction idiom is duplicated between the two kernels, which live
in different headers. I kept it inline rather than adding a shared device helper
because that is the exact code I benchmarked and validated; happy to factor it
into cpp/src/tsne/utils.cuh (which both already include) if reviewers prefer.

🤖 Generated with Claude Code

https://claude.ai/code/session_01HMTsqoAxrKxVQR6RfcUZhz

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 8da6db12-09e8-49f9-b648-b96aba7d6730

📥 Commits

Reviewing files that changed from the base of the PR and between b07de6e and 016e268.

📒 Files selected for processing (2)

• cpp/src/tsne/barnes_hut_kernels.cuh
• cpp/src/tsne/fft_kernels.cuh

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

📝 Summary

Summary by CodeRabbit

• Performance Improvements
  • Improved GPU-accelerated t-SNE computations by reducing synchronization and atomic update overhead.
  • Optimized force and normalization calculations for better efficiency across large datasets.
  • Improved handling of edge cases during parallel processing while preserving existing results and behavior.

Walkthrough

The t-SNE Barnes-Hut and FFT CUDA kernels now use warp-level reductions to aggregate normalization and force updates. The changes add CUDA utilities, preserve inactive lanes for warp operations, and reduce atomic updates for contiguous rows and warp participants.

Changes

t-SNE CUDA warp aggregation

Layer / File(s)	Summary
Barnes-Hut repulsion normalization cpp/src/tsne/barnes_hut_kernels.cuh	RepulsionKernel stores per-thread normalization values, reduces them within each warp, and performs one atomic update per warp.
Barnes-Hut attractive-force aggregation cpp/src/tsne/barnes_hut_kernels.cuh	attractive_kernel_bh keeps out-of-range lanes active with neutral values and aggregates equal-row runs before issuing force updates.
FFT force aggregation cpp/src/tsne/fft_kernels.cuh	compute_Pij_x_Qij_kernel aggregates contiguous COO row runs within each warp before issuing force atomics. Conditional Q and Qs handling remains in place.

Estimated code review effort: 4 (Complex) | ~45 minutes

Merge Risk: ⚪ Minimal · up to 016e2

The CUDA optimization reduces atomic contention while preserving reported t-SNE output quality and deterministic seeded behavior. No merge-blocking risk remains.

Suggested reviewers: divyegala, bdice

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely describes the main change: aggregating attractive-force atomic updates in the t-SNE kernels.
Description check	✅ Passed	The description is directly related to the changeset and explains the motivation, implementation, performance results, correctness validation, and deferred follow-up work.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check. Docstring coverage is scoped to functions touched by this diff. Analyzed 0 functions across 0…
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches 💡 1

🛠️ Fix failing CI checks 💡

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}