Jungfraujoch/image_analysis/image_preprocessing/ImagePreprocessorGPU.h

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

#pragma once

#include <memory>

#include "ImagePreprocessor.h"
#include "BSLZ4DecoderGPU.h"
#include "CBFDecoderGPU.h"
#include "../indexing/CUDAMemHelpers.h"
#include "../indexing/CudaSharedTables.h"

class ImagePreprocessorGPU : public ImagePreprocessor {
    std::shared_ptr<CudaStream> stream;
    const bool copy_image_to_host;
    int threads;
    int blocks;
    // Geometry-only, so one copy per GPU shared with every other engine on it (CudaSharedTables.h).
    std::shared_ptr<CudaDevicePtr<uint8_t>> gpu_mask;
    // Landing buffer for the HOST-upload path only. The device-decode path un-transposes straight
    // into the preprocessed image, so it never needs this - and at 4 bytes per pixel it is worth a
    // frame per worker, so it is allocated on first use rather than always.
    CudaDevicePtr<uint8_t> gpu_decompressed_image;
    CudaDevicePtr<ImageStatistics> gpu_stats;

    PinnedVector<ImageStatistics> cpu_stats;
    CudaRegisteredVector<uint8_t> input_reg;   // page-locks the caller's decompression buffer
    // The uncompressed-image counterpart of input_reg: the region PinInputRegion last page-locked.
    const void *pinned_input = nullptr;
    size_t pinned_input_bytes = 0;
    // Staging for a marCCD frame's stored 16-bit pixels (UINT16_WIDENED), allocated on first use.
    CudaHostPtr<uint8_t> host_stored;
    // The card reads host memory directly (gpu_direct_host_read()): images are read where they lie,
    // nothing is page-locked or staged; only the CBF device decoder still uploads its input.
    bool direct_host_read = false;

    std::vector<int32_t> cpu_image;

    // Built on first use: a decoder that can serve this engine's images, sized to the frame.
    std::unique_ptr<BSLZ4DecoderGPU> bslz4_decoder;
    std::unique_ptr<CBFDecoderGPU> cbf_decoder;
    // Which decoder the last AnalyzeCompressed used, for GetLastDecompressionTime_s.
    enum class LastDecode { None, BSLZ4, CBF } last_decode = LastDecode::None;

    // The file encodings of the miniCBF and marCCD readers (CompressionAlgorithmEnum.h): the stored
    // bytes are uploaded, decoded into gpu_decompressed_image as the int32 image the host would have
    // produced, and preprocessed by the same kernel as an uploaded int32 image.
    ImageStatistics AnalyzeStoredInt32(ImagePreprocessorBuffer &processed_image, const CompressedImage &image);

    template <class T> ImageStatistics Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *input, T err_value, T sat_value);
    // Preprocess the image at `input`: gpu_decompressed_image on the host-upload path, or the host
    // buffer itself where the card reads host memory directly. err_value and sat_value are compared
    // with the pixel's value, not its bits, so a narrower T can carry the int32 semantics of the image
    // it widens to.
    template <class T> ImageStatistics AnalyzeOnDevice(ImagePreprocessorBuffer &processed_image, const void *input,
                                                       int64_t err_value, int64_t sat_value);
    // Preprocess straight out of the bitshuffled bytes (the device-decode path). Same per-pixel
    // decision and same statistics as AnalyzeOnDevice, with the un-transpose folded in.
    template <class T, int ES> ImageStatistics UntransposeAndAnalyze(ImagePreprocessorBuffer &processed_image,
                                                                     const BSLZ4ShuffledImage &shuffled,
                                                                     T err_value, T sat_value);
public:
    // copy_image_to_host copies the preprocessed image back after every frame. It is only needed when
    // something on the CPU reads it - the GPU engines all work off the device buffer - and at 4 bytes
    // per pixel it is the single largest transfer in the pipeline, so the caller says whether it wants it.
    ImagePreprocessorGPU(const DiffractionExperiment &experiment, const PixelMask &mask, std::shared_ptr<CudaStream> stream,
                         bool copy_image_to_host = true);
    ~ImagePreprocessorGPU() override;
    ImageStatistics Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *decompressed_image, CompressedImageMode image_mode) override;
    bool AnalyzeCompressed(ImagePreprocessorBuffer &processed_image, const CompressedImage &image,
                           ImageStatistics &stats) override;
    [[nodiscard]] float GetLastDecompressionTime_s() const override;
    void PinInputBuffer(std::vector<uint8_t> &buffer, size_t size) override;
    void PinInputRegion(const void *ptr, size_t bytes) override;
};