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>

image_analysis/CMakeLists.txt

@@ -28,6 +28,8 @@ ADD_LIBRARY(JFJochImageAnalysis STATIC
         dark_mask_analysis/DarkMaskAnalysis.h
         beam_stop/ShadowFinder.cpp
         beam_stop/ShadowFinder.h
+        $<$<BOOL:${JFJOCH_USE_CUDA}>:beam_stop/ShadowAccumulatorGPU.cu>
+        beam_stop/ShadowAccumulatorGPU.h
         rotation_indexer/RotationIndexer.cpp
         rotation_indexer/RotationIndexer.h
         WriteReflections.cpp

image_analysis/beam_stop/ShadowAccumulatorGPU.cu

@@ -0,0 +1,186 @@
+// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
+// SPDX-License-Identifier: GPL-3.0-only
+
+#include "ShadowAccumulatorGPU.h"
+
+#include <limits>
+#include <type_traits>
+
+#include "../../common/CUDAWrapper.h"
+#include "../../common/JFJochException.h"
+
+inline void cuda_err(cudaError_t val) {
+    if (val != cudaSuccess)
+        throw JFJochException(JFJochExceptionCategory::GPUCUDAError, cudaGetErrorString(val));
+}
+
+namespace {
+
+// One frame folded into the projection. The sentinel test is predicated rather than branched: a
+// warp straddling a module gap then costs the same as one that does not, against a per-frame memory
+// bill of several hundred megabytes.
+//
+// The update is written exactly as the host writes it, including the "first value wins whatever it
+// is" rule for the maximum - a pixel that has never been counted holds 0, which a genuinely negative
+// maximum has to be allowed to replace. Each pixel is owned by one thread and the frames are
+// separate launches on one stream, so there is no race and no atomic.
+template <class T>
+__global__ void accumulate_kernel(const T *__restrict__ frames, int nframes, size_t frame_stride,
+                                  int64_t *__restrict__ max_value, int64_t *__restrict__ sum_value,
+                                  uint32_t *__restrict__ valid_count, size_t npixels, T masked) {
+    for (size_t i = blockIdx.x * static_cast<size_t>(blockDim.x) + threadIdx.x; i < npixels;
+         i += static_cast<size_t>(blockDim.x) * gridDim.x) {
+        int64_t mx = max_value[i];
+        int64_t sm = sum_value[i];
+        uint32_t c = valid_count[i];
+        for (int k = 0; k < nframes; k++) {
+            const T v = frames[k * frame_stride + i];
+            if (v == masked)
+                continue;
+            const int64_t vi = static_cast<int64_t>(v);
+            if (c == 0 || vi > mx)
+                mx = vi;
+            sm += vi;
+            c++;
+        }
+        max_value[i] = mx;
+        sum_value[i] = sm;
+        valid_count[i] = c;
+    }
+}
+
+template <class T>
+void launch(const uint8_t *raw, int nframes, size_t frame_bytes, int64_t *max_value,
+            int64_t *sum_value, uint32_t *valid_count, size_t npixels, int blocks,
+            cudaStream_t stream) {
+    T masked;
+    if constexpr (std::is_signed_v<T>)
+        masked = std::numeric_limits<T>::min();
+    else
+        masked = std::numeric_limits<T>::max();
+    accumulate_kernel<T><<<blocks, 256, 0, stream>>>(reinterpret_cast<const T *>(raw), nframes,
+                                                     frame_bytes / sizeof(T), max_value, sum_value,
+                                                     valid_count, npixels, masked);
+}
+
+}  // namespace
+
+ShadowAccumulatorGPU::ShadowAccumulatorGPU(size_t npixels)
+    : npixels(npixels),
+      stream(std::make_shared<CudaStream>()),
+      gpu_max(npixels),
+      gpu_sum(npixels),
+      gpu_count(npixels) {
+    cuda_err(cudaMemsetAsync(gpu_max, 0, sizeof(int64_t) * npixels, *stream));
+    cuda_err(cudaMemsetAsync(gpu_sum, 0, sizeof(int64_t) * npixels, *stream));
+    cuda_err(cudaMemsetAsync(gpu_count, 0, sizeof(uint32_t) * npixels, *stream));
+
+    int device = 0;
+    cuda_err(cudaGetDevice(&device));
+    cudaDeviceProp prop{};
+    cuda_err(cudaGetDeviceProperties(&prop, device));
+    blocks = 8 * prop.multiProcessorCount;
+
+    // Size the batch for the widest pixel type there is, so nothing has to be allocated once frames
+    // start arriving - a cudaMalloc then would stall every worker behind it.
+    raw_capacity = npixels * sizeof(uint32_t) * BATCH;
+    raw = CudaDevicePtr<uint8_t>(raw_capacity);
+    cuda_err(cudaStreamSynchronize(*stream));
+}
+
+bool ShadowAccumulatorGPU::Supports(const CompressedImage &image) {
+    if (!BSLZ4DecoderGPU::Supports(image))
+        return false;
+    switch (image.GetMode()) {
+        case CompressedImageMode::Int8:
+        case CompressedImageMode::Uint8:
+        case CompressedImageMode::Int16:
+        case CompressedImageMode::Uint16:
+        case CompressedImageMode::Int32:
+        case CompressedImageMode::Uint32:
+            return true;
+        default:
+            return false;
+    }
+}
+
+void ShadowAccumulatorGPU::EnsureRawCapacity(size_t bytes_per_frame) {
+    frame_bytes = bytes_per_frame;
+    if (bytes_per_frame * BATCH <= raw_capacity)
+        return;
+    // The constructor already sized this for the widest pixel type, so in practice this only runs
+    // if that guess was too small. cudaMalloc and cudaFree both synchronise the whole device, which
+    // is why it is never done per frame.
+    FoldPending();
+    cuda_err(cudaStreamSynchronize(*stream));
+    raw_capacity = bytes_per_frame * BATCH;
+    raw = CudaDevicePtr<uint8_t>(raw_capacity);
+}
+
+// Fold the frames decoded so far into the projection. One pass over the accumulator for the whole
+// batch rather than one per frame.
+void ShadowAccumulatorGPU::FoldPending() {
+    if (pending == 0)
+        return;
+    switch (pending_mode) {
+        case CompressedImageMode::Int8:
+            launch<int8_t>(raw, pending, frame_bytes, gpu_max, gpu_sum, gpu_count, npixels, blocks, *stream); break;
+        case CompressedImageMode::Uint8:
+            launch<uint8_t>(raw, pending, frame_bytes, gpu_max, gpu_sum, gpu_count, npixels, blocks, *stream); break;
+        case CompressedImageMode::Int16:
+            launch<int16_t>(raw, pending, frame_bytes, gpu_max, gpu_sum, gpu_count, npixels, blocks, *stream); break;
+        case CompressedImageMode::Uint16:
+            launch<uint16_t>(raw, pending, frame_bytes, gpu_max, gpu_sum, gpu_count, npixels, blocks, *stream); break;
+        case CompressedImageMode::Int32:
+            launch<int32_t>(raw, pending, frame_bytes, gpu_max, gpu_sum, gpu_count, npixels, blocks, *stream); break;
+        case CompressedImageMode::Uint32:
+            launch<uint32_t>(raw, pending, frame_bytes, gpu_max, gpu_sum, gpu_count, npixels, blocks, *stream); break;
+        default:
+            throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
+                                  "ShadowAccumulatorGPU: unsupported image mode");
+    }
+    pending = 0;
+}
+
+void ShadowAccumulatorGPU::Add(const CompressedImage &image) {
+    if (!Supports(image))
+        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
+                              "ShadowAccumulatorGPU: image cannot be decoded on the device");
+    if (static_cast<size_t>(image.GetWidth()) * image.GetHeight() != npixels)
+        throw JFJochException(JFJochExceptionCategory::InputParameterInvalid,
+                              "ShadowAccumulatorGPU: image size does not match the detector");
+
+    if (!decoder)
+        decoder = std::make_unique<BSLZ4DecoderGPU>(image.GetUncompressedSize(), stream);
+    EnsureRawCapacity(image.GetUncompressedSize());
+    // A batch holds one pixel type; a change of depth mid-run closes the batch first.
+    if (pending > 0 && image.GetMode() != pending_mode)
+        FoldPending();
+    pending_mode = image.GetMode();
+
+    decoder->Decode(image, raw.get() + static_cast<size_t>(pending) * frame_bytes);
+    // The decode has to be known good before its frame is counted, and the flag only arrives once
+    // the stream has drained.
+    cuda_err(cudaStreamSynchronize(*stream));
+    decoder->ThrowIfDecodeFailed();
+
+    pending++;
+    frames++;
+    if (pending == BATCH)
+        FoldPending();
+}
+
+void ShadowAccumulatorGPU::Download(std::vector<int64_t> &max_value, std::vector<int64_t> &sum_value,
+                                    std::vector<uint32_t> &valid_count) {
+    FoldPending();
+    max_value.resize(npixels);
+    sum_value.resize(npixels);
+    valid_count.resize(npixels);
+    cuda_err(cudaMemcpyAsync(max_value.data(), gpu_max.get(), sizeof(int64_t) * npixels,
+                             cudaMemcpyDeviceToHost, *stream));
+    cuda_err(cudaMemcpyAsync(sum_value.data(), gpu_sum.get(), sizeof(int64_t) * npixels,
+                             cudaMemcpyDeviceToHost, *stream));
+    cuda_err(cudaMemcpyAsync(valid_count.data(), gpu_count.get(), sizeof(uint32_t) * npixels,
+                             cudaMemcpyDeviceToHost, *stream));
+    cuda_err(cudaStreamSynchronize(*stream));
+}

