Jungfraujoch/image_analysis/geom_refinement/BeamCenterFFTGPU.h

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

#pragma once

// Included only where CUDA is known to be present, i.e. under JFJOCH_USE_CUDA - the same rule as
// FFTIndexerGPU.h. Deliberately free of CUDA headers so a host translation unit (the parity test)
// can construct the engine.

#include "BeamCenterFFTEngine.h"

// The beam-centre transforms on cuFFT. The capture is one call per run on an image of tens of
// megapixels, so plans and device buffers are created inside the call that needs them and freed
// when it returns: the card is shared with the analysis workers, and a plan held across the run
// would be several hundred megabytes doing nothing.
class BeamCenterFFTGPU : public BeamCenterFFTEngine {
public:
    BeamCenterConvSurfaces2D PointSurfaces(const std::vector<float> &a, const std::vector<float> &m,
                                           int64_t h, int64_t w) override;

    // The four convolutions are combined and searched where they already are, so the only thing
    // that crosses PCIe is the shortlist. On a 16 Mpixel detector that is 1.16 GB of surfaces, a
    // 1.4 s greedy search and a 0.4 s combination not done on the host - against 26 ms of
    // transform, which is what the host work had become.
    std::vector<BeamCenterFFTCandidate> PointShortlist(const std::vector<float> &a,
                                                       const std::vector<float> &m, int64_t h,
                                                       int64_t w,
                                                       const BeamCenterFFTSettings &settings,
                                                       double global_variance) override;

    // The scored surface itself, brought back. Nothing in the pipeline asks for it - the shortlist
    // is made on the device - but the parity test holds the device combination to
    // BeamCenterPointScore element by element, which is the sharpest comparison available once the
    // two paths no longer share that step.
    [[nodiscard]] std::vector<float> PointScoreSurface(const std::vector<float> &a,
                                                       const std::vector<float> &m, int64_t h,
                                                       int64_t w, float min_pair_fraction,
                                                       double global_variance);

    // Peak device memory the point surfaces need for an image of this size: buffers plus the cuFFT
    // plans' own workspace.
    [[nodiscard]] static size_t DeviceMemoryNeeded(int64_t width, int64_t height);
    // Whether that is at most half of the current device's TOTAL memory - the rest is the analysis
    // workers' beside it. Not what is free at the moment, so the same image on the same card always
    // takes the same path; a capture over it runs on the CPU.
    [[nodiscard]] static bool FitsInDeviceMemory(int64_t width, int64_t height);
};