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

Decode and accumulate the beam-stop projection on the GPU

Browse Source 
The pre-scan decompressed its frames on the host and folded them into a
per-pixel projection there. On a 16M-pixel detector that is 60 frames of 72 MB
to decompress and 20 bytes per pixel to read and write back per frame - about
40 GB of memory traffic - and it was the whole cost of the phase once the mask
was no longer the bottleneck.

Only the compressed chunk crosses PCIe now. BSLZ4DecoderGPU already exposes the
raw decoded bytes (Decode(), the path its own tests use), which is what this
needs: the projection is defined on the RAW STORED COUNTS with the pixel type's
sentinel skipped, not on the preprocessed image, so nothing here goes through
the preprocessor. Sums, maxima and counts are integers, so the device result is
identical to the host's rather than merely close.

Frames are folded in batches of four. The fold reads and writes the whole
accumulator whatever the batch holds, so per frame it was spending most of the
bandwidth on the accumulator rather than on the data; four is where that stops
mattering, and every frame beyond it is another full frame of device memory,
which costs more in cudaMalloc - device-synchronizing - than it saves.

The accumulator is built on a thread of its own. It allocates and clears
several hundred megabytes, and doing that in the constructor stalled the caller
before it had read its first frame.

Frames the device cannot take - anything but bitshuffle+LZ4 - still go to a host
shard, so a run mixing compressions needs no second code path, and a build
without CUDA is unchanged.

RotationScaleMergeGPU set the CUDA device in its constructor and never put it
back. CUDA's current device is per-thread, so that silently re-pinned the
calling thread for the rest of its life, and the destructor freed several
gigabytes against whatever device happened to be current by then - CudaDevicePtr
records no device of its own. Every entry point now sets the device on entry and
restores it on exit.

ParallelFor/ParallelChunks moved to common/ParallelFor.h; two files had copies
and a third wants them.

Measured on a 16M-pixel rotation dataset: pre-scan 4.78 s -> 2.37 s -> ~2.0 s,
shadow unchanged at 139126 pixels (22143 on a 2M-pixel dataset). Full 24-crystal
battery: same space group on all 24, none failed, 15m32s -> 14m49s.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
jungfrau
2026-08-23 08:19:20 -04:00
co-authored by Claude Opus 5
parent 20edf55063
commit ddf625d833
6 changed files with 419 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
image_analysis/beam_stop/ShadowFinder.cpp
+86 -1
View File
@@ -13,6 +13,10 @@
#include <queue>
#include <type_traits>
#include <spdlog/spdlog.h>
#include "../../common/CUDAWrapper.h"
#include "../../common/ParallelFor.h"
#include "../../common/JFJochException.h"
// A pixel is shadow when its background is below this fraction of the background it is
@@ -289,6 +293,14 @@ ShadowFinder::ShadowFinder(const DiffractionExperiment &experiment, const PixelM
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
"ShadowFinder: pixel mask does not match the detector");
SetShardCount(1);
#ifdef JFJOCH_USE_CUDA
if (get_gpu_count() > 0) {
const size_t npixels = static_cast<size_t>(width) * height;
gpu_pending = std::async(std::launch::async, [npixels] {
return std::make_unique<ShadowAccumulatorGPU>(npixels);
});
}
#endif
}
void ShadowFinder::SetShardCount(size_t n) {
@@ -304,7 +316,26 @@ void ShadowFinder::SetShardCount(size_t n) {
}
ShadowFinder::Projection ShadowFinder::Reduce() const {
if (shards.size() == 1)
#ifdef JFJOCH_USE_CUDA
// The device holds its own projection. Bring it back and let it take part in the fold below as
// one more shard; when every frame went to the GPU it is the whole answer.
Projection device;
if (Gpu() && gpu->GetFrameCount() > 0) {
gpu->Download(device.max_value, device.sum_value, device.valid_count);
device.frames = gpu->GetFrameCount();
bool host_empty = true;
for (const auto &p : shards)
host_empty = host_empty && (p.frames == 0);
if (host_empty)
return device;
}
#endif
if (shards.size() == 1
#ifdef JFJOCH_USE_CUDA
&& !(gpu && gpu->GetFrameCount() > 0)
#endif
)
return shards[0];
Projection out;
@@ -314,6 +345,9 @@ ShadowFinder::Projection ShadowFinder::Reduce() const {
out.valid_count.assign(npixels, 0);
for (const auto &p : shards)
out.frames += p.frames;
#ifdef JFJOCH_USE_CUDA
out.frames += device.frames;
#endif
// Each worker owns a slice of the pixels and folds every shard into it. The sums and counts are
// integers and a pixel is touched by one worker only, so the result is the same as folding them
@@ -328,6 +362,21 @@ ShadowFinder::Projection ShadowFinder::Reduce() const {
const size_t lo = t * chunk, hi = std::min(npixels, lo + chunk);
if (lo >= hi) break;
futures.emplace_back(std::async(std::launch::async, [&, lo, hi] {
#ifdef JFJOCH_USE_CUDA
const Projection *extra[1] = {&device};
for (const auto *pp : extra) {
const auto &p = *pp;
if (p.frames == 0) continue;
for (size_t i = lo; i < hi; i++) {
if (p.valid_count[i] == 0)
continue;
if (out.valid_count[i] == 0 || p.max_value[i] > out.max_value[i])
out.max_value[i] = p.max_value[i];
out.sum_value[i] += p.sum_value[i];
out.valid_count[i] += p.valid_count[i];
}
}
#endif
for (const auto &p : shards)
for (size_t i = lo; i < hi; i++) {
if (p.valid_count[i] == 0)
@@ -377,6 +426,21 @@ void ShadowFinder::AddImage(const DataMessage &data, std::vector<uint8_t> &buffe
throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
"ShadowFinder: shard out of range");
#ifdef JFJOCH_USE_CUDA
if (ShadowAccumulatorGPU::Supports(data.image) && Gpu()) {
// One device, so the frames queue here - but each is only a chunk upload plus two kernels,
// and nothing was decompressed on the host to get this far.
ShadowAccumulatorGPU *acc = Gpu();
std::unique_lock ul(gpu_mutex);
try {
acc->Add(data.image);
return;
} catch (const std::exception &e) {
spdlog::warn("Beam stop: GPU accumulate failed ({}), falling back to the host", e.what());
}
}
#endif
Projection &p = shards[shard];
const auto ptr = data.image.GetUncompressedPtr(buffer);
switch (data.image.GetMode()) {
@@ -392,10 +456,31 @@ void ShadowFinder::AddImage(const DataMessage &data, std::vector<uint8_t> &buffe
}
}
#ifdef JFJOCH_USE_CUDA
ShadowAccumulatorGPU *ShadowFinder::Gpu() const {
std::unique_lock ul(gpu_mutex);
if (gpu_pending.valid()) {
try {
gpu = gpu_pending.get();
} catch (const std::exception &e) {
// A GPU that cannot hold the projection is not a reason to fail: the host path gives
// the same answer, only slower.
spdlog::warn("Beam stop: GPU projection unavailable ({}), accumulating on the host",
e.what());
gpu.reset();
}
}
return gpu.get();
}
#endif
uint32_t ShadowFinder::GetFrameCount() const {
std::unique_lock ul(m);
uint32_t frames = 0;
for (const auto &p : shards) frames += p.frames;
#ifdef JFJOCH_USE_CUDA
if (gpu) frames += gpu->GetFrameCount();
#endif
return frames;
}