Jungfraujoch/image_analysis/image_preprocessing/ImagePreprocessorBuffer.h

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

#pragma once

#include <vector>
#include <cstdint>
#include <cstddef>
#include <climits>

#include "../../common/PanelBoxes.h"

class ImagePreprocessorBuffer {
protected:
    // The panels the image is made of (PanelBoxes); empty means one, the whole image.
    std::vector<PanelBox> panels;
    // The box the square about pixel (x, y) is clipped to: its panel's, or the image's.
    PanelBox TileBox(int32_t x, int32_t y, int width) const;

    // The image size, kept apart from `buffer` so size() still answers when the host copy was not
    // allocated at all (see the host_mirror constructor).
    const size_t npixels;
    std::vector<int32_t> buffer;

    // host_mirror = false leaves `buffer` empty. On the GPU path the preprocessed image lives on the
    // device and no CPU engine reads it back, so the four-bytes-per-pixel host copy would be
    // allocated, zeroed and page-locked for nothing - 72 MB per worker on a 16M-pixel detector.
    ImagePreprocessorBuffer(size_t npixels, bool host_mirror);
public:
    explicit ImagePreprocessorBuffer(size_t npixels);
    virtual ~ImagePreprocessorBuffer() = default;

    // Standard CPU operation. Defined here rather than in the .cpp: these are called once per pixel by
    // the CPU preprocessor, the CPU azimuthal integrator and the CPU spot finder, and the build has no
    // link-time optimisation - so out of line they are a real call per pixel, and they stop the callers
    // vectorising at all. Inlining them costs about a third of the per-image CPU work on every path
    // that does not have a GPU, and leaves the results bit-identical.
    std::vector<int32_t> &getBuffer() { return buffer; }
    const std::vector<int32_t> &getBuffer() const { return buffer; }
    int32_t &operator[](size_t i) { return buffer[i]; }
    const int32_t &operator[](size_t i) const { return buffer[i]; }
    size_t size() const { return npixels; }

    int32_t *data() { return buffer.data(); }
    const int32_t *data() const { return buffer.data(); }

    // GPU operations (overriden in ImagePreprocessorBufferGPU
    virtual int32_t *getGPUBuffer();
    virtual const int32_t *getGPUBuffer() const;

    // Values of the pixels with the given flat indices. Not just operator[] in a loop: on the GPU the
    // preprocessed image can live in device memory only, and then this gathers them there.
    virtual void Gather(const std::vector<uint32_t> &npixel, std::vector<int32_t> &values) const;

    // The (2*half+1)^2 pixels about each (x[i], y[i]) of the image, `width` pixels wide: square after
    // square, each row by row, with a pixel off the image - or off the panel (x[i], y[i]) is on - read as
    // INT32_MIN. Valid until the next call.
    virtual const int32_t *GatherTiles(const std::vector<int32_t> &x, const std::vector<int32_t> &y,
                                       int width, int half);
    // The panels of the image, for GatherTiles. Default: none, the whole image is one.
    void SetPanels(const std::vector<PanelBox> &in_panels) { panels = in_panels; }

    // The whole image on the host - on the GPU copied there first, where it was not already.
    virtual const int32_t *HostImage();
private:
    std::vector<int32_t> tiles;
};