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

#include <algorithm>
#include <type_traits>

#include "ImagePreprocessorGPU.h"
#include "BSLZ4DecodeWarp.h"
#include "../../common/JFJochException.h"

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

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);
    }
}

// The per-pixel decision preprocess_kernel makes, in a form the fused kernel can reuse so the two
// cannot drift apart. Priority: masked > error > saturated.
template<class T>
struct PreprocessAccum {
    unsigned long long masked = 0, saturated = 0, error = 0;
    long long max_v = INT64_MIN, min_v = INT64_MAX;

    __device__ __forceinline__ int32_t Apply(T v, bool is_masked, T sat_value, T err_value) {
        const bool is_err = (v == err_value);
        const bool is_sat = !is_err && (v >= sat_value);
        masked += is_masked;
        error += (!is_masked && is_err);
        saturated += (!is_masked && !is_err && is_sat);
        if (!(is_masked || is_sat || is_err)) {
            const int64_t val = (int64_t) v;
            if (val > max_v) max_v = val;
            if (val < min_v) min_v = val;
        }
        return is_masked ? INT32_MIN : is_err ? INT32_MIN : is_sat ? INT32_MAX : (int32_t) v;
    }
};

// Reduce a block's thread-local accumulators into the image-wide statistics. Every accumulator is an
// integer, so the result does not depend on the order the warps arrive in.
//
// The reduction is within the warp and then straight to global, rather than through a shared-memory
// staging area. It has to be: the fused kernel below fills its shared budget to the byte with the
// bitshuffle block it decodes, and forty-odd bytes of accumulators on top of that cost it a whole
// resident CUDA block per SM. A thread that saw no valid pixel still carries INT64_MIN/INT64_MAX,
// which is the identity for max/min, so it needs no guard.
template<class T>
__device__ __forceinline__ void FlushStats(PreprocessAccum<T> &l, ImageStatistics *stats) {
    for (int d = 16; d > 0; d >>= 1) {
        l.masked += __shfl_down_sync(0xffffffff, l.masked, d);
        l.saturated += __shfl_down_sync(0xffffffff, l.saturated, d);
        l.error += __shfl_down_sync(0xffffffff, l.error, d);
        l.max_v = max(l.max_v, __shfl_down_sync(0xffffffff, l.max_v, d));
        l.min_v = min(l.min_v, __shfl_down_sync(0xffffffff, l.min_v, d));
    }

    if ((threadIdx.x & 31) == 0) {
        atomicAdd(&stats->masked_pixel_count, l.masked);
        atomicAdd(&stats->saturated_pixel_count, l.saturated);
        atomicAdd(&stats->error_pixel_count, l.error);
        atomicMax((long long *) &stats->max_value, l.max_v);
        atomicMin((long long *) &stats->min_value, l.min_v);
    }
}

__device__ __forceinline__ uint64_t transpose8_fused(uint64_t x) {
    uint64_t t;
    t = (x ^ (x >> 7)) & 0x00aa00aa00aa00aaULL;  x = x ^ t ^ (t << 7);
    t = (x ^ (x >> 14)) & 0x0000cccc0000ccccULL; x = x ^ t ^ (t << 14);
    t = (x ^ (x >> 28)) & 0x00000000f0f0f0f0ULL; x = x ^ t ^ (t << 28);
    return x;
}

