image_preprocessing: fuse the bitshuffle inverse with preprocessing, and verify the decode

The device decoder was byte-exact on every valid input - 994 production-compressed
images, 927 hand-built LZ4 blocks covering engineered (offset, matchlen) pairs across
the overlap branch boundary, 18000 repeat decodes, sanitizer-clean - and an audit
against LZ4_decompress_generic could not construct a valid block it mis-decodes. What
it did not do was notice when the input was NOT valid, and that mattered more than it
looks: the decode buffers are reused frame to frame, so a block that stopped early left
the PREVIOUS image in place, and in the bitshuffled layout the untouched tail is the
most significant byte-plane. A corrupt chunk therefore did not look like a missing
corner. It looked like thousands of real pixels several powers of two too bright, fed
to spot finding with no diagnostic, where the host decoder had raised an error.

So the kernel now flags a block that fails to reach its declared length while consuming
exactly its payload, and the host turns that into an exception once the caller has
synchronised. Reads are clamped against the end of the payload as well as the output,
both length chains are bounded exactly as read_variable_length bounds them, the two
offset bytes are bounded, and LZ4's parsing restrictions are enforced. On the host side
a block size that is not a multiple of 8 elements is rejected (it made the un-transpose
read uninitialised shared memory), the block count is bounded by what the chunk could
hold before it becomes an allocation (twelve header bytes could demand hundreds of MB
of pinned memory, permanently, per worker), trailing bytes are rejected, and the stream
is synchronised before any throw that happens after work is queued. An image of fewer
than 8 elements is all verbatim tail and now decodes rather than throwing. When the
device route fails for any reason the host decoder gets its turn, so it costs speed
rather than the acquisition.

The lanes cooperate on the copies and a later match can read bytes another lane wrote,
which since Volta needs an explicit __syncwarp(); it worked only because ptxas happened
to reconverge at the post-dominator. The prototype's offset == 1 and power-of-two fast
paths are also restored - the shipped kernel ran a runtime modulo, an emulated 32-bit
division per output byte, on the path its own comment calls the common case.

The un-transpose is now fused with preprocessing. 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 emits 8 finished int32 pixels with the mask, the
error marker, the saturation cap and the statistics applied. The decompressed image is
never materialised: 0.623 -> 0.411 ms/frame at 18 Mpx, 0.523 -> 0.340 with 8 concurrent
workers. Staging nothing in shared memory also drops the 48 kB ceiling, which had made
any file whose bitshuffle blocks exceed it a hard failure; 64 kB blocks now decode.
gpu_compressed is sized from the chunk with grow-on-demand instead of from the
uncompressed size - it was reserving ~73 MB per worker to hold ~4 MB. Measured on a
1630x1553 uint32 rotation set at -N 32, peak GPU memory falls 3756 -> 3084 MiB; the
same model gives ~144 MB per worker on an 18 Mpx frame.

Decoding on the device also stopped reporting a decompression time, which blanked the
broker's compression plot trace and filled /entry/profiling/compressionTime with NaN.
The decoder brackets the decode with CUDA events and reports it again.

Tests: a differential fuzz suite against the CPU decoder - incompressible and highly
compressible data, engineered offsets, a size sweep hitting every rem%8 value twice,
all six element sizes, an 18 Mpx frame, decoder reuse, concurrency, hand-built LZ4
blocks across the overlap boundary, 26 foreign bitshuffle block sizes from 128 B to
64 kB, corrupt payloads and malformed containers, with a coverage report that proves
which LZ4 paths were reached rather than assuming it. Plus the fused path held byte for
byte against ImagePreprocessorCPU, statistics included, and against the host-upload
path on the same frame.

Battery: 37 crystals, every merged number identical to the host-decode run.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-08-03 14:13:55 +02:00
co-authored by Claude Opus 5
parent 47277674fa
commit bec7e2e922
12 changed files with 2023 additions and 92 deletions
@@ -16,6 +16,20 @@ 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.
struct BSLZ4ShuffledImage {
const uint8_t *shuffled = nullptr;
const BSLZ4BlockDesc *desc = nullptr;
int nblocks = 0;
uint32_t elem_size = 0;
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
};
// Decompress a bitshuffle+LZ4 image ON THE DEVICE, so the compressed bytes are what crosses PCIe.
//
// The idea - upload the compressed chunk and decode it on the GPU rather than decompressing on the
@@ -30,6 +44,14 @@ struct BSLZ4BlockDesc {
//
// Only BSHUF_LZ4 is handled. The zstd variants have no device decoder, so Supports() returns false
// and the caller decompresses on the host exactly as before.
//
// The container comes off the network or off disk, so it is not trusted. Everything the host can
// check cheaply is checked before any work is queued and throws; what only the kernel can see - a
// block that does not decode to its declared length, which is what a corrupt LZ4 payload looks like
// - raises a device-side flag that ThrowIfDecodeFailed() reports once the caller has synchronised.
// The CPU decoder makes exactly the same checks (LZ4_decompress_safe's length check plus the
// consumed-input check in JFJochDecompress.h), so a chunk either decodes identically on both or
// fails on both. It is never silently wrong on one and right on the other.
class BSLZ4DecoderGPU {
std::shared_ptr<CudaStream> stream;
@@ -37,19 +59,42 @@ class BSLZ4DecoderGPU {
CudaDevicePtr<uint8_t> gpu_shuffled; // LZ4 output, still bitshuffled
CudaDevicePtr<BSLZ4BlockDesc> gpu_desc;
CudaHostPtr<BSLZ4BlockDesc> host_desc; // pinned, so the descriptor upload is truly async
CudaDevicePtr<uint32_t> gpu_status; // set by the kernel when a block decodes short
CudaHostPtr<uint32_t> host_status;
// Bracket the decode so the time it takes can still be reported as decompression, which is what
// it is. Both are recorded on the decoder's stream and read after the caller synchronises.
CudaEvent decode_start;
CudaEvent decode_stop;
bool decode_timed = false;
size_t max_compressed_bytes = 0;
size_t max_uncompressed_bytes = 0;
size_t compressed_capacity = 0;
size_t max_blocks = 0;
void EnsureCompressedCapacity(size_t bytes);
void EnsureBlockCapacity(size_t nblocks);
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);
// Decode into gpu_out, which must hold image.GetUncompressedSize() bytes. Work is queued on the
// decoder's stream and the caller synchronises. Throws if the container is malformed - it comes
// off the network or off disk, so it is not trusted.
// 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);
// 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.
void Decode(const CompressedImage &image, uint8_t *gpu_out);
// Report a block that did not decode to its declared length. MUST be called after the caller has
// synchronised the stream; until then the flag has not arrived. Throws on failure.
void ThrowIfDecodeFailed();
// Device time spent decoding the last image, in seconds. Valid after the caller synchronises.
[[nodiscard]] float GetDecodeTime_s() const;
};