image_analysis/beam_stop/ShadowAccumulatorGPU.h

@@ -0,0 +1,66 @@
+// 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);
+};

image_analysis/beam_stop/ShadowFinder.cpp

@@ -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;
 }
 

image_analysis/beam_stop/ShadowFinder.h

@@ -4,6 +4,8 @@
 #pragma once
 
 #include <cstdint>
+#include <future>
+#include <memory>
 #include <mutex>
 #include <vector>
 
@@ -11,6 +13,9 @@
 #include "../../common/DiffractionExperiment.h"
 #include "../../common/JFJochMessages.h"
 #include "../../common/PixelMask.h"
+#ifdef JFJOCH_USE_CUDA
+#include "ShadowAccumulatorGPU.h"
+#endif
 
 // Finds the beam-stop shadow - the central disk and the holder arm - from a set of images,
 // mirroring the accumulate-then-finalize shape of DarkMaskAnalysis: feed frames with
@@ -53,6 +58,23 @@ class ShadowFinder {
     };
     std::vector<Projection> shards;
 
+#ifdef JFJOCH_USE_CUDA
+    // Present when a GPU is available. Frames it can decode are accumulated there instead of on the
+    // host - only the compressed chunk crosses PCIe - and its projection is folded in with the
+    // shards when the mask is read. Frames it cannot take (anything but bitshuffle+LZ4) still go to
+    // a host shard, so a run mixing compressions is handled without a second code path.
+    // Built on a thread of its own: it allocates and clears several hundred megabytes of device
+    // memory, and cudaMalloc synchronises the whole device, so doing it in the constructor would
+    // stall the caller before it has read its first frame. The first AddImage waits for it, by
+    // which time the reads have been running for a while.
+    mutable std::future<std::unique_ptr<ShadowAccumulatorGPU>> gpu_pending;
+    mutable std::unique_ptr<ShadowAccumulatorGPU> gpu;
+    mutable std::mutex gpu_mutex;
+
+    // The accumulator once its construction has finished, or null if there is none.
+    [[nodiscard]] ShadowAccumulatorGPU *Gpu() const;
+#endif
+
     template<class T> void Add(const T *ptr, Projection &p);
 
     // Sum the shards into one projection. max_value is only taken from a shard that actually

