Jungfraujoch/image_analysis/spot_finding/SpotExtractorGPU.h

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

#pragma once

// Device-side connected-component extraction for the GPU spot finders.
//
// The GPU finders flag strong pixels into a packed bit buffer ON THE DEVICE. Reading spots out of it
// used to mean copying that whole buffer back (2.26 MB per frame at 18 MP) and scanning it bit by bit
// on the host. This does the whole extraction where the data already is, so nothing about the image
// comes back - only the finished spot list, a few hundred entries.
//
// The algorithm is the sparse formulation the ACTS/traccc project settled on for the same problem
// (sparse silicon-detector hits): the strong pixels are compacted into a list that is sorted by flat
// index, each pixel finds its at most FOUR backward 8-neighbours by binary search in that list, and
// the resulting graph is labelled with a lock-free union-find. A dense image-wide labelling
// (Playne-equivalence, BUF/BKE, nppiLabelMarkers, cv::cuda::connectedComponents) would label 18
// million pixels to find five hundred.
//
// It reproduces the host StrongPixelSet::sparseccl EXACTLY, not just equivalently:
//   * both make a component's root its lowest list index, so both find the same roots;
//   * labels are handed out by a prefix sum over the roots in ascending order, which is the order the
//     host's second scan hands them out in, so the SPOT ORDER is identical;
//   * the centroid sums are accumulated per component in integers, term for term as
//     DiffractionSpot::AddPixel does them, so there is no rounding for the two compilers to disagree
//     about, and the order the device adds them in changes nothing.
// tests/SpotExtractorGPUParityTest.cpp holds the two to each other on realistic, occupancy-swept and
// pathological frames, and checks that repeating a frame gives byte-identical output.

#include <cstdint>
#include <memory>
#include <vector>

#include "../../common/DiffractionSpot.h"
#include "../../common/PanelBoxes.h"
#include "../indexing/CUDAMemHelpers.h"
#include "SpotFindingSettings.h"

// Per-component sums, in exactly the form DiffractionSpot holds them: x and y are sum(col*photons)
// and sum(line*photons), not a centroid.
struct SpotExtractorGPUSpot {
    int64_t x;
    int64_t y;
    int64_t photons;
    int64_t max_photons;
    int32_t pixel_count;
    // Longer side of the component's bounding box, for SpotShapeAccepted. It fits in what used to be
    // padding, so carrying it costs nothing.
    int32_t bbox_side;
    int32_t panel;   // PanelBox::panel of the component's pixels
};

class SpotExtractorGPU {
    std::shared_ptr<CudaStream> stream;
    const int32_t width;
    const size_t nwords;

    // Strong pixels this engine's buffers hold, and above which the extraction gives up on the frame -
    // StrongPixelLimit, so it follows the detector rather than standing at a constant. 116 bytes of
    // device memory apiece, 32 MB on an 18-megapixel detector.
    const uint32_t max_strong;
    // Spots copied back together with their count in one transfer. A frame with more than this many
    // surviving spots - far past anything indexable - simply takes a second copy.
    static constexpr uint32_t SPOT_PREFIX = 4096;

    int compact_blocks = 0;

    CudaDevicePtr<uint32_t> gpu_res_mask;      // packed, bit set = pixel excluded
    CudaDevicePtr<uint32_t> gpu_block_count;
    CudaDevicePtr<uint32_t> gpu_block_offset;
    CudaDevicePtr<uint32_t> gpu_nstrong;
    CudaDevicePtr<uint32_t> gpu_index;         // strong pixels, sorted by flat index
    CudaDevicePtr<int32_t>  gpu_value;
    CudaDevicePtr<uint32_t> gpu_parent;        // union-find parent
    CudaDevicePtr<uint32_t> gpu_root;
    CudaDevicePtr<uint32_t> gpu_label;         // compact label, indexed by root
    CudaDevicePtr<int4>     gpu_box;           // bounding box per component
    CudaDevicePtr<uint32_t> gpu_nlabel;        // components found
    CudaDevicePtr<SpotExtractorGPUSpot> gpu_spot;
    CudaDevicePtr<SpotExtractorGPUSpot> gpu_spot_out;
    CudaDevicePtr<uint32_t> gpu_nspot;
    CudaDevicePtr<PanelBox> gpu_panels;        // the panels components stay inside
    int npanels = 0;

    CudaHostPtr<uint32_t> host_nstrong;
    CudaHostPtr<uint32_t> host_nspot;
    CudaHostPtr<SpotExtractorGPUSpot> host_spot;      // SPOT_PREFIX entries, pinned
    std::vector<SpotExtractorGPUSpot> overflow_spot;  // only for a frame with more spots than that

public:
    SpotExtractorGPU(int32_t width, int32_t height, std::shared_ptr<CudaStream> stream);
    // Waits for the stream: its pooled buffers are freed next (see CudaDevicePtr).
    ~SpotExtractorGPU() { if (stream) cudaStreamSynchronize(*stream); }

    void SetResolutionMask(const std::vector<uint32_t> &packed_mask);
    // The panels of the image (PanelBoxes); default one, the whole image.
    void SetPanels(const std::vector<PanelBox> &panels);

    // gpu_strong is the finder's device bit buffer, gpu_image the preprocessed image it was built
    // from. Fills spots with every component of at most max-pix pixels, in the same order the host
    // extractor would.
    void Extract(const uint32_t *gpu_strong, const int32_t *gpu_image,
                 const SpotFindingSettings &settings, std::vector<DiffractionSpot> &spots);

    // Strong pixels the last Extract() saw, after the resolution mask. Reported whether or not the
    // frame was given up on, which is the point of it: a frame at or above max_strong yields no spots
    // at all, and this is what says so.
    [[nodiscard]] uint32_t StrongPixelCount() const { return *host_nstrong.get(); }
};