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

Merge branch 'gpu-spot-histclip' into rc173 (GPU adaptive spot finder: sigma clips from a per-ring histogram)

Browse Source
This commit is contained in:
leonarski_f
2026-09-27 21:39:12 +02:00
parent 03309032bd ccb38f469b
commit f4f9157cc0
3 changed files with 258 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/spot_finding/AdaptiveSpotFinderGPU.cu
+131 -21
View File
@@ -28,6 +28,43 @@ __device__ __forceinline__ void flush_ring(unsigned long long *s_sum, unsigned l
}
}
// The sigma-clip test, with both bounds as single fused multiply-adds. Written out rather than left to
// the compiler, which fuses `mean - clip_k * sigma` in one kernel and may not in another - and the
// histogram clip below must draw exactly the line the image pass draws.
__device__ __forceinline__ bool clip_keep(float fv, float mean, float sigma, float clip_k) {
const float lo = __fmaf_rn(-clip_k, sigma, mean);
const float hi = __fmaf_rn(clip_k, sigma, mean);
return !(fv < lo || fv > hi);
}
// One valid pixel of the plain pass into the per-ring value histogram, or - outside [0, HIST_VALUES) -
// onto the overflow list. Each warp holds 128 consecutive pixels, nearly all in one ring and at a few
// low values, so the warp's pixels are grouped by (ring, value) and each group is one atomic.
__device__ __forceinline__ void count_value(uint32_t *hist, int2 *overflow, uint32_t *overflow_n,
uint32_t overflow_cap, bool valid, int b, int32_t v) {
const bool in_hist = valid && v >= 0 && v < AdaptiveSpotFinderGPU::HIST_VALUES;
const bool in_overflow = valid && !in_hist;
const unsigned active = __activemask();
const int lane = threadIdx.x % 32;
const uint32_t key = in_hist ? b * AdaptiveSpotFinderGPU::HIST_VALUES + v : UINT32_MAX;
const unsigned same = __match_any_sync(active, key);
if (in_hist && lane == __ffs(same) - 1)
atomicAdd(&hist[key], __popc(same));
const unsigned ovf = __ballot_sync(active, in_overflow);
if (in_overflow) {
const int leader = __ffs(ovf) - 1;
uint32_t base = 0;
if (lane == leader)
base = atomicAdd(overflow_n, __popc(ovf));
base = __shfl_sync(ovf, base, leader);
const uint32_t i = base + __popc(ovf & ((1u << lane) - 1));
if (i < overflow_cap) // past the end the list is only counted, and the image pass clips
overflow[i] = make_int2(b, v);
}
}
// 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
@@ -42,8 +79,15 @@ __global__ void reduce_rings_shared(
bool accumulate_corrected,
unsigned long long *__restrict__ sum, unsigned long long *__restrict__ sum2, uint32_t *__restrict__ count,
float *__restrict__ sum_corr, float *__restrict__ sum2_corr,
uint32_t *__restrict__ hist, int2 *__restrict__ overflow, uint32_t *__restrict__ overflow_n,
uint32_t overflow_cap,
size_t npix, int nbins) {
// A sigma-clip pass reads the image only when the plain pass's overflow list did not fit;
// otherwise clip_rings_from_hist has done it from the histogram.
if (clip_k > 0.0f && *overflow_n <= overflow_cap)
return;
// The raw accumulators are INTEGERS, not floats. A preprocessed pixel is an exact int32 (the
// masked and saturated sentinels are skipped below), so v and v*v are exact in 64 bits and
// integer addition is associative - which makes the ring mean and sigma, and therefore the
@@ -103,15 +147,14 @@ __global__ void reduce_rings_shared(
#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;
}
const bool valid = v != INT32_MIN && v != INT32_MAX && b < nbins
&& (clip_k <= 0.0f || clip_keep(fv, mean[b], sigma[b], clip_k));
// The plain pass also records the values, so that the clip passes need not re-read the image.
if (accumulate_corrected)
count_value(hist, overflow, overflow_n, overflow_cap, valid, b, v);
if (!valid) 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);
@@ -137,15 +180,13 @@ __global__ void reduce_rings_shared(
// 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];
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;
}
const bool valid = v != INT32_MIN && v != INT32_MAX && b < nbins
&& (clip_k <= 0.0f || clip_keep(fv, mean[b], sigma[b], clip_k));
if (accumulate_corrected)
count_value(hist, overflow, overflow_n, overflow_cap, valid, b, v);
if (!valid) continue;
atomicAdd(&s_sum[b], static_cast<unsigned long long>(static_cast<long long>(v)));
atomicAdd(&s_sum2[b], static_cast<unsigned long long>(static_cast<long long>(v) * v));
atomicAdd(&s_count[b], 1u);
@@ -188,11 +229,7 @@ __global__ void reduce_rings_global(
const uint16_t b = pixel_to_bin[idx];
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 (clip_k > 0.0f && !clip_keep(fv, mean[b], sigma[b], clip_k)) continue;
atomicAdd(&sum[b], static_cast<unsigned long long>(static_cast<long long>(v)));
atomicAdd(&sum2[b], static_cast<unsigned long long>(static_cast<long long>(v) * v));
atomicAdd(&count[b], 1u);
@@ -204,6 +241,60 @@ __global__ void reduce_rings_global(
}
}
// A sigma-clip pass from the plain pass's per-ring value histogram and overflow list instead of the
// image, as AdaptiveSpotFinderCPU::AccumulateRings does it: each value of a ring meets the same test
// the pixels holding it would, and its pixels are added as a count. The sums are integers, so they
// are the image pass's exactly. One warp per ring; the overflow list is shared out across the grid.
// If the list did not fit, this does nothing and reduce_rings_shared reads the image instead.
__global__ void clip_rings_from_hist(const uint32_t *__restrict__ hist,
const int2 *__restrict__ overflow,
const uint32_t *__restrict__ overflow_n,
uint32_t overflow_cap,
const float *__restrict__ mean,
const float *__restrict__ sigma,
float clip_k,
unsigned long long *__restrict__ sum, unsigned long long *__restrict__ sum2,
uint32_t *__restrict__ count, int nbins) {
const uint32_t n_overflow = *overflow_n;
if (n_overflow > overflow_cap)
return;
const int lane = threadIdx.x % 32;
const int warp = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
const int nwarps = (gridDim.x * blockDim.x) / 32;
for (int b = warp; b < nbins; b += nwarps) {
unsigned long long r_sum = 0, r_sum2 = 0;
uint32_t r_count = 0;
for (int v = lane; v < AdaptiveSpotFinderGPU::HIST_VALUES; v += 32) {
const uint32_t n = hist[b * AdaptiveSpotFinderGPU::HIST_VALUES + v];
if (n == 0 || !clip_keep(static_cast<float>(v), mean[b], sigma[b], clip_k)) continue;
r_sum += static_cast<unsigned long long>(n) * v;
r_sum2 += static_cast<unsigned long long>(n) * static_cast<unsigned long long>(v * v);
r_count += n;
}
for (int offset = 16; offset > 0; offset /= 2) {
r_sum += __shfl_down_sync(0xffffffff, r_sum, offset);
r_sum2 += __shfl_down_sync(0xffffffff, r_sum2, offset);
r_count += __shfl_down_sync(0xffffffff, r_count, offset);
}
if (lane == 0 && r_count > 0) {
atomicAdd(&sum[b], r_sum);
atomicAdd(&sum2[b], r_sum2);
atomicAdd(&count[b], r_count);
}
}
for (uint32_t i = blockIdx.x * blockDim.x + threadIdx.x; i < n_overflow; i += gridDim.x * blockDim.x) {
const int b = overflow[i].x;
const int32_t v = overflow[i].y;
if (!clip_keep(static_cast<float>(v), mean[b], sigma[b], clip_k)) continue;
atomicAdd(&sum[b], static_cast<unsigned long long>(static_cast<long long>(v)));
atomicAdd(&sum2[b], static_cast<unsigned long long>(static_cast<long long>(v) * v));
atomicAdd(&count[b], 1u);
}
}
// Per-ring mean/sigma from the current raw accumulators. Rings with no pixels this pass keep their
// previous value (matches the CPU, which leaves ring_mean/ring_sigma untouched when the count is 0).
__global__ void finalize_rings(const unsigned long long *__restrict__ sum,
@@ -339,6 +430,12 @@ AdaptiveSpotFinderGPU::AdaptiveSpotFinderGPU(const AzimuthalIntegrationMapping &
if (use_shared) {
reduce_blocks_plain = blocks_per_sm(shared_plain);
reduce_blocks_clip = blocks_per_sm(shared_clip);
// The histogram clip needs the shared-memory path's small ring count: at the bin counts that
// fall back to global atomics it would be hundreds of MB.
overflow_cap = npix / 32;
gpu_hist = CudaDevicePtr<uint32_t>(static_cast<size_t>(nbins) * HIST_VALUES);
gpu_overflow = CudaDevicePtr<int2>(overflow_cap);
gpu_overflow_n = CudaDevicePtr<uint32_t>(1);
}
// Both tables are functions of the detector geometry alone, so they are uploaded once per GPU and
@@ -356,10 +453,17 @@ void AdaptiveSpotFinderGPU::ReducePass(const ImagePreprocessorBuffer &image, flo
if (use_shared) {
const size_t shared = accumulate_corrected ? shared_plain : shared_clip;
const int blocks = accumulate_corrected ? reduce_blocks_plain : reduce_blocks_clip;
if (!accumulate_corrected) {
clip_rings_from_hist<<<(nbins + 7) / 8, 256, 0, *stream>>>(
gpu_hist, gpu_overflow, gpu_overflow_n, overflow_cap, gpu_mean, gpu_sigma, clip_k,
gpu_sum, gpu_sum2, gpu_count, nbins);
cuda_err(cudaGetLastError());
}
// On a clip pass this returns at once unless the overflow list did not fit.
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);
gpu_hist, gpu_overflow, gpu_overflow_n, overflow_cap, npix, nbins);
cuda_err(cudaGetLastError());
} else {
reduce_rings_global<<<reduce_blocks, reduce_threads, 0, *stream>>>(
@@ -423,7 +527,8 @@ void AdaptiveSpotFinderGPU::Detect(const ImagePreprocessorBuffer &image,
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
"AdaptiveSpotFinderGPU::Detect: mismatch in pixel size");
// --- Stage A: robust per-ring background (one plain pass + two sigma-clip passes) ---
// --- Stage A: robust per-ring background (one plain pass + two sigma-clip passes, the clips from
// the plain pass's value histogram) ---
cuda_err(cudaMemsetAsync(gpu_sum, 0, sizeof(unsigned long long) * nbins, *stream));
cuda_err(cudaMemsetAsync(gpu_sum2, 0, sizeof(unsigned long long) * nbins, *stream));
cuda_err(cudaMemsetAsync(gpu_count, 0, sizeof(uint32_t) * nbins, *stream));
@@ -431,6 +536,10 @@ void AdaptiveSpotFinderGPU::Detect(const ImagePreprocessorBuffer &image,
cuda_err(cudaMemsetAsync(gpu_sigma, 0, sizeof(float) * nbins, *stream));
cuda_err(cudaMemsetAsync(gpu_sum_corr, 0, sizeof(float) * nbins, *stream));
cuda_err(cudaMemsetAsync(gpu_sum2_corr, 0, sizeof(float) * nbins, *stream));
if (use_shared) {
cuda_err(cudaMemsetAsync(gpu_hist, 0, sizeof(uint32_t) * nbins * HIST_VALUES, *stream));
cuda_err(cudaMemsetAsync(gpu_overflow_n, 0, sizeof(uint32_t), *stream));
}
ReducePass(image, 0.0f, true); // plain pass also fills the corrected profile accumulators
FinalizeStats();
@@ -455,6 +564,7 @@ void AdaptiveSpotFinderGPU::Detect(const ImagePreprocessorBuffer &image,
cuda_err(cudaMemcpyAsync(host_count.data(), gpu_count, sizeof(uint32_t) * nbins, cudaMemcpyDeviceToHost, *stream));
cuda_err(cudaStreamSynchronize(*stream));
// --- Stage B: per-ring threshold on the host (shared with the CPU finder) ---
ComputeThresholds(settings);
image_analysis/spot_finding/AdaptiveSpotFinderGPU.h
+11
View File
@@ -66,6 +66,15 @@ class AdaptiveSpotFinderGPU : public ImageSpotFinderGPU {
CudaDevicePtr<float> gpu_mean; // per-ring raw mean (clip predicate)
CudaDevicePtr<float> gpu_sigma; // per-ring raw sigma (clip predicate)
// The plain pass's valid pixels as a per-ring histogram of their values, and the (ring, value) of
// those outside [0, HIST_VALUES) - so the sigma-clip passes need not re-read the image. A list
// longer than overflow_cap is only counted, and the clip passes then read the image instead.
// Allocated only on the shared-memory path (few rings).
CudaDevicePtr<uint32_t> gpu_hist;
CudaDevicePtr<int2> gpu_overflow;
CudaDevicePtr<uint32_t> gpu_overflow_n;
uint32_t overflow_cap = 0;
// Corrected per-ring accumulators (plain first pass only) -> azimuthal-integration profile.
CudaDevicePtr<float> gpu_sum_corr;
CudaDevicePtr<float> gpu_sum2_corr;
@@ -109,6 +118,8 @@ class AdaptiveSpotFinderGPU : public ImageSpotFinderGPU {
void ComputeThresholds(const SpotFindingSettings &settings);
public:
static constexpr int32_t HIST_VALUES = 1024; // as AdaptiveSpotFinderCPU
AdaptiveSpotFinderGPU(const AzimuthalIntegrationMapping &mapping, std::shared_ptr<CudaStream> stream);
~AdaptiveSpotFinderGPU() override = default;
AdaptiveSpotFinderGPU(const AdaptiveSpotFinderGPU &) = delete;
tests/AdaptiveSpotFinderGPUTest.cpp
+116
View File
@@ -8,6 +8,7 @@
#include <algorithm>
#include <chrono>
#include <cmath>
#include "../common/AzimuthalIntegrationMapping.h"
#include "../common/AzimuthalIntegrationProfile.h"
@@ -15,6 +16,7 @@
#include "../image_analysis/spot_finding/AdaptiveSpotFinderCPU.h"
#include "../image_analysis/spot_finding/AdaptiveSpotFinderGPU.h"
#include "../image_analysis/spot_finding/ImageSpotFinderGPU.h"
#include "../image_analysis/spot_finding/AdaptiveThreshold.h"
#include "../image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.h"
namespace {
@@ -232,6 +234,120 @@ TEST_CASE("AdaptiveSpotFinderGPU_RingStatsMatchCPUExactly", "[AdaptiveSpotFinder
REQUIRE(SortedCoords(gpu_spots) == SortedCoords(cpu_spots));
}
// The per-ring background the GPU engine reported before its sigma clips were taken from a value
// histogram: a plain pass and two 3-sigma clip passes, each over every pixel of the image, with the
// clip bounds as the fused multiply-adds the image-pass kernel computes them with.
std::vector<float> ImagePassRingBackground(const ImagePreprocessorBuffer &image,
const AzimuthalIntegrationMapping &mapping) {
const auto &pixel_to_bin = mapping.GetPixelToBin();
const size_t nbins = mapping.GetBinNumber();
std::vector<float> mean(nbins, 0.0f), sigma(nbins, 0.0f);
std::vector<unsigned long long> sum(nbins), sum2(nbins);
std::vector<uint32_t> count(nbins);
for (int pass = 0; pass < 3; pass++) {
std::fill(sum.begin(), sum.end(), 0);
std::fill(sum2.begin(), sum2.end(), 0);
std::fill(count.begin(), count.end(), 0);
for (size_t i = 0; i < pixel_to_bin.size(); i++) {
const int32_t v = image[i];
const uint16_t b = pixel_to_bin[i];
if (v == INT32_MIN || v == INT32_MAX || b >= nbins) continue;
const float fv = static_cast<float>(v);
if (pass > 0 && (fv < std::fma(-3.0f, sigma[b], mean[b]) || fv > std::fma(3.0f, sigma[b], mean[b])))
continue;
sum[b] += static_cast<unsigned long long>(static_cast<long long>(v));
sum2[b] += static_cast<unsigned long long>(static_cast<long long>(v) * v);
count[b] += 1;
}
for (size_t b = 0; b < nbins; b++) {
if (count[b] == 0) continue;
const double m = static_cast<double>(static_cast<long long>(sum[b])) / count[b];
const double var = std::max(0.0, static_cast<double>(sum2[b]) / count[b] - m * m);
mean[b] = static_cast<float>(m);
sigma[b] = static_cast<float>(std::sqrt(var));
}
}
std::vector<float> bkg(nbins);
for (size_t b = 0; b < nbins; b++)
bkg[b] = (count[b] < adaptive_threshold::MIN_RING_PIXELS) ? NAN : mean[b];
return bkg;
}
// The sigma clips are taken from the plain pass's per-ring value histogram, with the values outside it
// on an overflow list - and, when that list does not fit, from the image after all. Every route has to
// give the ring statistics the image passes give, to the bit: they are integer sums of the same values.
TEST_CASE("AdaptiveSpotFinderGPU_HistogramClipMatchesImagePasses", "[AdaptiveSpotFinderGPU]") {
if (get_gpu_count() == 0) {
WARN("No CUDA GPU present. Skipping AdaptiveSpotFinderGPU_HistogramClipMatchesImagePasses");
return;
}
DiffractionExperiment x = MakeExperiment();
PixelMask pixel_mask(x);
AzimuthalIntegrationMapping mapping(x, pixel_mask);
const auto &pixel_to_bin = mapping.GetPixelToBin();
const size_t npix = x.GetPixelsNum();
const SpotFindingSettings settings = AdaptiveSettings();
auto stream = std::make_shared<CudaStream>();
AdaptiveSpotFinderGPU gpu(mapping, stream);
const auto check = [&](ImagePreprocessorBufferGPU &buffer) {
REQUIRE(cudaMemcpy(buffer.getGPUBuffer(), buffer.getBuffer().data(),
npix * sizeof(int32_t), cudaMemcpyHostToDevice) == cudaSuccess);
REQUIRE(cudaDeviceSynchronize() == cudaSuccess);
gpu.Run(buffer, settings);
const auto expected = ImagePassRingBackground(buffer, mapping);
const auto &got = gpu.GetRingBackground();
REQUIRE(got.size() == expected.size());
size_t differing = 0, rings = 0;
for (size_t b = 0; b < got.size(); b++) {
if (std::isnan(expected[b])) {
if (!std::isnan(got[b])) differing++;
} else {
rings++;
if (got[b] != expected[b]) differing++;
}
}
CHECK(rings > 10);
CHECK(differing == 0);
};
const auto count_outside = [&](const ImagePreprocessorBuffer &buffer) {
size_t n = 0;
for (size_t i = 0; i < npix; i++)
if (pixel_to_bin[i] < mapping.GetBinNumber() && buffer[i] != INT32_MIN && buffer[i] != INT32_MAX
&& (buffer[i] < 0 || buffer[i] >= AdaptiveSpotFinderGPU::HIST_VALUES))
n++;
return n;
};
ImagePreprocessorBufferGPU buffer(npix);
SECTION("Values beyond the histogram, on the overflow list") {
// Two rings far above the histogram's range and one straddling zero, next to rings inside it;
// with the bright blobs, a few percent of the pixels at most.
FillTestImage(buffer, x);
for (size_t i = 0; i < npix; i++) {
const uint16_t b = pixel_to_bin[i];
if (b == 20 || b == 21)
buffer[i] = 1020 + static_cast<int32_t>(i % 13); // across the top edge
else if (b == 30)
buffer[i] = -4 + static_cast<int32_t>(i % 9); // across zero
}
const size_t outside = count_outside(buffer);
REQUIRE(outside > 0);
REQUIRE(outside <= npix / 32);
check(buffer);
}
SECTION("Overflow list too long, clipped from the image") {
for (size_t i = 0; i < npix; i++)
buffer[i] = 1500 + static_cast<int32_t>((i * 7) % 61) + ((i % 997 == 0) ? 3000 : 0);
REQUIRE(count_outside(buffer) > npix / 32);
check(buffer);
}
}
TEST_CASE("AdaptiveSpotFinderGPU_Speed", "[AdaptiveSpotFinderGPU][.benchmark]") {
if (get_gpu_count() == 0) {
WARN("No CUDA GPU present. Skipping AdaptiveSpotFinderGPU_Speed");