diff --git a/docs/ACKNOWLEDGEMENT.md b/docs/ACKNOWLEDGEMENT.md index 56807b32..733552b5 100644 --- a/docs/ACKNOWLEDGEMENT.md +++ b/docs/ACKNOWLEDGEMENT.md @@ -7,4 +7,9 @@ The project is supported by : * ETH Domain via Open Research Data Contribute project (Jan - Dec 2023) * AMD University Program with donation of licenses of Ethernet IP cores and Vivado software +Decoding bitshuffle+LZ4 images on the GPU, rather than decompressing them on the host and uploading +the result, follows Jon Wright (ESRF): "Experiences with GPU decompression for bitshuffle + LZ4 +data", HDF5 User Group meeting (2021), and [bslz4decoders](https://github.com/jonwright/bslz4decoders). +The CUDA kernels in Jungfraujoch are its own, but the approach is his. + This software uses Viridis, Magma and Inferno colormaps from Matplotlib under its BSD-compatible license diff --git a/docs/CPU_DATA_ANALYSIS.md b/docs/CPU_DATA_ANALYSIS.md index a4c11bbc..50651916 100644 --- a/docs/CPU_DATA_ANALYSIS.md +++ b/docs/CPU_DATA_ANALYSIS.md @@ -30,8 +30,54 @@ The methods are inspired and reuising solutions implemented in: - S. French & K. Wilson, "On the treatment of negative intensity observations", *Acta Cryst.* **A34** (1978), 517-525 (Bayesian amplitude estimation from intensities). - A. T. Brünger, "Free R value: a novel statistical quantity for assessing the accuracy of crystal structures", *Nature* **355** (1992), 472-475 (R-free cross-validation). - M. Wojdyr, "GEMMI: A library for structural biology", *J. Open Source Softw.* **7** (2022), 4200 (model / structure-factor / map machinery used in §14). +- J. P. Wright, "Experiences with GPU decompression for bitshuffle + LZ4 data", HDF5 User Group meeting (2021), and [github.com/jonwright/bslz4decoders](https://github.com/jonwright/bslz4decoders) — decoding bitshuffle+LZ4 on the GPU instead of the host, which is the idea behind §0. (list is not exhaustive) +## 0. Getting the image onto the GPU: device-side bitshuffle+LZ4 decoding + +Images arrive bitshuffle+LZ4 compressed (HDF5 filter 32008), and everything from §1 onwards runs on +the GPU when one is present. The obvious arrangement — decompress on the host, upload the image — +turns out to be the most expensive part of the whole per-image loop. Profiling an 18 Mpx rotation +dataset showed the host-to-device copy occupying **78% of the loop** against **39% for all kernels +combined**: 3600 copies of 72.4 MB each, and running at only 12.5 GB/s of an available 27–28 +because the host-side decompression was itself saturating host memory bandwidth. + +So the compressed chunk is uploaded instead — about 4 MB rather than 72 MB, an 18× reduction — and +decoded on the device. This removes the transfer *and* the host decompression that was throttling +it. Measured on that dataset the run went from 49.4 s to 28.0 s with byte-identical output. + +**This approach is Jon Wright's (ESRF)** — see his HDF5 User Group talk "Experiences with GPU +decompression for bitshuffle + LZ4 data" (2021) and +[github.com/jonwright/bslz4decoders](https://github.com/jonwright/bslz4decoders). The kernels in +`image_analysis/image_preprocessing/BSLZ4DecoderGPU.cu` are our own, but the idea, and the +demonstration that it is worth doing at all, are his. + +Two kernels do the work, mirroring the CPU decoder exactly: + +1. **LZ4, one warp per bitshuffle block.** Blocks are independent, so the parallelism is across + them; within a warp every lane runs the same sequence parser over the same bytes (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 — the long zero runs of sparse detector data arrive here with `offset == 1`. One thread per + block instead, with each thread streaming its own 8 kB region, measured 13× slower. +2. **The bitshuffle inverse**, one CUDA block per bitshuffle block: bit un-transpose of each + byte-plane into shared memory, then interleave the planes back into elements so the final store + is coalesced. For 8-bit images (a real DECTRIS mode) there is a single plane and the interleave + degenerates to a copy, matching `bitshuf_decode_block`'s `elem_size == 1` branch. + +The block offsets inside the container can only be discovered by reading the block lengths in +order, so that scan stays on the host; it is a few hundred microseconds per frame against the +milliseconds it saves. + +**Only BSHUF_LZ4 is decoded on the device.** The zstd variants (`BSHUF_ZSTD`, `BSHUF_ZSTD_RLE`, +`BSHUF_ZSTD_RLE_HUFF`) have no device decoder, and neither does an uncompressed or float image; +for those `BSLZ4DecoderGPU::Supports()` returns false and the pipeline decompresses on the host and +uploads as before. The fallback is explicit rather than implicit, so an algorithm we cannot decode +on the device is a slower path and never a wrong answer. `tests/BSLZ4DecoderGPUTest.cpp` holds the +device decoder against the CPU one byte for byte, on data produced by the production compressor, +for every element size the detectors emit. + ## 1. Geometry, reciprocal-space mapping, and basic quantities ### 1.1 Coordinate conventions diff --git a/image_analysis/MXAnalysisWithoutFPGA.cpp b/image_analysis/MXAnalysisWithoutFPGA.cpp index e300a187..a21626dc 100644 --- a/image_analysis/MXAnalysisWithoutFPGA.cpp +++ b/image_analysis/MXAnalysisWithoutFPGA.cpp @@ -90,16 +90,30 @@ void MXAnalysisWithoutFPGA::Analyze(DataMessage &output, throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "Mismatch in pixel size"); + // Decompress on the device where the preprocessor can, so only the compressed chunk crosses PCIe + // and the host does no decompression at all. AnalyzeCompressed says whether it took the image; + // when it declines (a CPU preprocessor, or an algorithm with no device decoder) fall through to + // the host route unchanged. The two produce the same preprocessed image. const auto compression_start_time = std::chrono::steady_clock::now(); - const uint8_t *image_ptr = Decompress(output.image); + ImageStatistics ret{}; + const bool decoded_on_device = preprocessor->AnalyzeCompressed(*preprocessor_buffer, output.image, ret); const auto compression_end_time = std::chrono::steady_clock::now(); - if (output.image.GetCompressionAlgorithm() != CompressionAlgorithm::NO_COMPRESSION) - output.compression_time_s = std::chrono::duration(compression_end_time - compression_start_time).count(); - const auto preprocessing_start_time = std::chrono::steady_clock::now(); - auto ret = preprocessor->Analyze(*preprocessor_buffer, image_ptr, output.image.GetMode()); - const auto preprocessing_end_time = std::chrono::steady_clock::now(); - output.preprocessing_time_s = std::chrono::duration(preprocessing_end_time - preprocessing_start_time).count(); + if (!decoded_on_device) { + const uint8_t *image_ptr = Decompress(output.image); + const auto decompressed_time = std::chrono::steady_clock::now(); + if (output.image.GetCompressionAlgorithm() != CompressionAlgorithm::NO_COMPRESSION) + output.compression_time_s = std::chrono::duration(decompressed_time - compression_start_time).count(); + + const auto preprocessing_start_time = std::chrono::steady_clock::now(); + ret = preprocessor->Analyze(*preprocessor_buffer, image_ptr, output.image.GetMode()); + const auto preprocessing_end_time = std::chrono::steady_clock::now(); + output.preprocessing_time_s = std::chrono::duration(preprocessing_end_time - preprocessing_start_time).count(); + } else { + // Decode and preprocess are one device operation here, so they are reported together rather + // than split into a decompression time that no longer exists on the host. + output.preprocessing_time_s = std::chrono::duration(compression_end_time - compression_start_time).count(); + } // The fused GPU engine (rugnux offline, GPU, adaptive detection) produces the azimuthal profile as // a byproduct of spot finding, so the separate azint pass is skipped in that case and the profile is diff --git a/image_analysis/image_preprocessing/BSLZ4DecoderGPU.cu b/image_analysis/image_preprocessing/BSLZ4DecoderGPU.cu new file mode 100644 index 00000000..829d76c0 --- /dev/null +++ b/image_analysis/image_preprocessing/BSLZ4DecoderGPU.cu @@ -0,0 +1,227 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include + +#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 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(max_compressed_bytes); + gpu_shuffled = CudaDevicePtr(max_uncompressed_bytes); + gpu_desc = CudaDevicePtr(max_blocks); + host_desc = CudaHostPtr(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(max_blocks); + host_desc = CudaHostPtr(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(block_elems) : static_cast(last); + host_desc.get()[nblk++] = {static_cast(off), block_clen, static_cast(out_off), ne}; + off += block_clen; + out_off += static_cast(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(nblk); + lz4_decode_blocks<<<(nb * 32 + 255) / 256, 256, 0, *stream>>>( + gpu_compressed.get(), gpu_desc.get(), + gpu_shuffled.get(), nb, static_cast(elem_size)); + bitshuffle_untranspose<<>>( + gpu_shuffled.get(), gpu_desc.get(), + gpu_out, nb, static_cast(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)); + } +} diff --git a/image_analysis/image_preprocessing/BSLZ4DecoderGPU.h b/image_analysis/image_preprocessing/BSLZ4DecoderGPU.h new file mode 100644 index 00000000..d93b4008 --- /dev/null +++ b/image_analysis/image_preprocessing/BSLZ4DecoderGPU.h @@ -0,0 +1,55 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include + +#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 stream; + + CudaDevicePtr gpu_compressed; + CudaDevicePtr gpu_shuffled; // LZ4 output, still bitshuffled + CudaDevicePtr gpu_desc; + CudaHostPtr 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 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); +}; diff --git a/image_analysis/image_preprocessing/CMakeLists.txt b/image_analysis/image_preprocessing/CMakeLists.txt index 4c5a5b25..5ed65c05 100644 --- a/image_analysis/image_preprocessing/CMakeLists.txt +++ b/image_analysis/image_preprocessing/CMakeLists.txt @@ -10,6 +10,7 @@ IF (JFJOCH_CUDA_AVAILABLE) TARGET_SOURCES(JFJochImagePreprocessing PRIVATE ../indexing/CUDAMemHelpers.h ImagePreprocessorGPU.cu ImagePreprocessorGPU.h + BSLZ4DecoderGPU.cu BSLZ4DecoderGPU.h ImagePreprocessorBufferGPU.cu ImagePreprocessorBufferGPU.h) ENDIF() \ No newline at end of file diff --git a/image_analysis/image_preprocessing/ImagePreprocessor.h b/image_analysis/image_preprocessing/ImagePreprocessor.h index 3dd8af47..de30f73b 100644 --- a/image_analysis/image_preprocessing/ImagePreprocessor.h +++ b/image_analysis/image_preprocessing/ImagePreprocessor.h @@ -29,6 +29,16 @@ public: virtual ~ImagePreprocessor() = default; virtual ImageStatistics Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *decompressed_image, CompressedImageMode image_mode) = 0; + // Analyze straight from the COMPRESSED image, decompressing wherever the implementation prefers. + // The GPU implementation uploads the compressed chunk and decodes it on the device, so only a few + // MB cross PCIe instead of the whole frame and the host never decompresses at all. + // Returns false when this implementation cannot handle the image - the CPU preprocessor always, + // and the GPU one for any algorithm without a device decoder - and the caller then decompresses + // on the host and calls Analyze() as before. Keeping the fallback explicit means a format we + // cannot decode on the device is a slower path, never a wrong answer. + virtual bool AnalyzeCompressed(ImagePreprocessorBuffer &processed_image, const CompressedImage &image, + ImageStatistics &stats) { return false; } + // Resize the buffer an image will be decompressed into and page-lock it, so that the host->device // copy of Analyze() is a real DMA. Without page-locking the driver stages the copy through its own // pinned pool, which is a host-side copy on the calling thread: it does not overlap and it degrades diff --git a/image_analysis/image_preprocessing/ImagePreprocessorGPU.cu b/image_analysis/image_preprocessing/ImagePreprocessorGPU.cu index b8f74b84..31b25408 100644 --- a/image_analysis/image_preprocessing/ImagePreprocessorGPU.cu +++ b/image_analysis/image_preprocessing/ImagePreprocessorGPU.cu @@ -145,18 +145,58 @@ ImageStatistics ImagePreprocessorGPU::Analyze(ImagePreprocessorBuffer &processed } } +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) + bslz4_decoder = std::make_unique(npixels * sizeof(uint32_t), stream); + + // Straight into the buffer Analyze() would have filled by copying the decompressed frame across + // the bus; from here the two paths are the same code. + bslz4_decoder->Decode(image, gpu_decompressed_image.get()); + + switch (image.GetMode()) { + case CompressedImageMode::Int8: + stats = AnalyzeOnDevice(processed_image, INT8_MIN, INT8_MAX); return true; + case CompressedImageMode::Uint8: + stats = AnalyzeOnDevice(processed_image, UINT8_MAX, UINT8_MAX); return true; + case CompressedImageMode::Int16: + stats = AnalyzeOnDevice(processed_image, INT16_MIN, INT16_MAX); return true; + case CompressedImageMode::Uint16: + stats = AnalyzeOnDevice(processed_image, UINT16_MAX, UINT16_MAX); return true; + case CompressedImageMode::Int32: + stats = AnalyzeOnDevice(processed_image, INT32_MIN, INT32_MAX); return true; + case CompressedImageMode::Uint32: + stats = AnalyzeOnDevice(processed_image, UINT32_MAX, UINT32_MAX); return true; + default: + return false; // Supports() already excludes these; belt and braces + } +} + template ImageStatistics ImagePreprocessorGPU::Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *input, T err_value, T sat_value) { - if (sat_value > saturation_limit) - sat_value = static_cast(saturation_limit); - // 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. cudaMemcpyAsync(gpu_decompressed_image, input, npixels * sizeof(T), cudaMemcpyHostToDevice, *stream); + return AnalyzeOnDevice(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 +ImageStatistics ImagePreprocessorGPU::AnalyzeOnDevice(ImagePreprocessorBuffer &processed_image, + T err_value, T sat_value) { + if (sat_value > saturation_limit) + sat_value = static_cast(saturation_limit); cpu_stats[0] = ImageStatistics{.max_value = INT64_MIN, .min_value = INT64_MAX}; cudaMemcpyAsync(gpu_stats, cpu_stats.data(), sizeof(ImageStatistics), cudaMemcpyHostToDevice, *stream); diff --git a/image_analysis/image_preprocessing/ImagePreprocessorGPU.h b/image_analysis/image_preprocessing/ImagePreprocessorGPU.h index 9dbfb646..ec07458d 100644 --- a/image_analysis/image_preprocessing/ImagePreprocessorGPU.h +++ b/image_analysis/image_preprocessing/ImagePreprocessorGPU.h @@ -3,7 +3,10 @@ #pragma once +#include + #include "ImagePreprocessor.h" +#include "BSLZ4DecoderGPU.h" #include "../indexing/CUDAMemHelpers.h" #include "../indexing/CudaSharedTables.h" @@ -23,7 +26,12 @@ class ImagePreprocessorGPU : public ImagePreprocessor { std::vector cpu_image; + // Built on first use: a decoder that can serve this engine's images, sized to the frame. + std::unique_ptr bslz4_decoder; + template ImageStatistics Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *input, T err_value, T sat_value); + // Preprocess an image already sitting in gpu_decompressed_image, shared by both entry points. + template ImageStatistics AnalyzeOnDevice(ImagePreprocessorBuffer &processed_image, T err_value, T sat_value); public: // copy_image_to_host copies the preprocessed image back after every frame. It is only needed when // something on the CPU reads it - the GPU engines all work off the device buffer - and at 4 bytes @@ -31,6 +39,8 @@ public: ImagePreprocessorGPU(const DiffractionExperiment &experiment, const PixelMask &mask, std::shared_ptr stream, bool copy_image_to_host = true); ImageStatistics Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *decompressed_image, CompressedImageMode image_mode) override; + bool AnalyzeCompressed(ImagePreprocessorBuffer &processed_image, const CompressedImage &image, + ImageStatistics &stats) override; void PinInputBuffer(std::vector &buffer, size_t size) override; }; diff --git a/tests/BSLZ4DecoderGPUTest.cpp b/tests/BSLZ4DecoderGPUTest.cpp new file mode 100644 index 00000000..a02dda5b --- /dev/null +++ b/tests/BSLZ4DecoderGPUTest.cpp @@ -0,0 +1,159 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include +#include "../common/CUDAWrapper.h" + +#ifdef JFJOCH_USE_CUDA + +#include +#include +#include + +#include "../image_analysis/image_preprocessing/BSLZ4DecoderGPU.h" +#include "../compression/JFJochCompressor.h" +#include "../compression/JFJochDecompress.h" + +// The GPU decoder must agree with the CPU one BYTE FOR BYTE, on data produced by our own compressor, +// for every element size the detectors emit - including the 8-bit DECTRIS modes, which take a +// different branch in bitshuf_decode_block (bit un-transpose only, no byte interleave). +// +// Images are built to exercise what the LZ4 format actually does on detector data: long runs of a +// repeated byte (offset == 1 matches, the overlapping-match path), isolated bright pixels +// (literals), and a noisy region (short matches at assorted offsets). A uniformly random image +// would be almost incompressible and would never reach the match code at all. +namespace { + +template +std::vector MakeDetectorLikeImage(size_t npixels, uint32_t seed) { + std::mt19937 rng(seed); + std::vector img(npixels, 0); // sparse background: long zero runs + + // A band of low-level noise, so matches are short and offsets vary. + for (size_t i = npixels / 4; i < npixels / 2; i++) + img[i] = static_cast(rng() % 7); + + // Bright, isolated spots - these become literals. + for (size_t s = 0; s < 64; s++) { + const size_t c = rng() % npixels; + for (size_t d = 0; d < 9 && c + d < npixels; d++) + img[c + d] = static_cast(std::numeric_limits::max() / (2 + (d % 3))); + } + // A run of one repeated non-zero value, the classic offset==1 match. + for (size_t i = npixels * 3 / 4; i < npixels * 3 / 4 + 5000 && i < npixels; i++) + img[i] = static_cast(42); + return img; +} + +template +void RoundTrip(CompressedImageMode mode, size_t width, size_t height, uint32_t seed) { + const size_t npixels = width * height; + const auto original = MakeDetectorLikeImage(npixels, seed); + + // Compress with the production compressor, so the container is exactly what the pipeline reads. + JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_LZ4); + const std::vector compressed = compressor.Compress(original); + REQUIRE(!compressed.empty()); + + const CompressedImage image(compressed.data(), compressed.size(), width, height, mode, + CompressionAlgorithm::BSHUF_LZ4); + REQUIRE(BSLZ4DecoderGPU::Supports(image)); + REQUIRE(image.GetUncompressedSize() == npixels * sizeof(T)); + + // CPU reference: the same call the host path makes. + std::vector cpu_buffer; + const uint8_t *cpu_out = image.GetUncompressedPtr(cpu_buffer); + REQUIRE(std::memcmp(cpu_out, original.data(), npixels * sizeof(T)) == 0); + + auto stream = std::make_shared(); + BSLZ4DecoderGPU decoder(npixels * sizeof(uint32_t), stream); + CudaDevicePtr gpu_out(npixels * sizeof(T)); + decoder.Decode(image, gpu_out.get()); + REQUIRE(cudaStreamSynchronize(*stream) == cudaSuccess); + + std::vector gpu_result(npixels); + REQUIRE(cudaMemcpy(gpu_result.data(), gpu_out.get(), npixels * sizeof(T), + cudaMemcpyDeviceToHost) == cudaSuccess); + + REQUIRE(std::memcmp(gpu_result.data(), original.data(), npixels * sizeof(T)) == 0); +} + +} // namespace + +TEST_CASE("BSLZ4DecoderGPU_MatchesCPU_AllElementSizes", "[BSLZ4DecoderGPU]") { + if (get_gpu_count() == 0) + SKIP("No CUDA GPU present"); + + // Sizes chosen so the last block is partial and the leftover tail (the elements bitshuffle + // leaves uncompressed because they do not fill a multiple of 8) is non-empty on some of them. + RoundTrip(CompressedImageMode::Uint8, 1030, 517, 1); + RoundTrip(CompressedImageMode::Int8, 1030, 517, 2); + RoundTrip(CompressedImageMode::Uint16, 1030, 517, 3); + RoundTrip(CompressedImageMode::Int16, 1030, 517, 4); + RoundTrip(CompressedImageMode::Uint32, 1030, 517, 5); + RoundTrip(CompressedImageMode::Int32, 1030, 517, 6); +} + +TEST_CASE("BSLZ4DecoderGPU_MatchesCPU_LargeFrame", "[BSLZ4DecoderGPU]") { + if (get_gpu_count() == 0) + SKIP("No CUDA GPU present"); + + // Many blocks, so the per-block descriptor scan and the one-warp-per-block launch are exercised + // at a realistic scale rather than on a handful of blocks. + RoundTrip(CompressedImageMode::Uint32, 2068, 2162, 7); +} + +// A decoder that cannot handle an image must SAY so rather than produce something wrong: the caller +// relies on Supports() to decide whether the host route is needed. +TEST_CASE("BSLZ4DecoderGPU_DeclinesWhatItCannotDecode", "[BSLZ4DecoderGPU]") { + std::vector dummy(1024, 0); + const size_t w = 16, h = 16; + + CHECK_FALSE(BSLZ4DecoderGPU::Supports( + CompressedImage(dummy.data(), dummy.size(), w, h, CompressedImageMode::Uint32, + CompressionAlgorithm::BSHUF_ZSTD))); + CHECK_FALSE(BSLZ4DecoderGPU::Supports( + CompressedImage(dummy.data(), dummy.size(), w, h, CompressedImageMode::Uint32, + CompressionAlgorithm::BSHUF_ZSTD_RLE))); + CHECK_FALSE(BSLZ4DecoderGPU::Supports( + CompressedImage(dummy.data(), dummy.size(), w, h, CompressedImageMode::Uint32, + CompressionAlgorithm::BSHUF_ZSTD_RLE_HUFF))); + CHECK_FALSE(BSLZ4DecoderGPU::Supports( + CompressedImage(dummy.data(), dummy.size(), w, h, CompressedImageMode::Uint32, + CompressionAlgorithm::NO_COMPRESSION))); + CHECK_FALSE(BSLZ4DecoderGPU::Supports( + CompressedImage(dummy.data(), dummy.size(), w, h, CompressedImageMode::Float32, + CompressionAlgorithm::BSHUF_LZ4))); + + CHECK(BSLZ4DecoderGPU::Supports( + CompressedImage(dummy.data(), dummy.size(), w, h, CompressedImageMode::Uint32, + CompressionAlgorithm::BSHUF_LZ4))); +} + +// A malformed container comes off the network or off disk, so it must throw rather than run off +// the end of a buffer on the device. +TEST_CASE("BSLZ4DecoderGPU_RejectsMalformed", "[BSLZ4DecoderGPU]") { + if (get_gpu_count() == 0) + SKIP("No CUDA GPU present"); + + const size_t width = 128, height = 128, npixels = width * height; + const auto original = MakeDetectorLikeImage(npixels, 11); + JFJochBitShuffleCompressor compressor(CompressionAlgorithm::BSHUF_LZ4); + const std::vector compressed = compressor.Compress(original); + + auto stream = std::make_shared(); + BSLZ4DecoderGPU decoder(npixels * sizeof(uint32_t), stream); + CudaDevicePtr gpu_out(npixels * sizeof(uint32_t)); + + // Truncated mid-stream: the block header promises more than is there. + const CompressedImage truncated(compressed.data(), compressed.size() / 2, width, height, + CompressedImageMode::Uint32, CompressionAlgorithm::BSHUF_LZ4); + CHECK_THROWS(decoder.Decode(truncated, gpu_out.get())); + + // Shorter than the 12-byte container header. + const CompressedImage stub(compressed.data(), 8, width, height, + CompressedImageMode::Uint32, CompressionAlgorithm::BSHUF_LZ4); + CHECK_THROWS(decoder.Decode(stub, gpu_out.get())); +} + +#endif diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index e2ff8066..b5188c60 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -65,6 +65,7 @@ ADD_EXECUTABLE(jfjoch_test ResolutionShellsTest.cpp ImageSpotFinderCPUTest.cpp ImageSpotFinderGPUTest.cpp + BSLZ4DecoderGPUTest.cpp AdaptiveSpotFinderGPUTest.cpp SpotExtractorGPUParityTest.cpp CalcBraggPredictionTest.cpp