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();