leonarski_fandClaude Opus 5
2d3c39c9dd
The raw-bytes path assembled each element a byte at a time, which on a full frame cost about 4x against writing the 8 contiguous elements a thread owns through an element-typed pointer. They are 8*ES-byte aligned, so the compiler merges them. 72.4 MB frame: 1.524 -> 0.406 ms for upload plus both kernels. The test now also times the LZ4 pass on its own, so the bounds and validity checks in the hot loop can be costed rather than guessed at. They are free: 0.231 ms against 0.2297 ms measured for the kernel before any of them existed - the restored offset == 1 and power-of-two fast paths pay for them. compute-sanitizer memcheck reports no error over 400 single-bit-corrupted payloads and nine malformed containers. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
389 lines
20 KiB
Plaintext
389 lines
20 KiB
Plaintext
// 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 "BSLZ4DecoderGPU.h"
|
|
#include "../../common/JFJochException.h"
|
|
#include "../../compression/JFJochDecompress.h" // BSHUF_BLOCKED_MULT and the container layout
|
|
|
|
namespace {
|
|
void cuda_err(cudaError_t val) {
|
|
if (val != cudaSuccess)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
|
|
}
|
|
|
|
// One WARP per LZ4 block. Every lane runs the same sequence parser over the same bytes - a
|
|
// broadcast read, so no divergence - and the literal and match copies are split across the 32
|
|
// lanes so the stores coalesce. One thread per block instead has each thread streaming its own
|
|
// 8 kB region, which coalesces not at all and measured 13x slower.
|
|
//
|
|
// 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.
|
|
__global__ void lz4_decode_blocks(const uint8_t *__restrict__ src,
|
|
const BSLZ4BlockDesc *__restrict__ desc,
|
|
uint8_t *__restrict__ dst,
|
|
uint32_t *__restrict__ status,
|
|
int nblocks, uint32_t elem_size) {
|
|
const int lane = threadIdx.x & 31;
|
|
const int b = (blockIdx.x * blockDim.x + threadIdx.x) >> 5;
|
|
if (b >= nblocks) return;
|
|
|
|
const uint8_t *ip = src + desc[b].in_off;
|
|
const uint8_t *const iend = ip + desc[b].in_len;
|
|
uint8_t *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))
|
|
atomicExch(status, 1u);
|
|
}
|
|
|
|
__device__ __forceinline__ uint64_t transpose8(uint64_t x) {
|
|
uint64_t t;
|
|
t = (x ^ (x >> 7)) & 0x00aa00aa00aa00aaULL; x = x ^ t ^ (t << 7);
|
|
t = (x ^ (x >> 14)) & 0x0000cccc0000ccccULL; x = x ^ t ^ (t << 14);
|
|
t = (x ^ (x >> 28)) & 0x00000000f0f0f0f0ULL; x = x ^ t ^ (t << 28);
|
|
return x;
|
|
}
|
|
|
|
// The bitshuffle inverse, mirroring bitshuf_decode_block. One thread owns one group of 8 elements
|
|
// across EVERY byte-plane, so after transposing its 8 bytes out of each plane it holds all bytes
|
|
// of 8 complete elements and can write them straight out. That needs no staging buffer, which is
|
|
// what keeps the kernel free of the 48 kB dynamic-shared-memory ceiling a block size taken from
|
|
// the file header would otherwise run into.
|
|
template<int ES>
|
|
__global__ void bitshuffle_untranspose(const uint8_t *__restrict__ shuffled,
|
|
const BSLZ4BlockDesc *__restrict__ desc,
|
|
uint8_t *__restrict__ out, int nblocks) {
|
|
const int b = blockIdx.x;
|
|
if (b >= nblocks) return;
|
|
|
|
const uint32_t size = desc[b].nelem; // bytes per plane
|
|
const uint8_t *in = shuffled + desc[b].out_off;
|
|
uint8_t *dst = out + desc[b].out_off;
|
|
const uint32_t n = size / 8;
|
|
|
|
// The 8 elements a thread owns are contiguous and 8*ES-byte aligned, so they are assembled
|
|
// whole and written through an element-typed pointer. Storing them byte by byte instead
|
|
// costs about 4x on a full frame.
|
|
using UT = typename std::conditional<ES == 1, uint8_t,
|
|
typename std::conditional<ES == 2, uint16_t, uint32_t>::type>::type;
|
|
|
|
for (uint32_t i = threadIdx.x; i < n; i += blockDim.x) {
|
|
uint64_t x[ES];
|
|
#pragma unroll
|
|
for (int p = 0; p < ES; p++) {
|
|
const uint8_t *pin = in + p * size;
|
|
uint64_t a = 0;
|
|
#pragma unroll
|
|
for (int k = 0; k < 8; k++) a |= (uint64_t)pin[k * n + i] << (8 * k);
|
|
x[p] = transpose8(a);
|
|
}
|
|
UT *dstT = reinterpret_cast<UT *>(dst) + i * 8;
|
|
#pragma unroll
|
|
for (int k = 0; k < 8; k++) {
|
|
UT v = 0;
|
|
#pragma unroll
|
|
for (int p = 0; p < ES; p++) v |= (UT)((UT)((x[p] >> (8 * k)) & 0xff) << (8 * p));
|
|
dstT[k] = v;
|
|
}
|
|
}
|
|
}
|
|
|
|
uint64_t be64(const uint8_t *p) { uint64_t v = 0; for (int i = 0; i < 8; i++) v = (v << 8) | p[i]; return v; }
|
|
uint32_t be32(const uint8_t *p) { return ((uint32_t)p[0] << 24) | ((uint32_t)p[1] << 16) | ((uint32_t)p[2] << 8) | p[3]; }
|
|
|
|
size_t elem_size_of(CompressedImageMode mode) {
|
|
switch (mode) {
|
|
// 8-bit is a real DECTRIS mode. bitshuf_decode_block takes a separate branch for
|
|
// elem_size == 1 (bit un-transpose only, no byte interleave), and the kernel here
|
|
// reproduces that for free: with one plane the interleave step degenerates to a copy.
|
|
case CompressedImageMode::Int8:
|
|
case CompressedImageMode::Uint8: return 1;
|
|
case CompressedImageMode::Int16:
|
|
case CompressedImageMode::Uint16: return 2;
|
|
case CompressedImageMode::Int32:
|
|
case CompressedImageMode::Uint32: return 4;
|
|
default: return 0; // float and RGB modes are never bitshuffled by this pipeline
|
|
}
|
|
}
|
|
|
|
// The smallest a block can be on the wire: a 4-byte length plus at least one payload byte. Used
|
|
// to reject a header whose declared block size implies more blocks than the chunk could hold,
|
|
// before that count is turned into an allocation.
|
|
constexpr size_t MIN_BLOCK_BYTES_ON_WIRE = 5;
|
|
}
|
|
|
|
bool BSLZ4DecoderGPU::Supports(const CompressedImage &image) {
|
|
return image.GetCompressionAlgorithm() == CompressionAlgorithm::BSHUF_LZ4
|
|
&& elem_size_of(image.GetMode()) != 0;
|
|
}
|
|
|
|
BSLZ4DecoderGPU::BSLZ4DecoderGPU(size_t in_max_uncompressed_bytes, std::shared_ptr<CudaStream> in_stream)
|
|
: stream(std::move(in_stream)),
|
|
max_uncompressed_bytes(in_max_uncompressed_bytes) {
|
|
gpu_shuffled = CudaDevicePtr<uint8_t>(max_uncompressed_bytes);
|
|
gpu_status = CudaDevicePtr<uint32_t>(1);
|
|
host_status = CudaHostPtr<uint32_t>(1);
|
|
// The compressed buffer and the descriptors are grown to fit the first image instead of being
|
|
// sized for a worst case that no real frame reaches. A chunk is a few MB against an image of
|
|
// tens; sizing this from the UNCOMPRESSED size cost ~73 MB per worker to hold ~4 MB.
|
|
}
|
|
|
|
void BSLZ4DecoderGPU::EnsureCompressedCapacity(size_t bytes) {
|
|
if (bytes <= compressed_capacity)
|
|
return;
|
|
// A little slack, so a frame that compresses slightly worse than the last does not reallocate.
|
|
const size_t want = bytes + bytes / 4;
|
|
cuda_err(cudaStreamSynchronize(*stream)); // nothing may still be reading the old buffer
|
|
gpu_compressed = CudaDevicePtr<uint8_t>(want);
|
|
compressed_capacity = want;
|
|
}
|
|
|
|
void BSLZ4DecoderGPU::EnsureBlockCapacity(size_t nblocks) {
|
|
if (nblocks <= max_blocks)
|
|
return;
|
|
const size_t want = nblocks + nblocks / 4 + 16;
|
|
cuda_err(cudaStreamSynchronize(*stream)); // the previous descriptor upload must have landed
|
|
gpu_desc = CudaDevicePtr<BSLZ4BlockDesc>(want);
|
|
host_desc = CudaHostPtr<BSLZ4BlockDesc>(want);
|
|
max_blocks = want;
|
|
}
|
|
|
|
BSLZ4ShuffledImage BSLZ4DecoderGPU::DecodeShuffled(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());
|
|
const size_t total_bytes = image.GetUncompressedSize();
|
|
|
|
if (clen < 12)
|
|
throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 chunk shorter than its header");
|
|
if (total_bytes > max_uncompressed_bytes)
|
|
throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 image larger than the decoder was sized for");
|
|
if (be64(src) != total_bytes)
|
|
throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 header size does not match the image");
|
|
|
|
const uint32_t block_bytes = be32(src + 8);
|
|
if (block_bytes == 0 || block_bytes % elem_size != 0)
|
|
throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 block size invalid");
|
|
|
|
const size_t block_elems = block_bytes / elem_size;
|
|
// bitshuffle transposes 8 elements at a time and refuses a block that is not a multiple of 8;
|
|
// the host decoder rejects this too (JFJochDecompress.h). Without the check the un-transpose
|
|
// would silently drop the last size % 8 elements of every block.
|
|
if (block_elems % BSHUF_BLOCKED_MULT != 0)
|
|
throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 block size is not a multiple of 8 elements");
|
|
|
|
const size_t nelements = total_bytes / elem_size;
|
|
const size_t nfull = nelements / block_elems;
|
|
const size_t rem = nelements - nfull * block_elems;
|
|
const size_t last = rem - rem % BSHUF_BLOCKED_MULT;
|
|
const size_t leftover_bytes = (rem % BSHUF_BLOCKED_MULT) * elem_size;
|
|
|
|
// Walk the container to locate the blocks. Lengths are only knowable in order, so this scan is
|
|
// inherent to the format rather than an implementation choice.
|
|
// An image of fewer than 8 elements has no bitshuffle block at all - it is entirely the verbatim
|
|
// tail. The host decoder handles that, so handle it here rather than declining: nblocks is simply
|
|
// zero and only the tail is copied.
|
|
const size_t nblocks_needed = nfull + (last > 0 ? 1 : 0);
|
|
// Bound the block count by what the chunk could actually hold BEFORE it becomes an allocation:
|
|
// a header declaring a one-element block size would otherwise ask for hundreds of MB of pinned
|
|
// memory, and only then fail on the first block header.
|
|
if (nblocks_needed > (clen - 12) / MIN_BLOCK_BYTES_ON_WIRE)
|
|
throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 chunk too short for the blocks its header implies");
|
|
|
|
EnsureCompressedCapacity(clen);
|
|
EnsureBlockCapacity(nblocks_needed);
|
|
|
|
size_t nblk = 0, off = 12, out_off = 0;
|
|
for (size_t i = 0; i < nblocks_needed; i++) {
|
|
if (off + 4 > clen)
|
|
throw JFJochException(JFJochExceptionCategory::Compression, "truncated bslz4 block header");
|
|
const uint32_t block_clen = be32(src + off);
|
|
off += 4;
|
|
if (block_clen == 0 || off + block_clen > clen)
|
|
throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 block extends past the chunk");
|
|
const uint32_t ne = (i < nfull) ? static_cast<uint32_t>(block_elems) : static_cast<uint32_t>(last);
|
|
host_desc.get()[nblk++] = {static_cast<uint32_t>(off), block_clen, static_cast<uint32_t>(out_off), ne};
|
|
off += block_clen;
|
|
out_off += static_cast<size_t>(ne) * elem_size;
|
|
}
|
|
// The tail that bitshuffle leaves uncompressed and copies verbatim, and then nothing else: the
|
|
// host decoder requires the chunk to be consumed exactly, so require it here too rather than
|
|
// ignoring trailing bytes that indicate the container is not what it claims to be.
|
|
if (off + leftover_bytes > clen)
|
|
throw JFJochException(JFJochExceptionCategory::Compression, "truncated bslz4 leftover bytes");
|
|
if (off + leftover_bytes != clen)
|
|
throw JFJochException(JFJochExceptionCategory::Compression, "bslz4 chunk has trailing bytes after the last block");
|
|
|
|
// Everything that can be checked on the host has been checked; from here work is queued.
|
|
cuda_err(cudaEventRecord(decode_start, *stream));
|
|
host_status.get()[0] = 0;
|
|
cuda_err(cudaMemcpyAsync(gpu_status.get(), host_status.get(), sizeof(uint32_t),
|
|
cudaMemcpyHostToDevice, *stream));
|
|
cuda_err(cudaMemcpyAsync(gpu_compressed.get(), src, clen, cudaMemcpyHostToDevice, *stream));
|
|
|
|
const int nb = static_cast<int>(nblk);
|
|
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.desc = gpu_desc.get();
|
|
ret.nblocks = nb;
|
|
ret.elem_size = static_cast<uint32_t>(elem_size);
|
|
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;
|
|
}
|
|
|
|
void BSLZ4DecoderGPU::Decode(const CompressedImage &image, uint8_t *gpu_out) {
|
|
const BSLZ4ShuffledImage s = DecodeShuffled(image);
|
|
|
|
if (s.nblocks > 0) { // an image of fewer than 8 elements is all tail and has no block
|
|
switch (s.elem_size) {
|
|
case 1: bitshuffle_untranspose<1><<<s.nblocks, 256, 0, *stream>>>(s.shuffled, s.desc, gpu_out, s.nblocks); break;
|
|
case 2: bitshuffle_untranspose<2><<<s.nblocks, 256, 0, *stream>>>(s.shuffled, s.desc, gpu_out, s.nblocks); break;
|
|
default: bitshuffle_untranspose<4><<<s.nblocks, 256, 0, *stream>>>(s.shuffled, s.desc, gpu_out, s.nblocks); break;
|
|
}
|
|
cuda_err(cudaGetLastError());
|
|
}
|
|
|
|
// The verbatim tail is already on the device inside the uploaded chunk.
|
|
if (s.tail_elems > 0)
|
|
cuda_err(cudaMemcpyAsync(gpu_out + static_cast<size_t>(s.tail_elem0) * s.elem_size, s.tail_src,
|
|
static_cast<size_t>(s.tail_elems) * s.elem_size,
|
|
cudaMemcpyDeviceToDevice, *stream));
|
|
}
|
|
|
|
void BSLZ4DecoderGPU::ThrowIfDecodeFailed() {
|
|
if (host_status.get()[0] != 0)
|
|
throw JFJochException(JFJochExceptionCategory::Compression,
|
|
"bslz4 block did not decode to its declared length - the compressed data is corrupt");
|
|
}
|
|
|
|
float BSLZ4DecoderGPU::GetDecodeTime_s() const {
|
|
if (!decode_timed)
|
|
return 0.0f;
|
|
float ms = 0.0f;
|
|
if (cudaEventElapsedTime(&ms, decode_start, decode_stop) != cudaSuccess)
|
|
return 0.0f;
|
|
return ms * 1e-3f;
|
|
}
|