From 016e268042391b477d66373789aa086516345804 Mon Sep 17 00:00:00 2001 From: Max Buckley Date: Sat, 5 Sep 2026 17:38:18 +0200 Subject: [PATCH] Warp-aggregate the attractive-force atomics in t-SNE Both t-SNE attractive-force kernels apply one atomicAdd per non-zero to attr_forces[row], and RepulsionKernel applies one per body to the single Z_norm scalar. from_knn_symmetrize_matrix lays the symmetrized COO out row by row, so a warp normally spans one or two rows: all 32 lanes hit the same address and the adds serialize in L2. Sum each run of equal row indices inside the warp first and let the run's last lane issue one atomic per row, and fold Z_norm into one atomic per warp. The reduction is driven by lane position rather than by comparing row indices pairwise, so a row that appears in two separate runs within a warp is still counted once, and lanes past the end of the COO stay in the warp with a row index no edge can match rather than returning early. Measured on an RTX 5090 (sm_120, CUDA 13.2, driver 580.126.09) with the parameters cuml.TSNE configures by default, baseline and candidate runs interleaved, median of 3 rounds of 5: Barnes-Hut n= 20,000 604 ms -> 444 ms (1.36x) Barnes-Hut n= 50,000 760 ms -> 601 ms (1.27x) Barnes-Hut n=200,000 2105 ms -> 1551 ms (1.36x) FFT n= 20,000 392 ms -> 193 ms (2.03x) FFT n= 50,000 424 ms -> 268 ms (1.58x) FFT n=200,000 1699 ms -> 1185 ms (1.43x) attractive_kernel_bh alone goes from 225 us to 38 us per launch, and compute_Pij_x_Qij_kernel was 61% of unseeded FFT t-SNE GPU time before the change. Both kernels were already non-deterministic (atomicAdd ordering, atomicCAS tree construction), so no reproducibility contract changes. The seeded FFT path uses compute_Pij_x_Qij_deterministic_rows, which walks rows without atomics; it is untouched and stays bit-identical. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01HMTsqoAxrKxVQR6RfcUZhz --- cpp/src/tsne/barnes_hut_kernels.cuh | 93 ++++++++++++++++++++++------- cpp/src/tsne/fft_kernels.cuh | 70 ++++++++++++++++------ 2 files changed, 123 insertions(+), 40 deletions(-) diff --git a/cpp/src/tsne/barnes_hut_kernels.cuh b/cpp/src/tsne/barnes_hut_kernels.cuh index 9eeb3dba13..a363964028 100644 --- a/cpp/src/tsne/barnes_hut_kernels.cuh +++ b/cpp/src/tsne/barnes_hut_kernels.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2021-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -7,6 +7,7 @@ #include "utils.cuh" +#include #include #include @@ -600,6 +601,11 @@ CUML_KERNEL __launch_bounds__( // iterate over all bodies assigned to thread const auto MAX_SIZE = FOUR_NNODES + 4; + // Z_norm is one scalar for the whole grid, so accumulating it per body made + // every thread contend for the same address. Keep a per-thread partial and + // fold it down to one atomic per warp after the traversal. + value_t z_partial = 0.0f; + for (auto k = threadIdx.x + blockIdx.x * blockDim.x; k < N; k += blockDim.x * gridDim.x) { const auto i = sortd[k]; // get permuted/sorted index // cache position info @@ -660,8 +666,16 @@ CUML_KERNEL __launch_bounds__( // update velocity velxd[i] += vx; velyd[i] += vy; - atomicAdd(Z_norm, normsum); + z_partial += normsum; } + + // All threads reach this point together (the body loop has no early exit), + // so a full-warp reduction is safe here. +#pragma unroll + for (int offset = raft::WarpSize / 2; offset > 0; offset >>= 1) { + z_partial += __shfl_down_sync(0xffffffffu, z_partial, offset); + } + if (threadIdx.x % raft::WarpSize == 0 && z_partial != 0.0f) { atomicAdd(Z_norm, z_partial); } } /** @@ -680,27 +694,60 @@ CUML_KERNEL void attractive_kernel_bh(const value_t* restrict VAL, const value_t dof) { const auto index = (blockIdx.x * blockDim.x) + threadIdx.x; - if (index >= NNZ) return; - const auto i = ROW[index]; - const auto j = COL[index]; - - const value_t y1d = Y1[i] - Y1[j]; - const value_t y2d = Y2[i] - Y2[j]; - value_t dist = y1d * y1d + y2d * y2d; - // As a sum of squares, SED is mathematically >= 0. There might be a source of - // NaNs upstream though, so until we find and fix them, enforce that trait. - if (!(dist >= 0)) dist = 0.0f; - - const value_t P = VAL[index]; - const value_t Q = compute_q(dist, dof); - const value_t PQ = P * Q; - - // Apply forces - atomicAdd(&attract1[i], PQ * y1d); - atomicAdd(&attract2[i], PQ * y2d); - - if (Qs) { // when computing KL div - Qs[index] = Q; + // Lanes past the end stay in the warp (with a row index no real edge can + // match) so the warp-wide primitives below always see a full mask. + const bool active = index < NNZ; + const value_idx i = active ? ROW[index] : value_idx(-1); + + value_t force1 = 0; + value_t force2 = 0; + if (active) { + const auto j = COL[index]; + + const value_t y1d = Y1[i] - Y1[j]; + const value_t y2d = Y2[i] - Y2[j]; + value_t dist = y1d * y1d + y2d * y2d; + // As a sum of squares, SED is mathematically >= 0. There might be a source of + // NaNs upstream though, so until we find and fix them, enforce that trait. + if (!(dist >= 0)) dist = 0.0f; + + const value_t P = VAL[index]; + const value_t Q = compute_q(dist, dof); + const value_t PQ = P * Q; + + force1 = PQ * y1d; + force2 = PQ * y2d; + + if (Qs) { // when computing KL div + Qs[index] = Q; + } + } + + // Apply forces. The symmetrized COO is laid out row by row, so a warp + // normally covers one or two rows and the two atomicAdds below would + // serialize 32 deep on a single address. Sum each run of equal rows inside + // the warp first and let the run's last lane issue one atomic per row. The + // reduction is driven by lane position rather than by comparing row indices + // pairwise, so a row that appears in two separate runs is still counted once. + constexpr unsigned full_mask = 0xffffffffu; + const int lane = threadIdx.x % raft::WarpSize; + const value_idx prev = __shfl_up_sync(full_mask, i, 1); + const unsigned head_mask = __ballot_sync(full_mask, lane == 0 || prev != i); + // Lowest lane of this run: the highest run head at or before this lane. + const int run_start = raft::WarpSize - 1 - __clz(head_mask & ((2u << lane) - 1)); +#pragma unroll + for (int offset = 1; offset < raft::WarpSize; offset <<= 1) { + const value_t other1 = __shfl_up_sync(full_mask, force1, offset); + const value_t other2 = __shfl_up_sync(full_mask, force2, offset); + if (lane - offset >= run_start) { + force1 += other1; + force2 += other2; + } + } + const bool run_end = lane == raft::WarpSize - 1 || ((head_mask >> (lane + 1)) & 1u); + if (active && run_end) { + atomicAdd(&attract1[i], force1); + atomicAdd(&attract2[i], force2); } // TODO: Convert attractive forces to CSR format diff --git a/cpp/src/tsne/fft_kernels.cuh b/cpp/src/tsne/fft_kernels.cuh index 146b8cd5e2..1188c5537f 100644 --- a/cpp/src/tsne/fft_kernels.cuh +++ b/cpp/src/tsne/fft_kernels.cuh @@ -12,6 +12,8 @@ #pragma once +#include + #include namespace ML { @@ -441,29 +443,63 @@ CUML_KERNEL void compute_Pij_x_Qij_kernel(value_t* __restrict__ attr_forces, const value_t dof) { const value_idx TID = threadIdx.x + blockIdx.x * blockDim.x; - if (TID >= num_nonzero) return; - const value_idx i = coo_rows[TID]; - const value_idx j = coo_cols[TID]; + // Lanes past the end stay in the warp (with a row index no real edge can + // match) so the warp-wide primitives below always see a full mask. + const bool active = TID < num_nonzero; + const value_idx i = active ? coo_rows[TID] : value_idx(-1); - value_t ix = points[i]; - value_t iy = points[num_points + i]; - value_t jx = points[j]; - value_t jy = points[num_points + j]; + value_t x_force = 0; + value_t y_force = 0; + if (active) { + const value_idx j = coo_cols[TID]; - value_t dx = ix - jx; - value_t dy = iy - jy; + value_t ix = points[i]; + value_t iy = points[num_points + i]; + value_t jx = points[j]; + value_t jy = points[num_points + j]; - const value_t dist = (dx * dx) + (dy * dy); + value_t dx = ix - jx; + value_t dy = iy - jy; - const value_t P = pij[TID]; - const value_t Q = compute_q(dist, dof); - const value_t PQ = P * Q; + const value_t dist = (dx * dx) + (dy * dy); - atomicAdd(attr_forces + i, PQ * dx); - atomicAdd(attr_forces + num_points + i, PQ * dy); + const value_t P = pij[TID]; + const value_t Q = compute_q(dist, dof); + const value_t PQ = P * Q; - if (Qs) { // when computing KL div - Qs[TID] = Q; + x_force = PQ * dx; + y_force = PQ * dy; + + if (Qs) { // when computing KL div + Qs[TID] = Q; + } + } + + // The symmetrized COO is laid out row by row, so a warp normally covers one + // or two rows and the two atomicAdds below would serialize 32 deep on a + // single address. Sum each run of equal rows inside the warp first and let + // the run's last lane issue one atomic per row. The reduction is driven by + // lane position rather than by comparing row indices pairwise, so a row that + // appears in two separate runs is still counted once. + constexpr unsigned full_mask = 0xffffffffu; + const int lane = threadIdx.x % raft::WarpSize; + const value_idx prev = __shfl_up_sync(full_mask, i, 1); + const unsigned head_mask = __ballot_sync(full_mask, lane == 0 || prev != i); + // Lowest lane of this run: the highest run head at or before this lane. + const int run_start = raft::WarpSize - 1 - __clz(head_mask & ((2u << lane) - 1)); +#pragma unroll + for (int offset = 1; offset < raft::WarpSize; offset <<= 1) { + const value_t other_x = __shfl_up_sync(full_mask, x_force, offset); + const value_t other_y = __shfl_up_sync(full_mask, y_force, offset); + if (lane - offset >= run_start) { + x_force += other_x; + y_force += other_y; + } + } + const bool run_end = lane == raft::WarpSize - 1 || ((head_mask >> (lane + 1)) & 1u); + if (active && run_end) { + atomicAdd(attr_forces + i, x_force); + atomicAdd(attr_forces + num_points + i, y_force); } }