From c1b85c7e88f0c860bce478dcd70b33fb87ee476e Mon Sep 17 00:00:00 2001 From: jungfrau Date: Sat, 15 Aug 2026 22:15:00 -0400 Subject: [PATCH] Reduce within the warp before the Bragg integration atomics fit and boxsum were 80% of GPU time on a crowded crystal - 78 s of it. Neither was bandwidth- or occupancy-bound: both sat at about an eighth of the issue rate the card can sustain, stalled. What stalls them is the block-wide accumulations. Every one has all 128 lanes of the block adding into one shared address, and a shared-memory atomicAdd on a float or a 64-bit integer has no instruction on either Turing or Ada - it compiles to a compare-and-swap retry loop. So those 128 lanes serialise into 128 retries, eighteen times per thread in fit. Summing across the warp first and letting one lane do the atomic leaves four per block instead of 128. That is the whole story: the arithmetic below was worth 2%, the atomics 5.6x. The arithmetic is still worth having, and is what was expected to matter: - compute_shell ran on all 128 threads of a block for a value that belongs to the reflection. It is two software double-precision divisions, on a card whose double throughput is a thirty-second (a sixty-fourth on the production one) of its single. One thread does it now. - The Kabsch inner loop divided by the same weight three times; the compiler emits the whole correctly-rounded sequence each time. One reciprocal now. Likewise the two Gaussian widths and the profile normalisation, which are constant over a reflection's cells and were divided per cell. - boxsum read the pixel before deciding whether it wanted it. The window is the bounding box of an ellipse, so nearly half of it is neither the signal disk nor the background ring, and those slots were fetching a cache line for nothing. Measured: fit 50.8 s -> 9.1 s, boxsum 27.5 s -> 12.1 s. A crowded crystal 2m22s -> 1m58s, a 16M-pixel one 39.5 s -> 37.2 s, the whole battery 12m30s -> 11m35s. Same space group on all 24 crystals, none failed. The integer sums are unchanged - addition is associative. The float ones move in their last bits and become more reproducible, since a fixed shuffle tree replaces whatever order the atomics arrived in. Co-Authored-By: Claude Opus 5 (1M context) --- .../BraggIntegrationEngineGPU.cu | 95 ++++++++++++++----- 1 file changed, 72 insertions(+), 23 deletions(-) diff --git a/image_analysis/bragg_integration/BraggIntegrationEngineGPU.cu b/image_analysis/bragg_integration/BraggIntegrationEngineGPU.cu index 59d78686..c66d2156 100644 --- a/image_analysis/bragg_integration/BraggIntegrationEngineGPU.cu +++ b/image_analysis/bragg_integration/BraggIntegrationEngineGPU.cu @@ -69,6 +69,30 @@ __global__ void mark_mask(const float *px_x, const float *px_y, uint8_t *mask, u } } +// Sum one value across the warp so a single lane does the shared-memory atomic. +// +// Every block-wide accumulation here has all 128 lanes targeting one address, and a shared float or +// 64-bit atomicAdd has no native instruction on Turing OR Ada - it compiles to a compare-and-swap +// retry loop, so those 128 lanes serialise into 128 retries. Reducing within the warp first leaves +// four atomics per block instead of 128. Every call site below sits after a loop that all threads +// reach, so the full-warp mask is the right one. +// +// The integer sums are unchanged by this: addition is associative. The float and double ones change +// in their last bits, and become MORE reproducible - a fixed shuffle tree replaces whatever order +// the atomics happened to arrive in. +template +__device__ __forceinline__ T warp_sum(T v) { + #pragma unroll + for (int off = 16; off > 0; off >>= 1) + v += __shfl_down_sync(0xffffffffu, v, off); + return v; +} + +#define WARP_ATOMIC_ADD(dst, val) do { \ + const auto _w = warp_sum(val); \ + if ((threadIdx.x & 31u) == 0u) atomicAdd(&(dst), _w); \ + } while (0) + // --- Pass A box-sum: rough I / background / centroid / strong flag, one block per reflection. --- __global__ void boxsum(const float *px_x, const float *px_y, const float *dd, const int32_t *img, const uint8_t *mask, const uint32_t *owner, @@ -121,7 +145,9 @@ __global__ void boxsum(const float *px_x, const float *px_y, const float *dd, for (int t = threadIdx.x; t < area; t += blockDim.x) { const int x = x0 + t % bw, y = y0 + t / bw; const BraggStencilDist d = BraggStencilDistances(st, (float) x - cx, (float) y - cy); - const int32_t px = img[y * p.W + x]; + // The window is the bounding box of the outer ellipse, so nearly half of it is neither the + // signal disk nor the background ring. Reading the pixel only once it is known to be wanted + // keeps those slots from fetching a cache line for nothing. if (d.signal < p.r1_sq) { // A pixel a nearer neighbour owns carries that neighbour's flux; see the CPU engine. ++l_nd; @@ -130,22 +156,24 @@ __global__ void boxsum(const float *px_x, const float *px_y, const float *dd, else if (p.exclude_px) continue; } ++l_ni; + const int32_t px = img[y * p.W + x]; if (valid(px)) { l_Isum += px; l_Ix += (long long) x * px; l_Iy += (long long) y * px; ++l_niv; } } else if (d.inner >= p.r2_sq && d.outer < p.r3_sq) { if (mask[y * p.W + x]) continue; + const int32_t px = img[y * p.W + x]; if (!valid(px)) continue; l_bkg += (double) px; ++l_nb; } } - atomicAdd(&s_Isum, (unsigned long long) l_Isum); - atomicAdd(&s_Ix, (unsigned long long) l_Ix); - atomicAdd(&s_Iy, (unsigned long long) l_Iy); - atomicAdd(&s_ninner, l_ni); - atomicAdd(&s_ninner_valid, l_niv); - atomicAdd(&s_nbkg, l_nb); - atomicAdd(&s_bkgsum, l_bkg); - atomicAdd(&s_ndisk, l_nd); - atomicAdd(&s_nown, l_no); + WARP_ATOMIC_ADD(s_Isum, (unsigned long long) l_Isum); + WARP_ATOMIC_ADD(s_Ix, (unsigned long long) l_Ix); + WARP_ATOMIC_ADD(s_Iy, (unsigned long long) l_Iy); + WARP_ATOMIC_ADD(s_ninner, l_ni); + WARP_ATOMIC_ADD(s_ninner_valid, l_niv); + WARP_ATOMIC_ADD(s_nbkg, l_nb); + WARP_ATOMIC_ADD(s_bkgsum, l_bkg); + WARP_ATOMIC_ADD(s_ndisk, l_nd); + WARP_ATOMIC_ADD(s_nown, l_no); __syncthreads(); for (int t = threadIdx.x; t < p.rad_w; t += blockDim.x) { s_radv[t] = 0.0f; s_radn[t] = 0; } @@ -233,7 +261,7 @@ __global__ void boxsum(const float *px_x, const float *px_y, const float *dd, } } } - atomicAdd(&s_clipsum, c_l); atomicAdd(&s_clipn, cn_l); + WARP_ATOMIC_ADD(s_clipsum, c_l); WARP_ATOMIC_ADD(s_clipn, cn_l); } __syncthreads(); @@ -349,7 +377,13 @@ __global__ void learn_profile(const int32_t *img, const uint32_t *owner, if (i >= n || !ok_a[i] || !strong_a[i]) return; const float I = I_a[i]; if (!(I > 0.0f)) return; - const int cx = cx_a[i], cy = cy_a[i], sh = compute_shell(dd[i], invd2mm); + const int cx = cx_a[i], cy = cy_a[i]; + // The shell belongs to the reflection, not the pixel, so one thread works it out for the whole + // block: it is two software double-precision divisions, and all 128 threads were doing both. + __shared__ int s_sh; + if (threadIdx.x == 0) s_sh = compute_shell(dd[i], invd2mm); + __syncthreads(); + const int sh = s_sh; const float bkg = bkg_a[i]; const float rx = px_x[i] - p.beam_x, ry = px_y[i] - p.beam_y; const float Rpx = sqrtf(rx * rx + ry * ry); @@ -476,7 +510,12 @@ __global__ void fit(const int32_t *img, const uint32_t *owner, const float *px_x if (!ok_a[i]) { if (threadIdx.x == 0) ok_o[i] = 0; return; } const int cx = cx_a[i], cy = cy_a[i]; - const int sh = compute_shell(dd[i], invd2mm); + // As in learn_profile: one thread per block, not one per thread. Two double-precision divisions + // on a card whose double throughput is a sixty-fourth of its single is worth doing once. + __shared__ int s_sh; + if (threadIdx.x == 0) s_sh = compute_shell(dd[i], invd2mm); + __syncthreads(); + const int sh = s_sh; const bool use_shell = shell_n[sh] >= MIN_STRONG_PER_SHELL; // else fall back to the global profile const float bkg = bkg_a[i]; @@ -506,17 +545,23 @@ __global__ void fit(const int32_t *img, const uint32_t *owner, const float *px_x if (threadIdx.x == 0) { s_Rf = Rf; s_Gf = Gf; s_gs = 0.0f; } __syncthreads(); const float fx = px_x[i] - cx, fy = px_y[i] - cy; + // The two widths are the same for every cell of this reflection, so their reciprocals are + // taken once rather than per cell. + const float inv_2s2r = 0.5f / s2r, inv_2s2t = 0.5f / s2t; float l_gs = 0.0f; for (int k = threadIdx.x; k < Gf * Gf; k += blockDim.x) { const float ex = (k % Gf - Rf) - fx, ey = (k / Gf - Rf) - fy; const float rad = ex * ux + ey * uy, tn = -ex * uy + ey * ux; - const float g = expf(-rad * rad / (2.0f * s2r) - tn * tn / (2.0f * s2t)); + const float g = expf(-rad * rad * inv_2s2r - tn * tn * inv_2s2t); Pbuf[k] = g; l_gs += g; } - atomicAdd(&s_gs, l_gs); + WARP_ATOMIC_ADD(s_gs, l_gs); __syncthreads(); - const float gs = s_gs; - for (int k = threadIdx.x; k < Gf * Gf; k += blockDim.x) Pbuf[k] /= gs; + // Normalising by a reciprocal rather than dividing per cell: the divisor is the same for + // every cell of the reflection, and a division here is thirteen instructions of a + // twenty-instruction loop. + const float inv_gs = __frcp_rn(s_gs); + for (int k = threadIdx.x; k < Gf * Gf; k += blockDim.x) Pbuf[k] *= inv_gs; __syncthreads(); } @@ -558,8 +603,8 @@ __global__ void fit(const int32_t *img, const uint32_t *owner, const float *px_x if (own) l_pown += Pp; if (in_disk && (own || p.overlap != 2)) l_mread += Pp; } - atomicAdd(&s_pgrid, l_pgrid); atomicAdd(&s_pvalid, l_pvalid); atomicAdd(&s_pown, l_pown); - atomicAdd(&s_mall, l_mall); atomicAdd(&s_mread, l_mread); + WARP_ATOMIC_ADD(s_pgrid, l_pgrid); WARP_ATOMIC_ADD(s_pvalid, l_pvalid); WARP_ATOMIC_ADD(s_pown, l_pown); + WARP_ATOMIC_ADD(s_mall, l_mall); WARP_ATOMIC_ADD(s_mread, l_mread); atomicMax(&s_ppeak_i, __float_as_int(l_ppeak)); atomicMax(&s_plost_i, __float_as_int(l_plost)); __syncthreads(); if (s_pvalid < p.minpk * s_pgrid) { @@ -592,11 +637,15 @@ __global__ void fit(const int32_t *img, const uint32_t *owner, const float *px_x const int32_t px = img[y * p.W + x]; if (!valid(px)) continue; const float v = fmaxf(B + Ihere * Pp, (float) WEIGHT_VARIANCE_MIN_FRACTION * B); - l_num += Pp * ((float) px - bkg) / v; - l_den += Pp * Pp / v; - l_wsum += Pp / v; + // One reciprocal for the three weights. Written as three divisions the compiler emits + // the whole correctly-rounded sequence three times over - the same estimate, the same + // refinement, the same check - for a weight that only has to be a weight. + const float iv = __frcp_rn(v); + l_num += Pp * ((float) px - bkg) * iv; + l_den += Pp * Pp * iv; + l_wsum += Pp * iv; } - atomicAdd(&s_num, l_num); atomicAdd(&s_den, l_den); atomicAdd(&s_wsum, l_wsum); + WARP_ATOMIC_ADD(s_num, l_num); WARP_ATOMIC_ADD(s_den, l_den); WARP_ATOMIC_ADD(s_wsum, l_wsum); __syncthreads(); if (threadIdx.x == 0 && s_den > 0.0f) s_I = s_num / s_den; __syncthreads();