mx/Jungfraujoch
2
1
Fork 0
Code Issues Pull Requests Actions 4 Packages Projects Releases 70 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/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));