image_analysis/scale_merge/RotationScaleMergeGPU.cu

@@ -595,6 +595,7 @@ namespace {
 }
 
 struct RotationScaleMergeGPU::Impl {
+    int device = 0;   // the GPU this instance's buffers live on
     bool available = false;
     int n_obs = 0, n_frames = 0, n_groups = 0;
 
@@ -649,14 +650,44 @@ struct RotationScaleMergeGPU::Impl {
     CudaDevicePtr<int32_t> f_gperm, f_gstart, f_gcount;
 };
 
+// Set the device this instance's memory lives on for the duration of a call, and put the caller's
+// back afterwards. CUDA's current device is per-thread, so without this a single set_gpu() in the
+// constructor silently re-pins the calling thread for the rest of its life - and, worse, the
+// destructor would free several gigabytes against whatever device happened to be current then.
+// CudaDevicePtr records no device of its own, so every entry point needs this.
+namespace {
+struct DeviceGuard {
+    int prev = 0;
+    bool active = false;
+    explicit DeviceGuard(int device, bool enable) : active(enable) {
+        if (!active)
+            return;
+        cudaGetDevice(&prev);
+        if (prev != device)
+            set_gpu(device);
+    }
+    ~DeviceGuard() {
+        if (active)
+            set_gpu(prev);
+    }
+};
+}  // namespace
+
 RotationScaleMergeGPU::RotationScaleMergeGPU() : impl_(std::make_unique<Impl>()) {
     if (get_gpu_count() > 0) {
-        set_gpu(0);
+        // One instance, one device. It stays on device 0 for now - the merge is a single object and
+        // nothing else runs beside it - but it is recorded rather than assumed, so the guard below
+        // can put the caller's device back instead of leaving the thread moved.
+        impl_->device = 0;
+        DeviceGuard guard(impl_->device, true);
         impl_->available = true;
     }
 }
 
-RotationScaleMergeGPU::~RotationScaleMergeGPU() = default;
+RotationScaleMergeGPU::~RotationScaleMergeGPU() {
+    DeviceGuard guard(impl_->device, impl_->available);
+    impl_.reset();
+}
 
 bool RotationScaleMergeGPU::Available() const { return impl_->available; }
 
@@ -665,6 +696,7 @@ void RotationScaleMergeGPU::SetPartials(int n_obs, int n_frames,
                                         const float *partiality, const float *zeta, const uint8_t *on_ice,
                                         const int32_t *frame, const float *corr0,
                                         const int32_t *frame_start, const int32_t *frame_count) {
+    DeviceGuard guard(impl_->device, impl_->available);
     auto &d = *impl_;
     d.n_obs = n_obs;
     d.n_frames = n_frames;
@@ -683,6 +715,7 @@ void RotationScaleMergeGPU::SetPartials(int n_obs, int n_frames,
 void RotationScaleMergeGPU::SetGroups(int n_groups, const int32_t *group, const int32_t *group_perm,
                                       int n_group_perm, const int32_t *group_start,
                                       const int32_t *group_count) {
+    DeviceGuard guard(impl_->device, impl_->available);
     auto &d = *impl_;
     d.n_groups = n_groups;
     Upload(d.group, group, d.n_obs);
@@ -693,12 +726,14 @@ void RotationScaleMergeGPU::SetGroups(int n_groups, const int32_t *group, const
 }
 
 void RotationScaleMergeGPU::SetCorr(const float *corr) {
+    DeviceGuard guard(impl_->device, impl_->available);
     CudaCheck(cudaMemcpy(impl_->corr.get(), corr, size_t(impl_->n_obs) * sizeof(float),
                          cudaMemcpyHostToDevice), "upload corr");
 }
 
 void RotationScaleMergeGPU::ScalePartials(int iters, double robust_k, double min_partiality,
                                           bool /*has_d_min*/) {
+    DeviceGuard guard(impl_->device, impl_->available);
     auto &d = *impl_;
     CudaCheck(cudaMemset(d.scaled.get(), 0, size_t(d.n_frames) * sizeof(uint8_t)), "memset scaled");
     CudaCheck(cudaMemset(d.g.get(), 0, size_t(d.n_frames) * sizeof(double)), "memset g"); // unscaled g unused
@@ -721,11 +756,13 @@ void RotationScaleMergeGPU::ScalePartials(int iters, double robust_k, double min
 }
 
 void RotationScaleMergeGPU::GetCorr(float *corr_out) const {
+    DeviceGuard guard(impl_->device, impl_->available);
     CudaCheck(cudaMemcpy(corr_out, impl_->corr.get(), size_t(impl_->n_obs) * sizeof(float),
                          cudaMemcpyDeviceToHost), "download corr");
 }
 
 void RotationScaleMergeGPU::GetG(double *g_out, uint8_t *scaled_out) const {
+    DeviceGuard guard(impl_->device, impl_->available);
     CudaCheck(cudaMemcpy(g_out, impl_->g.get(), size_t(impl_->n_frames) * sizeof(double),
                          cudaMemcpyDeviceToHost), "download g");
     CudaCheck(cudaMemcpy(scaled_out, impl_->scaled.get(), size_t(impl_->n_frames) * sizeof(uint8_t),
@@ -733,6 +770,7 @@ void RotationScaleMergeGPU::GetG(double *g_out, uint8_t *scaled_out) const {
 }
 
 void RotationScaleMergeGPU::SetFrameCellOk(const uint8_t *frame_cell_ok) {
+    DeviceGuard guard(impl_->device, impl_->available);
     Upload(impl_->frame_cell_ok, frame_cell_ok, impl_->n_frames);
 }
 
@@ -741,6 +779,7 @@ void RotationScaleMergeGPU::SetFrameCellOk(const uint8_t *frame_cell_ok) {
 void RotationScaleMergeGPU::MergeEmSamples(bool for_search, double min_partiality,
                                            double *em_mean_out, int32_t *cnt_out, double *s2_out,
                                            double *I2_out, double *dev2_out, uint8_t *valid_out) {
+    DeviceGuard guard(impl_->device, impl_->available);
     auto &d = *impl_;
     const int ng = d.n_groups, nf = d.n_fulls;
     d.merge_for_search = for_search ? 1 : 0; d.merge_min_part = min_partiality;
@@ -783,6 +822,7 @@ void RotationScaleMergeGPU::MergeAccum(double error_model_a, double error_model_
                                        double *swI, double *sw, double *swIh0, double *swIh1,
                                        double *swh0, double *swh1, int32_t *nh0, int32_t *nh1, double *d_out,
                                        int32_t *rejected, uint8_t *rejected_obs) {
+    DeviceGuard guard(impl_->device, impl_->available);
     auto &d = *impl_;
     const int ng = d.n_groups;
     d.a_swI = CudaDevicePtr<double>(std::max(1, ng)); d.a_sw = CudaDevicePtr<double>(std::max(1, ng));
@@ -830,6 +870,7 @@ void RotationScaleMergeGPU::MergeAccum(double error_model_a, double error_model_
 
 void RotationScaleMergeGPU::MergeRmeas(const double *merged_I, double *absdev, double *sumI,
                                        int32_t *n, int32_t *nusable) {
+    DeviceGuard guard(impl_->device, impl_->available);
     auto &d = *impl_;
     const int ng = d.n_groups;
     Upload(d.merged_I, merged_I, ng);
@@ -856,6 +897,7 @@ void RotationScaleMergeGPU::MergeRmeas(const double *merged_I, double *absdev, d
 }
 
 void RotationScaleMergeGPU::SmoothCorr(const uint8_t *apply, const double *ratio) {
+    DeviceGuard guard(impl_->device, impl_->available);
     auto &d = *impl_;
     Upload(d.smooth_apply, apply, d.n_frames);
     Upload(d.smooth_ratio, ratio, d.n_frames);
@@ -867,6 +909,7 @@ void RotationScaleMergeGPU::SmoothCorr(const uint8_t *apply, const double *ratio
 }
 
 void RotationScaleMergeGPU::ComputePartialCC(double min_partiality, double *cc_out, int64_t *cc_n_out) {
+    DeviceGuard guard(impl_->device, impl_->available);
     auto &d = *impl_;
     const int grp_blocks = std::min(65535, (d.n_groups + BLK - 1) / BLK);
     // Post-smooth group means (reuse the scaling reduce; reads the resident, smoothed corr), then the
@@ -888,6 +931,7 @@ void RotationScaleMergeGPU::ComputePartialCC(double min_partiality, double *cc_o
 void RotationScaleMergeGPU::SetCombineInputs(const float *bkg, const float *var_bkg,
                                             const float *image_number, const float *d,
                                             const float *px, const float *py) {
+    DeviceGuard guard(impl_->device, impl_->available);
     auto &dd = *impl_;
     Upload(dd.bkg, bkg, dd.n_obs);
     Upload(dd.var_bkg, var_bkg, dd.n_obs);
@@ -900,6 +944,7 @@ void RotationScaleMergeGPU::SetCombineInputs(const float *bkg, const float *var_
 void RotationScaleMergeGPU::SetRawRuns(int n_runs, int n_perm, const int32_t *perm,
                                        const int32_t *rr_start, const int32_t *rr_count,
                                        const int32_t *rr_h, const int32_t *rr_k, const int32_t *rr_l) {
+    DeviceGuard guard(impl_->device, impl_->available);
     auto &d = *impl_;
     d.n_runs = n_runs;
     d.n_perm = n_perm;
@@ -917,6 +962,7 @@ void RotationScaleMergeGPU::SetRawRuns(int n_runs, int n_perm, const int32_t *pe
 
 int RotationScaleMergeGPU::Combine(const int32_t *rawrun_group, double min_partiality,
                                    double capture_uncertainty_coeff, double min_captured_fraction) {
+    DeviceGuard guard(impl_->device, impl_->available);
     auto &d = *impl_;
     CudaCheck(cudaMemcpy(d.rr_group.get(), rawrun_group, size_t(d.n_runs) * sizeof(int32_t),
                          cudaMemcpyHostToDevice), "upload rr_group");
@@ -982,6 +1028,7 @@ int RotationScaleMergeGPU::Combine(const int32_t *rawrun_group, double min_parti
 void RotationScaleMergeGPU::GetFulls(int32_t *h, int32_t *k, int32_t *l, float *I, float *sigma, float *d,
                                      float *image_number, int32_t *frame, uint8_t *on_ice,
                                      int32_t *group) const {
+    DeviceGuard guard(impl_->device, impl_->available);
     const auto &dd = *impl_;
     const size_t n = static_cast<size_t>(dd.n_fulls);
     if (n == 0) return;
@@ -997,6 +1044,7 @@ void RotationScaleMergeGPU::GetFulls(int32_t *h, int32_t *k, int32_t *l, float *
 }
 
 void RotationScaleMergeGPU::GetFullsKeys(int32_t *frame, int32_t *group) const {
+    DeviceGuard guard(impl_->device, impl_->available);
     const auto &d = *impl_;
     if (d.n_fulls == 0) return;
     const size_t bytes = size_t(d.n_fulls) * sizeof(int32_t);
@@ -1006,6 +1054,7 @@ void RotationScaleMergeGPU::GetFullsKeys(int32_t *frame, int32_t *group) const {
 
 void RotationScaleMergeGPU::SetFullsFrameCSR(const int32_t *frame_perm, int n_perm,
                                              const int32_t *frame_start, const int32_t *frame_count) {
+    DeviceGuard guard(impl_->device, impl_->available);
     auto &d = *impl_;
     Upload(d.f_frame_perm, frame_perm, n_perm);
     Upload(d.f_frame_start, frame_start, d.n_frames);
@@ -1014,6 +1063,7 @@ void RotationScaleMergeGPU::SetFullsFrameCSR(const int32_t *frame_perm, int n_pe
 
 void RotationScaleMergeGPU::SetFullsGroups(const int32_t *gperm, int n_gperm,
                                            const int32_t *gstart, const int32_t *gcount) {
+    DeviceGuard guard(impl_->device, impl_->available);
     auto &d = *impl_;
     Upload(d.f_gperm, gperm, n_gperm);
     Upload(d.f_gstart, gstart, d.n_groups);
@@ -1021,6 +1071,7 @@ void RotationScaleMergeGPU::SetFullsGroups(const int32_t *gperm, int n_gperm,
 }
 
 void RotationScaleMergeGPU::ScaleFulls(int iters, double robust_k, double min_partiality) {
+    DeviceGuard guard(impl_->device, impl_->available);
     auto &d = *impl_;
     const int nf = d.n_fulls;
     if (nf == 0) return;
@@ -1053,6 +1104,7 @@ void RotationScaleMergeGPU::ScaleFulls(int iters, double robust_k, double min_pa
 }
 
 void RotationScaleMergeGPU::GetFullsCorr(float *corr) const {
+    DeviceGuard guard(impl_->device, impl_->available);
     const auto &d = *impl_;
     if (d.n_fulls == 0) return;
     CudaCheck(cudaMemcpy(corr, d.f_corr.get(), size_t(d.n_fulls) * sizeof(float),
@@ -1060,6 +1112,7 @@ void RotationScaleMergeGPU::GetFullsCorr(float *corr) const {
 }
 
 void RotationScaleMergeGPU::GetFullsPxPy(float *px, float *py) const {
+    DeviceGuard guard(impl_->device, impl_->available);
     const auto &d = *impl_;
     if (d.n_fulls == 0) return;
     const size_t bytes = size_t(d.n_fulls) * sizeof(float);
@@ -1068,6 +1121,7 @@ void RotationScaleMergeGPU::GetFullsPxPy(float *px, float *py) const {
 }
 
 void RotationScaleMergeGPU::GetFullsVariance(float *var_bkg, float *var_per_I) const {
+    DeviceGuard guard(impl_->device, impl_->available);
     const auto &d = *impl_;
     if (d.n_fulls == 0) return;
     const size_t bytes = size_t(d.n_fulls) * sizeof(float);
@@ -1077,6 +1131,7 @@ void RotationScaleMergeGPU::GetFullsVariance(float *var_bkg, float *var_per_I) c
 }
 
 void RotationScaleMergeGPU::SetFullsCorr(const float *corr) {
+    DeviceGuard guard(impl_->device, impl_->available);
     auto &d = *impl_;
     if (d.n_fulls == 0) return;
     CudaCheck(cudaMemcpy(d.f_corr.get(), corr, size_t(d.n_fulls) * sizeof(float),