leonarski_fandClaude Opus 5
4b1c611bdf
Two leftovers from earlier fixes of the same shape. BraggIntegrationEngineGPU still read device 0's shared-memory size to decide whether its profile grid fits; workers are pinned round-robin across GPUs, so on a heterogeneous node that check can pass on a different card than the one the kernel launches on. SpotExtractorGPU still uploaded its default resolution mask with a pageable copy on the NULL stream, which is not ordered against the engine stream now that streams are created non-blocking. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
352 lines
17 KiB
Plaintext
352 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);
|
|
// On this engine's stream, then synchronised - the default stream is non-blocking, so a NULL-stream
|
|
// copy is not ordered against the kernels that read this mask.
|
|
cuda_err(cudaMemcpyAsync(gpu_res_mask, mask.data(), nwords * sizeof(uint32_t),
|
|
cudaMemcpyHostToDevice, *stream));
|
|
cuda_err(cudaStreamSynchronize(*stream));
|
|
}
|
|
|
|
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);
|
|
}
|