Jungfraujoch/image_analysis/image_preprocessing/BSLZ4DecoderGPU.h

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

#pragma once

#include <memory>

#include "../../common/CompressedImage.h"
#include "../indexing/CUDAMemHelpers.h"

// One bitshuffle block, located by the host scan and consumed by both kernels.
struct BSLZ4BlockDesc {
    uint32_t in_off;    // byte offset of the LZ4 payload within the chunk
    uint32_t in_len;    // compressed length
    uint32_t out_off;   // byte offset of this block's output in the image
    uint32_t nelem;     // elements in this block (the last one is usually shorter)
};

// One chunk, on the device, with everything the rest of the decode needs to finish the image. The
// tail is the handful of elements bitshuffle stores verbatim; it is already on the device inside the
// uploaded chunk, so it is handed over as a device pointer rather than copied again from the host.
//
// DecodeShuffled() fills in `shuffled` - the LZ4 output, still bitshuffled. UploadCompressed()
// leaves it null and hands over `compressed` and `status` instead, so the caller can run the LZ4
// pass itself.
struct BSLZ4ShuffledImage {
    const uint8_t *shuffled = nullptr;
    const uint8_t *compressed = nullptr;   // the uploaded chunk; desc[].in_off indexes into it
    const BSLZ4BlockDesc *desc = nullptr;
    uint32_t *status = nullptr;            // where a caller-run LZ4 pass flags a bad block
    int nblocks = 0;
    uint32_t elem_size = 0;
    uint32_t block_bytes = 0;              // uncompressed bytes in a full bitshuffle block
    const uint8_t *tail_src = nullptr;
    uint32_t tail_elems = 0;
    uint32_t tail_elem0 = 0;    // index of the first tail element in the image
};

// Decompress a bitshuffle+LZ4 image ON THE DEVICE, so the compressed bytes are what crosses PCIe.
//
// The idea - upload the compressed chunk and decode it on the GPU rather than decompressing on the
// host - is Jon Wright's (ESRF); see https://github.com/jonwright/bslz4decoders and his 2021 HDF5
// User Group talk "Experiences with GPU decompression for bitshuffle + LZ4 data". The kernels here
// are our own, but the approach, and the observation that it is worth doing at all, are his.
//
// Why it pays: a full 18 Mpx uint32 frame is 72 MB decompressed and about 4 MB compressed, and
// profiling showed the host-to-device copy owning ~78% of the per-image loop against ~39% for
// kernels. Decoding on the device removes both that transfer and the host-side decompression,
// whose memory traffic was itself holding the copy engine well below the link rate.
//
// Only BSHUF_LZ4 is handled. The zstd variants have no device decoder, so Supports() returns false
// and the caller decompresses on the host exactly as before.
//
// The container comes off the network or off disk, so it is not trusted. Everything the host can
// check cheaply is checked before any work is queued and throws; what only the kernel can see - a
// block that does not decode to its declared length, which is what a corrupt LZ4 payload looks like
// - raises a device-side flag that ThrowIfDecodeFailed() reports once the caller has synchronised.
// The CPU decoder makes exactly the same checks (LZ4_decompress_safe's length check plus the
// consumed-input check in JFJochDecompress.h), so a chunk either decodes identically on both or
// fails on both. It is never silently wrong on one and right on the other.
class BSLZ4DecoderGPU {
    std::shared_ptr<CudaStream> stream;

    CudaDevicePtr<uint8_t> gpu_compressed;
    CudaDevicePtr<uint8_t> gpu_shuffled;      // LZ4 output, still bitshuffled
    CudaDevicePtr<BSLZ4BlockDesc> gpu_desc;
    CudaHostPtr<BSLZ4BlockDesc> host_desc;    // pinned, so the descriptor upload is truly async
    CudaDevicePtr<uint32_t> gpu_status;       // set by the kernel when a block decodes short
    CudaHostPtr<uint32_t> host_status;

    // Bracket the decode so the time it takes can still be reported as decompression, which is what
    // it is. Both are recorded on the decoder's stream and read after the caller synchronises.
    CudaEvent decode_start;
    CudaEvent decode_stop;
    bool decode_timed = false;

    size_t max_uncompressed_bytes = 0;
    size_t compressed_capacity = 0;
    size_t max_blocks = 0;

    void EnsureCompressedCapacity(size_t bytes);
    void EnsureUncompressedCapacity(size_t bytes);
    void EnsureBlockCapacity(size_t nblocks);

    // Scan the container and upload it, stopping short of decoding the blocks.
    BSLZ4ShuffledImage PrepareChunk(const CompressedImage &image);

public:
    BSLZ4DecoderGPU(size_t max_uncompressed_bytes, std::shared_ptr<CudaStream> stream);
    // Waits for the stream: its pooled buffers are freed next (see CudaDevicePtr).
    ~BSLZ4DecoderGPU() { if (stream) cudaStreamSynchronize(*stream); }

    // True when this image can be decoded on the device. Everything else must go the host route.
    static bool Supports(const CompressedImage &image);

    // The uncompressed size of one bitshuffle block, straight out of the chunk header, so a caller
    // that wants to decode the blocks in shared memory can size that memory before it commits to
    // the route. Zero when the chunk is too short to hold a header, which the decode then reports.
    static uint32_t BlockBytes(const CompressedImage &image);

    // Locate the blocks, upload the chunk, and run the LZ4 pass. The result is still bitshuffled -
    // the caller finishes it, either with Decode()'s un-transpose or by fusing the un-transpose into
    // its own kernel. Work is queued on the decoder's stream and the caller synchronises.
    BSLZ4ShuffledImage DecodeShuffled(const CompressedImage &image);

    // Upload the chunk and locate its blocks, and stop there. For a caller that runs the LZ4 pass
    // in its OWN kernel, decoding each block into shared memory and consuming it there, so the
    // bitshuffled bytes never reach device memory. Such a caller must call QueueDecodeStatus()
    // once that kernel is queued.
    BSLZ4ShuffledImage UploadCompressed(const CompressedImage &image);

    // Queue the device-side failure flag back to the host. DecodeShuffled() does this itself; a
    // caller that decodes the blocks in its own kernel does it after queueing that kernel, or the
    // flag ThrowIfDecodeFailed() reads is the one from before the decode.
    void QueueDecodeStatus();

    // Decode into gpu_out, which must hold image.GetUncompressedSize() bytes. The plain raw-bytes
    // path: DecodeShuffled() plus the un-transpose. Used by the tests and by any caller that wants
    // the decompressed image rather than a preprocessed one.
    void Decode(const CompressedImage &image, uint8_t *gpu_out);

    // Report a block that did not decode to its declared length. MUST be called after the caller has
    // synchronised the stream; until then the flag has not arrived. Throws on failure.
    void ThrowIfDecodeFailed();

    // Device time spent decoding the last image, in seconds. Valid after the caller synchronises.
    [[nodiscard]] float GetDecodeTime_s() const;
};