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
103 lines
5.1 KiB
C++
103 lines
5.1 KiB
C++
// SPDX-FileCopyrightText: 2026 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
|
|
// SPDX-License-Identifier: GPL-3.0-only
|
|
|
|
#pragma once
|
|
|
|
// Device-side connected-component extraction for the GPU spot finders.
|
|
//
|
|
// The GPU finders flag strong pixels into a packed bit buffer ON THE DEVICE. Reading spots out of it
|
|
// used to mean copying that whole buffer back (2.26 MB per frame at 18 MP) and scanning it bit by bit
|
|
// on the host. This does the whole extraction where the data already is, so nothing about the image
|
|
// comes back - only the finished spot list, a few hundred entries.
|
|
//
|
|
// The algorithm is the sparse formulation the ACTS/traccc project settled on for the same problem
|
|
// (sparse silicon-detector hits): the strong pixels are compacted into a list that is sorted by flat
|
|
// index, each pixel finds its at most FOUR backward 8-neighbours by binary search in that list, and
|
|
// the resulting graph is labelled with a lock-free union-find. A dense image-wide labelling
|
|
// (Playne-equivalence, BUF/BKE, nppiLabelMarkers, cv::cuda::connectedComponents) would label 18
|
|
// million pixels to find five hundred.
|
|
//
|
|
// It reproduces the host StrongPixelSet::sparseccl EXACTLY, not just equivalently:
|
|
// * both make a component's root its lowest list index, so both find the same roots;
|
|
// * labels are handed out by a prefix sum over the roots in ascending order, which is the order the
|
|
// host's second scan hands them out in, so the SPOT ORDER is identical;
|
|
// * the centroid sums are accumulated per component in ascending list order, in integers, term for
|
|
// term as DiffractionSpot::AddPixel does them, so there is no rounding for the two compilers to
|
|
// disagree about.
|
|
// tests/SpotExtractorGPUParityTest.cpp holds the two to each other on realistic, occupancy-swept and
|
|
// pathological frames, and checks that repeating a frame gives byte-identical output.
|
|
|
|
#include <cstdint>
|
|
#include <memory>
|
|
#include <vector>
|
|
|
|
#include "../../common/DiffractionSpot.h"
|
|
#include "../indexing/CUDAMemHelpers.h"
|
|
#include "SpotFindingSettings.h"
|
|
|
|
// Per-component sums, in exactly the form DiffractionSpot holds them: x and y are sum(col*photons)
|
|
// and sum(line*photons), not a centroid.
|
|
struct SpotExtractorGPUSpot {
|
|
int64_t x;
|
|
int64_t y;
|
|
int64_t photons;
|
|
int64_t max_photons;
|
|
int32_t pixel_count;
|
|
// Longer side of the component's bounding box, for SpotShapeAccepted. It fits in what used to be
|
|
// padding, so carrying it costs nothing.
|
|
int32_t bbox_side;
|
|
};
|
|
|
|
class SpotExtractorGPU {
|
|
std::shared_ptr<CudaStream> stream;
|
|
const int32_t width;
|
|
const size_t nwords;
|
|
|
|
// Strong pixels this engine's buffers hold, and above which the extraction gives up on the frame -
|
|
// StrongPixelLimit, so it follows the detector rather than standing at a constant. 104 bytes of
|
|
// device memory apiece, 29 MB on an 18-megapixel detector.
|
|
const uint32_t max_strong;
|
|
// Spots copied back together with their count in one transfer. A frame with more than this many
|
|
// surviving spots - far past anything indexable - simply takes a second copy.
|
|
static constexpr uint32_t SPOT_PREFIX = 4096;
|
|
|
|
int compact_blocks = 0;
|
|
|
|
CudaDevicePtr<uint32_t> gpu_res_mask; // packed, bit set = pixel excluded
|
|
CudaDevicePtr<uint32_t> gpu_block_count;
|
|
CudaDevicePtr<uint32_t> gpu_block_offset;
|
|
CudaDevicePtr<uint32_t> gpu_nstrong;
|
|
CudaDevicePtr<uint32_t> gpu_index; // strong pixels, sorted by flat index
|
|
CudaDevicePtr<int32_t> gpu_value;
|
|
CudaDevicePtr<uint32_t> gpu_parent; // union-find parent
|
|
CudaDevicePtr<uint32_t> gpu_root;
|
|
CudaDevicePtr<uint32_t> gpu_label; // compact label, indexed by root
|
|
CudaDevicePtr<int32_t> gpu_count; // pixels per component
|
|
CudaDevicePtr<SpotExtractorGPUSpot> gpu_spot;
|
|
CudaDevicePtr<SpotExtractorGPUSpot> gpu_spot_out;
|
|
CudaDevicePtr<uint32_t> gpu_nspot;
|
|
|
|
CudaHostPtr<uint32_t> host_nstrong;
|
|
CudaHostPtr<uint32_t> host_nspot;
|
|
CudaHostPtr<SpotExtractorGPUSpot> host_spot; // SPOT_PREFIX entries, pinned
|
|
std::vector<SpotExtractorGPUSpot> overflow_spot; // only for a frame with more spots than that
|
|
|
|
public:
|
|
SpotExtractorGPU(int32_t width, int32_t height, std::shared_ptr<CudaStream> stream);
|
|
// Waits for the stream: its pooled buffers are freed next (see CudaDevicePtr).
|
|
~SpotExtractorGPU() { if (stream) cudaStreamSynchronize(*stream); }
|
|
|
|
void SetResolutionMask(const std::vector<uint32_t> &packed_mask);
|
|
|
|
// gpu_strong is the finder's device bit buffer, gpu_image the preprocessed image it was built
|
|
// from. Fills spots with every component of at most max-pix pixels, in the same order the host
|
|
// extractor would.
|
|
void Extract(const uint32_t *gpu_strong, const int32_t *gpu_image,
|
|
const SpotFindingSettings &settings, std::vector<DiffractionSpot> &spots);
|
|
|
|
// Strong pixels the last Extract() saw, after the resolution mask. Reported whether or not the
|
|
// frame was given up on, which is the point of it: a frame at or above max_strong yields no spots
|
|
// at all, and this is what says so.
|
|
[[nodiscard]] uint32_t StrongPixelCount() const { return *host_nstrong.get(); }
|
|
};
|