Files
Jungfraujoch/image_analysis/image_preprocessing/ImagePreprocessorGPU.cu
T
leonarski_fandjungfrau 4dc2534dbf
Build Packages / build:rpm (rocky9_sls9) (push) Successful in 18m57s
Build Packages / Unit tests (push) Skipped
Build Packages / build:windows:nocuda (push) Successful in 16m55s
Build Packages / build:windows:cuda (push) Successful in 18m48s
Build Packages / build:viewer-tgz:cpu (push) Successful in 13m10s
Build Packages / build:viewer-tgz:cuda (push) Successful in 14m45s
Build Packages / build:rpm (rocky8_nocuda) (push) Successful in 22m23s
Build Packages / build:rpm (rocky9_nocuda) (push) Successful in 20m12s
Build Packages / build:rpm (ubuntu2204_nocuda) (push) Successful in 23m7s
Build Packages / build:rpm (ubuntu2404_nocuda) (push) Successful in 20m43s
Build Packages / build:rpm (rocky8_sls9) (push) Successful in 23m9s
Build Packages / XDS test (durin plugin) (push) Successful in 12m26s
Build Packages / build:rpm (rocky9) (push) Successful in 24m58s
Build Packages / Generate python client (push) Successful in 50s
Build Packages / build:rpm (ubuntu2404) (push) Successful in 23m20s
Build Packages / Create release (push) Skipped
Build Packages / XDS test (JFJoch plugin) (push) Successful in 12m37s
Build Packages / build:rpm (rocky8) (push) Successful in 27m58s
Build Packages / build:rpm (ubuntu2204) (push) Successful in 25m38s
Build Packages / Build documentation (push) Successful in 59s
Build Packages / DIALS test (push) Successful in 23m16s
Build Packages / XDS test (neggia plugin) (push) Successful in 6m38s
v1.0.0.rc-162 (#72)
**Files written by Jungfraujoch now import correctly in DIALS, XDS and pyFAI.** A tilted detector, a grid scan, a still recorded at a goniometer position, and saturated or unreadable pixels were each described in a way that a third-party program acted on wrongly. If you process Jungfraujoch data outside Jungfraujoch, prefer this release to any earlier one.

* HDF5: the detector tilt (`rot1`/`rot2`/`rot3`) is exported correctly in the NXmx transformation chain; untilted geometries are unaffected.
* HDF5: a still recorded at a goniometer position is no longer read back as a single image, and a grid scan records a stationary spindle so a program that requires a rotation axis can open it.
* HDF5: the sample transformation chain is written in mounting order, with a Smargon head position told apart from the spindle, one entry per image, `module_offset` as a float unit vector, and `offset_units` on every offset.
* HDF5: saturated, underloaded and unreadable pixels are described so a downstream program masks them - `saturation_value`, `underload_value`, `error_value` and `bit_depth_readout` are written correctly, and a data file missing next to a VDS master reads as the error marker rather than as zero counts.
* HDF5: the rotation axis is read back under whatever name it carries, and `mirror_y` records whether the assembled image is mirrored in Y relative to the detector's raw readout.
* A grid scan and a goniometer axis can both be set; they are no longer alternatives.
* `images_per_file` is chosen from the acquisition when it is not given: a rotation sweep of at most 20000 images goes into a single data file, a grid scan splits on whole fast-axis rows, and stills and serial keep 1000.
* The writer refuses a stream whose start message declares a different pixel format than its images carry, and a DECTRIS detector sending signed images is no longer declared unsigned.
* The image stream can carry the sample transformation chain (`transformations`, in the END message); a producer that does not send it gets the same chain built by the writer.
* rugnux: fixing the space group with `-S` no longer prevents the lattice from being found - a lattice indexed in a different setting is reindexed into that group's own setting, and a run whose crystal does not have that group's lattice stops and names the cell it indexed as, rather than reporting statistics that cannot describe it.
* rugnux: the per-image resolution estimate now predicts the resolution the merged data reach rather than the highest-resolution spot found, and is reported as `SPOT_RESOLUTION_ESTIMATE`.
* rugnux: two runs of the same command on the same images produce the same merged intensities; the azimuthal profile written alongside them is not yet reproducible in the same way.
* rugnux: the offline lattice refinement is bounded by iterations rather than by a wall clock, so a loaded machine can no longer refine to a different lattice; a live acquisition keeps its real-time bound.
* rugnux: the detector-frame modulation correction is fitted on a grid spanning the detector, so whether it is applied no longer depends on how far integration reached.
* rugnux: the geometry pre-pass no longer writes `<prefix>_01.mtz`, `_01.cif`, `_01.hkl` and `_01_image.dat`; the refined second pass writes those files under `<prefix>`, and that is the result to use.
* rugnux: `_process.h5` describes the pixel format of the images it links to, and is written on a thread of its own.
* rugnux: the detector geometry is also logged in XDS's convention (`ORGX`/`ORGY`, detector axis vectors, rotation axis), so it can be compared with an XDS refinement.
* rugnux: an image integrated in pyFAI through the `.poni` file written by `--mode calibration` comes out with the correct azimuth, and the file declares pyFAI's `orientation`, which needs pyFAI 2024.01 or newer. Radial integration is unchanged.
* rugnux: a rotation run is substantially faster throughout - beam-stop detection, first-pass indexing, geometry refinement, integration, scaling and merging - and observations outside the scaling resolution range are dropped as they are ingested. The refined geometry, the space group chosen and the merged statistics are unchanged.
* Faster spot finding and indexing, on the broker as well as in rugnux; the spots found and the lattices indexed are unchanged.
* A run reserves substantially less GPU memory: nothing is allocated for buffers that are never read, and a worker builds only the engines it uses.
* rugnux: with `-N` left at its default the per-image loop of `--mode mx` uses at most 16 workers per GPU, rather than one per hardware thread; an explicit `-N` is obeyed as given.
* CUDA 12 builds now contain device code for Volta, so the RHEL 8 packages and the portable Linux `.tgz` run on a V100; the CUDA 13 artefacts (RHEL 9, Ubuntu, Windows) remain Turing and newer.
* The build resolves a single Eigen for the whole project, and refuses to configure if Ceres picks up a different one; a build that mixed two Eigen versions was undefined behaviour and crashed at -O2.
* Documentation: a security page, and the supported GPU generations and minimum NVIDIA driver version of every released artefact.

**Breaking change to OpenAPI** - regenerate the client (`jfjoch-client` 1.0.0-rc.162, `frontend/src/client`):
* `dataset_settings.images_per_file` is no longer `default: 1000` and no longer accepts `0`; it is optional, and its minimum is 1. A client sending `0` (previously "one file for the whole run") is now rejected - omit the field instead, which for a rotation sweep gives the same single file.
* `file_writer_format` now defaults to `NXmxVDS`, matching the server's own default and the layout recommended for DIALS, XDS and CrystFEL. A generated client that fills in schema defaults and does not set the format explicitly will write VDS masters where it previously wrote legacy ones; set `NXmxLegacy` explicitly to keep them.

---------

Co-authored-by: jungfrau <jungfrau@mx-aare-test.psi.ch>
Reviewed-on: #72
Co-authored-by: Filip Leonarski <filip.leonarski@psi.ch>
2026-08-25 08:21:39 +02:00

525 lines
25 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;
}
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];
}