mx/Jungfraujoch
2
1
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 68 Wiki Activity

Read four pixels at a time in the ring reduction

Browse Source 
reduce_rings_shared was 69% of all GPU kernel time - 116.8 s of a 70 s run
across four cards. It is not bandwidth bound: flag_strong streams the same two
arrays through the same grid-stride loop and reaches 196 GB/s, while this
reached 30. The difference is the shared-memory atomics. Lanes in a warp read
consecutive pixels along a detector row, a ring is a few pixels wide, so most of
a warp lands in a handful of rings and the atomics to each one serialise.

Two changes.

The block reads four pixels per thread as one 16-byte and one 8-byte
transaction, and merges the ones that fall in the same ring in registers before
touching shared memory. Consecutive pixels usually DO share a ring, so this is
where the win is: a run costs one set of atomics instead of one per pixel.
npix is not guaranteed to be a multiple of four - it is width x height on the
converted path, and detectors are not obliged to be even - so the vector loop
stops short and a scalar loop finishes the remainder. Reading past the end would
not fault, which is worse than if it did: it would fold uninitialised device
memory into the accumulators and move the detection threshold in a way that does
not reproduce.

And the grid is sized from the occupancy the device reports, per pass. The two
passes have different shared footprints - the first carries the corrected rings
as well - so they do not fit the same number of blocks, and a grid sized for one
left the other running a second wave at a quarter occupancy. The comment that
justified the old grid reasoned from 1536 threads per SM, which is an Ada
number; the card it ran on holds 1024.

The run totals are still exactly what they were. The accumulators are unsigned
64-bit, so summing a run in a register and adding it once is the same value as
adding each pixel separately - addition mod 2^64 is associative, overflow
included - which is what keeps the ring statistics, and therefore the detection
threshold, independent of how the work was grouped. That is the property the
integer accumulators exist for. (The run accumulators are unsigned for the same
reason: signed overflow would be undefined, and four squares of a large pixel
value reach 2^64.) The corrected float sums, which feed the reported profile
rather than any decision, change in their last bits as they already did between
runs.

