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

// 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.
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

    size_t max_compressed_bytes = 0;
    size_t max_uncompressed_bytes = 0;
    size_t max_blocks = 0;

public:
    BSLZ4DecoderGPU(size_t max_uncompressed_bytes, std::shared_ptr<CudaStream> stream);

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

    // Decode into gpu_out, which must hold image.GetUncompressedSize() bytes. Work is queued on the
    // decoder's stream and the caller synchronises. Throws if the container is malformed - it comes
    // off the network or off disk, so it is not trusted.
    void Decode(const CompressedImage &image, uint8_t *gpu_out);
};