Jungfraujoch/image_analysis/image_preprocessing/ImagePreprocessorGPU.cu

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

#include "ImagePreprocessorGPU.h"

template<class T>
__global__ void preprocess_kernel(
    const T *__restrict__ input,
    const uint8_t *__restrict__ mask,
    int32_t *__restrict__ output,
    ImageStatistics *__restrict__ stats,
    T saturation_limit,
    T err_value,
    int npixels) {
    // Shared block accumulators
    __shared__ unsigned long long s_masked;
    __shared__ unsigned long long s_saturated;
    __shared__ unsigned long long s_error;
    __shared__ long long s_max;
    __shared__ long long s_min;

    if (threadIdx.x == 0) {
        s_masked = 0;
        s_saturated = 0;
        s_error = 0;
        s_max = INT64_MIN;
        s_min = INT64_MAX;
    }
    __syncthreads();

    // Thread-local accumulators
    unsigned long long local_masked = 0;
    unsigned long long local_saturated = 0;
    unsigned long long local_error = 0;
    long long local_max = INT64_MIN;
    long long local_min = INT64_MAX;

    for (int i = blockIdx.x * blockDim.x + threadIdx.x;
         i < npixels;
         i += blockDim.x * gridDim.x) {
        T v = input[i];

        bool is_masked = mask[i];
        // Error/invalid marker = the pixel type's extreme value (0xFFFFFFFF for EIGER uint32); tested
        // before saturation, since for unsigned types the marker also exceeds saturation_limit (which is
        // clipped to the HDF5 saturation_value). Priority: masked > error > saturated.
        bool is_err = (v == err_value);
        bool is_sat = !is_err && (v >= saturation_limit);

        bool valid = !(is_masked || is_sat || is_err);

        // Output
        output[i] =
                is_masked ? INT32_MIN : is_err ? INT32_MIN : is_sat ? INT32_MAX : (int32_t) v;

        // Counters
        local_masked += is_masked;
        local_error += (!is_masked && is_err);
        local_saturated += (!is_masked && !is_err && is_sat);

        // Min/max only for valid
        if (valid) {
            int64_t val = (int64_t) v;
            if (val > local_max) local_max = val;
            if (val < local_min) local_min = val;
        }
    }

    // Reduce to shared memory
    atomicAdd(&s_masked, local_masked);
    atomicAdd(&s_saturated, local_saturated);
    atomicAdd(&s_error, local_error);

    if (local_min <= local_max) {
        atomicMax((long long *) &s_max, (long long) local_max);
        atomicMin((long long *) &s_min, (long long) local_min);
    }

    __syncthreads();

    // One thread writes block result
    if (threadIdx.x == 0) {
        atomicAdd(&stats->masked_pixel_count, s_masked);
        atomicAdd(&stats->saturated_pixel_count, s_saturated);
        atomicAdd(&stats->error_pixel_count, s_error);

        atomicMax((long long *) &stats->max_value, (long long) s_max);
        atomicMin((long long *) &stats->min_value, (long long) s_min);
    }
}

ImagePreprocessorGPU::ImagePreprocessorGPU(const DiffractionExperiment &experiment, const PixelMask &mask,
                                           std::shared_ptr<CudaStream> stream, bool copy_image_to_host)
    : ImagePreprocessor(experiment),
      stream(stream),
      copy_image_to_host(copy_image_to_host),
      gpu_decompressed_image(npixels * sizeof(uint32_t)),  // Overshoot - if input image is 1- or 2-byte, then it is still fine, while memory loss is minimal
      gpu_stats(1),
      cpu_stats(1),
      cpu_stats_reg(cpu_stats) {
    // Setup mask. The same for every worker, so it is uploaded once per GPU and shared; keyed on the
    // PixelMask's own vector, which the derived table is a pure function of.
    std::vector<uint8_t> mask_vec(npixels);
    for (int i = 0; i < npixels; i++)
        mask_vec[i] = (mask.GetMask().at(i) != 0);
    gpu_mask = SharedDeviceTable(mask.GetMask().data(), npixels, mask_vec.data(), *stream);

    // Setup GPU settings. The current device, not device 0: workers are pinned round-robin across GPUs,
    // so device 0's SM count can belong to a different card than the one these kernels launch on.
    int device = 0;
    cudaGetDevice(&device);
    cudaDeviceProp prop{};
    cudaGetDeviceProperties(&prop, device);

    threads = 128;
    blocks = 4 * prop.multiProcessorCount;
}

void ImagePreprocessorGPU::PinInputBuffer(std::vector<uint8_t> &buffer, size_t size) {
    if (buffer.size() == size)
        return;
    // Unregister before the resize, which can move the buffer.
    input_reg.unregister();
    buffer.resize(size);
    input_reg.rebind(buffer);
}

