Jungfraujoch/image_analysis/spot_finding/SpotExtractorGPU.cu

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

// Sparse connected-component labelling following 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 adapts on the host.
//   https://github.com/acts-project/traccc
//   (c) 2021-2025 CERN for the benefit of the ACTS project
//   Mozilla Public License Version 2.0
// The kernels below are this project's own - the algorithm is traccc's. Cited in
// docs/ACKNOWLEDGEMENT.md: P. Gessinger et al., "traccc: GPU track reconstruction library for HEP
// experiments" (2025), arXiv:2505.22822.

#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 warp per block emits its range, 32 words at a time: each lane takes a word, a warp prefix sum
// over the words' bit counts gives every lane where its pixels go, and the lane writes them in
// increasing order - so the list comes out in exactly the order a single thread walking the range
// would write it. Strong pixels are ~1e-4 of the image, so most passes find nothing at all; doing the
// walk on one thread left the other 31 idle and the reads of the range one after another.
__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) {
    const int lane = threadIdx.x % 32;
    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 base = w0; base < w1; base += 32) {
        const size_t w = base + lane;
        uint32_t word = (w < w1) ? (strong[w] & ~res_mask[w]) : 0;
        const uint32_t n = __popc(word);
        uint32_t before = n;                      // inclusive prefix sum of n over the lanes
        for (int offset = 1; offset < 32; offset *= 2) {
            const uint32_t t = __shfl_up_sync(0xffffffff, before, offset);
            if (lane >= offset) before += t;
        }
        const uint32_t total = __shfl_sync(0xffffffff, before, 31);
        uint32_t p = pos + before - n;
        while (word != 0) {
            const uint32_t flat = static_cast<uint32_t>(w * 32 + (__ffs(word) - 1));
            word &= word - 1;
            if (p < capacity) {
                out_index[p] = flat;
                out_value[p] = image[flat];
            }
            ++p;
        }
        pos += total;
    }
}

// --- 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 three launches and still one host synchronisation: hand out
// labels and clear each component's sums (one block), add every pixel to its component (all blocks),
// then filter by max-pix and shape and compact (one block) - each step order-preserving or made of
// integer sums, so the result does not depend on the order anything runs in.
__global__ void label_components(const uint32_t *__restrict__ root, uint32_t *__restrict__ label,
                                 SpotExtractorGPUSpot *__restrict__ scratch, uint32_t *__restrict__ nlabel_out,
                                 uint32_t *__restrict__ nout, const uint32_t *__restrict__ nstrong,
                                 uint32_t capacity) {
    const int n = static_cast<int>(min(*nstrong, capacity));
    if (threadIdx.x == 0) {
        *nout = 0;
        *nlabel_out = 0;
    }
    __syncthreads();
    // The same give-up the host makes at StrongPixelLimit - 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);

    // 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]);
    if (t == 0) *nlabel_out = nlabel;

    for (int l = t; l < nlabel; l += nthreads) {
        SpotExtractorGPUSpot &s = scratch[l];
        s.x = 0;
        s.y = 0;
        s.photons = 0;
        s.max_photons = LLONG_MIN;
        s.pixel_count = 0;
        s.bbox_side = 0;
    }
}

// The bounding box of each component (x, y = lowest and highest column, z, w = lowest and highest
// line), kept aside from the sums while they are being accumulated.
__global__ void init_boxes(int4 *__restrict__ box, const uint32_t *__restrict__ nlabel) {
    const int n = static_cast<int>(*nlabel);
    for (int l = blockIdx.x * blockDim.x + threadIdx.x; l < n; l += blockDim.x * gridDim.x)
        box[l] = make_int4(INT_MAX, INT_MIN, INT_MAX, INT_MIN);
}

// Every strong pixel into its component: the member count, the sums exactly as DiffractionSpot::AddPixel
// forms them (integers, so the order the atomics arrive in does not matter and host and device agree by
// construction) and the bounding box.
__global__ void sum_components(const uint32_t *__restrict__ index, const int32_t *__restrict__ value,
                               const uint32_t *__restrict__ root, const uint32_t *__restrict__ label,
                               SpotExtractorGPUSpot *__restrict__ scratch, int4 *__restrict__ box,
                               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 l = label[root[i]];
        const long long counts = value[i];
        const int col = static_cast<int>(index[i] % width), line = static_cast<int>(index[i] / width);
        SpotExtractorGPUSpot &s = scratch[l];
        atomicAdd(reinterpret_cast<unsigned long long *>(&s.x), static_cast<unsigned long long>(col * counts));
        atomicAdd(reinterpret_cast<unsigned long long *>(&s.y), static_cast<unsigned long long>(line * counts));
        atomicAdd(reinterpret_cast<unsigned long long *>(&s.photons), static_cast<unsigned long long>(counts));
        atomicMax(reinterpret_cast<long long *>(&s.max_photons), counts);
        atomicAdd(&s.pixel_count, 1);
        atomicMin(&box[l].x, col);
        atomicMax(&box[l].y, col);
        atomicMin(&box[l].z, line);
        atomicMax(&box[l].w, line);
    }
}

