Jungfraujoch/image_analysis/image_preprocessing/ImagePreprocessorBufferGPU.cu

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

#include "ImagePreprocessorBufferGPU.h"

static inline void cuda_err(cudaError_t val) {
    if (val != cudaSuccess)
        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
}

__global__ void gather_kernel(const int32_t *__restrict__ image,
                              const uint32_t *__restrict__ npixel,
                              int32_t *__restrict__ values,
                              int count) {
    for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < count; i += blockDim.x * gridDim.x)
        values[i] = image[npixel[i]];
}

ImagePreprocessorBufferGPU::ImagePreprocessorBufferGPU(size_t npixel, bool host_mirror)
    : ImagePreprocessorBuffer(npixel, host_mirror),
      gpu_image(npixel),
      // A no-op when the mirror was not allocated: CudaRegisteredVector skips an empty vector.
      buffer_reg(buffer),
      max_gather(StrongPixelLimit(npixel)),
      gpu_gather_index(max_gather),
      gpu_gather_value(max_gather) {
}

int32_t *ImagePreprocessorBufferGPU::getGPUBuffer() {
    return gpu_image;
}

const int32_t *ImagePreprocessorBufferGPU::getGPUBuffer() const {
    return gpu_image;
}

void ImagePreprocessorBufferGPU::Gather(const std::vector<uint32_t> &npixel, std::vector<int32_t> &values) const {
    values.resize(npixel.size());
    if (npixel.empty())
        return;

    const int count = static_cast<int>(npixel.size());
    cudaMemcpyAsync(gpu_gather_index.get(), npixel.data(), count * sizeof(uint32_t),
                    cudaMemcpyHostToDevice, gather_stream);
    gather_kernel<<<(count + 255) / 256, 256, 0, gather_stream>>>(
        gpu_image.get(), gpu_gather_index.get(), gpu_gather_value.get(), count);
    cuda_err(cudaGetLastError());
    cudaMemcpyAsync(values.data(), gpu_gather_value.get(), count * sizeof(int32_t),
                    cudaMemcpyDeviceToHost, gather_stream);
    cudaStreamSynchronize(gather_stream);
}