diff --git a/image_analysis/CMakeLists.txt b/image_analysis/CMakeLists.txt index 0ce15825..342a7e76 100644 --- a/image_analysis/CMakeLists.txt +++ b/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 + $<$:beam_stop/ShadowAccumulatorGPU.cu> + beam_stop/ShadowAccumulatorGPU.h rotation_indexer/RotationIndexer.cpp rotation_indexer/RotationIndexer.h WriteReflections.cpp diff --git a/image_analysis/beam_stop/ShadowAccumulatorGPU.cu b/image_analysis/beam_stop/ShadowAccumulatorGPU.cu new file mode 100644 index 00000000..efbb9130 --- /dev/null +++ b/image_analysis/beam_stop/ShadowAccumulatorGPU.cu @@ -0,0 +1,186 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#include "ShadowAccumulatorGPU.h" + +#include +#include + +#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 +__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(blockDim.x) + threadIdx.x; i < npixels; + i += static_cast(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(v); + if (c == 0 || vi > mx) + mx = vi; + sm += vi; + c++; + } + max_value[i] = mx; + sum_value[i] = sm; + valid_count[i] = c; + } +} + +template +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) + masked = std::numeric_limits::min(); + else + masked = std::numeric_limits::max(); + accumulate_kernel<<>>(reinterpret_cast(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()), + 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(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(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(raw, pending, frame_bytes, gpu_max, gpu_sum, gpu_count, npixels, blocks, *stream); break; + case CompressedImageMode::Uint8: + launch(raw, pending, frame_bytes, gpu_max, gpu_sum, gpu_count, npixels, blocks, *stream); break; + case CompressedImageMode::Int16: + launch(raw, pending, frame_bytes, gpu_max, gpu_sum, gpu_count, npixels, blocks, *stream); break; + case CompressedImageMode::Uint16: + launch(raw, pending, frame_bytes, gpu_max, gpu_sum, gpu_count, npixels, blocks, *stream); break; + case CompressedImageMode::Int32: + launch(raw, pending, frame_bytes, gpu_max, gpu_sum, gpu_count, npixels, blocks, *stream); break; + case CompressedImageMode::Uint32: + launch(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(image.GetWidth()) * image.GetHeight() != npixels) + throw JFJochException(JFJochExceptionCategory::InputParameterInvalid, + "ShadowAccumulatorGPU: image size does not match the detector"); + + if (!decoder) + decoder = std::make_unique(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(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 &max_value, std::vector &sum_value, + std::vector &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)); +} diff --git a/image_analysis/beam_stop/ShadowAccumulatorGPU.h b/image_analysis/beam_stop/ShadowAccumulatorGPU.h new file mode 100644 index 00000000..c45a424c --- /dev/null +++ b/image_analysis/beam_stop/ShadowAccumulatorGPU.h @@ -0,0 +1,66 @@ +// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include +#include +#include + +#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 stream; + std::unique_ptr 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 raw; // BATCH decoded frames, still in their stored type + CudaDevicePtr gpu_max; + CudaDevicePtr gpu_sum; + CudaDevicePtr 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 &max_value, std::vector &sum_value, + std::vector &valid_count); +}; diff --git a/image_analysis/beam_stop/ShadowFinder.cpp b/image_analysis/beam_stop/ShadowFinder.cpp index 29eb36bc..660831fc 100644 --- a/image_analysis/beam_stop/ShadowFinder.cpp +++ b/image_analysis/beam_stop/ShadowFinder.cpp @@ -13,6 +13,10 @@ #include #include +#include + +#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(width) * height; + gpu_pending = std::async(std::launch::async, [npixels] { + return std::make_unique(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 &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 &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; } diff --git a/image_analysis/beam_stop/ShadowFinder.h b/image_analysis/beam_stop/ShadowFinder.h index 4899efb3..4bcb9930 100644 --- a/image_analysis/beam_stop/ShadowFinder.h +++ b/image_analysis/beam_stop/ShadowFinder.h @@ -4,6 +4,8 @@ #pragma once #include +#include +#include #include #include @@ -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 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> gpu_pending; + mutable std::unique_ptr 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 void Add(const T *ptr, Projection &p); // Sum the shards into one projection. max_value is only taken from a shard that actually diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu index 4c3166c8..28ea5c52 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.cu +++ b/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 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()) { 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(std::max(1, ng)); d.a_sw = CudaDevicePtr(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(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),