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

CBFDecoderGPU: walk short parts from shared memory

Browse Source 
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) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01EizBKhTqrAYago9KJG3dA3
This commit is contained in:
leonarski_f
2026-10-13 02:53:52 +02:00
co-authored by Claude Opus 5.5
parent dfd69fcb91
commit 11052bc376
2 changed files with 116 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/image_preprocessing/CBFDecoderGPU.cu
+109 -37
View File
@@ -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<int8_t>(d[p]);
if (b != -128) {
delta = static_cast<uint32_t>(static_cast<int32_t>(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<uint8_t>(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<size_t>(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<size_t>(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<size_t>(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<int32_t>(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<size_t>(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<int32_t>(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<uint8_t>(compressed_capacity);
gpu_compressed = CudaDevicePtr<uint8_t>(compressed_capacity);
gpu_part_maps = CudaDevicePtr<CBFTileMap>(tile_capacity * PARTS_PER_TILE);
gpu_maps = CudaDevicePtr<CBFTileMap>(tile_capacity);
gpu_entries = CudaDevicePtr<CBFTileMap>(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<int>((ntiles + threads - 1) / threads);
tile_map_kernel<<<blocks, threads, 0, *stream>>>(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<<<blocks, threads, 0, *stream>>>(gpu_compressed, size, ntiles, gpu_entries, gpu_maps,
gpu_out, static_cast<uint32_t>(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<uint32_t>(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));