// The un-transpose and the preprocessing of ONE bitshuffle block, mirroring bitshuf_decode_block.
// One thread owns one group of 8 elements across every byte-plane, so once it has transposed its 8
// bytes out of each plane it holds 8 complete elements and can emit 8 finished int32 pixels - the
// decompressed image never has to exist in device memory at all. That removes a full-frame buffer
// per worker and a full-frame write plus read from the pipeline.
//
// `in` is the block's bitshuffled bytes: device memory when the LZ4 pass ran in its own kernel, and
// shared memory when it ran in this one. A generic pointer covers both, so the two routes cannot
// produce different pixels from the same block.
template<class T, int ES>
__device__ __forceinline__ void untranspose_preprocess_block(const uint8_t *in, uint32_t size, uint32_t elem0,
                                                             const uint8_t *__restrict__ mask,
                                                             int32_t *__restrict__ out,
                                                             T sat_value, T err_value,
                                                             PreprocessAccum<T> &l) {
    // The bytes are assembled in the unsigned counterpart of T - shifting a byte into the top of a
    // signed type overflows it - and converted back at the end, which C++20 defines as two's
    // complement reinterpretation. That is exactly what the byte order in the file means.
    using U = typename std::make_unsigned<T>::type;

    const uint32_t n = size / 8;
    for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
        uint64_t x[ES];
        #pragma unroll
        for (int p = 0; p < ES; p++) {
            const uint8_t *pin = in + p * size;
            uint64_t a = 0;
            #pragma unroll
            for (int k = 0; k < 8; k++) a |= (uint64_t)pin[k * n + i] << (8 * k);
            x[p] = transpose8_fused(a);
        }
        int32_t o[8];
        #pragma unroll
        for (int k = 0; k < 8; k++) {
            U uv = 0;
            #pragma unroll
            for (int p = 0; p < ES; p++) uv |= (U)((U)((x[p] >> (8 * k)) & 0xff) << (8 * p));
            o[k] = l.Apply((T) uv, mask[elem0 + i * 8 + k] != 0, sat_value, err_value);
        }
        // elem0 is a multiple of 8 (bitshuffle blocks are), so this is 32-byte aligned.
        int4 *dst = reinterpret_cast<int4 *>(out + elem0 + i * 8);
        dst[0] = make_int4(o[0], o[1], o[2], o[3]);
        dst[1] = make_int4(o[4], o[5], o[6], o[7]);
    }
}

// The handful of elements bitshuffle stores verbatim rather than in a block. They are already on the
// device inside the uploaded chunk, so they only need the per-pixel decision, not the un-transpose.
template<class T, int ES>
__device__ __forceinline__ void preprocess_tail(const uint8_t *__restrict__ tail_src, uint32_t tail_elems,
                                                uint32_t tail_elem0, const uint8_t *__restrict__ mask,
                                                int32_t *__restrict__ out, T sat_value, T err_value,
                                                PreprocessAccum<T> &l) {
    using U = typename std::make_unsigned<T>::type;
    if (threadIdx.x < tail_elems) {
        U uv = 0;
        #pragma unroll
        for (int p = 0; p < ES; p++)
            uv |= (U)((U)tail_src[threadIdx.x * ES + p] << (8 * p));
        const T v = (T) uv;
        out[tail_elem0 + threadIdx.x] = l.Apply(v, mask[tail_elem0 + threadIdx.x] != 0, sat_value, err_value);
    }
}

// One CUDA block per bitshuffle block, over blocks the LZ4 kernel has already decoded. The last CUDA
// block (blockIdx.x == nblocks) finishes the verbatim tail.
template<class T, int ES>
__global__ __launch_bounds__(256) void untranspose_preprocess_kernel(
        const uint8_t *__restrict__ shuffled,
        const BSLZ4BlockDesc *__restrict__ desc,
        const uint8_t *__restrict__ mask,
        int32_t *__restrict__ out,
        ImageStatistics *__restrict__ stats,
        T sat_value, T err_value, int nblocks,
        const uint8_t *__restrict__ tail_src, uint32_t tail_elems, uint32_t tail_elem0) {
    PreprocessAccum<T> l;

    if (blockIdx.x == nblocks)
        preprocess_tail<T, ES>(tail_src, tail_elems, tail_elem0, mask, out, sat_value, err_value, l);
    else
        untranspose_preprocess_block<T, ES>(shuffled + desc[blockIdx.x].out_off, desc[blockIdx.x].nelem,
                                            desc[blockIdx.x].out_off / ES, mask, out,
                                            sat_value, err_value, l);
    FlushStats<T>(l, stats);
}

