A pooled CudaDevicePtr frees on the thread's allocation stream, not on the engine's stream, and the
pool may hand the memory to another engine - or, past its release threshold, unmap it - as soon as
that free is reached, which on an idle allocation stream is at once. An engine destroyed with work
still queued (FFTIndexerGPU after SearchCap's last DirectionsChanged upload, a spot finder between
DetectAt and Extract, a shadow accumulator after a pending fold, any engine on an exception path)
thus had kernels or copies writing memory that was someone else's or no longer mapped. Now:
- CudaStream synchronises before cudaStreamDestroy (destructor and move-assignment), which covers
engines whose own stream is declared after their buffers (FFTIndexerGPU, the gather buffer);
- every engine holding pooled buffers and a stream (shared or own, declared before the buffers)
synchronises it in its destructor; BraggIntegrationEngineGPU also before EnsureCapacity
reallocates, where a Run that threw leaves work queued.
A GPU failure that is handled no longer hides a lost context: ShadowFinder, BeamCenterFFT, the
rigid-body pool and model validation call cuda_throw_if_context_lost() before cuda_clear_error(),
as the device-decode fallbacks already did; RotationScaleMergeGPU's Alloc does so before waiting up
to ten minutes for GPU work beside it and then reporting a lost device as out of memory; and a
failed cudaMalloc says why. BeamCenterFFT logs the failure it used to drop silently, the
speculative geometry probe logs the exception it swallowed (its GPU fault was otherwise reported by
the merge beside it, under the merge's name), and RotationScaleMergeGPU's DeviceGuard no longer
throws from its destructor.
Only synchronisation and error paths change: p.hkl md5 and the MTZ data (gemmi) are identical to
the b530c2d full battery on myob_x10sa, cytc_x10sa, 8a1a, 9gdj, 11if, kdp_x10sa_20keV and 6z9g.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi
556 lines
27 KiB
Plaintext
556 lines
27 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 "ImagePreprocessorGPU.h"
|
|
#include "BSLZ4DecodeWarp.h"
|
|
#include "../../common/JFJochException.h"
|
|
|
|
namespace {
|
|
void cuda_err(cudaError_t val) {
|
|
if (val != cudaSuccess)
|
|
throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
|
|
}
|
|
}
|
|
|
|
template<class T>
|
|
__global__ void preprocess_kernel(
|
|
const T *__restrict__ input,
|
|
const uint8_t *__restrict__ mask,
|
|
int32_t *__restrict__ output,
|
|
ImageStatistics *__restrict__ stats,
|
|
T saturation_limit,
|
|
T err_value,
|
|
int npixels) {
|
|
// Shared block accumulators
|
|
__shared__ unsigned long long s_masked;
|
|
__shared__ unsigned long long s_saturated;
|
|
__shared__ unsigned long long s_error;
|
|
__shared__ long long s_max;
|
|
__shared__ long long s_min;
|
|
|
|
if (threadIdx.x == 0) {
|
|
s_masked = 0;
|
|
s_saturated = 0;
|
|
s_error = 0;
|
|
s_max = INT64_MIN;
|
|
s_min = INT64_MAX;
|
|
}
|
|
__syncthreads();
|
|
|
|
// Thread-local accumulators
|
|
unsigned long long local_masked = 0;
|
|
unsigned long long local_saturated = 0;
|
|
unsigned long long local_error = 0;
|
|
long long local_max = INT64_MIN;
|
|
long long local_min = INT64_MAX;
|
|
|
|
for (int i = blockIdx.x * blockDim.x + threadIdx.x;
|
|
i < npixels;
|
|
i += blockDim.x * gridDim.x) {
|
|
T v = input[i];
|
|
|
|
bool is_masked = mask[i];
|
|
// Error/invalid marker = the pixel type's extreme value (0xFFFFFFFF for EIGER uint32); tested
|
|
// before saturation, since for unsigned types the marker also exceeds saturation_limit (which is
|
|
// clipped to the HDF5 saturation_value). Priority: masked > error > saturated.
|
|
bool is_err = (v == err_value);
|
|
bool is_sat = !is_err && (v >= saturation_limit);
|
|
|
|
bool valid = !(is_masked || is_sat || is_err);
|
|
|
|
// Output
|
|
output[i] =
|
|
is_masked ? INT32_MIN : is_err ? INT32_MIN : is_sat ? INT32_MAX : (int32_t) v;
|
|
|
|
// Counters
|
|
local_masked += is_masked;
|
|
local_error += (!is_masked && is_err);
|
|
local_saturated += (!is_masked && !is_err && is_sat);
|
|
|
|
// Min/max only for valid
|
|
if (valid) {
|
|
int64_t val = (int64_t) v;
|
|
if (val > local_max) local_max = val;
|
|
if (val < local_min) local_min = val;
|
|
}
|
|
}
|
|
|
|
// Reduce to shared memory
|
|
atomicAdd(&s_masked, local_masked);
|
|
atomicAdd(&s_saturated, local_saturated);
|
|
atomicAdd(&s_error, local_error);
|
|
|
|
if (local_min <= local_max) {
|
|
atomicMax((long long *) &s_max, (long long) local_max);
|
|
atomicMin((long long *) &s_min, (long long) local_min);
|
|
}
|
|
|
|
__syncthreads();
|
|
|
|
// One thread writes block result
|
|
if (threadIdx.x == 0) {
|
|
atomicAdd(&stats->masked_pixel_count, s_masked);
|
|
atomicAdd(&stats->saturated_pixel_count, s_saturated);
|
|
atomicAdd(&stats->error_pixel_count, s_error);
|
|
|
|
atomicMax((long long *) &stats->max_value, (long long) s_max);
|
|
atomicMin((long long *) &stats->min_value, (long long) s_min);
|
|
}
|
|
}
|
|
|
|
// The per-pixel decision preprocess_kernel makes, in a form the fused kernel can reuse so the two
|
|
// cannot drift apart. Priority: masked > error > saturated.
|
|
template<class T>
|
|
struct PreprocessAccum {
|
|
unsigned long long masked = 0, saturated = 0, error = 0;
|
|
long long max_v = INT64_MIN, min_v = INT64_MAX;
|
|
|
|
__device__ __forceinline__ int32_t Apply(T v, bool is_masked, T sat_value, T err_value) {
|
|
const bool is_err = (v == err_value);
|
|
const bool is_sat = !is_err && (v >= sat_value);
|
|
masked += is_masked;
|
|
error += (!is_masked && is_err);
|
|
saturated += (!is_masked && !is_err && is_sat);
|
|
if (!(is_masked || is_sat || is_err)) {
|
|
const int64_t val = (int64_t) v;
|
|
if (val > max_v) max_v = val;
|
|
if (val < min_v) min_v = val;
|
|
}
|
|
return is_masked ? INT32_MIN : is_err ? INT32_MIN : is_sat ? INT32_MAX : (int32_t) v;
|
|
}
|
|
};
|
|
|
|
// Reduce a block's thread-local accumulators into the image-wide statistics. Every accumulator is an
|
|
// integer, so the result does not depend on the order the warps arrive in.
|
|
//
|
|
// The reduction is within the warp and then straight to global, rather than through a shared-memory
|
|
// staging area. It has to be: the fused kernel below fills its shared budget to the byte with the
|
|
// bitshuffle block it decodes, and forty-odd bytes of accumulators on top of that cost it a whole
|
|
// resident CUDA block per SM. A thread that saw no valid pixel still carries INT64_MIN/INT64_MAX,
|
|
// which is the identity for max/min, so it needs no guard.
|
|
template<class T>
|
|
__device__ __forceinline__ void FlushStats(PreprocessAccum<T> &l, ImageStatistics *stats) {
|
|
for (int d = 16; d > 0; d >>= 1) {
|
|
l.masked += __shfl_down_sync(0xffffffff, l.masked, d);
|
|
l.saturated += __shfl_down_sync(0xffffffff, l.saturated, d);
|
|
l.error += __shfl_down_sync(0xffffffff, l.error, d);
|
|
l.max_v = max(l.max_v, __shfl_down_sync(0xffffffff, l.max_v, d));
|
|
l.min_v = min(l.min_v, __shfl_down_sync(0xffffffff, l.min_v, d));
|
|
}
|
|
|
|
if ((threadIdx.x & 31) == 0) {
|
|
atomicAdd(&stats->masked_pixel_count, l.masked);
|
|
atomicAdd(&stats->saturated_pixel_count, l.saturated);
|
|
atomicAdd(&stats->error_pixel_count, l.error);
|
|
atomicMax((long long *) &stats->max_value, l.max_v);
|
|
atomicMin((long long *) &stats->min_value, l.min_v);
|
|
}
|
|
}
|
|
|
|
__device__ __forceinline__ uint64_t transpose8_fused(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 un-transpose and the preprocessing of ONE bitshuffle block, mirroring bitshuf_decode_block.
|
|
// 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 can emit 8 finished int32 pixels - the
|
|
// decompressed image never has to exist in device memory at all. That removes a full-frame buffer
|
|
// per worker and a full-frame write plus read from the pipeline.
|
|
//
|
|
// `in` is the block's bitshuffled bytes: device memory when the LZ4 pass ran in its own kernel, and
|
|
// shared memory when it ran in this one. A generic pointer covers both, so the two routes cannot
|
|
// produce different pixels from the same block.
|
|
template<class T, int ES>
|
|
__device__ __forceinline__ void untranspose_preprocess_block(const uint8_t *in, uint32_t size, uint32_t elem0,
|
|
const uint8_t *__restrict__ mask,
|
|
int32_t *__restrict__ out,
|
|
T sat_value, T err_value,
|
|
PreprocessAccum<T> &l) {
|
|
// The bytes are assembled in the unsigned counterpart of T - shifting a byte into the top of a
|
|
// signed type overflows it - and converted back at the end, which C++20 defines as two's
|
|
// complement reinterpretation. That is exactly what the byte order in the file means.
|
|
using U = typename std::make_unsigned<T>::type;
|
|
|
|
const uint32_t n = size / 8;
|
|
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_fused(a);
|
|
}
|
|
int32_t o[8];
|
|
#pragma unroll
|
|
for (int k = 0; k < 8; k++) {
|
|
U uv = 0;
|
|
#pragma unroll
|
|
for (int p = 0; p < ES; p++) uv |= (U)((U)((x[p] >> (8 * k)) & 0xff) << (8 * p));
|
|
o[k] = l.Apply((T) uv, mask[elem0 + i * 8 + k] != 0, sat_value, err_value);
|
|
}
|
|
// elem0 is a multiple of 8 (bitshuffle blocks are), so this is 32-byte aligned.
|
|
int4 *dst = reinterpret_cast<int4 *>(out + elem0 + i * 8);
|
|
dst[0] = make_int4(o[0], o[1], o[2], o[3]);
|
|
dst[1] = make_int4(o[4], o[5], o[6], o[7]);
|
|
}
|
|
}
|
|
|
|
// The handful of elements bitshuffle stores verbatim rather than in a block. They are already on the
|
|
// device inside the uploaded chunk, so they only need the per-pixel decision, not the un-transpose.
|
|
template<class T, int ES>
|
|
__device__ __forceinline__ void preprocess_tail(const uint8_t *__restrict__ tail_src, uint32_t tail_elems,
|
|
uint32_t tail_elem0, const uint8_t *__restrict__ mask,
|
|
int32_t *__restrict__ out, T sat_value, T err_value,
|
|
PreprocessAccum<T> &l) {
|
|
using U = typename std::make_unsigned<T>::type;
|
|
if (threadIdx.x < tail_elems) {
|
|
U uv = 0;
|
|
#pragma unroll
|
|
for (int p = 0; p < ES; p++)
|
|
uv |= (U)((U)tail_src[threadIdx.x * ES + p] << (8 * p));
|
|
const T v = (T) uv;
|
|
out[tail_elem0 + threadIdx.x] = l.Apply(v, mask[tail_elem0 + threadIdx.x] != 0, sat_value, err_value);
|
|
}
|
|
}
|
|
|
|
// One CUDA block per bitshuffle block, over blocks the LZ4 kernel has already decoded. The last CUDA
|
|
// block (blockIdx.x == nblocks) finishes the verbatim tail.
|
|
template<class T, int ES>
|
|
__global__ __launch_bounds__(256) void untranspose_preprocess_kernel(
|
|
const uint8_t *__restrict__ shuffled,
|
|
const BSLZ4BlockDesc *__restrict__ desc,
|
|
const uint8_t *__restrict__ mask,
|
|
int32_t *__restrict__ out,
|
|
ImageStatistics *__restrict__ stats,
|
|
T sat_value, T err_value, int nblocks,
|
|
const uint8_t *__restrict__ tail_src, uint32_t tail_elems, uint32_t tail_elem0) {
|
|
PreprocessAccum<T> l;
|
|
|
|
if (blockIdx.x == nblocks)
|
|
preprocess_tail<T, ES>(tail_src, tail_elems, tail_elem0, mask, out, sat_value, err_value, l);
|
|
else
|
|
untranspose_preprocess_block<T, ES>(shuffled + desc[blockIdx.x].out_off, desc[blockIdx.x].nelem,
|
|
desc[blockIdx.x].out_off / ES, mask, out,
|
|
sat_value, err_value, l);
|
|
FlushStats<T>(l, stats);
|
|
}
|
|
|
|
// The same thing with the LZ4 decode folded in as well: one CUDA block owns one bitshuffle block
|
|
// from the compressed payload all the way to finished pixels. Its first warp decodes the payload
|
|
// into a shared-memory buffer the size of one block, and then the whole CUDA block un-transposes and
|
|
// preprocesses out of that buffer. Nothing of the block reaches device memory but the pixels.
|
|
//
|
|
// The saved bandwidth is the smaller half of it - 72 MB of bitshuffled bytes per 18 Mpx frame stop
|
|
// being written and read back. What the shared buffer is really for is the LZ4 copy loop: a match
|
|
// sources bytes that other lanes of the warp wrote a few sequences earlier, so every copy step is a
|
|
// dependent round trip to wherever the output lives, and there are a few hundred of them per block.
|
|
// In device memory that round trip is hundreds of cycles, which is why the standalone decode runs at
|
|
// a fraction of the streaming rate this hardware reaches on a plain pass over an image; in shared
|
|
// memory it is tens of cycles.
|
|
//
|
|
// It is paid for in residency. A whole bitshuffle block of shared memory per CUDA block means an SM
|
|
// holds only as many concurrent decoders as its shared memory divides into: four on a T4 at the
|
|
// 16 kB blocks 32-bit detector data comes in, against the thirty-two warps the standalone kernel
|
|
// keeps in flight. So this trades decode parallelism away for decode latency, and which way that
|
|
// comes out is a measurement rather than an argument.
|
|
template<class T, int ES>
|
|
__global__ void decode_untranspose_preprocess_kernel(
|
|
const uint8_t *__restrict__ src,
|
|
const BSLZ4BlockDesc *__restrict__ desc,
|
|
uint32_t *__restrict__ status,
|
|
const uint8_t *__restrict__ mask,
|
|
int32_t *__restrict__ out,
|
|
ImageStatistics *__restrict__ stats,
|
|
T sat_value, T err_value, int nblocks,
|
|
const uint8_t *__restrict__ tail_src, uint32_t tail_elems, uint32_t tail_elem0) {
|
|
extern __shared__ uint8_t s_shuffled[];
|
|
PreprocessAccum<T> l;
|
|
|
|
if (blockIdx.x == nblocks) {
|
|
preprocess_tail<T, ES>(tail_src, tail_elems, tail_elem0, mask, out, sat_value, err_value, l);
|
|
FlushStats<T>(l, stats);
|
|
return;
|
|
}
|
|
|
|
// An LZ4 match never reaches back past the start of its own bitshuffle block - the decoder
|
|
// rejects an offset larger than what the block has written - so one block's worth of shared
|
|
// memory is the whole of what the decode can address.
|
|
const uint32_t size = desc[blockIdx.x].nelem; // bytes per plane
|
|
if (threadIdx.x < 32) {
|
|
const uint8_t *const ip = src + desc[blockIdx.x].in_off;
|
|
const bool ok = lz4_decode_block_warp(ip, ip + desc[blockIdx.x].in_len,
|
|
s_shuffled, s_shuffled + size * ES, threadIdx.x);
|
|
if (threadIdx.x == 0 && !ok)
|
|
atomicExch(status, 1u);
|
|
}
|
|
__syncthreads();
|
|
|
|
untranspose_preprocess_block<T, ES>(s_shuffled, size, desc[blockIdx.x].out_off / ES, mask, out,
|
|
sat_value, err_value, l);
|
|
FlushStats<T>(l, stats);
|
|
}
|
|
|
|
// Largest bitshuffle block the fused kernel takes. A CUDA block holds one whole block in shared
|
|
// memory, so this figure is directly what limits residency: an SM's shared memory divided by it is
|
|
// how many blocks the SM can decode at once - four on a T4 at 16 kB - and past that there is too
|
|
// little parallelism left to cover even a shared-memory latency. Both writers this pipeline reads
|
|
// stay at or under it: our own compressor targets 16 kB of block whatever the pixel depth, and the
|
|
// EIGER files use 4096-element blocks, which is 16 kB at 32 bits and less below that. Anything
|
|
// larger takes the LZ4 kernel and the un-transposing preprocessor instead.
|
|
constexpr uint32_t FUSED_MAX_BLOCK_BYTES = 16384;
|
|
|
|
ImagePreprocessorGPU::ImagePreprocessorGPU(const DiffractionExperiment &experiment, const PixelMask &mask,
|
|
std::shared_ptr<CudaStream> stream, bool copy_image_to_host)
|
|
: ImagePreprocessor(experiment),
|
|
stream(stream),
|
|
copy_image_to_host(copy_image_to_host),
|
|
gpu_stats(1),
|
|
cpu_stats(1),
|
|
cpu_stats_reg(cpu_stats) {
|
|
// Setup mask. The byte-per-pixel form and its checksum come from the PixelMask, which derives them
|
|
// whenever the mask changes: they are the same for every worker, and deriving them walks every
|
|
// pixel of the detector - 18 million of them on a 16 Mpx one, per engine, with an engine built per
|
|
// worker per pass. The table is then uploaded once per GPU and shared by the engines on it.
|
|
gpu_mask = SharedDeviceTable(mask.GetBinaryMask().data(), npixels, mask.GetBinaryMask().data(),
|
|
mask.GetBinaryMaskChecksum(), *stream);
|
|
|
|
// Setup GPU settings. The current device, not device 0: workers are pinned round-robin across GPUs,
|
|
// so device 0's SM count can belong to a different card than the one these kernels launch on.
|
|
int device = 0;
|
|
cuda_err(cudaGetDevice(&device));
|
|
cudaDeviceProp prop{};
|
|
cuda_err(cudaGetDeviceProperties(&prop, device));
|
|
|
|
threads = 128;
|
|
blocks = 4 * prop.multiProcessorCount;
|
|
}
|
|
|
|
float ImagePreprocessorGPU::GetLastDecompressionTime_s() const {
|
|
return bslz4_decoder ? bslz4_decoder->GetDecodeTime_s() : 0.0f;
|
|
}
|
|
|
|
ImagePreprocessorGPU::~ImagePreprocessorGPU() {
|
|
// Unchecked, as everywhere else in this teardown path: a destructor is noexcept and the region
|
|
// may already be gone with the device. The stream is waited for first: its pooled buffers are
|
|
// freed next (see CudaDevicePtr), and the region may still be the source of a copy.
|
|
if (stream)
|
|
cudaStreamSynchronize(*stream);
|
|
if (pinned_input)
|
|
cudaHostUnregister(const_cast<void *>(pinned_input));
|
|
}
|
|
|
|
void ImagePreprocessorGPU::PinInputRegion(const void *ptr, size_t bytes) {
|
|
if (ptr == pinned_input && bytes == pinned_input_bytes)
|
|
return;
|
|
if (pinned_input)
|
|
cuda_err(cudaHostUnregister(const_cast<void *>(pinned_input)));
|
|
pinned_input = nullptr;
|
|
pinned_input_bytes = 0;
|
|
// The buffer may already be page-locked by whoever owns it - the receiver registers its image
|
|
// buffer once and hands out ranges inside it - and registering a range twice is an error rather
|
|
// than a no-op. There is nothing to do in that case: the upload is already a DMA, and the
|
|
// region is not ours to unregister later, so it is deliberately not remembered.
|
|
const auto err = cudaHostRegister(const_cast<void *>(ptr), bytes, cudaHostRegisterDefault);
|
|
if (err == cudaErrorHostMemoryAlreadyRegistered) {
|
|
cudaGetLastError();
|
|
return;
|
|
}
|
|
cuda_err(err);
|
|
pinned_input = ptr;
|
|
pinned_input_bytes = bytes;
|
|
}
|
|
|
|
void ImagePreprocessorGPU::PinInputBuffer(std::vector<uint8_t> &buffer, size_t size) {
|
|
if (buffer.size() == size)
|
|
return;
|
|
// Unregister before the resize, which can move the buffer.
|
|
input_reg.unregister();
|
|
buffer.resize(size);
|
|
input_reg.rebind(buffer);
|
|
}
|
|
|
|
ImageStatistics ImagePreprocessorGPU::Analyze(ImagePreprocessorBuffer &processed_image, const uint8_t *image_ptr,
|
|
CompressedImageMode image_mode) {
|
|
switch (image_mode) {
|
|
case CompressedImageMode::Int8:
|
|
return Analyze<int8_t>(processed_image, image_ptr, INT8_MIN, INT8_MAX);
|
|
case CompressedImageMode::Int16:
|
|
return Analyze<int16_t>(processed_image, image_ptr, INT16_MIN, INT16_MAX);
|
|
case CompressedImageMode::Int32:
|
|
return Analyze<int32_t>(processed_image, image_ptr, INT32_MIN, INT32_MAX);
|
|
case CompressedImageMode::Uint8:
|
|
return Analyze<uint8_t>(processed_image, image_ptr, UINT8_MAX, UINT8_MAX);
|
|
case CompressedImageMode::Uint16:
|
|
return Analyze<uint16_t>(processed_image, image_ptr, UINT16_MAX, UINT16_MAX);
|
|
case CompressedImageMode::Uint32:
|
|
return Analyze<uint32_t>(processed_image, image_ptr, UINT32_MAX, UINT32_MAX);
|
|
default:
|
|
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, "RGB/float mode not supported");
|
|
}
|
|
}
|
|
|
|
bool ImagePreprocessorGPU::AnalyzeCompressed(ImagePreprocessorBuffer &processed_image,
|
|
const CompressedImage &image,
|
|
ImageStatistics &stats) {
|
|
if (!BSLZ4DecoderGPU::Supports(image))
|
|
return false; // caller decompresses on the host and uses Analyze()
|
|
if (image.GetUncompressedSize() != npixels * image.GetByteDepth())
|
|
return false;
|
|
|
|
if (!bslz4_decoder)
|
|
// Sized for the depth this image actually has, not for the widest one there could be: on
|
|
// 16-bit data the difference is half of a full frame per worker thread.
|
|
bslz4_decoder = std::make_unique<BSLZ4DecoderGPU>(image.GetUncompressedSize(), stream);
|
|
|
|
// Everything from the compressed chunk to finished int32 pixels in ONE kernel, as long as a
|
|
// bitshuffle block fits the shared-memory buffer it decodes into. Neither the bitshuffled bytes
|
|
// nor the decompressed image is ever materialised. A block too large for that budget takes the
|
|
// older route: the LZ4 kernel writes the shuffled image, and the un-transposing preprocessor
|
|
// reads it back.
|
|
const uint32_t block_bytes = BSLZ4DecoderGPU::BlockBytes(image);
|
|
const BSLZ4ShuffledImage shuffled =
|
|
(block_bytes > 0 && block_bytes <= FUSED_MAX_BLOCK_BYTES)
|
|
? bslz4_decoder->UploadCompressed(image)
|
|
: bslz4_decoder->DecodeShuffled(image);
|
|
|
|
switch (image.GetMode()) {
|
|
case CompressedImageMode::Int8:
|
|
stats = UntransposeAndAnalyze<int8_t, 1>(processed_image, shuffled, INT8_MIN, INT8_MAX); return true;
|
|
case CompressedImageMode::Uint8:
|
|
stats = UntransposeAndAnalyze<uint8_t, 1>(processed_image, shuffled, UINT8_MAX, UINT8_MAX); return true;
|
|
case CompressedImageMode::Int16:
|
|
stats = UntransposeAndAnalyze<int16_t, 2>(processed_image, shuffled, INT16_MIN, INT16_MAX); return true;
|
|
case CompressedImageMode::Uint16:
|
|
stats = UntransposeAndAnalyze<uint16_t, 2>(processed_image, shuffled, UINT16_MAX, UINT16_MAX); return true;
|
|
case CompressedImageMode::Int32:
|
|
stats = UntransposeAndAnalyze<int32_t, 4>(processed_image, shuffled, INT32_MIN, INT32_MAX); return true;
|
|
case CompressedImageMode::Uint32:
|
|
stats = UntransposeAndAnalyze<uint32_t, 4>(processed_image, shuffled, UINT32_MAX, UINT32_MAX); return true;
|
|
default:
|
|
return false; // Supports() already excludes these; belt and braces
|
|
}
|
|
}
|
|
|
|
// The device-decode counterpart of AnalyzeOnDevice: same per-pixel decision, same statistics, but
|
|
// fed from the compressed chunk rather than from a decompressed image. Which kernel does it is what
|
|
// UploadCompressed() left behind - a shuffled image to read, or a chunk still to decode.
|
|
template<class T, int ES>
|
|
ImageStatistics ImagePreprocessorGPU::UntransposeAndAnalyze(ImagePreprocessorBuffer &processed_image,
|
|
const BSLZ4ShuffledImage &shuffled,
|
|
T err_value, T sat_value) {
|
|
if (sat_value > saturation_limit)
|
|
sat_value = static_cast<T>(saturation_limit);
|
|
|
|
cpu_stats[0] = ImageStatistics{.max_value = INT64_MIN, .min_value = INT64_MAX};
|
|
cuda_err(cudaMemcpyAsync(gpu_stats, cpu_stats.data(), sizeof(ImageStatistics), cudaMemcpyHostToDevice, *stream));
|
|
|
|
// One CUDA block per bitshuffle block, plus one for the verbatim tail when there is one.
|
|
const int nb = shuffled.nblocks + (shuffled.tail_elems > 0 ? 1 : 0);
|
|
if (shuffled.shuffled) {
|
|
untranspose_preprocess_kernel<T, ES> <<< nb, 256, 0, *stream >>>(
|
|
shuffled.shuffled,
|
|
shuffled.desc,
|
|
gpu_mask->get(),
|
|
processed_image.getGPUBuffer(),
|
|
gpu_stats,
|
|
sat_value,
|
|
err_value,
|
|
shuffled.nblocks,
|
|
shuffled.tail_src,
|
|
shuffled.tail_elems,
|
|
shuffled.tail_elem0);
|
|
cuda_err(cudaGetLastError());
|
|
} else {
|
|
// One warp per 2 kB of bitshuffle block. Shared memory is what caps how many CUDA blocks an
|
|
// SM can hold, and a T4 has 2 kB of it per warp slot (64 kB against 32 warps), so this is
|
|
// the shape that fills the SM instead of leaving warp slots no resident block can claim.
|
|
const int threads = std::clamp<int>(shuffled.block_bytes / 64, 32, 1024);
|
|
decode_untranspose_preprocess_kernel<T, ES> <<< nb, threads, shuffled.block_bytes, *stream >>>(
|
|
shuffled.compressed,
|
|
shuffled.desc,
|
|
shuffled.status,
|
|
gpu_mask->get(),
|
|
processed_image.getGPUBuffer(),
|
|
gpu_stats,
|
|
sat_value,
|
|
err_value,
|
|
shuffled.nblocks,
|
|
shuffled.tail_src,
|
|
shuffled.tail_elems,
|
|
shuffled.tail_elem0);
|
|
cuda_err(cudaGetLastError());
|
|
bslz4_decoder->QueueDecodeStatus();
|
|
}
|
|
|
|
if (copy_image_to_host)
|
|
cuda_err(cudaMemcpyAsync(processed_image.data(), processed_image.getGPUBuffer(), npixels * sizeof(int32_t), cudaMemcpyDeviceToHost, *stream));
|
|
cuda_err(cudaMemcpyAsync(cpu_stats.data(), gpu_stats, sizeof(ImageStatistics), cudaMemcpyDeviceToHost, *stream));
|
|
|
|
// Check the synchronise: this path decodes bytes we did not produce, and cpu_stats still holds the
|
|
// sentinel the host wrote above, so an unchecked failure here returns it as if it were a real
|
|
// measurement and the caller's fallback to the host decoder never fires.
|
|
cuda_err(cudaStreamSynchronize(*stream));
|
|
|
|
// Only now can the device tell us whether every block actually decoded.
|
|
bslz4_decoder->ThrowIfDecodeFailed();
|
|
|
|
return cpu_stats[0];
|
|
}
|
|
|
|
template<class T>
|
|
ImageStatistics ImagePreprocessorGPU::Analyze(ImagePreprocessorBuffer &processed_image,
|
|
const uint8_t *input,
|
|
T err_value,
|
|
T sat_value) {
|
|
// Allocated here rather than in the constructor: only the host-upload path needs it, and the
|
|
// device-decode path - which is what BSHUF_LZ4 images take - never touches it. Overshoot to
|
|
// 4 bytes per pixel so a 1- or 2-byte image fits the same buffer.
|
|
if (!gpu_decompressed_image.get())
|
|
gpu_decompressed_image = CudaDevicePtr<uint8_t>(npixels * sizeof(uint32_t));
|
|
|
|
// On this engine's own stream, not the NULL stream: a NULL-stream copy implicitly synchronises with
|
|
// every blocking stream in the process, which serialised all workers behind whichever one was
|
|
// uploading. The stream is synchronised at the end of this function, so the ordering is unchanged.
|
|
cuda_err(cudaMemcpyAsync(gpu_decompressed_image, input, npixels * sizeof(T), cudaMemcpyHostToDevice, *stream));
|
|
return AnalyzeOnDevice<T>(processed_image, err_value, sat_value);
|
|
}
|
|
|
|
// Everything after the image is on the device, shared by the host-upload and the device-decode
|
|
// entry points so the two cannot drift apart.
|
|
template<class T>
|
|
ImageStatistics ImagePreprocessorGPU::AnalyzeOnDevice(ImagePreprocessorBuffer &processed_image,
|
|
T err_value, T sat_value) {
|
|
if (sat_value > saturation_limit)
|
|
sat_value = static_cast<T>(saturation_limit);
|
|
|
|
cpu_stats[0] = ImageStatistics{.max_value = INT64_MIN, .min_value = INT64_MAX};
|
|
cuda_err(cudaMemcpyAsync(gpu_stats, cpu_stats.data(), sizeof(ImageStatistics), cudaMemcpyHostToDevice, *stream));
|
|
preprocess_kernel<T> <<< blocks, threads, 0, *stream >>>(
|
|
reinterpret_cast<const T *>(gpu_decompressed_image.get()),
|
|
gpu_mask->get(),
|
|
processed_image.getGPUBuffer(),
|
|
gpu_stats,
|
|
sat_value,
|
|
err_value,
|
|
npixels);
|
|
cuda_err(cudaGetLastError());
|
|
// The preprocessed image is 4 bytes per pixel - by far the largest transfer here - and every GPU
|
|
// engine reads it straight from the device buffer, so it only comes back when a CPU engine needs it.
|
|
if (copy_image_to_host)
|
|
cuda_err(cudaMemcpyAsync(processed_image.data(), processed_image.getGPUBuffer(), npixels * sizeof(int32_t), cudaMemcpyDeviceToHost, *stream));
|
|
cuda_err(cudaMemcpyAsync(cpu_stats.data(), gpu_stats, sizeof(ImageStatistics), cudaMemcpyDeviceToHost, *stream));
|
|
|
|
cuda_err(cudaStreamSynchronize(*stream));
|
|
|
|
return cpu_stats[0];
|
|
}
|