Measured on a 16M-pixel rotation dataset: the kernel 116.8 s -> 19.1 s (6.1x),
no longer the largest; the whole run 70 s -> 39.5 s. Full 24-crystal battery:
same space group on all 24, none failed, 15m32s -> 12m47s.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
jungfrau
2026-08-15 19:21:24 -04:00
co-authored by Claude Opus 5
parent 56512414aa
commit eb634400ea
2 changed files with 104 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/spot_finding/AdaptiveSpotFinderGPU.cu
+100 -8
View File
@@ -12,6 +12,22 @@ inline void cuda_err(cudaError_t val) {
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
}
// One thread's merged contribution to a single ring, pushed to the block's shared accumulators.
__device__ __forceinline__ void flush_ring(unsigned long long *s_sum, unsigned long long *s_sum2,
uint32_t *s_count, float *s_sum_corr, float *s_sum2_corr,
bool accumulate_corrected, int b,
unsigned long long r_sum, unsigned long long r_sum2,
uint32_t r_count, float r_sum_corr, float r_sum2_corr) {
if (r_count == 0) return; // also covers the initial "no ring yet"
atomicAdd(&s_sum[b], r_sum);
atomicAdd(&s_sum2[b], r_sum2);
atomicAdd(&s_count[b], r_count);
if (accumulate_corrected) {
atomicAdd(&s_sum_corr[b], r_sum_corr);
atomicAdd(&s_sum2_corr[b], r_sum2_corr);
}
}
// One ring reduction, staging per-ring sums in shared memory (fast path). Shared layout:
// [ sum(float) | sum2(float) | count(uint32) | sum_corr(float) | sum2_corr(float) ] x nbins
// The corrected arrays exist only when accumulate_corrected is true (the plain first pass); on the
@@ -55,7 +71,71 @@ __global__ void reduce_rings_shared(
}
__syncthreads();
for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < npix; idx += blockDim.x * gridDim.x) {
// Four pixels per thread, read as one 16-byte and one 8-byte transaction instead of four of
// each. The image and both tables come straight from cudaMalloc, which aligns to at least the
// 16 bytes int4/float4 want, and no caller offsets them. npix need not be a multiple of four -
// the vector loop stops short and the leftovers are done one at a time below, so nothing is read
// past the end (which would fold uninitialised device memory into the accumulators).
const size_t stride = static_cast<size_t>(blockDim.x) * gridDim.x;
const size_t nquad = npix / 4;
for (size_t q = blockIdx.x * blockDim.x + threadIdx.x; q < nquad; q += stride) {
const int4 v4 = reinterpret_cast<const int4 *>(image)[q];
const ushort4 b4 = reinterpret_cast<const ushort4 *>(pixel_to_bin)[q];
float4 c4 = make_float4(0.0f, 0.0f, 0.0f, 0.0f);
if (accumulate_corrected)
c4 = reinterpret_cast<const float4 *>(corrections)[q];
const int32_t vq[4] = {v4.x, v4.y, v4.z, v4.w};
const uint16_t bq[4] = {b4.x, b4.y, b4.z, b4.w};
const float cq[4] = {c4.x, c4.y, c4.z, c4.w};
// A ring is several pixels wide, so consecutive pixels along a row usually fall in the same
// one. Carry a running total for the ring in registers and push it to shared memory only
// when the ring changes - one set of atomics for the run instead of one per pixel, which is
// what this kernel is actually limited by. A pixel dropped by one of the tests below does
// not end a run; it simply contributes nothing.
int r_b = -1;
unsigned long long r_sum = 0, r_sum2 = 0;
uint32_t r_count = 0;
float r_sum_corr = 0.0f, r_sum2_corr = 0.0f;
#pragma unroll
for (int k = 0; k < 4; k++) {
const int32_t v = vq[k];
if (v == INT32_MIN || v == INT32_MAX) continue;
const int b = bq[k];
if (b >= nbins) continue;
const float fv = static_cast<float>(v);
if (clip_k > 0.0f) {
const float lo = mean[b] - clip_k * sigma[b];
const float hi = mean[b] + clip_k * sigma[b];
if (fv < lo || fv > hi) continue;
}
if (b != r_b) {
flush_ring(s_sum, s_sum2, s_count, s_sum_corr, s_sum2_corr, accumulate_corrected,
r_b, r_sum, r_sum2, r_count, r_sum_corr, r_sum2_corr);
r_b = b;
r_sum = 0; r_sum2 = 0; r_count = 0;
r_sum_corr = 0.0f; r_sum2_corr = 0.0f;
}
// Unsigned, so a run that overflows wraps exactly as the shared accumulator would have:
// addition mod 2^64 is associative, which is what keeps the regrouping bit-identical.
r_sum += static_cast<unsigned long long>(static_cast<long long>(v));
r_sum2 += static_cast<unsigned long long>(static_cast<long long>(v) * v);
r_count += 1u;
if (accumulate_corrected) {
const float cv = fv * cq[k];
r_sum_corr += cv;
r_sum2_corr += cv * cv;
}
}
flush_ring(s_sum, s_sum2, s_count, s_sum_corr, s_sum2_corr, accumulate_corrected,
r_b, r_sum, r_sum2, r_count, r_sum_corr, r_sum2_corr);
}
// The last npix % 4 pixels, one per thread.
for (size_t idx = 4 * nquad + blockIdx.x * blockDim.x + threadIdx.x; idx < npix; idx += stride) {
const int32_t v = image[idx];
if (v == INT32_MIN || v == INT32_MAX) continue;
const uint16_t b = pixel_to_bin[idx];
@@ -207,12 +287,6 @@ AdaptiveSpotFinderGPU::AdaptiveSpotFinderGPU(const AzimuthalIntegrationMapping &
cuda_err(cudaGetDevice(&device));
cudaDeviceProp prop{};
cuda_err(cudaGetDeviceProperties(&prop, device));
// Eight blocks per SM, not four. Both kernels are grid-stride loops, so any grid is correct and
// a device that cannot co-schedule eight simply queues the rest - but four left only 512 of the
// 1536 threads an SM can hold resident (33%), and reduce_rings_shared is bound by shared-memory
// atomic replay rather than by bandwidth, which is exactly the case that wants more resident
// warps to hide the serialisation. The per-block histogram is nbins * 20 B (~9.6 kB at the
// default 0.01 1/A spacing), so eight blocks fit in an SM's shared memory with room to spare.
reduce_blocks = 8 * prop.multiProcessorCount;
// flag_strong stays at four: it is bandwidth-shaped rather than atomic-bound, and eight
// measured no better (181 vs 175 us/launch).
@@ -222,6 +296,23 @@ AdaptiveSpotFinderGPU::AdaptiveSpotFinderGPU(const AzimuthalIntegrationMapping &
shared_clip = static_cast<size_t>(nbins) * (2 * sizeof(unsigned long long) + sizeof(uint32_t));
use_shared = (shared_plain < prop.sharedMemPerBlock);
// Launch as many blocks as the device can actually hold resident, and ask it rather than guess.
// The two passes have different shared-memory footprints, so they do not fit the same number of
// blocks: the plain pass carries the corrected rings as well, and a grid sized for the clip pass
// left it running a second wave at a quarter of the occupancy. The kernel is a grid-stride loop,
// so any grid is correct - but a block that is not resident is a wave, not parallelism. This also
// tracks nbins, which is data-driven: a fine q spacing shrinks the number of blocks that fit.
const auto blocks_per_sm = [&](size_t shared) {
int bpsm = 0;
cuda_err(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&bpsm, reduce_rings_shared,
reduce_threads, shared));
return std::max(1, bpsm) * prop.multiProcessorCount;
};
if (use_shared) {
reduce_blocks_plain = blocks_per_sm(shared_plain);
reduce_blocks_clip = blocks_per_sm(shared_clip);
}
// Both tables are functions of the detector geometry alone, so they are uploaded once per GPU and
// shared: the azimuthal-integration engine in the same worker reads the very same two arrays.
gpu_pixel_to_bin = SharedDeviceTable(mapping.GetPixelToBin().data(), npix,
@@ -234,7 +325,8 @@ void AdaptiveSpotFinderGPU::ReducePass(const ImagePreprocessorBuffer &image, flo
bool accumulate_corrected) {
if (use_shared) {
const size_t shared = accumulate_corrected ? shared_plain : shared_clip;
reduce_rings_shared<<<reduce_blocks, reduce_threads, shared, *stream>>>(
const int blocks = accumulate_corrected ? reduce_blocks_plain : reduce_blocks_clip;
reduce_rings_shared<<<blocks, reduce_threads, shared, *stream>>>(
gpu_pixel_to_bin->get(), gpu_corrections->get(), image.getGPUBuffer(), gpu_mean, gpu_sigma,
clip_k, accumulate_corrected, gpu_sum, gpu_sum2, gpu_count, gpu_sum_corr, gpu_sum2_corr,
npix, nbins);
image_analysis/spot_finding/AdaptiveSpotFinderGPU.h
+4 -2
View File
@@ -40,8 +40,10 @@ class AdaptiveSpotFinderGPU : public ImageSpotFinder {
const int nbins;
const size_t npix;
int reduce_threads = 128;
int reduce_blocks = 0;
int reduce_threads = 256;
int reduce_blocks = 0; // global-atomics fallback
int reduce_blocks_plain = 0; // as many blocks as actually fit, per shared-memory footprint
int reduce_blocks_clip = 0;
int flag_threads = 256;
int flag_blocks = 0;
size_t shared_plain = 0; // per-block shared bytes for the plain pass (raw + corrected rings)