leonarski_fandClaude Opus 5.5
09c41d067d
The GPU engines are built per worker per step, and each took its page-locked host buffers with cudaHostAlloc / cudaHostRegister and gave them back with cudaFreeHost / cudaHostUnregister. Those calls map and lock pages under the process's memory-map lock, so all workers queue behind each other: on the 5080 workstation a myoglobin run made ~3400 of them (6-22 thread-s, depending on the load), 9gdj ~10400 (124 thread-s); on a four-socket host they cost 20-30 ms each. - CudaHostPtr takes its memory from a pool: a freed buffer is kept in a power-of-two size class for the next engine, up to 512 MiB in all; a buffer above 64 MiB is never kept. Rugnux::Run hands the pool back at its end. - PinnedVector (std::vector with the pool's allocator) replaces the registered vectors of AdaptiveSpotFinderGPU and ImagePreprocessorGPU; the Bragg predictors download into a pooled pinned buffer and copy into their output, instead of registering it. - BSLZ4DecoderGPU copies the compressed chunk into a pooled pinned buffer before the upload, so it is one DMA instead of the driver's pageable staging (9gdj: 163 GB pageable H2D, 52 s of copy time, -> 147 GB pinned, 13 s). Measured with nsys on the workstation: myob 3398 -> 1396 calls (674 of them the release at the end), cytc 2986 -> 1270, 9gdj 10428 -> 1748 (124 -> 3.7 thread-s). Output identical (p.hkl md5 on myob, cytc, 9gdj; CI tier battery). Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3
675 lines
35 KiB
Plaintext
675 lines
35 KiB
Plaintext
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include "AdaptiveSpotFinderGPU.h"
|
|
#include "AdaptiveThreshold.h"
|
|
#include "../../common/JFJochException.h"
|
|
|
|
namespace {
|
|
|
|
inline void cuda_err(cudaError_t val) {
|
|
if (val != cudaSuccess)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
|
|
}
|
|
|
|
// Fixed-point scale of the corrected (profile) sums. A corrected value is a pixel times a float
|
|
// correction, so unlike the raw sums it has no exact integer form, and a float atomicAdd adds the
|
|
// values in whatever order the threads and blocks happen to arrive - the profile, and the background
|
|
// estimate read from it, then moved in their last bits between two runs of the same command. Each
|
|
// value is quantised on its own, and the sum of the quanta is integer addition, which is associative:
|
|
// the result no longer depends on the order, nor on how the grid was sized for this GPU. 2^12 puts
|
|
// the quantum at 2.4e-4 of a count, and leaves room for a ring summing to ~2e15 in either sum.
|
|
constexpr float CORR_FIXED = 4096.0f; // 2^12
|
|
|
|
// Signed, carried as two's complement like the raw sums.
|
|
__device__ __forceinline__ unsigned long long corr_fixed(float v) {
|
|
return static_cast<unsigned long long>(static_cast<long long>(llrintf(v * CORR_FIXED)));
|
|
}
|
|
|
|
// 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, unsigned long long *s_sum_corr,
|
|
unsigned long long *s_sum2_corr,
|
|
bool accumulate_corrected, int b,
|
|
unsigned long long r_sum, unsigned long long r_sum2,
|
|
uint32_t r_count, unsigned long long r_sum_corr,
|
|
unsigned long long 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);
|
|
}
|
|
}
|
|
|
|
// 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);
|
|
}
|
|
}
|
|
|
|
// A pixel value as an unsigned key with the same order, the masked sentinel INT32_MIN mapping to 0 -
|
|
// so a zeroed word maximum is "every pixel masked" and atomicMax on unsigned keeps the order.
|
|
__device__ __forceinline__ uint32_t order_key(int32_t v) {
|
|
return static_cast<uint32_t>(v) ^ 0x80000000u;
|
|
}
|
|
|
|
// The maximum of a quad's key over the eight lanes that hold one 32-pixel word (quads 8w .. 8w + 7;
|
|
// a warp starts on a multiple of 32 quads, so the eight are lanes 8j .. 8j + 7), written by the first
|
|
// of them. Only the warp's last pass can have lanes past the end, and those are a suffix: a partner
|
|
// outside the active mask is simply not taken. A word whose eight quads are all here has no other
|
|
// writer, so it is stored outright; only the last word of the image, which may also hold the pixels
|
|
// past the last quad, is merged with an atomic (an atomic per word measured ~10 % of the kernel).
|
|
__device__ __forceinline__ void note_word_max(uint32_t *word_max, size_t q, uint32_t key) {
|
|
const unsigned active = __activemask();
|
|
const int lane = threadIdx.x % 32;
|
|
for (int offset = 1; offset < 8; offset *= 2) {
|
|
const uint32_t other = __shfl_xor_sync(active, key, offset);
|
|
if (active & (1u << (lane ^ offset)))
|
|
key = max(key, other);
|
|
}
|
|
if (lane % 8 == 0) {
|
|
if (((active >> lane) & 0xFFu) == 0xFFu)
|
|
word_max[q / 8] = key;
|
|
else
|
|
atomicMax(&word_max[q / 8], key);
|
|
}
|
|
}
|
|
|
|
// One ring reduction, staging per-ring sums in shared memory (fast path). Shared layout:
|
|
// [ sum | sum2 | sum_corr | sum2_corr | count(uint32) ] x nbins
|
|
// The corrected arrays exist only when accumulate_corrected is true (the plain first pass); on the
|
|
// sigma-clip passes they are left out and the count follows sum2.
|
|
__global__ void reduce_rings_shared(
|
|
const uint16_t *__restrict__ pixel_to_bin,
|
|
const float *__restrict__ corrections,
|
|
const int32_t *__restrict__ image,
|
|
const float *__restrict__ mean,
|
|
const float *__restrict__ sigma,
|
|
float clip_k,
|
|
bool accumulate_corrected,
|
|
unsigned long long *__restrict__ sum, unsigned long long *__restrict__ sum2, uint32_t *__restrict__ count,
|
|
unsigned long long *__restrict__ sum_corr, unsigned long long *__restrict__ sum2_corr,
|
|
uint32_t *__restrict__ hist, int2 *__restrict__ overflow, uint32_t *__restrict__ overflow_n,
|
|
uint32_t overflow_cap, uint32_t *__restrict__ word_max,
|
|
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
|
|
// detection threshold, independent of the order the atomics happen to arrive in. With float
|
|
// accumulators the threshold moved in its last bits between runs, and because detection compares
|
|
// an integer pixel value against it, a threshold crossing an integer flipped every pixel of that
|
|
// value in the ring at once.
|
|
// The CORRECTED sums have no exact integer form, so they are summed in fixed point (CORR_FIXED).
|
|
extern __shared__ unsigned long long sh[];
|
|
unsigned long long *s_sum = sh; // signed value carried as two's complement
|
|
unsigned long long *s_sum2 = &s_sum[nbins];
|
|
unsigned long long *s_sum_corr = &s_sum2[nbins];
|
|
unsigned long long *s_sum2_corr = &s_sum_corr[nbins];
|
|
uint32_t *s_count = reinterpret_cast<uint32_t *>(accumulate_corrected ? &s_sum2_corr[nbins]
|
|
: &s_sum2[nbins]);
|
|
|
|
for (int i = threadIdx.x; i < nbins; i += blockDim.x) {
|
|
s_sum[i] = 0;
|
|
s_sum2[i] = 0;
|
|
s_count[i] = 0;
|
|
if (accumulate_corrected) {
|
|
s_sum_corr[i] = 0;
|
|
s_sum2_corr[i] = 0;
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
// 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};
|
|
|
|
// The plain pass also notes the largest value of every 32-pixel word, for flag_strong.
|
|
if (accumulate_corrected)
|
|
note_word_max(word_max, q, max(max(order_key(v4.x), order_key(v4.y)),
|
|
max(order_key(v4.z), order_key(v4.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;
|
|
unsigned long long r_sum_corr = 0, r_sum2_corr = 0;
|
|
|
|
#pragma unroll
|
|
for (int k = 0; k < 4; k++) {
|
|
const int32_t v = vq[k];
|
|
const int b = bq[k];
|
|
const float fv = static_cast<float>(v);
|
|
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);
|
|
r_b = b;
|
|
r_sum = 0; r_sum2 = 0; r_count = 0;
|
|
r_sum_corr = 0; r_sum2_corr = 0;
|
|
}
|
|
// 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 += corr_fixed(cv);
|
|
r_sum2_corr += corr_fixed(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];
|
|
const uint16_t b = pixel_to_bin[idx];
|
|
const float fv = static_cast<float>(v);
|
|
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);
|
|
atomicMax(&word_max[idx / 32], order_key(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);
|
|
if (accumulate_corrected) {
|
|
const float cv = fv * corrections[idx];
|
|
atomicAdd(&s_sum_corr[b], corr_fixed(cv));
|
|
atomicAdd(&s_sum2_corr[b], corr_fixed(cv * cv));
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
for (int i = threadIdx.x; i < nbins; i += blockDim.x) {
|
|
atomicAdd(&sum[i], s_sum[i]);
|
|
atomicAdd(&sum2[i], s_sum2[i]);
|
|
atomicAdd(&count[i], s_count[i]);
|
|
if (accumulate_corrected) {
|
|
atomicAdd(&sum_corr[i], s_sum_corr[i]);
|
|
atomicAdd(&sum2_corr[i], s_sum2_corr[i]);
|
|
}
|
|
}
|
|
}
|
|
|
|
// Same reduction with direct global atomics (used only when nbins is too large to stage in shared
|
|
// memory - a rare, high-bin-count configuration).
|
|
__global__ void reduce_rings_global(
|
|
const uint16_t *__restrict__ pixel_to_bin,
|
|
const float *__restrict__ corrections,
|
|
const int32_t *__restrict__ image,
|
|
const float *__restrict__ mean,
|
|
const float *__restrict__ sigma,
|
|
float clip_k,
|
|
bool accumulate_corrected,
|
|
unsigned long long *__restrict__ sum, unsigned long long *__restrict__ sum2, uint32_t *__restrict__ count,
|
|
unsigned long long *__restrict__ sum_corr, unsigned long long *__restrict__ sum2_corr,
|
|
size_t npix, int nbins) {
|
|
|
|
for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < npix; idx += blockDim.x * gridDim.x) {
|
|
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 && !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);
|
|
if (accumulate_corrected) {
|
|
const float cv = fv * corrections[idx];
|
|
atomicAdd(&sum_corr[b], corr_fixed(cv));
|
|
atomicAdd(&sum2_corr[b], corr_fixed(cv * cv));
|
|
}
|
|
}
|
|
}
|
|
|
|
// 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,
|
|
const unsigned long long *__restrict__ sum2,
|
|
const uint32_t *__restrict__ count,
|
|
float *__restrict__ mean, float *__restrict__ sigma, int nbins) {
|
|
for (int b = blockIdx.x * blockDim.x + threadIdx.x; b < nbins; b += blockDim.x * gridDim.x) {
|
|
if (count[b] > 0) {
|
|
// In double, then rounded to float for the clip predicate - the same two steps, in the same
|
|
// order and the same types, as AdaptiveSpotFinderCPU::AccumulateRings.
|
|
const double m = static_cast<double>(static_cast<long long>(sum[b])) / count[b];
|
|
const double var = fmax(0.0, static_cast<double>(sum2[b]) / count[b] - m * m);
|
|
mean[b] = static_cast<float>(m);
|
|
sigma[b] = static_cast<float>(sqrt(var));
|
|
}
|
|
}
|
|
}
|
|
|
|
// The range of rings the pixels of each 32-pixel word fall in (x = lowest, y = highest; x > y when
|
|
// none of them is in a ring). A function of the geometry alone, filled once.
|
|
__global__ void word_ring_range(const uint16_t *__restrict__ pixel_to_bin, ushort2 *__restrict__ range,
|
|
size_t npix, int nbins) {
|
|
const size_t nwords = (npix + 31) / 32;
|
|
for (size_t w = blockIdx.x * blockDim.x + threadIdx.x; w < nwords; w += static_cast<size_t>(blockDim.x) * gridDim.x) {
|
|
int lo = nbins, hi = -1;
|
|
for (size_t idx = 32 * w; idx < min(32 * w + 32, npix); idx++) {
|
|
const int b = pixel_to_bin[idx];
|
|
if (b < nbins) {
|
|
lo = min(lo, b);
|
|
hi = max(hi, b);
|
|
}
|
|
}
|
|
range[w] = (hi < 0) ? make_ushort2(1, 0) : make_ushort2(lo, hi);
|
|
}
|
|
}
|
|
|
|
// Flag strong pixels (value >= ring threshold, or saturated) into the packed bit buffer. Mirrors
|
|
// AdaptiveSpotFinderCPU Stage C exactly.
|
|
// One thread per 32-pixel word, writing the whole word, so the buffer needs neither clearing nor
|
|
// atomics. The image is read only for the words that can hold a strong pixel: the plain pass noted
|
|
// each word's largest value, and a word whose largest value is below the lowest threshold of the
|
|
// rings it touches has none. Strong pixels are sparse, so that is a few thousand words of the image
|
|
// instead of all of it. The test reads the same float comparison as the per-pixel one below, and
|
|
// the conversion to float keeps the order, so it never drops a word the per-pixel test would flag.
|
|
__global__ void flag_strong(const int32_t *__restrict__ image,
|
|
const uint16_t *__restrict__ pixel_to_bin,
|
|
const uint32_t *__restrict__ word_max,
|
|
const ushort2 *__restrict__ word_range,
|
|
const float *__restrict__ thr,
|
|
uint32_t *__restrict__ strong,
|
|
size_t npix, int nbins) {
|
|
const size_t nwords = (npix + 31) / 32;
|
|
for (size_t w = blockIdx.x * blockDim.x + threadIdx.x; w < nwords; w += static_cast<size_t>(blockDim.x) * gridDim.x) {
|
|
const int32_t vmax = static_cast<int32_t>(word_max[w] ^ 0x80000000u);
|
|
bool check = (vmax == INT32_MAX);
|
|
if (!check && vmax != INT32_MIN) {
|
|
float lowest = INFINITY;
|
|
for (int b = word_range[w].x; b <= word_range[w].y; b++)
|
|
lowest = fminf(lowest, thr[b]);
|
|
check = static_cast<float>(vmax) >= lowest;
|
|
}
|
|
uint32_t bits = 0;
|
|
if (check) {
|
|
for (size_t idx = 32 * w; idx < min(32 * w + 32, npix); idx++) {
|
|
const int32_t v = image[idx];
|
|
bool s = false;
|
|
if (v == INT32_MAX) {
|
|
s = true;
|
|
} else if (v != INT32_MIN) {
|
|
const uint16_t b = pixel_to_bin[idx];
|
|
if (b < nbins && static_cast<float>(v) >= thr[b])
|
|
s = true;
|
|
}
|
|
if (s)
|
|
bits |= 1u << (idx % 32);
|
|
}
|
|
}
|
|
strong[w] = bits;
|
|
}
|
|
}
|
|
|
|
} // namespace
|
|
|
|
AdaptiveSpotFinderGPU::AdaptiveSpotFinderGPU(const AzimuthalIntegrationMapping &in_mapping,
|
|
std::shared_ptr<CudaStream> in_stream)
|
|
: ImageSpotFinderGPU(static_cast<int32_t>(in_mapping.GetWidth()),
|
|
static_cast<int32_t>(in_mapping.GetHeight()), std::move(in_stream)),
|
|
mapping(in_mapping),
|
|
nbins(in_mapping.GetBinNumber()),
|
|
npix(in_mapping.GetPixelToBin().size()),
|
|
gpu_sum(nbins),
|
|
gpu_sum2(nbins),
|
|
gpu_count(nbins),
|
|
gpu_mean(nbins),
|
|
gpu_sigma(nbins),
|
|
gpu_sum_corr(nbins),
|
|
gpu_sum2_corr(nbins),
|
|
gpu_thr(nbins),
|
|
gpu_ring(OutputSize()),
|
|
gpu_word_max(OutputSize()),
|
|
gpu_word_range(OutputSize()),
|
|
host_sum(nbins),
|
|
host_sum2(nbins),
|
|
host_count(nbins),
|
|
host_bkg(nbins, NAN),
|
|
prof_sum(nbins),
|
|
prof_sum2(nbins),
|
|
prof_count(nbins),
|
|
last_profile(in_mapping) {
|
|
|
|
// The current device, not device 0: callers round-robin engines across GPUs, so device 0's shared
|
|
// memory and SM count can belong to a different card than the one these kernels launch on.
|
|
int device = 0;
|
|
cuda_err(cudaGetDevice(&device));
|
|
cudaDeviceProp prop{};
|
|
cuda_err(cudaGetDeviceProperties(&prop, device));
|
|
reduce_blocks = 8 * prop.multiProcessorCount;
|
|
// flag_strong: a thread per word. Its few loads per word depend on one another (the word's maximum,
|
|
// its ring range, their thresholds), so it wants every word in flight at once rather than a
|
|
// grid-stride walk.
|
|
flag_blocks = static_cast<int>((OutputSize() + flag_threads - 1) / flag_threads);
|
|
|
|
shared_plain = static_cast<size_t>(nbins) * (4 * sizeof(unsigned long long) + sizeof(uint32_t));
|
|
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);
|
|
// 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
|
|
// shared: the azimuthal-integration engine in the same worker reads the very same two arrays.
|
|
gpu_pixel_to_bin = SharedDeviceTable(mapping.GetPixelToBin().data(), npix,
|
|
mapping.GetPixelToBin().data(),
|
|
mapping.GetPixelToBinChecksum(), *stream);
|
|
gpu_corrections = SharedDeviceTable(mapping.Corrections().data(), npix,
|
|
mapping.Corrections().data(),
|
|
mapping.GetCorrectionsChecksum(), *stream);
|
|
word_ring_range<<<flag_blocks, flag_threads, 0, *stream>>>(gpu_pixel_to_bin->get(), gpu_word_range,
|
|
npix, nbins);
|
|
cuda_err(cudaGetLastError());
|
|
}
|
|
|
|
void AdaptiveSpotFinderGPU::ReducePass(const ImagePreprocessorBuffer &image, float clip_k,
|
|
bool accumulate_corrected) {
|
|
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,
|
|
gpu_hist, gpu_overflow, gpu_overflow_n, overflow_cap, gpu_word_max, npix, nbins);
|
|
cuda_err(cudaGetLastError());
|
|
} else {
|
|
reduce_rings_global<<<reduce_blocks, reduce_threads, 0, *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);
|
|
cuda_err(cudaGetLastError());
|
|
}
|
|
}
|
|
|
|
void AdaptiveSpotFinderGPU::FinalizeStats() {
|
|
const int threads = 128;
|
|
const int blocks = (nbins + threads - 1) / threads;
|
|
finalize_rings<<<blocks, threads, 0, *stream>>>(gpu_sum, gpu_sum2, gpu_count, gpu_mean, gpu_sigma, nbins);
|
|
cuda_err(cudaGetLastError());
|
|
}
|
|
|
|
// Host reproduction of AdaptiveSpotFinderCPU Stage B, from the clipped raw per-ring stats.
|
|
void AdaptiveSpotFinderGPU::ComputeThresholds(const SpotFindingSettings &settings) {
|
|
int64_t n_total = 0;
|
|
double g_sum = 0.0, g_sum2 = 0.0;
|
|
for (int b = 0; b < nbins; ++b) {
|
|
n_total += host_count[b];
|
|
g_sum += static_cast<double>(static_cast<int64_t>(host_sum[b]));
|
|
g_sum2 += static_cast<double>(host_sum2[b]);
|
|
}
|
|
if (n_total == 0) {
|
|
host_thr.clear();
|
|
std::fill(host_bkg.begin(), host_bkg.end(), NAN);
|
|
return;
|
|
}
|
|
|
|
const double E = std::max(1.0f, settings.false_pixels_per_frame);
|
|
double p = E / static_cast<double>(n_total);
|
|
p = std::min(std::max(p, 1e-9), 0.1);
|
|
const float z = static_cast<float>(adaptive_threshold::NormalQuantile(1.0 - p));
|
|
|
|
const double g_mean = g_sum / n_total;
|
|
const double g_sigma = std::sqrt(std::max(0.0, g_sum2 / n_total - g_mean * g_mean));
|
|
const float g_thr = adaptive_threshold::RingThreshold(static_cast<float>(g_mean),
|
|
static_cast<float>(g_sigma), p, z);
|
|
|
|
host_thr.assign(nbins, 0.0f);
|
|
for (int b = 0; b < nbins; ++b) {
|
|
if (host_count[b] < adaptive_threshold::MIN_RING_PIXELS) {
|
|
host_thr[b] = g_thr;
|
|
host_bkg[b] = NAN;
|
|
} else {
|
|
const double m = static_cast<double>(static_cast<int64_t>(host_sum[b])) / host_count[b];
|
|
const double var = std::max(0.0, static_cast<double>(host_sum2[b]) / host_count[b] - m * m);
|
|
host_thr[b] = adaptive_threshold::RingThreshold(static_cast<float>(m),
|
|
static_cast<float>(std::sqrt(var)), p, z);
|
|
host_bkg[b] = static_cast<float>(m);
|
|
}
|
|
}
|
|
}
|
|
|
|
void AdaptiveSpotFinderGPU::Detect(const ImagePreprocessorBuffer &image,
|
|
const SpotFindingSettings &settings) {
|
|
if (image.size() != npix)
|
|
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
|
|
"AdaptiveSpotFinderGPU::Detect: mismatch in pixel size");
|
|
|
|
// --- 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));
|
|
cuda_err(cudaMemsetAsync(gpu_mean, 0, sizeof(float) * nbins, *stream));
|
|
cuda_err(cudaMemsetAsync(gpu_sigma, 0, sizeof(float) * nbins, *stream));
|
|
cuda_err(cudaMemsetAsync(gpu_sum_corr, 0, sizeof(unsigned long long) * nbins, *stream));
|
|
cuda_err(cudaMemsetAsync(gpu_sum2_corr, 0, sizeof(unsigned long long) * 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));
|
|
cuda_err(cudaMemsetAsync(gpu_word_max, 0, OutputByteSize(), *stream));
|
|
} else {
|
|
// The global-atomics pass notes no word maxima: all ones makes flag_strong read every word.
|
|
cuda_err(cudaMemsetAsync(gpu_word_max, 0xFF, OutputByteSize(), *stream));
|
|
}
|
|
|
|
ReducePass(image, 0.0f, true); // plain pass also fills the corrected profile accumulators
|
|
FinalizeStats();
|
|
|
|
// Snapshot the plain corrected profile (and its pixel count) before the raw accumulators are
|
|
// re-zeroed for the sigma-clip passes.
|
|
cuda_err(cudaMemcpyAsync(prof_sum.data(), gpu_sum_corr, sizeof(unsigned long long) * nbins, cudaMemcpyDeviceToHost, *stream));
|
|
cuda_err(cudaMemcpyAsync(prof_sum2.data(), gpu_sum2_corr, sizeof(unsigned long long) * nbins, cudaMemcpyDeviceToHost, *stream));
|
|
cuda_err(cudaMemcpyAsync(prof_count.data(), gpu_count, sizeof(uint32_t) * nbins, cudaMemcpyDeviceToHost, *stream));
|
|
|
|
for (int pass = 0; pass < 2; ++pass) {
|
|
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));
|
|
ReducePass(image, 3.0f, false);
|
|
FinalizeStats();
|
|
}
|
|
|
|
// Snapshot the clipped raw stats that drive the threshold.
|
|
cuda_err(cudaMemcpyAsync(host_sum.data(), gpu_sum, sizeof(unsigned long long) * nbins, cudaMemcpyDeviceToHost, *stream));
|
|
cuda_err(cudaMemcpyAsync(host_sum2.data(), gpu_sum2, sizeof(unsigned long long) * nbins, cudaMemcpyDeviceToHost, *stream));
|
|
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);
|
|
|
|
// The profile is a byproduct even when the frame has no valid pixels for detection.
|
|
std::vector<float> sum_corr(nbins), sum2_corr(nbins);
|
|
for (int b = 0; b < nbins; ++b) {
|
|
sum_corr[b] = static_cast<float>(static_cast<double>(static_cast<long long>(prof_sum[b])) / CORR_FIXED);
|
|
sum2_corr[b] = static_cast<float>(static_cast<double>(static_cast<long long>(prof_sum2[b])) / CORR_FIXED);
|
|
}
|
|
last_profile.Clear(mapping);
|
|
last_profile.Add(sum_corr, sum2_corr, std::vector<uint32_t>(prof_count.begin(), prof_count.end()));
|
|
|
|
if (host_thr.empty()) {
|
|
// Nothing valid to threshold against: leave no strong pixels for the extractor to build on.
|
|
cuda_err(cudaMemsetAsync(gpu_out_1, 0, OutputByteSize(), *stream));
|
|
cuda_err(cudaStreamSynchronize(*stream));
|
|
return;
|
|
}
|
|
|
|
// --- Stage C: the ring threshold, intersected with the classic local-box SNR test ---
|
|
cuda_err(cudaMemcpyAsync(gpu_thr, host_thr.data(), sizeof(float) * nbins, cudaMemcpyHostToDevice, *stream));
|
|
flag_strong<<<flag_blocks, flag_threads, 0, *stream>>>(
|
|
image.getGPUBuffer(), gpu_pixel_to_bin->get(), gpu_word_max, gpu_word_range, gpu_thr, gpu_ring,
|
|
npix, nbins);
|
|
cuda_err(cudaGetLastError());
|
|
|
|
if (settings.signal_to_noise_threshold <= 0.0f) {
|
|
// No local test asked for: the ring threshold alone decides, as the fixed photon floor
|
|
// alone would in the classic finder.
|
|
cuda_err(cudaMemcpyAsync(gpu_out_1, gpu_ring, OutputByteSize(), cudaMemcpyDeviceToDevice, *stream));
|
|
cuda_err(cudaStreamSynchronize(*stream));
|
|
return;
|
|
}
|
|
|
|
// The ring threshold IS the photon floor here, so the local pass must not apply another one.
|
|
SpotFindingSettings local = settings;
|
|
local.photon_count_threshold = 0;
|
|
// Only the ring pixels survive the intersection, so the local test's second pass is evaluated at
|
|
// those alone (bit-identical to the dense pass followed by out &= ring).
|
|
DetectAt(image, local, gpu_ring);
|
|
// The bit buffer stays on the device and ExtractComponents reads it there, on this same stream,
|
|
// so the ordering already guarantees the detection has finished. Waiting here only idled the host.
|
|
}
|