// Size and shape filter, compacted by a prefix sum so the surviving spots keep their order. The test
// is SpotShapeAccepted written out - the constants come in as arguments rather than being included
// here, so there is exactly one definition of them.
__global__ void filter_components(SpotExtractorGPUSpot *__restrict__ scratch, const int4 *__restrict__ box,
                                  SpotExtractorGPUSpot *__restrict__ out, uint32_t *__restrict__ nout,
                                  const uint32_t *__restrict__ nlabel_in,
                                  int max_pix, int shape_free_pix, int min_fill_percent) {
    const int nlabel = static_cast<int>(*nlabel_in);
    if (nlabel == 0) return;

    __shared__ uint32_t shared[FINISH_THREADS];
    const int t = threadIdx.x, nthreads = blockDim.x;
    auto keep = [&](const SpotExtractorGPUSpot &s) {
        if (s.pixel_count > max_pix) return false;
        if (s.pixel_count <= shape_free_pix) return true;
        return static_cast<long long>(s.pixel_count) * 100
               >= static_cast<long long>(min_fill_percent) * s.bbox_side * s.bbox_side;
    };
    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++) {
        scratch[i].bbox_side = max(box[i].y - box[i].x, box[i].w - box[i].z) + 1;
        nkeep += keep(scratch[i]) ? 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 (keep(scratch[i])) 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),
      max_strong(StrongPixelLimit(static_cast<size_t>(in_width) * in_height)),
      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_box(max_strong),
      gpu_nlabel(1),
      gpu_spot(max_strong),
      gpu_spot_out(max_strong),
      gpu_nspot(1),
      host_nstrong(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);
    cuda_err(cudaGetLastError());
    scan_block_counts<<<1, FINISH_THREADS, 0, *stream>>>(gpu_block_count, gpu_block_offset, gpu_nstrong,
                                                          compact_blocks);
    cuda_err(cudaGetLastError());
    scatter_bits<<<compact_blocks, 32, 0, *stream>>>(gpu_strong, gpu_res_mask, gpu_block_offset, gpu_image,
                                                      gpu_index, gpu_value, nwords, max_strong);
    cuda_err(cudaGetLastError());
    // 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);
    cuda_err(cudaGetLastError());
    union_neighbours<<<512, THREADS, 0, *stream>>>(gpu_index, gpu_parent, gpu_nstrong, max_strong, width);
    cuda_err(cudaGetLastError());
    resolve_roots<<<512, THREADS, 0, *stream>>>(gpu_parent, gpu_root, gpu_nstrong, max_strong);
    cuda_err(cudaGetLastError());
    label_components<<<1, FINISH_THREADS, 0, *stream>>>(gpu_root, gpu_label, gpu_spot, gpu_nlabel, gpu_nspot,
                                                         gpu_nstrong, max_strong);
    cuda_err(cudaGetLastError());
    init_boxes<<<512, THREADS, 0, *stream>>>(gpu_box, gpu_nlabel);
    cuda_err(cudaGetLastError());
    sum_components<<<512, THREADS, 0, *stream>>>(gpu_index, gpu_value, gpu_root, gpu_label, gpu_spot, gpu_box,
                                                 gpu_nstrong, max_strong, width);
    cuda_err(cudaGetLastError());
    filter_components<<<1, FINISH_THREADS, 0, *stream>>>(gpu_spot, gpu_box, gpu_spot_out, gpu_nspot, gpu_nlabel,
                                                          max_pix, static_cast<int>(SPOT_SHAPE_FREE_PIXELS),
                                                          static_cast<int>(SPOT_MIN_FILL_PERCENT));
    cuda_err(cudaGetLastError());
    // Rides along with the spot count on the frame's one synchronisation, so knowing how many strong
    // pixels there were costs nothing.
    cuda_err(cudaMemcpyAsync(host_nstrong, gpu_nstrong, sizeof(uint32_t), cudaMemcpyDeviceToHost, *stream));
    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);
}