diff --git a/image_analysis/image_preprocessing/BSLZ4DecodeWarp.h b/image_analysis/image_preprocessing/BSLZ4DecodeWarp.h new file mode 100644 index 00000000..9948fcb9 --- /dev/null +++ b/image_analysis/image_preprocessing/BSLZ4DecodeWarp.h @@ -0,0 +1,128 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include + +// The LZ4 block parser, as a device function, so the standalone decode kernel and the fused +// decode+un-transpose+preprocess kernel run exactly the same code over the same bytes. The two +// differ only in where the output lands - device memory for the first, a shared-memory staging +// buffer for the second - and a generic pointer covers both, so there is one parser and no way for +// the two paths to disagree on what a chunk decodes to. +// +// CUDA only: include it from a .cu, never from a header a .cpp sees. + +// 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. +// +// Because the lanes cooperate on the copies, a match can source bytes that OTHER lanes wrote in +// an earlier sequence. Since Volta that needs an explicit __syncwarp() - implicit reconvergence +// is not part of the programming model - so there is one after every copy loop. The full mask is +// correct: the early return and every break test warp-uniform values, so lanes never diverge +// permanently. +// +// Bounds: every read is clamped against iend and every write against oend, so a malformed or +// corrupt payload cannot walk off either buffer. It can still stop early, which leaves the block +// short; that is what the false return reports. +__device__ __forceinline__ bool lz4_decode_block_warp(const uint8_t *ip, const uint8_t *const iend, + uint8_t *const obase, uint8_t *const oend, + int lane) { + uint8_t *op = obase; + bool malformed = false; + + while (ip < iend) { + const uint32_t token = *ip++; + uint32_t litlen = token >> 4; + if (litlen == 15) { + // read_variable_length(&ip, iend - RUN_MASK, initial_check=1) in the reference: the + // chain may not start within, nor run into, the last RUN_MASK (15) input bytes. A + // valid stream never does - the literals it counts have to follow it - so a chain + // that reaches there is corruption, and this is the only place it shows up. + if ((size_t)(iend - ip) <= 15) { malformed = true; break; } + uint32_t s; + do { + s = *ip++; + litlen += s; + if ((size_t)(iend - ip) < 15) { malformed = true; break; } + } while (s == 255); + if (malformed) break; + } + if (litlen) { + // Clamped by the INPUT as well as the output: a corrupt litlen must not read past the + // end of this block's payload or write past the end of the block. Clamping keeps the + // kernel in bounds; needing to clamp at all means the stream is not decodable, which + // is what the reference reports as an error, so record it. + if (litlen > (uint32_t)(oend - op) || litlen > (uint32_t)(iend - ip)) + malformed = true; + const uint32_t n = min(min(litlen, (uint32_t)(oend - op)), (uint32_t)(iend - ip)); + for (uint32_t i = lane; i < n; i += 32) op[i] = ip[i]; + __syncwarp(); + op += n; ip += litlen; + } + + // LZ4's parsing restrictions: an encoder may not leave a match within MFLIMIT (12) bytes + // of the end of the block, nor fewer than 2+1+LASTLITERALS (8) input bytes after a + // literal run that is not the last one. So once either limit is reached this can ONLY be + // the final sequence, and the final sequence must consume the payload exactly. The + // reference applies this whether or not the run was empty, which is why the test sits + // outside the copy - a zero-length literal run near the end is just as illegal. + if ((size_t)(oend - op) < 12 || (size_t)(iend - ip) < 8) { + malformed = (ip != iend) || (op != oend); + break; // necessarily EOF + } + if (iend - ip < 2) 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) { + // read_variable_length(&ip, iend - LASTLITERALS + 1, initial_check=0): bounded by the + // last 4 input bytes rather than 15, and with no check before the first read. + uint32_t s; + do { + s = *ip++; + matchlen += s; + if ((size_t)(iend - ip) < 4) { malformed = true; break; } + } while (s == 255); + if (malformed) break; + } + matchlen += 4; // minmatch + + if (offset == 0 || offset > (uint32_t)(op - obase)) { malformed = true; break; } + const uint8_t *mp = op - offset; + // A match may reach the end of the block but never past it - the reference treats an + // overrun as an error rather than truncating, and so must this. + if (matchlen > (uint32_t)(oend - op)) + malformed = true; + 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, and a runtime + // modulo is an emulated division, so the two cheap cases are peeled off first. + if (offset == 1) { + const uint8_t v = mp[0]; + for (uint32_t i = lane; i < n; i += 32) op[i] = v; + } else if ((offset & (offset - 1)) == 0) { + const uint32_t m = offset - 1; + for (uint32_t i = lane; i < n; i += 32) op[i] = mp[i & m]; + } else { + for (uint32_t i = lane; i < n; i += 32) op[i] = mp[i % offset]; + } + } + __syncwarp(); + op += n; + } + + // A block must decode to exactly its declared length AND consume exactly its payload. Both + // are conditions LZ4_decompress_safe reports to the host path, and both are needed: a corrupt + // stream can land on the right output length while leaving input over, or run its input out + // early. Either way the bytes are not the ones that were compressed. + return !malformed && op == oend && ip == iend; +} diff --git a/image_analysis/image_preprocessing/BSLZ4DecoderGPU.cu b/image_analysis/image_preprocessing/BSLZ4DecoderGPU.cu index 29cd70d0..1bbe3ba3 100644 --- a/image_analysis/image_preprocessing/BSLZ4DecoderGPU.cu +++ b/image_analysis/image_preprocessing/BSLZ4DecoderGPU.cu @@ -5,6 +5,7 @@ #include #include "BSLZ4DecoderGPU.h" +#include "BSLZ4DecodeWarp.h" #include "../../common/JFJochException.h" #include "../../compression/JFJochDecompress.h" // BSHUF_BLOCKED_MULT and the container layout @@ -14,20 +15,8 @@ namespace { 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. - // - // Because the lanes cooperate on the copies, a match can source bytes that OTHER lanes wrote in - // an earlier sequence. Since Volta that needs an explicit __syncwarp() - implicit reconvergence - // is not part of the programming model - so there is one after every copy loop. The full mask is - // correct: the early return and every break test warp-uniform values, so lanes never diverge - // permanently. - // - // Bounds: every read is clamped against iend and every write against oend, so a malformed or - // corrupt payload cannot walk off either buffer. It can still stop early, which leaves the block - // short; lane 0 flags that at the end and the host turns it into an exception. + // One CUDA block per 256/32 = 8 LZ4 blocks; the parser itself is lz4_decode_block_warp, shared + // with the fused decode+preprocess kernel so the two cannot decode a chunk differently. __global__ void lz4_decode_blocks(const uint8_t *__restrict__ src, const BSLZ4BlockDesc *__restrict__ desc, uint8_t *__restrict__ dst, @@ -37,105 +26,11 @@ namespace { 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; + const uint8_t *const ip = src + desc[b].in_off; uint8_t *const obase = dst + desc[b].out_off; - uint8_t *op = obase; - uint8_t *const oend = obase + desc[b].nelem * elem_size; - bool malformed = false; - - while (ip < iend) { - const uint32_t token = *ip++; - uint32_t litlen = token >> 4; - if (litlen == 15) { - // read_variable_length(&ip, iend - RUN_MASK, initial_check=1) in the reference: the - // chain may not start within, nor run into, the last RUN_MASK (15) input bytes. A - // valid stream never does - the literals it counts have to follow it - so a chain - // that reaches there is corruption, and this is the only place it shows up. - if ((size_t)(iend - ip) <= 15) { malformed = true; break; } - uint32_t s; - do { - s = *ip++; - litlen += s; - if ((size_t)(iend - ip) < 15) { malformed = true; break; } - } while (s == 255); - if (malformed) break; - } - if (litlen) { - // Clamped by the INPUT as well as the output: a corrupt litlen must not read past the - // end of this block's payload or write past the end of the block. Clamping keeps the - // kernel in bounds; needing to clamp at all means the stream is not decodable, which - // is what the reference reports as an error, so record it. - if (litlen > (uint32_t)(oend - op) || litlen > (uint32_t)(iend - ip)) - malformed = true; - const uint32_t n = min(min(litlen, (uint32_t)(oend - op)), (uint32_t)(iend - ip)); - for (uint32_t i = lane; i < n; i += 32) op[i] = ip[i]; - __syncwarp(); - op += n; ip += litlen; - } - - // LZ4's parsing restrictions: an encoder may not leave a match within MFLIMIT (12) bytes - // of the end of the block, nor fewer than 2+1+LASTLITERALS (8) input bytes after a - // literal run that is not the last one. So once either limit is reached this can ONLY be - // the final sequence, and the final sequence must consume the payload exactly. The - // reference applies this whether or not the run was empty, which is why the test sits - // outside the copy - a zero-length literal run near the end is just as illegal. - if ((size_t)(oend - op) < 12 || (size_t)(iend - ip) < 8) { - malformed = (ip != iend) || (op != oend); - break; // necessarily EOF - } - if (iend - ip < 2) 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) { - // read_variable_length(&ip, iend - LASTLITERALS + 1, initial_check=0): bounded by the - // last 4 input bytes rather than 15, and with no check before the first read. - uint32_t s; - do { - s = *ip++; - matchlen += s; - if ((size_t)(iend - ip) < 4) { malformed = true; break; } - } while (s == 255); - if (malformed) break; - } - matchlen += 4; // minmatch - - if (offset == 0 || offset > (uint32_t)(op - obase)) { malformed = true; break; } - const uint8_t *mp = op - offset; - // A match may reach the end of the block but never past it - the reference treats an - // overrun as an error rather than truncating, and so must this. - if (matchlen > (uint32_t)(oend - op)) - malformed = true; - 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, and a runtime - // modulo is an emulated division, so the two cheap cases are peeled off first. - if (offset == 1) { - const uint8_t v = mp[0]; - for (uint32_t i = lane; i < n; i += 32) op[i] = v; - } else if ((offset & (offset - 1)) == 0) { - const uint32_t m = offset - 1; - for (uint32_t i = lane; i < n; i += 32) op[i] = mp[i & m]; - } else { - for (uint32_t i = lane; i < n; i += 32) op[i] = mp[i % offset]; - } - } - __syncwarp(); - op += n; - } - - // A block must decode to exactly its declared length AND consume exactly its payload. Both - // are conditions LZ4_decompress_safe reports to the host path, and both are needed: a corrupt - // stream can land on the right output length while leaving input over, or run its input out - // early. Either way the bytes are not the ones that were compressed. - if (lane == 0 && (malformed || op != oend || ip != iend)) + const bool ok = lz4_decode_block_warp(ip, ip + desc[b].in_len, + obase, obase + desc[b].nelem * elem_size, lane); + if (lane == 0 && !ok) atomicExch(status, 1u); } @@ -263,7 +158,17 @@ void BSLZ4DecoderGPU::EnsureBlockCapacity(size_t nblocks) { max_blocks = want; } -BSLZ4ShuffledImage BSLZ4DecoderGPU::DecodeShuffled(const CompressedImage &image) { +uint32_t BSLZ4DecoderGPU::BlockBytes(const CompressedImage &image) { + if (image.GetCompressedSize() < 12) + return 0; // not a chunk at all; the caller's decode reports it properly + return be32(image.GetCompressed() + 8); +} + +// Everything up to and including getting the chunk onto the device: the container scan, which is the +// only part that has to happen on the host, the capacity checks, and the uploads. What decodes the +// blocks is left to the caller of this - either the LZ4 kernel below or the fused kernel in the +// preprocessor - because that is the whole difference between the two routes. +BSLZ4ShuffledImage BSLZ4DecoderGPU::PrepareChunk(const CompressedImage &image) { const uint8_t *src = image.GetCompressed(); const size_t clen = image.GetCompressedSize(); const size_t elem_size = elem_size_of(image.GetMode()); @@ -336,34 +241,55 @@ BSLZ4ShuffledImage BSLZ4DecoderGPU::DecodeShuffled(const CompressedImage &image) cuda_err(cudaMemcpyAsync(gpu_compressed.get(), src, clen, cudaMemcpyHostToDevice, *stream)); const int nb = static_cast(nblk); - if (nb > 0) { + if (nb > 0) cuda_err(cudaMemcpyAsync(gpu_desc.get(), host_desc.get(), nblk * sizeof(BSLZ4BlockDesc), cudaMemcpyHostToDevice, *stream)); - lz4_decode_blocks<<<(nb * 32 + 255) / 256, 256, 0, *stream>>>( - gpu_compressed.get(), gpu_desc.get(), gpu_shuffled.get(), gpu_status.get(), - nb, static_cast(elem_size)); - cuda_err(cudaGetLastError()); - } - cuda_err(cudaMemcpyAsync(host_status.get(), gpu_status.get(), sizeof(uint32_t), - cudaMemcpyDeviceToHost, *stream)); - // Stop the clock here rather than after the un-transpose: getting the chunk onto the device and - // LZ4-decoding it is the part that replaced the host decompression, and it is the same work on - // both the raw-bytes path and the fused one, where the un-transpose is inseparable from - // preprocessing and is reported with it. - cuda_err(cudaEventRecord(decode_stop, *stream)); - decode_timed = true; BSLZ4ShuffledImage ret; - ret.shuffled = gpu_shuffled.get(); + ret.compressed = gpu_compressed.get(); ret.desc = gpu_desc.get(); + ret.status = gpu_status.get(); ret.nblocks = nb; ret.elem_size = static_cast(elem_size); + ret.block_bytes = block_bytes; ret.tail_elems = static_cast(leftover_bytes / elem_size); ret.tail_elem0 = static_cast(out_off / elem_size); ret.tail_src = leftover_bytes > 0 ? gpu_compressed.get() + off : nullptr; return ret; } +BSLZ4ShuffledImage BSLZ4DecoderGPU::DecodeShuffled(const CompressedImage &image) { + BSLZ4ShuffledImage ret = PrepareChunk(image); + if (ret.nblocks > 0) { + lz4_decode_blocks<<<(ret.nblocks * 32 + 255) / 256, 256, 0, *stream>>>( + ret.compressed, ret.desc, gpu_shuffled.get(), ret.status, + ret.nblocks, ret.elem_size); + cuda_err(cudaGetLastError()); + } + ret.shuffled = gpu_shuffled.get(); + QueueDecodeStatus(); + // Stop the clock here rather than after the un-transpose: getting the chunk onto the device and + // LZ4-decoding it is the part that replaced the host decompression. + cuda_err(cudaEventRecord(decode_stop, *stream)); + decode_timed = true; + return ret; +} + +BSLZ4ShuffledImage BSLZ4DecoderGPU::UploadCompressed(const CompressedImage &image) { + BSLZ4ShuffledImage ret = PrepareChunk(image); + // All this route decodes on its own is the PCIe upload, so that is what the clock brackets. The + // LZ4 pass happens inside the caller's kernel, inseparably from the un-transpose and the + // preprocessing, and is reported with them. + cuda_err(cudaEventRecord(decode_stop, *stream)); + decode_timed = true; + return ret; +} + +void BSLZ4DecoderGPU::QueueDecodeStatus() { + cuda_err(cudaMemcpyAsync(host_status.get(), gpu_status.get(), sizeof(uint32_t), + cudaMemcpyDeviceToHost, *stream)); +} + void BSLZ4DecoderGPU::Decode(const CompressedImage &image, uint8_t *gpu_out) { const BSLZ4ShuffledImage s = DecodeShuffled(image); diff --git a/image_analysis/image_preprocessing/BSLZ4DecoderGPU.h b/image_analysis/image_preprocessing/BSLZ4DecoderGPU.h index 98da7b5f..acfac8b0 100644 --- a/image_analysis/image_preprocessing/BSLZ4DecoderGPU.h +++ b/image_analysis/image_preprocessing/BSLZ4DecoderGPU.h @@ -16,15 +16,21 @@ struct BSLZ4BlockDesc { uint32_t nelem; // elements in this block (the last one is usually shorter) }; -// What DecodeShuffled() leaves on the device: the LZ4 output, still bitshuffled, plus everything the -// un-transpose 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. +// 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 @@ -76,17 +82,36 @@ class BSLZ4DecoderGPU { 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 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. diff --git a/image_analysis/image_preprocessing/CMakeLists.txt b/image_analysis/image_preprocessing/CMakeLists.txt index 5ed65c05..0b958acc 100644 --- a/image_analysis/image_preprocessing/CMakeLists.txt +++ b/image_analysis/image_preprocessing/CMakeLists.txt @@ -10,7 +10,7 @@ IF (JFJOCH_CUDA_AVAILABLE) TARGET_SOURCES(JFJochImagePreprocessing PRIVATE ../indexing/CUDAMemHelpers.h ImagePreprocessorGPU.cu ImagePreprocessorGPU.h - BSLZ4DecoderGPU.cu BSLZ4DecoderGPU.h + BSLZ4DecoderGPU.cu BSLZ4DecoderGPU.h BSLZ4DecodeWarp.h ImagePreprocessorBufferGPU.cu ImagePreprocessorBufferGPU.h) ENDIF() \ No newline at end of file diff --git a/image_analysis/image_preprocessing/ImagePreprocessorGPU.cu b/image_analysis/image_preprocessing/ImagePreprocessorGPU.cu index 19b86924..5fc0e46f 100644 --- a/image_analysis/image_preprocessing/ImagePreprocessorGPU.cu +++ b/image_analysis/image_preprocessing/ImagePreprocessorGPU.cu @@ -1,9 +1,11 @@ // SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute // SPDX-License-Identifier: GPL-3.0-only +#include #include #include "ImagePreprocessorGPU.h" +#include "BSLZ4DecodeWarp.h" #include "../../common/JFJochException.h" namespace { @@ -122,31 +124,29 @@ struct PreprocessAccum { }; // 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 blocks arrive in. +// 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 __device__ __forceinline__ void FlushStats(PreprocessAccum &l, ImageStatistics *stats) { - __shared__ unsigned long long s_masked, s_saturated, s_error; - __shared__ long long s_max, s_min; - if (threadIdx.x == 0) { - s_masked = 0; s_saturated = 0; s_error = 0; s_max = INT64_MIN; s_min = INT64_MAX; + 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)); } - __syncthreads(); - atomicAdd(&s_masked, l.masked); - atomicAdd(&s_saturated, l.saturated); - atomicAdd(&s_error, l.error); - if (l.min_v <= l.max_v) { - atomicMax(&s_max, l.max_v); - atomicMin(&s_min, l.min_v); - } - __syncthreads(); - - 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, s_max); - atomicMin((long long *) &stats->min_value, s_min); + 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); } } @@ -158,49 +158,27 @@ __device__ __forceinline__ uint64_t transpose8_fused(uint64_t x) { return x; } -// The bitshuffle inverse and the preprocessing in ONE pass. 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. +// 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. // -// The last CUDA block (blockIdx.x == nblocks) finishes the handful of elements bitshuffle stores -// verbatim; they are already on the device inside the uploaded chunk. +// `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 -__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 l; - +__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 &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::type; - if (blockIdx.x == nblocks) { - 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); - } - FlushStats(l, stats); - return; - } - - const int b = blockIdx.x; - const uint32_t size = desc[b].nelem; // bytes per plane - const uint8_t *in = shuffled + desc[b].out_off; - const uint32_t elem0 = desc[b].out_off / ES; // first pixel of this block const uint32_t n = size / 8; - for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) { uint64_t x[ES]; #pragma unroll @@ -224,9 +202,112 @@ __global__ __launch_bounds__(256) void untranspose_preprocess_kernel( 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 +__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 &l) { + using U = typename std::make_unsigned::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 +__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 l; + + if (blockIdx.x == nblocks) + preprocess_tail(tail_src, tail_elems, tail_elem0, mask, out, sat_value, err_value, l); + else + untranspose_preprocess_block(shuffled + desc[blockIdx.x].out_off, desc[blockIdx.x].nelem, + desc[blockIdx.x].out_off / ES, mask, out, + sat_value, err_value, l); FlushStats(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 +__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 l; + + if (blockIdx.x == nblocks) { + preprocess_tail(tail_src, tail_elems, tail_elem0, mask, out, sat_value, err_value, l); + FlushStats(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(s_shuffled, size, desc[blockIdx.x].out_off / ES, mask, out, + sat_value, err_value, l); + FlushStats(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 stream, bool copy_image_to_host) : ImagePreprocessor(experiment), @@ -299,10 +380,16 @@ bool ImagePreprocessorGPU::AnalyzeCompressed(ImagePreprocessorBuffer &processed_ // 16-bit data the difference is half of a full frame per worker thread. bslz4_decoder = std::make_unique(image.GetUncompressedSize(), stream); - // LZ4 on the device, then ONE kernel that un-transposes the bitshuffle blocks and preprocesses - // them as it goes. The decompressed image is never materialised: the fused kernel reads the - // shuffled bytes and writes finished int32 pixels. - const BSLZ4ShuffledImage shuffled = bslz4_decoder->DecodeShuffled(image); + // 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: @@ -323,7 +410,8 @@ bool ImagePreprocessorGPU::AnalyzeCompressed(ImagePreprocessorBuffer &processed_ } // The device-decode counterpart of AnalyzeOnDevice: same per-pixel decision, same statistics, but -// fed from the bitshuffled bytes rather than from a decompressed image. +// 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 ImageStatistics ImagePreprocessorGPU::UntransposeAndAnalyze(ImagePreprocessorBuffer &processed_image, const BSLZ4ShuffledImage &shuffled, @@ -336,19 +424,41 @@ ImageStatistics ImagePreprocessorGPU::UntransposeAndAnalyze(ImagePreprocessorBuf // 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); - untranspose_preprocess_kernel <<< 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()); + if (shuffled.shuffled) { + untranspose_preprocess_kernel <<< 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(shuffled.block_bytes / 64, 32, 1024); + decode_untranspose_preprocess_kernel <<< 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));