From 03dee2dda7d4e84759188a9708710457afc8218c Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Sat, 3 Oct 2026 11:43:12 +0200 Subject: [PATCH] AzIntEngineGPU: four pixels per thread, one shared add per ring run The standalone GPU azimuthal integration did three shared-memory atomics per pixel on the same few ring addresses, which is what it was limited by. It now reads four pixels per thread as vector loads and keeps a running total per ring, flushed when the ring changes - the scheme the adaptive finder's ring pass (reduce_rings_shared) already uses. The npix % 4 leftovers are done one at a time. Used wherever the fused adaptive engine is not (fixed-threshold spot finding, the broker's non-adaptive path). Measured on a 16 Mpx sweep (1800 frames, --no-adaptive-spots, RTX 5080): 843 -> 295 us per call (min 621 -> 196 us). Not bit-identical, and the old kernel was not either: float atomics arrive in any order, so two runs of the OLD kernel already differ by up to 1.7e-6 relative in the per-frame profile; new vs old differs by up to 1.9e-6, the same order. Per-ring pixel counts are identical. Default rugnux runs do not reach this kernel (p.mtz md5 unchanged on three sets); on the fixed-threshold path p_unmerged.mtz is md5-identical to the old kernel's. New test: GPU vs CPU engine on a pixel count that is not a multiple of four, with masked and saturated pixels. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01K5K8jvPPbmCrbqnWkddTuB --- image_analysis/azint/AzIntEngineGPU.cu | 57 ++++++++++++++++++++++---- tests/AdaptiveSpotFinderGPUTest.cpp | 41 ++++++++++++++++++ 2 files changed, 89 insertions(+), 9 deletions(-) diff --git a/image_analysis/azint/AzIntEngineGPU.cu b/image_analysis/azint/AzIntEngineGPU.cu index b35adaa95..7de941260 100644 --- a/image_analysis/azint/AzIntEngineGPU.cu +++ b/image_analysis/azint/AzIntEngineGPU.cu @@ -8,6 +8,15 @@ inline void cuda_err(cudaError_t val) { throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val)); } +// Pushes one ring run's totals to the shared accumulators; nothing for an empty run. +__device__ __forceinline__ void flush_azim_run(float *s_sum, float *s_sum2, uint32_t *s_count, + int b, float r_sum, float r_sum2, uint32_t r_count) { + 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); +} + __global__ void gpu_azim_shared( const uint16_t *__restrict__ pixel_to_bin, @@ -33,19 +42,49 @@ void gpu_azim_shared( __syncthreads(); - for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; - idx < num_pixels; - idx += blockDim.x * gridDim.x) { - uint16_t bin = pixel_to_bin[idx]; + // Four pixels per thread, read as vector loads, and a running total per ring pushed to shared + // memory only when the ring changes: consecutive pixels along a row mostly share a ring, and an + // atomic per pixel on the same few addresses is what this kernel was limited by. The same scheme + // as the adaptive spot finder's ring pass (reduce_rings_shared). The buffers come straight from + // cudaMalloc, aligned for int4/float4; the npix % 4 leftovers are done one at a time below. + const size_t stride = static_cast(blockDim.x) * gridDim.x; + const size_t nquad = num_pixels / 4; + for (size_t q = blockIdx.x * blockDim.x + threadIdx.x; q < nquad; q += stride) { + const int4 v4 = reinterpret_cast(input_buffer)[q]; + const ushort4 b4 = reinterpret_cast(pixel_to_bin)[q]; + const float4 c4 = reinterpret_cast(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}; - int32_t v = input_buffer[idx]; - bool valid = (v != INT32_MIN) & (v != INT32_MAX); + int r_b = -1; + float r_sum = 0.0f, r_sum2 = 0.0f; + uint32_t r_count = 0; + #pragma unroll + for (int k = 0; k < 4; k++) { + const int32_t v = vq[k]; + const int b = bq[k]; + if (v == INT32_MIN || v == INT32_MAX || b >= azint_bins) continue; + if (b != r_b) { + flush_azim_run(s_sum, s_sum2, s_count, r_b, r_sum, r_sum2, r_count); + r_b = b; + r_sum = 0.0f; r_sum2 = 0.0f; r_count = 0; + } + const float val = static_cast(v) * cq[k]; + r_sum += val; + r_sum2 += val * val; + r_count += 1; + } + flush_azim_run(s_sum, s_sum2, s_count, r_b, r_sum, r_sum2, r_count); + } - if (bin < azint_bins && valid) { + for (size_t idx = 4 * nquad + blockIdx.x * blockDim.x + threadIdx.x; idx < num_pixels; idx += stride) { + const uint16_t bin = pixel_to_bin[idx]; + const int32_t v = input_buffer[idx]; + if (bin < azint_bins && v != INT32_MIN && v != INT32_MAX) { const float val = static_cast(v) * corrections[idx]; - const float val2 = val * val; atomicAdd(&s_sum[bin], val); - atomicAdd(&s_sum2[bin], val2); + atomicAdd(&s_sum2[bin], val * val); atomicAdd(&s_count[bin], 1); } } diff --git a/tests/AdaptiveSpotFinderGPUTest.cpp b/tests/AdaptiveSpotFinderGPUTest.cpp index 87d9d9e44..efee443fd 100644 --- a/tests/AdaptiveSpotFinderGPUTest.cpp +++ b/tests/AdaptiveSpotFinderGPUTest.cpp @@ -12,6 +12,7 @@ #include "../common/AzimuthalIntegrationMapping.h" #include "../common/AzimuthalIntegrationProfile.h" +#include "../image_analysis/azint/AzIntEngineCPU.h" #include "../image_analysis/azint/AzIntEngineGPU.h" #include "../image_analysis/spot_finding/AdaptiveSpotFinderCPU.h" #include "../image_analysis/spot_finding/AdaptiveSpotFinderGPU.h" @@ -161,6 +162,46 @@ TEST_CASE("AdaptiveSpotFinderGPU_AzimuthalIntegration", "[AdaptiveSpotFinderGPU] } } +// The GPU azimuthal integration against the CPU one, on a pixel count that is not a multiple of four (the +// kernel reads four pixels at a time and does the rest one by one) and with masked and saturated pixels +// in it. The per-ring pixel counts are integers and must agree exactly; the float sums only to rounding. +TEST_CASE("AzIntEngineGPU_MatchesCPU", "[AdaptiveSpotFinderGPU]") { + if (get_gpu_count() == 0) { + WARN("No CUDA GPU present. Skipping AzIntEngineGPU_MatchesCPU"); + return; + } + + DiffractionExperiment x(DetDECTRIS(1031, 1063, "Test", {})); + x.DetectorDistance_mm(80).BeamX_pxl(515).BeamY_pxl(530); + x.QSpacingForAzimInt_recipA(0.05).QRangeForAzimInt_recipA(0.05, 5.0); + REQUIRE(x.GetPixelsNum() % 4 != 0); + PixelMask pixel_mask(x); + AzimuthalIntegrationMapping mapping(x, pixel_mask); + + ImagePreprocessorBufferGPU buffer(x.GetPixelsNum()); + for (size_t i = 0; i < x.GetPixelsNum(); i++) + buffer[i] = (i % 997 == 0) ? INT32_MIN : (i % 1009 == 0) ? INT32_MAX + : 8 + static_cast((i * 7919) % 23); + REQUIRE(cudaMemcpy(buffer.getGPUBuffer(), buffer.getBuffer().data(), + x.GetPixelsNum() * sizeof(int32_t), cudaMemcpyHostToDevice) == cudaSuccess); + REQUIRE(cudaDeviceSynchronize() == cudaSuccess); + + AzimuthalIntegrationProfile cpu_profile(mapping), gpu_profile(mapping); + AzIntEngineCPU(mapping).Run(buffer, cpu_profile); + AzIntEngineGPU(mapping, std::make_shared()).Run(buffer, gpu_profile); + + REQUIRE(gpu_profile.GetPixelCount() == cpu_profile.GetPixelCount()); + const auto ref = cpu_profile.GetResult(); + const auto got = gpu_profile.GetResult(); + REQUIRE(ref.size() == got.size()); + for (size_t b = 0; b < ref.size(); b++) { + if (std::isnan(ref[b])) + CHECK(std::isnan(got[b])); + else + CHECK(got[b] == Catch::Approx(ref[b]).epsilon(1e-5)); + } +} + // The ring sums are built by atomics, which arrive in an arbitrary order, so the same frame has to be // re-run to show the engine agrees with itself: detection is a hard "value >= threshold" on integer // counts, and a threshold that wobbles between runs flips pixels on the boundary and with them the size