Jungfraujoch/image_analysis/spot_finding/ImageSpotFinderGPU.h

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

#pragma once

#include <vector>

#include "SpotFindingSettings.h"
#include "ImageSpotFinder.h"
#include "SpotExtractorGPU.h"
#include "../indexing/CUDAMemHelpers.h"

class ImageSpotFinderGPU : public ImageSpotFinder {
protected:
    // Protected rather than private because AdaptiveSpotFinderGPU derives from this engine: it is
    // this same local-box detection with the fixed photon floor replaced by a per-resolution-ring
    // one, so it reuses the stream, the bit buffers and the extractor rather than owning a second
    // set of them.
    std::shared_ptr<CudaStream> stream;

    CudaDevicePtr<uint32_t> gpu_out_0;
    CudaDevicePtr<uint32_t> gpu_out_1;   // holds the strong-pixel bits after Detect()
    SpotExtractorGPU extractor;
    CudaDevicePtr<PanelBox> gpu_panels;   // ImageSpotFinder::panels on the device

private:
    const int numberOfCudaThreads = 128;    // #threads per block of analyze_pixel (one warp per 32 columns)
    const int numberOfWaves = 32;           // #row bands of analyze_pixel
    const int sparseRowsPerWave = 16;       // rows per band of analyze_pixel when only near candidates
    const int windowSizeLimit = 32;         // limit on the window size (2nby+1, 2nbx+1): a warp holds 2 x 32 columns
    int candidateBlocks = 0;                // grid of analyze_candidates (256 threads per block)

    void UploadPanels();

    void RunDetect(const ImagePreprocessorBuffer &image, const SpotFindingSettings &settings,
                   const uint32_t *gpu_candidates);
protected:
    // Detect, but with the result wanted only at the candidate pixels (device bit buffer, the layout
    // of the output): gpu_out_1 ends up holding (Detect's result & candidates), and the second pass is
    // evaluated at the candidates alone. Asynchronous - it does not wait for the stream.
    void DetectAt(const ImagePreprocessorBuffer &image, const SpotFindingSettings &settings,
                  const uint32_t *gpu_candidates);
public:
    ImageSpotFinderGPU(int32_t width, int32_t height, std::shared_ptr<CudaStream> stream);
    // Waits for the stream: its pooled buffers are freed next (see CudaDevicePtr).
    ~ImageSpotFinderGPU() override { if (stream) cudaStreamSynchronize(*stream); }

    void Detect(const ImagePreprocessorBuffer &image, const SpotFindingSettings &settings) override;
    void SetResolutionMaskBits(const std::vector<uint32_t> &packed_mask) override;
    void SetPanels(const DetectorModules &modules) override;
    [[nodiscard]] uint32_t StrongPixelCount() const override { return extractor.StrongPixelCount(); }
    const std::vector<DiffractionSpot> &ExtractComponents(const ImagePreprocessorBuffer &image,
                                                          const SpotFindingSettings &settings) override;
};