leonarski_fandClaude Opus 5
8f1b0b2281
The photon-weighted position sums were floats, so the centroid's last bit depended on the build rather than on the data: gcc contracts the multiply-add in AddPixel into an FMA under -march=x86-64-v3 and cannot at the baseline, and MSVC does not contract at all under /fp:precise. The GPU extractor had to match with __fmaf_rn, and the parity test still needed a two-ulp slack for hosts that do not fuse. Column, line and the per-pixel count are all integral, so the sums are exact in int64 and both implementations reach the same bits with nothing to match. The parity test now demands exact equality unconditionally and gets it, including on a baseline build. ConvertToImageCoordinates keeps the sums integral too: the raw -> image map is a signed axis swap plus an integer translation, so it is applied to the sums instead of to the centroid. Drops the SpotToSave constructor, which had no callers and could not have been converted without quantising the stored centroid. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
348 lines
17 KiB
Plaintext
348 lines
17 KiB
Plaintext
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
// Sparse connected-component labelling adapted from the design the ACTS/traccc project arrived at
|
|
// for sparse silicon-detector clusterization (backward-neighbour graph over a sorted hit list, then
|
|
// a parallel union-find), which is itself the GPU counterpart of the SparseCCL that
|
|
// StrongPixelSet.cpp uses on the host.
|
|
|
|
#include <climits>
|
|
|
|
#include "SpotExtractorGPU.h"
|
|
#include "../../common/JFJochException.h"
|
|
|
|
namespace {
|
|
|
|
inline void cuda_err(cudaError_t val) {
|
|
if (val != cudaSuccess)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
|
|
}
|
|
|
|
constexpr int THREADS = 256;
|
|
constexpr int FINISH_THREADS = 1024;
|
|
|
|
// --- compaction: packed bit buffer -> strong-pixel list, sorted by flat index -------------------
|
|
// Each block owns a CONTIGUOUS range of words. Pass 1 counts its strong bits, a single-block scan
|
|
// turns the counts into offsets, and pass 2 walks the same range in increasing order and writes at
|
|
// that offset. No atomics anywhere, which is what keeps the output sorted.
|
|
__global__ void count_bits(const uint32_t *__restrict__ strong, const uint32_t *__restrict__ res_mask,
|
|
uint32_t *__restrict__ block_count, size_t nwords) {
|
|
const size_t per_block = (nwords + gridDim.x - 1) / gridDim.x;
|
|
const size_t w0 = static_cast<size_t>(blockIdx.x) * per_block;
|
|
const size_t w1 = min(w0 + per_block, nwords);
|
|
uint32_t local = 0;
|
|
for (size_t w = w0 + threadIdx.x; w < w1; w += blockDim.x)
|
|
local += __popc(strong[w] & ~res_mask[w]);
|
|
|
|
__shared__ uint32_t block_total;
|
|
if (threadIdx.x == 0) block_total = 0;
|
|
__syncthreads();
|
|
atomicAdd(&block_total, local);
|
|
__syncthreads();
|
|
if (threadIdx.x == 0) block_count[blockIdx.x] = block_total;
|
|
}
|
|
|
|
__global__ void scan_block_counts(const uint32_t *__restrict__ in, uint32_t *__restrict__ out,
|
|
uint32_t *__restrict__ total, int n) {
|
|
__shared__ uint32_t shared[FINISH_THREADS];
|
|
const int t = threadIdx.x, nthreads = blockDim.x;
|
|
const int chunk = (n + nthreads - 1) / nthreads;
|
|
const int lo = min(t * chunk, n), hi = min(lo + chunk, n);
|
|
uint32_t sum = 0;
|
|
for (int i = lo; i < hi; i++) sum += in[i];
|
|
shared[t] = sum;
|
|
__syncthreads();
|
|
for (int d = 1; d < nthreads; d <<= 1) {
|
|
const uint32_t v = (t >= d) ? shared[t - d] : 0u;
|
|
__syncthreads();
|
|
shared[t] += v;
|
|
__syncthreads();
|
|
}
|
|
uint32_t acc = shared[t] - sum;
|
|
for (int i = lo; i < hi; i++) { out[i] = acc; acc += in[i]; }
|
|
if (t == nthreads - 1) *total = shared[t];
|
|
}
|
|
|
|
// One thread per block emits its range. Strong pixels are ~1e-4 of the image, so a block's range
|
|
// holds a handful of them and serial emission is both trivially ordered and fast; the parallelism
|
|
// comes from the block count.
|
|
__global__ void scatter_bits(const uint32_t *__restrict__ strong, const uint32_t *__restrict__ res_mask,
|
|
const uint32_t *__restrict__ block_offset, const int32_t *__restrict__ image,
|
|
uint32_t *__restrict__ out_index, int32_t *__restrict__ out_value,
|
|
size_t nwords, uint32_t capacity) {
|
|
if (threadIdx.x != 0) return;
|
|
const size_t per_block = (nwords + gridDim.x - 1) / gridDim.x;
|
|
const size_t w0 = static_cast<size_t>(blockIdx.x) * per_block;
|
|
const size_t w1 = min(w0 + per_block, nwords);
|
|
uint32_t pos = block_offset[blockIdx.x];
|
|
for (size_t w = w0; w < w1; ++w) {
|
|
uint32_t word = strong[w] & ~res_mask[w];
|
|
while (word != 0) {
|
|
const uint32_t flat = static_cast<uint32_t>(w * 32 + (__ffs(word) - 1));
|
|
word &= word - 1;
|
|
if (pos < capacity) {
|
|
out_index[pos] = flat;
|
|
out_value[pos] = image[flat];
|
|
}
|
|
++pos;
|
|
}
|
|
}
|
|
}
|
|
|
|
// --- connected components ----------------------------------------------------------------------
|
|
__device__ __forceinline__ int lower_bound_device(const uint32_t *a, int n, uint32_t key) {
|
|
int lo = 0, hi = n;
|
|
while (lo < hi) {
|
|
const int mid = (lo + hi) >> 1;
|
|
if (a[mid] < key) lo = mid + 1; else hi = mid;
|
|
}
|
|
return lo;
|
|
}
|
|
|
|
// Walk to the root, halving the path on the way. The plain store is safe: a parent only ever
|
|
// decreases and the grandparent is always still an ancestor, so a concurrent writer can only make
|
|
// the chain shorter. Without halving, a long thin feature - a diffraction ring is exactly one -
|
|
// builds a parent chain as long as the feature itself and every later merge walks all of it.
|
|
__device__ __forceinline__ uint32_t find_root(uint32_t *parent, uint32_t a) {
|
|
uint32_t p = parent[a];
|
|
while (p != a) {
|
|
const uint32_t gp = parent[p];
|
|
if (gp == p) return p;
|
|
parent[a] = gp;
|
|
a = gp;
|
|
p = parent[a];
|
|
}
|
|
return a;
|
|
}
|
|
|
|
// Lock-free union. Terminates because max(a,b) strictly decreases; converges on the component's
|
|
// LOWEST index as its root, which is what the host's make_union does too.
|
|
__device__ __forceinline__ void merge_roots(uint32_t *parent, uint32_t a, uint32_t b) {
|
|
a = find_root(parent, a);
|
|
b = find_root(parent, b);
|
|
while (a != b) {
|
|
if (a < b) { const uint32_t t = a; a = b; b = t; }
|
|
const uint32_t old = atomicMin(&parent[a], b);
|
|
if (old == a) return; // a was a root and now points at b: done
|
|
a = find_root(parent, old); // someone re-parented a; carry on from its root
|
|
b = find_root(parent, b);
|
|
}
|
|
}
|
|
|
|
__global__ void init_parent(uint32_t *__restrict__ parent, const uint32_t *__restrict__ nstrong,
|
|
uint32_t capacity) {
|
|
const int n = static_cast<int>(min(*nstrong, capacity));
|
|
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x)
|
|
parent[i] = i;
|
|
}
|
|
|
|
// The list is sorted by flat index, so a pixel's 8-neighbours that come EARLIER in it are exactly
|
|
// four: (line, col-1), (line-1, col-1), (line-1, col) and (line-1, col+1). Each is one binary
|
|
// search away, which is what makes the sparse formulation cheap.
|
|
__global__ void union_neighbours(const uint32_t *__restrict__ index, uint32_t *__restrict__ parent,
|
|
const uint32_t *__restrict__ nstrong, uint32_t capacity, int width) {
|
|
const int n = static_cast<int>(min(*nstrong, capacity));
|
|
if (static_cast<uint32_t>(n) >= capacity) return;
|
|
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x) {
|
|
const uint32_t flat = index[i];
|
|
const uint32_t col = flat % width;
|
|
const uint32_t line = flat / width;
|
|
uint32_t candidate[4];
|
|
int ncandidate = 0;
|
|
if (col > 0) candidate[ncandidate++] = flat - 1;
|
|
if (line > 0 && col > 0) candidate[ncandidate++] = flat - width - 1;
|
|
if (line > 0) candidate[ncandidate++] = flat - width;
|
|
if (line > 0 && col + 1 < static_cast<uint32_t>(width)) candidate[ncandidate++] = flat - width + 1;
|
|
for (int k = 0; k < ncandidate; k++) {
|
|
const int j = lower_bound_device(index, i, candidate[k]);
|
|
if (j < i && index[j] == candidate[k]) merge_roots(parent, static_cast<uint32_t>(i), static_cast<uint32_t>(j));
|
|
}
|
|
}
|
|
}
|
|
|
|
__global__ void resolve_roots(uint32_t *__restrict__ parent, uint32_t *__restrict__ root,
|
|
const uint32_t *__restrict__ nstrong, uint32_t capacity) {
|
|
const int n = static_cast<int>(min(*nstrong, capacity));
|
|
for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x)
|
|
root[i] = find_root(parent, static_cast<uint32_t>(i));
|
|
}
|
|
|
|
// Everything after the labelling in ONE block, so a frame needs a single host synchronisation:
|
|
// hand out labels, count the members of each component, sum the surviving ones, filter by max-pix
|
|
// and compact - all of it order-preserving.
|
|
__global__ void finish_components(const uint32_t *__restrict__ index, const int32_t *__restrict__ value,
|
|
const uint32_t *__restrict__ root, uint32_t *__restrict__ label,
|
|
int32_t *__restrict__ count, SpotExtractorGPUSpot *__restrict__ scratch,
|
|
SpotExtractorGPUSpot *__restrict__ out, uint32_t *__restrict__ nout,
|
|
const uint32_t *__restrict__ nstrong, uint32_t capacity,
|
|
int width, int max_pix) {
|
|
const int n = static_cast<int>(min(*nstrong, capacity));
|
|
if (threadIdx.x == 0) *nout = 0;
|
|
__syncthreads();
|
|
// Same give-up as StrongPixelSet::FindComponentsImage - except that here the count is known
|
|
// before a single pixel has been written anywhere, so the frame costs nothing to reject.
|
|
if (n == 0 || static_cast<uint32_t>(n) >= capacity) return;
|
|
|
|
__shared__ uint32_t shared[FINISH_THREADS];
|
|
const int t = threadIdx.x, nthreads = blockDim.x;
|
|
const int chunk = (n + nthreads - 1) / nthreads;
|
|
const int lo = min(t * chunk, n), hi = min(lo + chunk, n);
|
|
|
|
// 1) labels, by a prefix sum over the roots in ascending order - the order the host's second
|
|
// scan hands them out in, which is what makes the spot ORDER identical.
|
|
uint32_t nroot = 0;
|
|
for (int i = lo; i < hi; i++) nroot += (root[i] == static_cast<uint32_t>(i)) ? 1u : 0u;
|
|
shared[t] = nroot;
|
|
__syncthreads();
|
|
for (int d = 1; d < nthreads; d <<= 1) {
|
|
const uint32_t v = (t >= d) ? shared[t - d] : 0u;
|
|
__syncthreads();
|
|
shared[t] += v;
|
|
__syncthreads();
|
|
}
|
|
uint32_t next_label = shared[t] - nroot;
|
|
for (int i = lo; i < hi; i++)
|
|
if (root[i] == static_cast<uint32_t>(i)) label[i] = next_label++;
|
|
const int nlabel = static_cast<int>(shared[nthreads - 1]);
|
|
__syncthreads();
|
|
|
|
// 2) member counts
|
|
for (int i = t; i < nlabel; i += nthreads) count[i] = 0;
|
|
__syncthreads();
|
|
for (int i = t; i < n; i += nthreads) atomicAdd(&count[label[root[i]]], 1);
|
|
__syncthreads();
|
|
|
|
// 3) sums, one thread per component, walking its members in ascending list order so the
|
|
// accumulation matches DiffractionSpot::AddPixel term for term. A component bigger than
|
|
// max-pix is thrown away below, so it is not summed - which is also what keeps a whole lit
|
|
// module or diffraction ring from turning into one thread walking tens of thousands of
|
|
// entries.
|
|
for (int i = t; i < n; i += nthreads) {
|
|
if (root[i] != static_cast<uint32_t>(i)) continue;
|
|
const uint32_t l = label[i];
|
|
const int want = count[l];
|
|
scratch[l].pixel_count = want;
|
|
if (want > max_pix) continue;
|
|
long long x = 0, y = 0;
|
|
long long photons = 0, max_photons = LLONG_MIN;
|
|
int found = 0;
|
|
for (int j = i; j < n && found < want; j++) {
|
|
if (root[j] != static_cast<uint32_t>(i)) continue;
|
|
const long long counts = value[j];
|
|
// Integers, exactly as DiffractionSpot::AddPixel does them, so host and device agree by
|
|
// construction - no rounding mode to match and nothing for either compiler to contract.
|
|
x += static_cast<long long>(index[j] % width) * counts;
|
|
y += static_cast<long long>(index[j] / width) * counts;
|
|
photons += counts;
|
|
max_photons = max(max_photons, counts);
|
|
found++;
|
|
}
|
|
scratch[l].x = x;
|
|
scratch[l].y = y;
|
|
scratch[l].photons = photons;
|
|
scratch[l].max_photons = max_photons;
|
|
}
|
|
__syncthreads();
|
|
|
|
// 4) max-pix filter, compacted by another prefix sum so the surviving spots keep their order
|
|
const int label_chunk = (nlabel + nthreads - 1) / nthreads;
|
|
const int label_lo = min(t * label_chunk, nlabel), label_hi = min(label_lo + label_chunk, nlabel);
|
|
uint32_t nkeep = 0;
|
|
for (int i = label_lo; i < label_hi; i++) nkeep += (scratch[i].pixel_count <= max_pix) ? 1u : 0u;
|
|
shared[t] = nkeep;
|
|
__syncthreads();
|
|
for (int d = 1; d < nthreads; d <<= 1) {
|
|
const uint32_t v = (t >= d) ? shared[t - d] : 0u;
|
|
__syncthreads();
|
|
shared[t] += v;
|
|
__syncthreads();
|
|
}
|
|
uint32_t pos = shared[t] - nkeep;
|
|
for (int i = label_lo; i < label_hi; i++)
|
|
if (scratch[i].pixel_count <= max_pix) out[pos++] = scratch[i];
|
|
if (t == nthreads - 1) *nout = shared[t];
|
|
}
|
|
|
|
} // namespace
|
|
|
|
SpotExtractorGPU::SpotExtractorGPU(int32_t in_width, int32_t in_height, std::shared_ptr<CudaStream> in_stream)
|
|
: stream(std::move(in_stream)),
|
|
width(in_width),
|
|
nwords((static_cast<size_t>(in_width) * in_height + 31) / 32),
|
|
gpu_res_mask(nwords),
|
|
gpu_nstrong(1),
|
|
gpu_index(MAX_STRONG),
|
|
gpu_value(MAX_STRONG),
|
|
gpu_parent(MAX_STRONG),
|
|
gpu_root(MAX_STRONG),
|
|
gpu_label(MAX_STRONG),
|
|
gpu_count(MAX_STRONG),
|
|
gpu_spot(MAX_STRONG),
|
|
gpu_spot_out(MAX_STRONG),
|
|
gpu_nspot(1),
|
|
host_nspot(1),
|
|
host_spot(SPOT_PREFIX) {
|
|
// One block per few hundred words: enough blocks to fill the device, few enough that the serial
|
|
// emission inside a block stays short even when a whole detector row lights up.
|
|
compact_blocks = static_cast<int>((nwords + 255) / 256);
|
|
if (compact_blocks > 4096) compact_blocks = 4096;
|
|
if (compact_blocks < 1) compact_blocks = 1;
|
|
gpu_block_count = CudaDevicePtr<uint32_t>(compact_blocks);
|
|
gpu_block_offset = CudaDevicePtr<uint32_t>(compact_blocks);
|
|
|
|
// Nothing excluded except the padding bits of the last word - the same starting point as
|
|
// ImageSpotFinder's own mask, so the two agree even when no resolution mask is ever set.
|
|
std::vector<uint32_t> mask(nwords, 0);
|
|
const size_t npixel = static_cast<size_t>(in_width) * in_height;
|
|
if (npixel % 32 != 0)
|
|
mask.back() = ~((1u << (npixel % 32)) - 1u);
|
|
cuda_err(cudaMemcpy(gpu_res_mask, mask.data(), nwords * sizeof(uint32_t), cudaMemcpyHostToDevice));
|
|
}
|
|
|
|
void SpotExtractorGPU::SetResolutionMask(const std::vector<uint32_t> &packed_mask) {
|
|
if (packed_mask.size() != nwords)
|
|
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
|
|
"SpotExtractorGPU::SetResolutionMask: mask size mismatch");
|
|
cuda_err(cudaMemcpyAsync(gpu_res_mask, packed_mask.data(), nwords * sizeof(uint32_t),
|
|
cudaMemcpyHostToDevice, *stream));
|
|
cuda_err(cudaStreamSynchronize(*stream));
|
|
}
|
|
|
|
void SpotExtractorGPU::Extract(const uint32_t *gpu_strong, const int32_t *gpu_image,
|
|
const SpotFindingSettings &settings, std::vector<DiffractionSpot> &spots) {
|
|
const int max_pix = static_cast<int>(settings.max_pix_per_spot);
|
|
|
|
count_bits<<<compact_blocks, THREADS, 0, *stream>>>(gpu_strong, gpu_res_mask, gpu_block_count, nwords);
|
|
scan_block_counts<<<1, FINISH_THREADS, 0, *stream>>>(gpu_block_count, gpu_block_offset, gpu_nstrong,
|
|
compact_blocks);
|
|
scatter_bits<<<compact_blocks, 32, 0, *stream>>>(gpu_strong, gpu_res_mask, gpu_block_offset, gpu_image,
|
|
gpu_index, gpu_value, nwords, MAX_STRONG);
|
|
// Fixed grids reading the strong-pixel count from device memory: the host never learns it, so it
|
|
// never has to synchronise in the middle of the frame.
|
|
init_parent<<<512, THREADS, 0, *stream>>>(gpu_parent, gpu_nstrong, MAX_STRONG);
|
|
union_neighbours<<<512, THREADS, 0, *stream>>>(gpu_index, gpu_parent, gpu_nstrong, MAX_STRONG, width);
|
|
resolve_roots<<<512, THREADS, 0, *stream>>>(gpu_parent, gpu_root, gpu_nstrong, MAX_STRONG);
|
|
finish_components<<<1, FINISH_THREADS, 0, *stream>>>(gpu_index, gpu_value, gpu_root, gpu_label,
|
|
gpu_count, gpu_spot, gpu_spot_out, gpu_nspot,
|
|
gpu_nstrong, MAX_STRONG, width, max_pix);
|
|
cuda_err(cudaMemcpyAsync(host_nspot, gpu_nspot, sizeof(uint32_t), cudaMemcpyDeviceToHost, *stream));
|
|
cuda_err(cudaMemcpyAsync(host_spot, gpu_spot_out, SPOT_PREFIX * sizeof(SpotExtractorGPUSpot),
|
|
cudaMemcpyDeviceToHost, *stream));
|
|
cuda_err(cudaStreamSynchronize(*stream)); // the only synchronisation in the frame
|
|
|
|
const uint32_t nspot = *host_nspot.get();
|
|
const SpotExtractorGPUSpot *s = host_spot.get();
|
|
if (nspot > SPOT_PREFIX) {
|
|
overflow_spot.resize(nspot);
|
|
cuda_err(cudaMemcpyAsync(overflow_spot.data(), gpu_spot_out, nspot * sizeof(SpotExtractorGPUSpot),
|
|
cudaMemcpyDeviceToHost, *stream));
|
|
cuda_err(cudaStreamSynchronize(*stream));
|
|
s = overflow_spot.data();
|
|
}
|
|
|
|
spots.clear();
|
|
spots.reserve(nspot);
|
|
for (uint32_t i = 0; i < nspot; i++)
|
|
spots.emplace_back(s[i].x, s[i].y, s[i].pixel_count, s[i].photons, s[i].max_photons);
|
|
}
|