Jungfraujoch/image_analysis/spot_finding/AdaptiveSpotFinderGPU.h

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

#pragma once

// GPU adaptive spot finder that FUSES azimuthal integration and spot finding into one image pass.
//
// The CPU adaptive finder (AdaptiveSpotFinderCPU) and the azimuthal integrator both bin every pixel
// into resolution rings and reduce (sum / sum^2 / count). Today azint runs on the GPU while the
// adaptive finder re-does the identical per-ring reduction on the HOST - a wasted second pass over a
// ~10 MP image. This engine does the ring reduction on the GPU and drives BOTH products from it:
//   - the azimuthal-integration profile (mean intensity per ring, in flat-field-corrected space), and
//   - the per-ring background (mean, sigma, peak-excluded via two sigma-clip passes) that sets the
//     self-calibrating spot-detection threshold (in raw photon counts).
// It then flags strong pixels (value >= ring threshold) into a packed bit buffer and hands that
// buffer - still on the device - to SpotExtractorGPU, which builds the spots there.
//
// Numerically it reproduces AdaptiveSpotFinderCPU: the same three-pass robust background, the same
// per-ring threshold formula (shared via AdaptiveThreshold.h, computed on the host once per frame),
// and the same raw-count detection test. The only differences from the CPU are those inherent to a
// GPU reduction (float per-ring accumulation in atomic order vs the CPU's serial double sums), which
// shift a handful of borderline pixels at most. The corrected sums for the azint profile are
// accumulated in the SAME plain first pass, so one reduction feeds both products.

#include <memory>
#include <vector>

#include "ImageSpotFinder.h"
#include "SpotExtractorGPU.h"
#include "SpotFindingSettings.h"
#include "../../common/AzimuthalIntegrationProfile.h"
#include "../../common/AzimuthalIntegrationMapping.h"
#include "../indexing/CUDAMemHelpers.h"
#include "../indexing/CudaSharedTables.h"

class AdaptiveSpotFinderGPU : public ImageSpotFinder {
    const AzimuthalIntegrationMapping &mapping;
    std::shared_ptr<CudaStream> stream;

    const int nbins;
    const size_t npix;

    int reduce_threads = 128;
    int reduce_blocks = 0;
    int flag_threads = 256;
    int flag_blocks = 0;
    size_t shared_plain = 0;   // per-block shared bytes for the plain pass (raw + corrected rings)
    size_t shared_clip = 0;    // per-block shared bytes for a sigma-clip pass (raw rings only)
    bool use_shared = true;    // false -> nbins too large for shared memory, use the global-atomics kernel

    // Static mapping inputs: geometry-only, so one copy per GPU shared with every other engine on it
    // (see CudaSharedTables.h) rather than one copy per worker thread.
    std::shared_ptr<CudaDevicePtr<uint16_t>> gpu_pixel_to_bin;
    std::shared_ptr<CudaDevicePtr<float>>    gpu_corrections;

    // Raw per-ring accumulators (re-zeroed each pass) + derived stats used to clip and threshold.
    // double, like the CPU engine's ring accumulators: the ring sigma is the cancelling difference
    // sum2/n - m^2, and the block atomics that fill these arrive in an arbitrary order.
    CudaDevicePtr<unsigned long long> gpu_sum;
    CudaDevicePtr<unsigned long long> gpu_sum2;
    CudaDevicePtr<uint32_t> gpu_count;
    CudaDevicePtr<float>    gpu_mean;    // per-ring raw mean  (clip predicate)
    CudaDevicePtr<float>    gpu_sigma;   // per-ring raw sigma (clip predicate)

    // Corrected per-ring accumulators (plain first pass only) -> azimuthal-integration profile.
    CudaDevicePtr<float>    gpu_sum_corr;
    CudaDevicePtr<float>    gpu_sum2_corr;

    // Per-ring detection threshold (host-computed, uploaded) and the strong-pixel bit buffer.
    CudaDevicePtr<float>    gpu_thr;
    CudaDevicePtr<uint32_t> gpu_strong;

    // Host mirrors of the small per-ring transfers.
    std::vector<unsigned long long> host_sum;    // clipped raw sum   } input to the host threshold computation
    std::vector<unsigned long long> host_sum2;   // clipped raw sum^2 }  (exact integers - see the kernel)
    std::vector<uint32_t> host_count;  // clipped raw count }
    std::vector<float>    host_thr;    // per-ring threshold (empty -> frame had no valid pixels)
    std::vector<float>    host_bkg;    // clipped per-ring mean, NaN where the ring is too sparse to trust
    std::vector<float>    prof_sum;    // plain corrected sum   } azimuthal-integration profile
    std::vector<float>    prof_sum2;   // plain corrected sum^2 }
    std::vector<uint32_t> prof_count;  // plain pixel count     }

    // Every per-ring array above is a device-to-host copy once per frame. A D2H copy into PAGEABLE
    // memory blocks the host until it completes, whatever stream it was issued on - which would stall
    // Detect() between the plain pass and the clip passes, with the device then idle while the host
    // enqueues them. Pinning the destinations makes the copies genuinely asynchronous, as the
    // azimuthal-integration engine already does with its own.
    CudaRegisteredVector<unsigned long long> host_sum_reg;
    CudaRegisteredVector<unsigned long long> host_sum2_reg;
    CudaRegisteredVector<uint32_t> host_count_reg;
    CudaRegisteredVector<float>    prof_sum_reg;
    CudaRegisteredVector<float>    prof_sum2_reg;
    CudaRegisteredVector<uint32_t> prof_count_reg;

    SpotExtractorGPU extractor;   // builds the spots from gpu_strong without it leaving the device

    AzimuthalIntegrationProfile last_profile;   // filled every Run(), retrievable via GetProfile()

    // One reduction pass over the image into the raw accumulators. clip_k <= 0 -> plain pass (all
    // valid pixels); clip_k > 0 -> keep only pixels within clip_k sigma of the current gpu_mean.
    // accumulate_corrected additionally fills gpu_sum_corr/gpu_sum2_corr for the profile (plain pass).
    void ReducePass(const ImagePreprocessorBuffer &image, float clip_k, bool accumulate_corrected);
    // Finalize gpu_mean/gpu_sigma from the current raw accumulators (per ring).
    void FinalizeStats();
    // Host: per-ring threshold from the clipped raw stats and the single knob E (false pixels/frame).
    void ComputeThresholds(const SpotFindingSettings &settings);

public:
    AdaptiveSpotFinderGPU(const AzimuthalIntegrationMapping &mapping, std::shared_ptr<CudaStream> stream);
    ~AdaptiveSpotFinderGPU() override = default;
    AdaptiveSpotFinderGPU(const AdaptiveSpotFinderGPU &) = delete;
    AdaptiveSpotFinderGPU &operator=(const AdaptiveSpotFinderGPU &) = delete;

    void Detect(const ImagePreprocessorBuffer &image, const SpotFindingSettings &settings) override;
    void SetResolutionMask(const std::vector<bool> &mask) override;
    const std::vector<DiffractionSpot> &ExtractComponents(const ImagePreprocessorBuffer &image,
                                                          const SpotFindingSettings &settings) override;

    // The azimuthal profile computed as a byproduct of the last Detect() - lets this engine replace the
    // separate azint pass in the analysis pipeline.
    [[nodiscard]] const AzimuthalIntegrationProfile &GetProfile() const { return last_profile; }
    [[nodiscard]] const std::vector<float> &GetRingBackground() const override { return host_bkg; }
};