image_preprocessing: decode bitshuffle+LZ4 on the GPU

The pipeline decompressed each image on the host and uploaded the result. On
an 18 Mpx rotation dataset that made the host-to-device copy the bottleneck of
the whole per-image loop: nsys puts the copies at 78% of the loop against 39%
for every kernel combined - 3600 transfers of 72.4 MB - and they ran at only
12.5 GB/s of an available 27-28 because the host-side decompression was itself
saturating host memory bandwidth. The GPU was mostly waiting.

So the compressed chunk goes across instead, about 4 MB rather than 72 MB, and
is decoded on the device. That removes the transfer and the host decompression
that was throttling it, in one change. Measured on an idle machine, a run goes
from 45.11 s to 24.97 s - 1.81x - with the merged output unchanged.

THE APPROACH IS JON WRIGHT'S (ESRF): "Experiences with GPU decompression for
bitshuffle + LZ4 data", HDF5 User Group 2021, and github.com/jonwright/
bslz4decoders. The kernels here are ours, but the idea and the demonstration
that it is worth doing are his. Cited in docs/ACKNOWLEDGEMENT.md and in the new
section 0 of docs/CPU_DATA_ANALYSIS.md.

Two kernels mirror the CPU decoder. LZ4 runs one WARP per bitshuffle block:
every lane parses the same sequence stream (a broadcast read, no divergence)
and the literal and match copies are split across the 32 lanes so the stores
coalesce; an overlapping match is treated as a pattern of period offset sourced
from bytes that already precede the write position, which keeps it parallel
rather than a serial byte loop. One thread per block instead measured 13x
slower. The bitshuffle inverse then un-transposes each byte-plane through
shared memory and interleaves the planes back into elements.

Only BSHUF_LZ4 is decoded on the device. The zstd variants have no device
decoder, and neither has an uncompressed or float image; Supports() returns
false for those and the caller decompresses on the host exactly as before. The
fallback is explicit, so a format we cannot decode on the device is a slower
path and never a wrong answer.

Tests hold the device decoder against the CPU one byte for byte, on data from
the production compressor, for every element size the detectors emit -
including the 8-bit DECTRIS modes, which take bitshuf_decode_block's separate
elem_size == 1 branch - plus a many-block frame, the formats it must decline,
and malformed containers, which must throw rather than run off a buffer.

Battery: 37 crystals, no failures, identical to the host-decode run.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

image_analysis/image_preprocessing/BSLZ4DecoderGPU.cu

