Files
Jungfraujoch/image_analysis/beam_stop/ShadowAccumulatorGPU.h
T
jungfrauandClaude Opus 5 6368c00173 Decode and accumulate the beam-stop projection on the GPU
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>
2026-08-15 18:38:26 -04:00

67 lines
3.0 KiB
C++

// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only
#pragma once
#include <cstdint>
#include <memory>
#include <vector>
#include "../../common/CompressedImage.h"
#include "../image_preprocessing/BSLZ4DecoderGPU.h"
#include "../indexing/CUDAMemHelpers.h"
// The beam-stop projection accumulated on the device: only the compressed chunk crosses PCIe, and
// both the decode and the per-pixel maximum / sum / count run on the GPU. The projection comes back
// once, when the mask is read.
//
// This accumulates the RAW STORED COUNTS, with the pixel type's sentinel skipped, exactly as the
// host ShadowFinder does - NOT the preprocessed image. The two are not the same thing and the
// shadow is defined on the raw one. Sums and counts are integers, so the device result is identical
// to the host's rather than merely close.
//
// Only bitshuffle+LZ4 images are handled; anything else is left to the host path.
class ShadowAccumulatorGPU {
const size_t npixels;
std::shared_ptr<CudaStream> stream;
std::unique_ptr<BSLZ4DecoderGPU> decoder;
// Frames are decoded into a batch and folded into the projection together. The fold reads and
// writes 20 bytes per pixel whatever the batch holds, so doing it once per frame would spend
// most of the bandwidth on the accumulator rather than on the data. Four is where that stops
// mattering - the fold is then a few percent of the phase - and every frame beyond it is a
// further full frame of device memory, which costs more in cudaMalloc (device-synchronizing)
// than it saves.
static constexpr int BATCH = 4;
CudaDevicePtr<uint8_t> raw; // BATCH decoded frames, still in their stored type
CudaDevicePtr<int64_t> gpu_max;
CudaDevicePtr<int64_t> gpu_sum;
CudaDevicePtr<uint32_t> gpu_count;
size_t frame_bytes = 0; // stride of one frame within `raw`
size_t raw_capacity = 0; // bytes actually allocated for `raw`
int pending = 0; // frames decoded but not yet folded in
CompressedImageMode pending_mode = CompressedImageMode::Uint32;
uint32_t frames = 0;
int blocks = 0;
void EnsureRawCapacity(size_t bytes_per_frame);
void FoldPending();
public:
explicit ShadowAccumulatorGPU(size_t npixels);
// True when this image can be decoded and accumulated on the device.
static bool Supports(const CompressedImage &image);
// Decode and accumulate one frame. Throws if the image is not one this can take (ask Supports
// first) or if the chunk does not decode.
void Add(const CompressedImage &image);
[[nodiscard]] uint32_t GetFrameCount() const { return frames; }
// Bring the projection back to the host, folding in whatever the last batch still holds. Cheap
// to call once; it moves 20 bytes per pixel.
void Download(std::vector<int64_t> &max_value, std::vector<int64_t> &sum_value,
std::vector<uint32_t> &valid_count);
};