leonarski_fandClaude Opus 5.5
11052bc376
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
294 lines
13 KiB
Plaintext
294 lines
13 KiB
Plaintext
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#include <algorithm>
|
|
#include <cstring>
|
|
#include <string>
|
|
|
|
#include <cub/cub.cuh>
|
|
|
|
#include "CBFDecoderGPU.h"
|
|
#include "../../common/JFJochException.h"
|
|
|
|
namespace {
|
|
void cuda_err(cudaError_t val) {
|
|
if (val != cudaSuccess)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
|
|
}
|
|
|
|
// 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 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));
|
|
return 1;
|
|
}
|
|
if (p + 3 > size) return 0;
|
|
const int16_t v16 = static_cast<int16_t>(d[p + 1] | (d[p + 2] << 8));
|
|
if (v16 != -32768) {
|
|
delta = static_cast<uint32_t>(static_cast<int32_t>(v16));
|
|
return 3;
|
|
}
|
|
if (p + 7 > size) return 0;
|
|
const uint32_t v32 = d[p + 3] | (d[p + 4] << 8) | (d[p + 5] << 16) | (static_cast<uint32_t>(d[p + 6]) << 24);
|
|
if (v32 != 0x80000000u) {
|
|
delta = v32;
|
|
return 7;
|
|
}
|
|
if (p + 15 > size) return 0;
|
|
// An int64 delta: only its low 32 bits reach an int32 pixel.
|
|
delta = d[p + 7] | (d[p + 8] << 8) | (d[p + 9] << 16) | (static_cast<uint32_t>(d[p + 10]) << 24);
|
|
return 15;
|
|
}
|
|
|
|
// One step of a walk at position pos: false when the stream has ended there.
|
|
__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, pos, delta);
|
|
if (len == 0) return false;
|
|
pos += len;
|
|
count++;
|
|
sum += delta;
|
|
return true;
|
|
}
|
|
|
|
__device__ inline uint8_t ExitOf(size_t pos, size_t tile_end, size_t size, bool ended) {
|
|
return (ended || pos >= size) ? CBFTileMap::DEAD : static_cast<uint8_t>(pos - tile_end);
|
|
}
|
|
|
|
// 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) {
|
|
__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 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, pos0, count0, sum0);
|
|
m.count[0] = count0;
|
|
m.sum[0] = sum0;
|
|
m.exit[0] = ExitOf(pos0, end, size, ended0);
|
|
|
|
// The walks from the other entries. Each is stepped beside the first walk, whichever is behind
|
|
// moving, until both stand on the same byte - from there they are one walk, and the rest is
|
|
// the first walk's remainder.
|
|
for (int e = 1; e < CBFTileMap::DEAD; e++) {
|
|
size_t a = start, b = start + e;
|
|
uint32_t count_a = 0, sum_a = 0, count_b = 0, sum_b = 0;
|
|
bool a_ended = false, b_ended = false, merged = false;
|
|
while (b < limit && !b_ended) {
|
|
if (a == b && !a_ended) {
|
|
merged = true;
|
|
break;
|
|
}
|
|
if (a < b && !a_ended)
|
|
a_ended = !Step(d, a, count_a, sum_a);
|
|
else
|
|
b_ended = !Step(d, b, count_b, sum_b);
|
|
}
|
|
if (merged) {
|
|
m.count[e] = count_b + (count0 - count_a);
|
|
m.sum[e] = sum_b + (sum0 - sum_a);
|
|
m.exit[e] = m.exit[0];
|
|
} else {
|
|
m.count[e] = count_b;
|
|
m.sum[e] = sum_b;
|
|
m.exit[e] = ExitOf(b, end, size, b_ended);
|
|
}
|
|
}
|
|
m.count[CBFTileMap::DEAD] = 0;
|
|
m.sum[CBFTileMap::DEAD] = 0;
|
|
m.exit[CBFTileMap::DEAD] = CBFTileMap::DEAD;
|
|
maps[t] = m;
|
|
}
|
|
|
|
// a, then b. The scan's look-back also composes the uninitialised state of tiles before the first
|
|
// one and discards the result, so an exit is masked to the table rather than trusted to be in it.
|
|
struct ComposeTileMaps {
|
|
__host__ __device__ CBFTileMap operator()(const CBFTileMap &a, const CBFTileMap &b) const {
|
|
CBFTileMap r;
|
|
for (int e = 0; e < 16; e++) {
|
|
const int mid = a.exit[e] & 15;
|
|
r.exit[e] = b.exit[mid];
|
|
r.count[e] = a.count[e] + b.count[mid];
|
|
r.sum[e] = a.sum[e] + b.sum[mid];
|
|
}
|
|
return r;
|
|
}
|
|
};
|
|
|
|
__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;
|
|
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];
|
|
}
|
|
}
|
|
|
|
CBFDecoderGPU::CBFDecoderGPU(std::shared_ptr<CudaStream> in_stream)
|
|
: stream(std::move(in_stream)), gpu_total(1), host_total(1) {}
|
|
|
|
void CBFDecoderGPU::Decode(const CompressedImage &image, int32_t *gpu_out) {
|
|
const size_t size = image.GetCompressedSize();
|
|
last_npixels = image.GetWidth() * image.GetHeight();
|
|
if (size == 0)
|
|
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.
|
|
cuda_err(cudaStreamSynchronize(*stream));
|
|
compressed_capacity = size + size / 4;
|
|
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{};
|
|
cuda_err(cub::DeviceScan::ExclusiveScan(nullptr, scan_temp_bytes, gpu_maps.get(), gpu_entries.get(),
|
|
ComposeTileMaps(), identity, tile_capacity, *stream));
|
|
gpu_scan_temp = CudaDevicePtr<uint8_t>(scan_temp_bytes);
|
|
}
|
|
|
|
// The upload waits for the previous frame's to have left the staging buffer.
|
|
cuda_err(cudaStreamSynchronize(*stream));
|
|
std::memcpy(host_compressed.get(), image.GetCompressed(), size);
|
|
cuda_err(cudaEventRecord(decode_start, *stream));
|
|
cuda_err(cudaMemcpyAsync(gpu_compressed, host_compressed.get(), size, cudaMemcpyHostToDevice, *stream));
|
|
|
|
const int threads = 128;
|
|
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.
|
|
CBFTileMap identity{};
|
|
for (int e = 0; e < 16; e++)
|
|
identity.exit[e] = static_cast<uint8_t>(e);
|
|
size_t temp_bytes = scan_temp_bytes;
|
|
cuda_err(cub::DeviceScan::ExclusiveScan(gpu_scan_temp.get(), temp_bytes, gpu_maps.get(), gpu_entries.get(),
|
|
ComposeTileMaps(), identity, ntiles, *stream));
|
|
|
|
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));
|
|
}
|
|
|
|
void CBFDecoderGPU::ThrowIfDecodeFailed() const {
|
|
const size_t total = host_total.get()[0];
|
|
if (total < last_npixels)
|
|
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
|
|
"miniCBF byte-offset stream ended after " + std::to_string(total)
|
|
+ " of " + std::to_string(last_npixels) + " pixels");
|
|
}
|
|
|
|
float CBFDecoderGPU::GetDecodeTime_s() const {
|
|
float ms = 0.0f;
|
|
if (cudaEventElapsedTime(&ms, decode_start, decode_stop) != cudaSuccess)
|
|
return 0.0f;
|
|
return ms * 1e-3f;
|
|
}
|