From 11052bc376371ce95c3f2dc604b1b0e8f667ef16 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Tue, 13 Oct 2026 02:53:52 +0200 Subject: [PATCH] CBFDecoderGPU: walk short parts from shared memory The first device decoder walked 256-byte tiles, one thread each, straight from global memory: on a PILATUS 6M frame (~6 MB of stream) that is ~24k threads walking alone, and the two walks took ~0.5 ms of the card per frame - with the CBF decode gone from the workers the card is what bounds the 7qij loops (kernels 67-78 % busy). Now each thread walks a 68-byte part that its block has read into shared memory together (68, not 64, so the threads of a warp do not read the same banks), the scan runs over tiles of 8 parts whose maps are composed first, and the write kernel stages a block's pixels in shared memory and writes them out coalesced. Measured on 36 real frames of 9 CBF sets (nsys, average per frame): 0.41 ms against ~0.6 ms before. Output unchanged: the same tests, and the same 36 frames identical to the host decoder. Co-Authored-By: Claude Opus 5.5 (1M context) Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3 --- .../image_preprocessing/CBFDecoderGPU.cu | 146 +++++++++++++----- .../image_preprocessing/CBFDecoderGPU.h | 13 +- 2 files changed, 116 insertions(+), 43 deletions(-) diff --git a/image_analysis/image_preprocessing/CBFDecoderGPU.cu b/image_analysis/image_preprocessing/CBFDecoderGPU.cu index 13cf93845..e404afb12 100644 --- a/image_analysis/image_preprocessing/CBFDecoderGPU.cu +++ b/image_analysis/image_preprocessing/CBFDecoderGPU.cu @@ -16,12 +16,40 @@ namespace { throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val)); } - constexpr size_t TILE_BYTES = 256; + // The stream is walked in parts of PART_BYTES, one thread each, so that a frame of a few MB is many + // short walks. The scan that strings them together runs over tiles of PARTS_PER_TILE parts: each + // of its elements is a whole 16-entry map, and on parts that scan cost more than the walks. + constexpr size_t PART_BYTES = 68; + constexpr size_t PARTS_PER_TILE = 8; + constexpr size_t TILE_BYTES = PART_BYTES * PARTS_PER_TILE; + // Threads per block of the part kernels. A block reads its parts' bytes into shared memory together, + // coalesced, and walks them there; the write kernel stages its pixels the same way - at most one + // pixel per byte, so PART_BYTES per thread - and writes them out together. + constexpr int BLOCK_THREADS = 128; + // A token starting in a block's last part can run 14 bytes past it. + constexpr size_t BLOCK_BYTES = PART_BYTES * BLOCK_THREADS + 16; + + // The bytes of the stream from `first` on, as a block holds them in shared memory. + struct StagedBytes { + const uint8_t *bytes; + size_t first; + size_t size; // of the whole stream + __device__ uint32_t operator[](size_t p) const { return bytes[p - first]; } + }; + + __device__ inline StagedBytes StageBlock(const uint8_t *__restrict__ d, size_t size, uint8_t *staged) { + const size_t first = blockIdx.x * PART_BYTES * BLOCK_THREADS; + for (size_t i = threadIdx.x; i < BLOCK_BYTES; i += BLOCK_THREADS) + staged[i] = first + i < size ? d[first + i] : 0; + __syncthreads(); + return {staged, first, size}; + } // The token starting at byte p: its length, and its delta (low 32 bits) in `delta`. 0 when the // token runs past the end of the stream - where the host decoder stops too. // Following Bernstein & Hammersley (2006) Int. Tables Cryst. G, 37-43 - __device__ inline int TokenAt(const uint8_t *__restrict__ d, size_t size, size_t p, uint32_t &delta) { + __device__ inline int TokenAt(const StagedBytes &d, size_t p, uint32_t &delta) { + const size_t size = d.size; const int8_t b = static_cast(d[p]); if (b != -128) { delta = static_cast(static_cast(b)); @@ -46,10 +74,9 @@ namespace { } // One step of a walk at position pos: false when the stream has ended there. - __device__ inline bool Step(const uint8_t *__restrict__ d, size_t size, size_t &pos, - uint32_t &count, uint32_t &sum) { + __device__ inline bool Step(const StagedBytes &d, size_t &pos, uint32_t &count, uint32_t &sum) { uint32_t delta = 0; - const int len = TokenAt(d, size, pos, delta); + const int len = TokenAt(d, pos, delta); if (len == 0) return false; pos += len; count++; @@ -61,21 +88,24 @@ namespace { return (ended || pos >= size) ? CBFTileMap::DEAD : static_cast(pos - tile_end); } - __global__ void tile_map_kernel(const uint8_t *__restrict__ d, size_t size, size_t ntiles, + // The map of one part. A part wholly past the end of the stream maps every entry to DEAD. + __global__ void part_map_kernel(const uint8_t *__restrict__ stream, size_t size, size_t nparts, CBFTileMap *__restrict__ maps) { - const size_t t = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; - if (t >= ntiles) return; - const size_t start = t * TILE_BYTES; - const size_t end = start + TILE_BYTES; + __shared__ uint8_t staged[BLOCK_BYTES]; + const StagedBytes d = StageBlock(stream, size, staged); + const size_t t = blockIdx.x * static_cast(BLOCK_THREADS) + threadIdx.x; + if (t >= nparts) return; + const size_t start = t * PART_BYTES; + const size_t end = start + PART_BYTES; const size_t limit = min(end, size); CBFTileMap m; - // The walk from the tile's first byte, in full. + // The walk from the part's first byte, in full. size_t pos0 = start; uint32_t count0 = 0, sum0 = 0; bool ended0 = false; while (pos0 < limit && !ended0) - ended0 = !Step(d, size, pos0, count0, sum0); + ended0 = !Step(d, pos0, count0, sum0); m.count[0] = count0; m.sum[0] = sum0; m.exit[0] = ExitOf(pos0, end, size, ended0); @@ -93,9 +123,9 @@ namespace { break; } if (a < b && !a_ended) - a_ended = !Step(d, size, a, count_a, sum_a); + a_ended = !Step(d, a, count_a, sum_a); else - b_ended = !Step(d, size, b, count_b, sum_b); + b_ended = !Step(d, b, count_b, sum_b); } if (merged) { m.count[e] = count_b + (count0 - count_a); @@ -128,28 +158,65 @@ namespace { } }; - __global__ void tile_write_kernel(const uint8_t *__restrict__ d, size_t size, size_t ntiles, - const CBFTileMap *__restrict__ entries, const CBFTileMap *__restrict__ maps, - int32_t *__restrict__ out, uint32_t npixels, uint32_t *__restrict__ total) { + __global__ void tile_map_kernel(const CBFTileMap *__restrict__ part_maps, size_t ntiles, + CBFTileMap *__restrict__ tile_maps) { const size_t t = blockIdx.x * static_cast(blockDim.x) + threadIdx.x; if (t >= ntiles) return; - const CBFTileMap &in = entries[t]; - if (t == ntiles - 1) - *total = in.count[0] + maps[t].count[in.exit[0]]; - if (in.exit[0] == CBFTileMap::DEAD) - return; - const size_t end = min((t + 1) * TILE_BYTES, size); - size_t pos = t * TILE_BYTES + in.exit[0]; - uint32_t pixel = in.count[0]; - uint32_t value = in.sum[0]; - while (pos < end && pixel < npixels) { - uint32_t delta = 0; - const int len = TokenAt(d, size, pos, delta); - if (len == 0) break; - pos += len; - value += delta; - out[pixel++] = static_cast(value); + CBFTileMap m = part_maps[t * PARTS_PER_TILE]; + for (size_t k = 1; k < PARTS_PER_TILE; k++) + m = ComposeTileMaps()(m, part_maps[t * PARTS_PER_TILE + k]); + tile_maps[t] = m; + } + + __global__ void part_write_kernel(const uint8_t *__restrict__ stream, size_t size, size_t nparts, + const CBFTileMap *__restrict__ tile_entries, + const CBFTileMap *__restrict__ part_maps, + int32_t *__restrict__ out, uint32_t npixels, uint32_t *__restrict__ total) { + __shared__ uint8_t staged_bytes[BLOCK_BYTES]; + __shared__ int32_t staged[PART_BYTES * BLOCK_THREADS]; + __shared__ uint32_t block_first; + __shared__ uint32_t block_end; + const StagedBytes d = StageBlock(stream, size, staged_bytes); + const size_t p = blockIdx.x * static_cast(BLOCK_THREADS) + threadIdx.x; + + // The state this part is entered in: its tile's, carried through the parts before it in the tile. + uint32_t state = CBFTileMap::DEAD, pixel = 0, value = 0; + if (p < nparts) { + const size_t tile = p / PARTS_PER_TILE; + state = tile_entries[tile].exit[0]; + pixel = tile_entries[tile].count[0]; + value = tile_entries[tile].sum[0]; + for (size_t k = tile * PARTS_PER_TILE; k < p; k++) { + pixel += part_maps[k].count[state]; + value += part_maps[k].sum[state]; + state = part_maps[k].exit[state]; + } + if (p == nparts - 1) + *total = pixel + part_maps[p].count[state]; } + // The block's pixels are contiguous, from the first one of its first part. + if (threadIdx.x == 0) { + block_first = pixel; + block_end = pixel; + } + __syncthreads(); + + if (p < nparts && state != CBFTileMap::DEAD) { + const size_t end = min((p + 1) * PART_BYTES, size); + size_t pos = p * PART_BYTES + state; + while (pos < end && pixel < npixels) { + uint32_t delta = 0; + const int len = TokenAt(d, pos, delta); + if (len == 0) break; + pos += len; + value += delta; + staged[pixel++ - block_first] = static_cast(value); + } + atomicMax(&block_end, min(pixel, npixels)); + } + __syncthreads(); + for (uint32_t i = threadIdx.x; block_first + i < block_end; i += BLOCK_THREADS) + out[block_first + i] = staged[i]; } } @@ -163,6 +230,7 @@ void CBFDecoderGPU::Decode(const CompressedImage &image, int32_t *gpu_out) { throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "miniCBF byte-offset stream ended after 0 of " + std::to_string(last_npixels) + " pixels"); const size_t ntiles = (size + TILE_BYTES - 1) / TILE_BYTES; + const size_t nparts = ntiles * PARTS_PER_TILE; if (size > compressed_capacity || ntiles > tile_capacity) { // A frame's stream is a little longer or shorter than the last one's: grow with slack. @@ -171,6 +239,7 @@ void CBFDecoderGPU::Decode(const CompressedImage &image, int32_t *gpu_out) { tile_capacity = (compressed_capacity + TILE_BYTES - 1) / TILE_BYTES; host_compressed = CudaHostPtr(compressed_capacity); gpu_compressed = CudaDevicePtr(compressed_capacity); + gpu_part_maps = CudaDevicePtr(tile_capacity * PARTS_PER_TILE); gpu_maps = CudaDevicePtr(tile_capacity); gpu_entries = CudaDevicePtr(tile_capacity); CBFTileMap identity{}; @@ -186,8 +255,10 @@ void CBFDecoderGPU::Decode(const CompressedImage &image, int32_t *gpu_out) { cuda_err(cudaMemcpyAsync(gpu_compressed, host_compressed.get(), size, cudaMemcpyHostToDevice, *stream)); const int threads = 128; - const int blocks = static_cast((ntiles + threads - 1) / threads); - tile_map_kernel<<>>(gpu_compressed, size, ntiles, gpu_maps); + part_map_kernel<<<(nparts + BLOCK_THREADS - 1) / BLOCK_THREADS, BLOCK_THREADS, 0, *stream>>>( + gpu_compressed, size, nparts, gpu_part_maps); + cuda_err(cudaGetLastError()); + tile_map_kernel<<<(ntiles + threads - 1) / threads, threads, 0, *stream>>>(gpu_part_maps, ntiles, gpu_maps); cuda_err(cudaGetLastError()); // Entering tile 0 at its first byte: every entry maps to itself, nothing emitted yet. @@ -198,8 +269,9 @@ void CBFDecoderGPU::Decode(const CompressedImage &image, int32_t *gpu_out) { cuda_err(cub::DeviceScan::ExclusiveScan(gpu_scan_temp.get(), temp_bytes, gpu_maps.get(), gpu_entries.get(), ComposeTileMaps(), identity, ntiles, *stream)); - tile_write_kernel<<>>(gpu_compressed, size, ntiles, gpu_entries, gpu_maps, - gpu_out, static_cast(last_npixels), gpu_total); + part_write_kernel<<<(nparts + BLOCK_THREADS - 1) / BLOCK_THREADS, BLOCK_THREADS, 0, *stream>>>( + gpu_compressed, size, nparts, gpu_entries, gpu_part_maps, gpu_out, + static_cast(last_npixels), gpu_total); cuda_err(cudaGetLastError()); cuda_err(cudaEventRecord(decode_stop, *stream)); cuda_err(cudaMemcpyAsync(host_total.get(), gpu_total.get(), sizeof(uint32_t), cudaMemcpyDeviceToHost, *stream)); diff --git a/image_analysis/image_preprocessing/CBFDecoderGPU.h b/image_analysis/image_preprocessing/CBFDecoderGPU.h index b264f5090..d07a1a31c 100644 --- a/image_analysis/image_preprocessing/CBFDecoderGPU.h +++ b/image_analysis/image_preprocessing/CBFDecoderGPU.h @@ -23,12 +23,12 @@ struct CBFTileMap { // int32 pixels, and the same refusal of a stream that ends before the image does. // // The format is a running sum of deltas whose tokens are 1, 3, 7 or 15 bytes long, so where a token -// starts depends on every token before it. The stream is cut into tiles; for each tile and each of the -// 15 offsets a decode can enter it at, one thread works out where it leaves, how many pixels it emits -// and what they sum to (the walks from different offsets meet within a few bytes, after which they -// are the same walk). Composing those maps across tiles is a scan, which gives every tile its true -// entry offset, its first pixel and the running value there; a second walk of each tile then writes -// the pixels. +// starts depends on every token before it. The stream is cut into short parts; for each part and each +// of the 15 offsets a decode can enter it at, one thread works out where it leaves, how many pixels it +// emits and what they sum to (the walks from different offsets meet within a few bytes, after which +// they are the same walk). Composing those maps is a scan - over tiles of a few parts, whose maps are +// the parts' composed - which gives every part its true entry offset, its first pixel and the running +// value there; a second walk of each part then writes the pixels. class CBFDecoderGPU { std::shared_ptr stream; @@ -37,6 +37,7 @@ class CBFDecoderGPU { CudaDevicePtr gpu_compressed; size_t compressed_capacity = 0; + CudaDevicePtr gpu_part_maps; CudaDevicePtr gpu_maps; CudaDevicePtr gpu_entries; // exclusive scan of gpu_maps: the state each tile is entered in CudaDevicePtr gpu_scan_temp;