leonarski_fandClaude Opus 5.5
672e182d6a
A pooled CudaDevicePtr frees on the thread's allocation stream, not on the engine's stream, and the
pool may hand the memory to another engine - or, past its release threshold, unmap it - as soon as
that free is reached, which on an idle allocation stream is at once. An engine destroyed with work
still queued (FFTIndexerGPU after SearchCap's last DirectionsChanged upload, a spot finder between
DetectAt and Extract, a shadow accumulator after a pending fold, any engine on an exception path)
thus had kernels or copies writing memory that was someone else's or no longer mapped. Now:
- CudaStream synchronises before cudaStreamDestroy (destructor and move-assignment), which covers
engines whose own stream is declared after their buffers (FFTIndexerGPU, the gather buffer);
- every engine holding pooled buffers and a stream (shared or own, declared before the buffers)
synchronises it in its destructor; BraggIntegrationEngineGPU also before EnsureCapacity
reallocates, where a Run that threw leaves work queued.
A GPU failure that is handled no longer hides a lost context: ShadowFinder, BeamCenterFFT, the
rigid-body pool and model validation call cuda_throw_if_context_lost() before cuda_clear_error(),
as the device-decode fallbacks already did; RotationScaleMergeGPU's Alloc does so before waiting up
to ten minutes for GPU work beside it and then reporting a lost device as out of memory; and a
failed cudaMalloc says why. BeamCenterFFT logs the failure it used to drop silently, the
speculative geometry probe logs the exception it swallowed (its GPU fault was otherwise reported by
the merge beside it, under the merge's name), and RotationScaleMergeGPU's DeviceGuard no longer
throws from its destructor.
Only synchronisation and error paths change: p.hkl md5 and the MTZ data (gemmi) are identical to
the b530c2d full battery on myob_x10sa, cytc_x10sa, 8a1a, 9gdj, 11if, kdp_x10sa_20keV and 6z9g.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi
100 lines
5.8 KiB
C++
100 lines
5.8 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 "BraggIntegrationEngine.h"
|
|
#include "../indexing/CUDAMemHelpers.h"
|
|
#include "../../common/PixelMask.h"
|
|
|
|
// CUDA engine: reproduces BraggIntegrationEngineCPU up to floating-point precision. Each stage is a
|
|
// kernel with one CUDA block per reflection cooperating over the small window via shared-memory
|
|
// reductions (the natural mapping for thousands of independent, tiny per-spot integrations).
|
|
//
|
|
// Pipeline (profile modes): reset -> mark_mask -> boxsum -> learn_profile -> build_profiles -> fit
|
|
// (the resolution shell is computed inline, so there is no separate shell pass). BoxSum mode stops
|
|
// after boxsum (that pass is the BraggIntegrate2D box integrator and the seed of the profile fit).
|
|
// The preprocessed image already lives on the device (ImagePreprocessorBufferGPU::getGPUBuffer());
|
|
// only the per-frame predicted centres are uploaded.
|
|
class BraggIntegrationEngineGPU : public BraggIntegrationEngine {
|
|
std::shared_ptr<CudaStream> stream;
|
|
int threads;
|
|
size_t fit_shared_bytes;
|
|
int rad_w = 0; // radial-background window of boxsum, in bins of one pixel
|
|
size_t boxsum_shared_bytes = 0;
|
|
|
|
size_t capacity = 0; // per-reflection device/host arrays hold at least this many reflections
|
|
|
|
// Whether d_mask / d_owner may still carry the marks of an earlier image. Run() clears what it
|
|
// marked before it returns when that is cheaper than clearing the frame, and this is then false in
|
|
// the steady state; it is true before the first call, after one that threw part-way through, and
|
|
// whenever the marks covered enough of the frame that clearing all of it was the cheaper choice.
|
|
bool dirty = true;
|
|
size_t mask_box_px = 0; // pixels one reflection's mark_mask box covers, at the widest aperture
|
|
|
|
// --- per-reflection device arrays (grown by EnsureCapacity) ---
|
|
CudaDevicePtr<float> d_px_x, d_px_y, d_d;
|
|
CudaDevicePtr<uint8_t> d_mark; // the reflection marks its signal region in d_mask
|
|
CudaDevicePtr<int> d_cx, d_cy;
|
|
CudaDevicePtr<float> d_I, d_sigma, d_bkg, d_bkg_var, d_var_bkg, d_obs_x, d_obs_y;
|
|
CudaDevicePtr<float> d_isum; // box-sum raw sum, for the radial correction
|
|
CudaDevicePtr<int> d_ninner, d_rbin, d_kbin;
|
|
// The run's pixel mask, one byte per pixel (PixelMask::GetBinaryMask), shared with the
|
|
// preprocessor's copy. See BraggIntegrationEngineCPU::static_mask.
|
|
std::shared_ptr<CudaDevicePtr<uint8_t>> d_static_mask;
|
|
CudaDevicePtr<uint8_t> d_ok, d_strong, d_has_obs, d_overloaded;
|
|
|
|
// --- radial background curvature correction (see BraggIntegrationEngine) ---
|
|
int n_rad = 0; // radial bins, 0 when the correction is off
|
|
CudaDevicePtr<unsigned long long> d_rad_sum; // integer pixel sums, see boxsum
|
|
CudaDevicePtr<float> d_k_diff;
|
|
CudaDevicePtr<int> d_rad_cnt;
|
|
|
|
// --- fixed-size device arrays ---
|
|
// The learning/fit math is single precision: FP64 is heavily throttled on consumer GPUs and the
|
|
// extraction is Poisson-noise limited, so float reproduces the double CPU path to ~1e-4.
|
|
CudaDevicePtr<uint8_t> d_mask; // per-pixel inner-stencil reflection mask
|
|
// Per-pixel (distance, reflection) key naming the nearest predicted centre; allocated only when
|
|
// an overlap treatment is on, so the default path costs no extra device memory.
|
|
CudaDevicePtr<uint32_t> d_owner;
|
|
// Fixed-point (see PROFILE_FIXED): a float atomicAdd here made the profile depend on the order
|
|
// the blocks arrived in, and with it every intensity fitted through it.
|
|
CudaDevicePtr<unsigned long long> d_shell_grid, d_global_grid; // learned profile accumulators (N_SHELL*GG, GG)
|
|
CudaDevicePtr<float> d_shell_P, d_global_P; // normalised profiles (empirical mode)
|
|
CudaDevicePtr<unsigned long long> d_mom; // learned 2nd moments, 3 per shell + global
|
|
CudaDevicePtr<float> d_sigma2_r, d_sigma2_t; // radial/tangential widths, N_SHELL + global
|
|
CudaDevicePtr<int> d_shell_n, d_global_n;
|
|
CudaDevicePtr<unsigned long long> d_invd2; // [min,max] inv-d^2 as monotonic bit patterns
|
|
// BraggIntegrationCounts, accumulated on the device so an image costs no transfer; brought back
|
|
// only when Counts() is asked for.
|
|
CudaDevicePtr<unsigned long long> d_counts;
|
|
|
|
// --- host staging (copied back once per frame) ---
|
|
// Pinned, like every other engine's staging: a copy out of pageable memory does not return until the
|
|
// driver has staged it through a bounce buffer, so eight of them in a row are eight serialised
|
|
// round-trips rather than eight queued transfers.
|
|
CudaHostPtr<float> h_px_x, h_px_y, h_d;
|
|
CudaHostPtr<uint8_t> h_mark;
|
|
CudaHostPtr<float> h_I, h_sigma, h_bkg, h_var_bkg, h_obs_x, h_obs_y;
|
|
CudaHostPtr<uint8_t> h_ok, h_has_obs, h_overloaded;
|
|
|
|
void EnsureCapacity(size_t n);
|
|
|
|
public:
|
|
BraggIntegrationEngineGPU(const DiffractionExperiment &experiment, std::shared_ptr<CudaStream> stream,
|
|
const PixelMask &mask);
|
|
// Waits for the stream: its pooled buffers are freed next (see CudaDevicePtr).
|
|
~BraggIntegrationEngineGPU() override { if (stream) cudaStreamSynchronize(*stream); }
|
|
std::vector<Reflection> Run(const ImagePreprocessorBuffer &image,
|
|
const std::vector<Reflection> &predicted, size_t npredicted,
|
|
int64_t image_number) override;
|
|
|
|
// Brings the two device counters back before answering. Synchronises the stream, so ask once a
|
|
// pass rather than once an image.
|
|
[[nodiscard]] BraggIntegrationCounts Counts() const override;
|
|
};
|