@@ -0,0 +1,227 @@
+// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
+// SPDX-License-Identifier: GPL-3.0-only
+
+#include <algorithm>
+
+#include "BSLZ4DecoderGPU.h"
+#include "../../common/JFJochException.h"
+#include "../../compression/JFJochDecompress.h"   // BSHUF_BLOCKED_MULT and the container layout
+
+namespace {
+    void cuda_err(cudaError_t val) {
+        if (val != cudaSuccess)
+            throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
+    }
+
+    // One WARP per LZ4 block. Every lane runs the same sequence parser over the same bytes - a
+    // broadcast read, so no divergence - and the literal and match copies are split across the 32
+    // lanes so the stores coalesce. One thread per block instead has each thread streaming its own
+    // 8 kB region, which coalesces not at all and measured 13x slower.
+    __global__ void lz4_decode_blocks(const uint8_t *__restrict__ src,
+                                      const BSLZ4BlockDesc *__restrict__ desc,
+                                      uint8_t *__restrict__ dst,
+                                      int nblocks, uint32_t elem_size) {
+        const int lane = threadIdx.x & 31;
+        const int b = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
+        if (b >= nblocks) return;
+
+        const uint8_t *ip = src + desc[b].in_off;
+        const uint8_t *const iend = ip + desc[b].in_len;
+        uint8_t *op = dst + desc[b].out_off;
+        uint8_t *const oend = op + desc[b].nelem * elem_size;
+
+        while (ip < iend) {
+            const uint32_t token = *ip++;
+            uint32_t litlen = token >> 4;
+            if (litlen == 15) {
+                uint32_t s;
+                do { s = *ip++; litlen += s; } while (s == 255 && ip < iend);
+            }
+            if (litlen) {
+                const uint32_t n = min(litlen, (uint32_t)(oend - op));
+                for (uint32_t i = lane; i < n; i += 32) op[i] = ip[i];
+                op += n; ip += litlen;
+            }
+            if (ip >= iend) break;   // last sequence carries literals only
+
+            const uint32_t offset = (uint32_t)ip[0] | ((uint32_t)ip[1] << 8);
+            ip += 2;
+            uint32_t matchlen = token & 0x0F;
+            if (matchlen == 15) {
+                uint32_t s;
+                do { s = *ip++; matchlen += s; } while (s == 255 && ip < iend);
+            }
+            matchlen += 4;   // minmatch
+
+            if (offset == 0 || offset > (uint32_t)(op - (dst + desc[b].out_off))) break;  // malformed
+            const uint8_t *mp = op - offset;
+            const uint32_t n = min(matchlen, (uint32_t)(oend - op));
+            if (offset >= matchlen) {
+                for (uint32_t i = lane; i < n; i += 32) op[i] = mp[i];
+            } else {
+                // An overlapping match is a pattern of period `offset`. mp[0..offset-1] all lie
+                // before op and are already final, so each output byte can be sourced from them
+                // independently - which keeps this parallel rather than a serial byte loop. Long
+                // zero runs in sparse detector data arrive here with offset == 1.
+                for (uint32_t i = lane; i < n; i += 32) op[i] = mp[i % offset];
+            }
+            op += n;
+        }
+    }
+
+    __device__ __forceinline__ uint64_t transpose8(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 bitshuffle inverse, one CUDA block per bitshuffle block, mirroring bitshuf_decode_block:
+    // un-transpose the bits of each byte-plane, then interleave the planes back into elements. The
+    // planes are staged in shared memory so the final store is coalesced.
+    extern __shared__ uint8_t smem[];
+    __global__ void bitshuffle_untranspose(const uint8_t *__restrict__ shuffled,
+                                           const BSLZ4BlockDesc *__restrict__ desc,
+                                           uint8_t *__restrict__ out,
+                                           int nblocks, uint32_t elem_size) {
+        const int b = blockIdx.x;
+        if (b >= nblocks) return;
+
+        const uint32_t size = desc[b].nelem;             // bytes per plane
+        const uint8_t *in = shuffled + desc[b].out_off;
+        uint8_t *dst = out + desc[b].out_off;
+        const uint32_t n = size / 8;
+
+        for (uint32_t p = 0; p < elem_size; p++) {
+            const uint8_t *pin = in + p * size;
+            uint8_t *pout = smem + p * size;
+            for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
+                uint64_t a = 0;
+                #pragma unroll
+                for (int k = 0; k < 8; k++) a |= (uint64_t)pin[k * n + i] << (8 * k);
+                const uint64_t x = transpose8(a);
+                #pragma unroll
+                for (int k = 0; k < 8; k++) pout[i * 8 + k] = (uint8_t)(x >> (8 * k));
+            }
+        }
+        __syncthreads();
+
+        for (uint32_t i = threadIdx.x; i < size; i += blockDim.x)
+            for (uint32_t j = 0; j < elem_size; j++)
+                dst[i * elem_size + j] = smem[j * size + i];
+    }
+
+    uint64_t be64(const uint8_t *p) { uint64_t v = 0; for (int i = 0; i < 8; i++) v = (v << 8) | p[i]; return v; }
+    uint32_t be32(const uint8_t *p) { return ((uint32_t)p[0] << 24) | ((uint32_t)p[1] << 16) | ((uint32_t)p[2] << 8) | p[3]; }
+
+    size_t elem_size_of(CompressedImageMode mode) {
+        switch (mode) {
+            // 8-bit is a real DECTRIS mode. bitshuf_decode_block takes a separate branch for
+            // elem_size == 1 (bit un-transpose only, no byte interleave), and the kernel here
+            // reproduces that for free: with one plane the interleave step degenerates to a copy.
+            case CompressedImageMode::Int8:
+            case CompressedImageMode::Uint8: return 1;
+            case CompressedImageMode::Int16:
+            case CompressedImageMode::Uint16: return 2;
+            case CompressedImageMode::Int32:
+            case CompressedImageMode::Uint32: return 4;
+            default: return 0;   // float and RGB modes are never bitshuffled by this pipeline
+        }
+    }
+}
+
+bool BSLZ4DecoderGPU::Supports(const CompressedImage &image) {
+    return image.GetCompressionAlgorithm() == CompressionAlgorithm::BSHUF_LZ4
+           && elem_size_of(image.GetMode()) != 0;
+}
+
+BSLZ4DecoderGPU::BSLZ4DecoderGPU(size_t in_max_uncompressed_bytes, std::shared_ptr<CudaStream> in_stream)
+    : stream(std::move(in_stream)),
+      max_uncompressed_bytes(in_max_uncompressed_bytes) {
+    // The compressed chunk is smaller than the image in every case worth having, but a pathological
+    // frame can expand slightly, so allow headroom rather than risk a per-frame reallocation.
+    max_compressed_bytes = max_uncompressed_bytes + max_uncompressed_bytes / 16 + 4096;
+    // Enough descriptors for the block size the compressor actually uses (8 kB); Decode() grows
+    // these if a stream turns up with smaller blocks. Sizing for the format's theoretical minimum
+    // block instead would allocate megabytes of pinned memory that no real file needs.
+    max_blocks = max_uncompressed_bytes / 8192 + 2;
+
+    gpu_compressed = CudaDevicePtr<uint8_t>(max_compressed_bytes);
+    gpu_shuffled = CudaDevicePtr<uint8_t>(max_uncompressed_bytes);
+    gpu_desc = CudaDevicePtr<BSLZ4BlockDesc>(max_blocks);
+    host_desc = CudaHostPtr<BSLZ4BlockDesc>(max_blocks);
+}
+
+void BSLZ4DecoderGPU::Decode(const CompressedImage &image, uint8_t *gpu_out) {
+    const uint8_t *src = image.GetCompressed();
+    const size_t clen = image.GetCompressedSize();
+    const size_t elem_size = elem_size_of(image.GetMode());
+    const size_t total_bytes = image.GetUncompressedSize();
+
+    if (clen < 12)
+        throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 chunk shorter than its header");
+    if (total_bytes > max_uncompressed_bytes || clen > max_compressed_bytes)
+        throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 image larger than the decoder was sized for");
+    if (be64(src) != total_bytes)
+        throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 header size does not match the image");
+
+    const uint32_t block_bytes = be32(src + 8);
+    if (block_bytes == 0 || block_bytes % elem_size != 0)
+        throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 block size invalid");
+
+    const size_t block_elems = block_bytes / elem_size;
+    const size_t nelements = total_bytes / elem_size;
+    const size_t nfull = nelements / block_elems;
+    const size_t rem = nelements - nfull * block_elems;
+    const size_t last = rem - rem % BSHUF_BLOCKED_MULT;
+    const size_t leftover_bytes = (rem % BSHUF_BLOCKED_MULT) * elem_size;
+
+    // Walk the container to locate the blocks. Lengths are only knowable in order, so this scan is
+    // inherent to the format rather than an implementation choice.
+    const size_t nblocks_needed = nfull + (last > 0 ? 1 : 0);
+    if (nblocks_needed > max_blocks) {   // a stream with smaller blocks than the compressor emits
+        max_blocks = nblocks_needed;
+        gpu_desc = CudaDevicePtr<BSLZ4BlockDesc>(max_blocks);
+        host_desc = CudaHostPtr<BSLZ4BlockDesc>(max_blocks);
+    }
+
+    size_t nblk = 0, off = 12, out_off = 0;
+    for (size_t i = 0; i < nfull + (last > 0 ? 1 : 0); i++) {
+        if (off + 4 > clen)
+            throw JFJochException(JFJochExceptionCategory::Compression, "truncated bslz4 block header");
+        const uint32_t block_clen = be32(src + off);
+        off += 4;
+        if (block_clen == 0 || off + block_clen > clen)
+            throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 block extends past the chunk");
+        if (nblk >= max_blocks)
+            throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 chunk has more blocks than expected");
+        const uint32_t ne = (i < nfull) ? static_cast<uint32_t>(block_elems) : static_cast<uint32_t>(last);
+        host_desc.get()[nblk++] = {static_cast<uint32_t>(off), block_clen, static_cast<uint32_t>(out_off), ne};
+        off += block_clen;
+        out_off += static_cast<size_t>(ne) * elem_size;
+    }
+    if (nblk == 0)
+        throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 chunk contains no blocks");
+
+    cuda_err(cudaMemcpyAsync(gpu_compressed.get(), src, clen, cudaMemcpyHostToDevice, *stream));
+    cuda_err(cudaMemcpyAsync(gpu_desc.get(), host_desc.get(), nblk * sizeof(BSLZ4BlockDesc),
+                             cudaMemcpyHostToDevice, *stream));
+
+    const int nb = static_cast<int>(nblk);
+    lz4_decode_blocks<<<(nb * 32 + 255) / 256, 256, 0, *stream>>>(
+            gpu_compressed.get(), gpu_desc.get(),
+            gpu_shuffled.get(), nb, static_cast<uint32_t>(elem_size));
+    bitshuffle_untranspose<<<nb, 256, block_bytes, *stream>>>(
+            gpu_shuffled.get(), gpu_desc.get(),
+            gpu_out, nb, static_cast<uint32_t>(elem_size));
+    cuda_err(cudaGetLastError());
+
+    // The tail that bitshuffle leaves uncompressed and copies verbatim.
+    if (leftover_bytes > 0) {
+        if (off + leftover_bytes > clen)
+            throw JFJochException(JFJochExceptionCategory::Compression, "truncated bslz4 leftover bytes");
+        cuda_err(cudaMemcpyAsync(gpu_out + out_off, src + off, leftover_bytes,
+                                 cudaMemcpyHostToDevice, *stream));
+    }
+}