mx/Jungfraujoch
2
1
Fork 0
Code Issues Pull Requests 1 Actions 2 Packages Projects Releases 71 Wiki Activity

Decode a compressed frame into shared memory and preprocess it there

Browse Source 
The image loop on a 16 Mpx detector is two thirds of the run and both cards are busy for essentially
all of it, so card time removed is wall time removed. Of the six milliseconds a frame costs, two and
a half were spent decompressing it - and not because the card was short of bandwidth. The LZ4 pass
moved 53 GB/s where the strong-pixel flagger, reading the same image and the same bin table, gets
276. It is latency, not bandwidth: the copy loop moves 32 bytes per warp iteration with a syncwarp
after each one, and for a match copy the source and the destination both derive from the same
pointer, so nothing pipelines. The warp spends its time waiting for global memory, one dependent
round trip at a time.

So decode where the waiting is cheap. One CUDA block now owns one bitshuffle block: its first warp
decodes the payload into shared memory, and the whole block then un-transposes and preprocesses out
of shared and writes finished pixels. A shared round trip is tens of cycles rather than hundreds,
and the 72 MB shuffled intermediate never reaches DRAM at all - the pair of kernels moved about 238
MB a frame and the fused one moves 93.

The parser is lifted into a device function that both kernels call over the same bytes, so the
standalone path and the fused one cannot decode a chunk differently. The statistics reduction had to
change with it: 48 bytes of static shared on top of a full bitshuffle block costs a whole resident
block per multiprocessor, so the counts now reduce through a warp shuffle and one integer atomic per
warp. Blocks larger than 16 kB keep the two-kernel path, and the beam stop's own decoder is
untouched.

What this costs is decoder parallelism: a block that holds 16 kB of shared is one of four resident
per multiprocessor on this card, where the old kernel fitted thirty-two warps each decoding on its
own. The trade is favourable here and should be better on the production cards, which have half
again as much shared memory per multiprocessor. Measured on a 16 Mpx rotation set at the production
GPU count, with the indexing work of the next commit: 37.2 s -> 32.9 s, and the merged output is
byte-identical.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_011n8riB6X59oRjkrSHzNPAU
This commit is contained in:
jungfrau
2026-08-23 14:03:40 -04:00
co-authored by Claude Opus 5
parent f09fe4e2c5
commit cf2336e523
5 changed files with 396 additions and 207 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/image_preprocessing/BSLZ4DecodeWarp.h
+128
View File
@@ -0,0 +1,128 @@
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#pragma once
#include <cstdint>
// 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;
}
image_analysis/image_preprocessing/BSLZ4DecoderGPU.cu
+54 -128
View File
@@ -5,6 +5,7 @@
#include <type_traits>
#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<int>(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<uint32_t>(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<uint32_t>(elem_size);
ret.block_bytes = block_bytes;
ret.tail_elems = static_cast<uint32_t>(leftover_bytes / elem_size);
ret.tail_elem0 = static_cast<uint32_t>(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);
image_analysis/image_preprocessing/BSLZ4DecoderGPU.h
+29 -4
View File
@@ -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<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);
// 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.
image_analysis/image_preprocessing/CMakeLists.txt
+1 -1
View File
@@ -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()
image_analysis/image_preprocessing/ImagePreprocessorGPU.cu
+184 -74
View File
@@ -1,9 +1,11 @@
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#include <algorithm>
#include <type_traits>
#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<class T>
__device__ __forceinline__ void FlushStats(PreprocessAccum<T> &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<class T, int ES>
__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<T> 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<T> &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<T>::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<T>(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<class T, int ES>
__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<T> &l) {
using U = typename std::make_unsigned<T>::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<class T, int ES>
__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<T> l;
if (blockIdx.x == nblocks)
preprocess_tail<T, ES>(tail_src, tail_elems, tail_elem0, mask, out, sat_value, err_value, l);
else
untranspose_preprocess_block<T, ES>(shuffled + desc[blockIdx.x].out_off, desc[blockIdx.x].nelem,
desc[blockIdx.x].out_off / ES, mask, out,
sat_value, err_value, l);
FlushStats<T>(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<class T, int ES>
__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<T> l;
if (blockIdx.x == nblocks) {
preprocess_tail<T, ES>(tail_src, tail_elems, tail_elem0, mask, out, sat_value, err_value, l);
FlushStats<T>(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<T, ES>(s_shuffled, size, desc[blockIdx.x].out_off / ES, mask, out,
sat_value, err_value, l);
FlushStats<T>(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<CudaStream> 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<BSLZ4DecoderGPU>(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<class T, int ES>
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<T, ES> <<< 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<T, ES> <<< 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<int>(shuffled.block_bytes / 64, 32, 1024);
decode_untranspose_preprocess_kernel<T, ES> <<< 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));