// The same thing with the LZ4 decode folded in as well: one CUDA block owns one bitshuffle block
// from the compressed payload all the way to finished pixels. Its first warp decodes the payload
// into a shared-memory buffer the size of one block, and then the whole CUDA block un-transposes and
// preprocesses out of that buffer. Nothing of the block reaches device memory but the pixels.
//
// The saved bandwidth is the smaller half of it - 72 MB of bitshuffled bytes per 18 Mpx frame stop
// being written and read back. What the shared buffer is really for is the LZ4 copy loop: a match
// sources bytes that other lanes of the warp wrote a few sequences earlier, so every copy step is a
// dependent round trip to wherever the output lives, and there are a few hundred of them per block.
// In device memory that round trip is hundreds of cycles, which is why the standalone decode runs at
// a fraction of the streaming rate this hardware reaches on a plain pass over an image; in shared
// memory it is tens of cycles.
//
// It is paid for in residency. A whole bitshuffle block of shared memory per CUDA block means an SM
// holds only as many concurrent decoders as its shared memory divides into: four on a T4 at the
// 16 kB blocks 32-bit detector data comes in, against the thirty-two warps the standalone kernel
// keeps in flight. So this trades decode parallelism away for decode latency, and which way that
// comes out is a measurement rather than an argument.
template<class T, int ES>
__global__ void decode_untranspose_preprocess_kernel(
        const uint8_t *__restrict__ src,
        const BSLZ4BlockDesc *__restrict__ desc,
        uint32_t *__restrict__ status,
        const uint8_t *__restrict__ mask,
        int32_t *__restrict__ out,
        ImageStatistics *__restrict__ stats,
        T sat_value, T err_value, int nblocks,
        const uint8_t *__restrict__ tail_src, uint32_t tail_elems, uint32_t tail_elem0) {
    extern __shared__ uint8_t s_shuffled[];
    PreprocessAccum<T> l;

    if (blockIdx.x == nblocks) {
        preprocess_tail<T, ES>(tail_src, tail_elems, tail_elem0, mask, out, sat_value, err_value, l);
        FlushStats<T>(l, stats);
        return;
    }

    // An LZ4 match never reaches back past the start of its own bitshuffle block - the decoder
    // rejects an offset larger than what the block has written - so one block's worth of shared
    // memory is the whole of what the decode can address.
    const uint32_t size = desc[blockIdx.x].nelem;         // bytes per plane
    if (threadIdx.x < 32) {
        const uint8_t *const ip = src + desc[blockIdx.x].in_off;
        const bool ok = lz4_decode_block_warp(ip, ip + desc[blockIdx.x].in_len,
                                              s_shuffled, s_shuffled + size * ES, threadIdx.x);
        if (threadIdx.x == 0 && !ok)
            atomicExch(status, 1u);
    }
    __syncthreads();

    untranspose_preprocess_block<T, ES>(s_shuffled, size, desc[blockIdx.x].out_off / ES, mask, out,
                                        sat_value, err_value, l);
    FlushStats<T>(l, stats);
}