ImageStatistics ImagePreprocessorGPU::Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *image_ptr,
                                              CompressedImageMode image_mode) {
    switch (image_mode) {
        case CompressedImageMode::Int8:
            return Analyze<int8_t>(processed_image, image_ptr, INT8_MIN, INT8_MAX);
        case CompressedImageMode::Int16:
            return Analyze<int16_t>(processed_image, image_ptr, INT16_MIN, INT16_MAX);
        case CompressedImageMode::Int32:
            return Analyze<int32_t>(processed_image, image_ptr, INT32_MIN, INT32_MAX);
        case CompressedImageMode::Uint8:
            return Analyze<uint8_t>(processed_image, image_ptr, UINT8_MAX, UINT8_MAX);
        case CompressedImageMode::Uint16:
            return Analyze<uint16_t>(processed_image, image_ptr, UINT16_MAX, UINT16_MAX);
        case CompressedImageMode::Uint32:
            return Analyze<uint32_t>(processed_image, image_ptr, UINT32_MAX, UINT32_MAX);
        default:
            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "RGB/float mode not supported");
    }
}

bool ImagePreprocessorGPU::AnalyzeCompressed(ImagePreprocessorBuffer &processed_image,
                                             const CompressedImage &image,
                                             ImageStatistics &stats) {
    if (!BSLZ4DecoderGPU::Supports(image))
        return false;   // caller decompresses on the host and uses Analyze()
    if (image.GetUncompressedSize() != npixels * image.GetByteDepth())
        return false;

    if (!bslz4_decoder)
        bslz4_decoder = std::make_unique<BSLZ4DecoderGPU>(npixels * sizeof(uint32_t), stream);

    // Straight into the buffer Analyze() would have filled by copying the decompressed frame across
    // the bus; from here the two paths are the same code.
    bslz4_decoder->Decode(image, gpu_decompressed_image.get());

    switch (image.GetMode()) {
        case CompressedImageMode::Int8:
            stats = AnalyzeOnDevice<int8_t>(processed_image, INT8_MIN, INT8_MAX); return true;
        case CompressedImageMode::Uint8:
            stats = AnalyzeOnDevice<uint8_t>(processed_image, UINT8_MAX, UINT8_MAX); return true;
        case CompressedImageMode::Int16:
            stats = AnalyzeOnDevice<int16_t>(processed_image, INT16_MIN, INT16_MAX); return true;
        case CompressedImageMode::Uint16:
            stats = AnalyzeOnDevice<uint16_t>(processed_image, UINT16_MAX, UINT16_MAX); return true;
        case CompressedImageMode::Int32:
            stats = AnalyzeOnDevice<int32_t>(processed_image, INT32_MIN, INT32_MAX); return true;
        case CompressedImageMode::Uint32:
            stats = AnalyzeOnDevice<uint32_t>(processed_image, UINT32_MAX, UINT32_MAX); return true;
        default:
            return false;   // Supports() already excludes these; belt and braces
    }
}

template<class T>
ImageStatistics ImagePreprocessorGPU::Analyze(ImagePreprocessorBuffer &processed_image,
                                              const uint8_t *input,
                                              T err_value,
                                              T sat_value) {
    // On this engine's own stream, not the NULL stream: a NULL-stream copy implicitly synchronises with
    // every blocking stream in the process, which serialised all workers behind whichever one was
    // uploading. The stream is synchronised at the end of this function, so the ordering is unchanged.
    cudaMemcpyAsync(gpu_decompressed_image, input, npixels * sizeof(T), cudaMemcpyHostToDevice, *stream);
    return AnalyzeOnDevice<T>(processed_image, err_value, sat_value);
}

// Everything after the image is on the device, shared by the host-upload and the device-decode
// entry points so the two cannot drift apart.
template<class T>
ImageStatistics ImagePreprocessorGPU::AnalyzeOnDevice(ImagePreprocessorBuffer &processed_image,
                                                      T err_value, T sat_value) {
    if (sat_value > saturation_limit)
        sat_value = static_cast<T>(saturation_limit);

    cpu_stats[0] = ImageStatistics{.max_value = INT64_MIN, .min_value = INT64_MAX};
    cudaMemcpyAsync(gpu_stats, cpu_stats.data(), sizeof(ImageStatistics), cudaMemcpyHostToDevice, *stream);
    preprocess_kernel<T> <<< blocks, threads, 0, *stream >>>(
        reinterpret_cast<const T *>(gpu_decompressed_image.get()),
        gpu_mask->get(),
        processed_image.getGPUBuffer(),
        gpu_stats,
        sat_value,
        err_value,
        npixels);
    // The preprocessed image is 4 bytes per pixel - by far the largest transfer here - and every GPU
    // engine reads it straight from the device buffer, so it only comes back when a CPU engine needs it.
    if (copy_image_to_host)
        cudaMemcpyAsync(processed_image.data(), processed_image.getGPUBuffer(), npixels * sizeof(int32_t), cudaMemcpyDeviceToHost, *stream);
    cudaMemcpyAsync(cpu_stats.data(), gpu_stats, sizeof(ImageStatistics), cudaMemcpyDeviceToHost, *stream);

    cudaStreamSynchronize(*stream);

    return cpu_stats[0];
}