Jungfraujoch/rugnux/HotPixelsGPU.h

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

#pragma once

#include <cstdint>
#include <memory>
#include <mutex>
#include <vector>

#include "../image_analysis/indexing/CUDAMemHelpers.h"

// The device half of HotPixelFinder, for frames already preprocessed on the GPU: the per-frame order
// statistics (each ring-sector's median, each ring's median and median absolute deviation) and the
// per-pixel sums run where the image already is, and only the per-key statistics - tens of thousands
// of numbers - come to the host, which turns them into levels and thresholds with the very code the
// host path uses. Every statistic is an exact order statistic of integers and every sum an integer,
// so the sums, and with them the mask, are identical to what HotPixelFinder::AddImage produces.
class HotPixelFinderGPU {
    const size_t npixels;
    const size_t nkeys;
    const int nrings;
    const int sectors;

    CudaDevicePtr<int32_t> key;               // ring * sectors + sector of each pixel, -1 masked
    CudaDevicePtr<uint32_t> pixels_by_key;    // the unmasked pixels, grouped by key
    CudaDevicePtr<uint32_t> key_begin;        // where each key's pixels start in pixels_by_key

    // The per-pixel sums, exactly those of HotPixelFinder.
    CudaDevicePtr<uint16_t> n_lit, n_error, n_error_ring_ok;
    CudaDevicePtr<int64_t> sum_value, error_level_sum;
    // Frames are selected on their workers' streams in parallel, but each pixel's sums are plain
    // read-modify-writes, so one frame at a time adds to them.
    std::mutex accumulate_mutex;

public:
    // One worker's buffers, on the stream its frames are preprocessed on.
    struct Frame {
        explicit Frame(std::shared_ptr<CudaStream> stream) : stream(std::move(stream)) {}
        std::shared_ptr<CudaStream> stream;
        CudaDevicePtr<uint32_t> count;           // valid pixels per key
        CudaDevicePtr<int32_t> sector_median;    // per key
        CudaDevicePtr<int32_t> ring_median, ring_mad;
        CudaDevicePtr<int32_t> level;
        CudaDevicePtr<float> threshold;
        CudaDevicePtr<char> ring_ok;
    };

    HotPixelFinderGPU(const int32_t *key, size_t npixels, const std::vector<uint32_t> &key_begin, int nrings,
                      int sectors);

    // The lower median of the valid values of each key (count[k] of them, 0 where there are none), and
    // of each ring the median and the lower median of the absolute deviations from it.
    void Statistics(const int32_t *device_image, Frame &frame, std::vector<uint32_t> &count,
                    std::vector<int32_t> &sector_median, std::vector<int32_t> &ring_median,
                    std::vector<int32_t> &ring_mad);

    // Add the frame to the per-pixel sums, with each key's level and lit threshold and each ring's
    // verdict on whether it has a level at all - as HotPixelFinder::AddImage does.
    void Accumulate(const int32_t *device_image, Frame &frame, const std::vector<int32_t> &level,
                    const std::vector<float> &threshold, const std::vector<char> &ring_ok);

    // The per-pixel sums, npixels each.
    void Download(uint16_t *n_lit, uint16_t *n_error, int64_t *sum_value, uint16_t *n_error_ring_ok,
                  int64_t *error_level_sum);
};