// Largest bitshuffle block the fused kernel takes. A CUDA block holds one whole block in shared
// memory, so this figure is directly what limits residency: an SM's shared memory divided by it is
// how many blocks the SM can decode at once - four on a T4 at 16 kB - and past that there is too
// little parallelism left to cover even a shared-memory latency. Both writers this pipeline reads
// stay at or under it: our own compressor targets 16 kB of block whatever the pixel depth, and the
// EIGER files use 4096-element blocks, which is 16 kB at 32 bits and less below that. Anything
// larger takes the LZ4 kernel and the un-transposing preprocessor instead.
constexpr uint32_t FUSED_MAX_BLOCK_BYTES = 16384;

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_stats(1),
      cpu_stats(1),
      cpu_stats_reg(cpu_stats) {
    // Setup mask. The byte-per-pixel form and its checksum come from the PixelMask, which derives them
    // whenever the mask changes: they are the same for every worker, and deriving them walks every
    // pixel of the detector - 18 million of them on a 16 Mpx one, per engine, with an engine built per
    // worker per pass. The table is then uploaded once per GPU and shared by the engines on it.
    gpu_mask = SharedDeviceTable(mask.GetBinaryMask().data(), npixels, mask.GetBinaryMask().data(),
                                 mask.GetBinaryMaskChecksum(), *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;
    cuda_err(cudaGetDevice(&device));
    cudaDeviceProp prop{};
    cuda_err(cudaGetDeviceProperties(&prop, device));

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

float ImagePreprocessorGPU::GetLastDecompressionTime_s() const {
    return bslz4_decoder ? bslz4_decoder->GetDecodeTime_s() : 0.0f;
}

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)
        // Sized for the depth this image actually has, not for the widest one there could be: on
        // 16-bit data the difference is half of a full frame per worker thread.
        bslz4_decoder = std::make_unique<BSLZ4DecoderGPU>(image.GetUncompressedSize(), stream);

    // Everything from the compressed chunk to finished int32 pixels in ONE kernel, as long as a
    // bitshuffle block fits the shared-memory buffer it decodes into. Neither the bitshuffled bytes
    // nor the decompressed image is ever materialised. A block too large for that budget takes the
    // older route: the LZ4 kernel writes the shuffled image, and the un-transposing preprocessor
    // reads it back.
    const uint32_t block_bytes = BSLZ4DecoderGPU::BlockBytes(image);
    const BSLZ4ShuffledImage shuffled =
            (block_bytes > 0 && block_bytes <= FUSED_MAX_BLOCK_BYTES)
            ? bslz4_decoder->UploadCompressed(image)
            : bslz4_decoder->DecodeShuffled(image);

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

// The device-decode counterpart of AnalyzeOnDevice: same per-pixel decision, same statistics, but
// fed from the compressed chunk rather than from a decompressed image. Which kernel does it is what
// UploadCompressed() left behind - a shuffled image to read, or a chunk still to decode.
template<class T, int ES>
ImageStatistics ImagePreprocessorGPU::UntransposeAndAnalyze(ImagePreprocessorBuffer &processed_image,
                                                            const BSLZ4ShuffledImage &shuffled,
                                                            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};
    cuda_err(cudaMemcpyAsync(gpu_stats, cpu_stats.data(), sizeof(ImageStatistics), cudaMemcpyHostToDevice, *stream));

    // One CUDA block per bitshuffle block, plus one for the verbatim tail when there is one.
    const int nb = shuffled.nblocks + (shuffled.tail_elems > 0 ? 1 : 0);
    if (shuffled.shuffled) {
        untranspose_preprocess_kernel<T, ES> <<< nb, 256, 0, *stream >>>(
            shuffled.shuffled,
            shuffled.desc,
            gpu_mask->get(),
            processed_image.getGPUBuffer(),
            gpu_stats,
            sat_value,
            err_value,
            shuffled.nblocks,
            shuffled.tail_src,
            shuffled.tail_elems,
            shuffled.tail_elem0);
        cuda_err(cudaGetLastError());
    } else {
        // One warp per 2 kB of bitshuffle block. Shared memory is what caps how many CUDA blocks an
        // SM can hold, and a T4 has 2 kB of it per warp slot (64 kB against 32 warps), so this is
        // the shape that fills the SM instead of leaving warp slots no resident block can claim.
        const int threads = std::clamp<int>(shuffled.block_bytes / 64, 32, 1024);
        decode_untranspose_preprocess_kernel<T, ES> <<< nb, threads, shuffled.block_bytes, *stream >>>(
            shuffled.compressed,
            shuffled.desc,
            shuffled.status,
            gpu_mask->get(),
            processed_image.getGPUBuffer(),
            gpu_stats,
            sat_value,
            err_value,
            shuffled.nblocks,
            shuffled.tail_src,
            shuffled.tail_elems,
            shuffled.tail_elem0);
        cuda_err(cudaGetLastError());
        bslz4_decoder->QueueDecodeStatus();
    }

    if (copy_image_to_host)
        cuda_err(cudaMemcpyAsync(processed_image.data(), processed_image.getGPUBuffer(), npixels * sizeof(int32_t), cudaMemcpyDeviceToHost, *stream));
    cuda_err(cudaMemcpyAsync(cpu_stats.data(), gpu_stats, sizeof(ImageStatistics), cudaMemcpyDeviceToHost, *stream));

    // Check the synchronise: this path decodes bytes we did not produce, and cpu_stats still holds the
    // sentinel the host wrote above, so an unchecked failure here returns it as if it were a real
    // measurement and the caller's fallback to the host decoder never fires.
    cuda_err(cudaStreamSynchronize(*stream));

    // Only now can the device tell us whether every block actually decoded.
    bslz4_decoder->ThrowIfDecodeFailed();

    return cpu_stats[0];
}

template<class T>
ImageStatistics ImagePreprocessorGPU::Analyze(ImagePreprocessorBuffer &processed_image,
                                              const uint8_t *input,
                                              T err_value,
                                              T sat_value) {
    // Allocated here rather than in the constructor: only the host-upload path needs it, and the
    // device-decode path - which is what BSHUF_LZ4 images take - never touches it. Overshoot to
    // 4 bytes per pixel so a 1- or 2-byte image fits the same buffer.
    if (!gpu_decompressed_image.get())
        gpu_decompressed_image = CudaDevicePtr<uint8_t>(npixels * sizeof(uint32_t));

    // 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.
    cuda_err(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};
    cuda_err(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);
    cuda_err(cudaGetLastError());
    // 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)
        cuda_err(cudaMemcpyAsync(processed_image.data(), processed_image.getGPUBuffer(), npixels * sizeof(int32_t), cudaMemcpyDeviceToHost, *stream));
    cuda_err(cudaMemcpyAsync(cpu_stats.data(), gpu_stats, sizeof(ImageStatistics), cudaMemcpyDeviceToHost, *stream));

    cuda_err(cudaStreamSynchronize(*stream));

    return cpu_stats[0];
}