From debd73c0bd06e1fa0fd57c8b7a0a1f00bfa70f4f Mon Sep 17 00:00:00 2001 From: jungfrau Date: Mon, 17 Aug 2026 18:06:03 -0400 Subject: [PATCH] Find the first pass's spots on every worker The first pass of rotation indexing feeds two sampling schemes and a validation set, and it found their spots one frame at a time on one CUDA stream. It reads 271 to 402 frames depending on the crystal, once per pass, and there are two passes. On a 16M-pixel dataset that is 6.7 s of a 30 s run at a mean occupancy of 1.9 threads out of 48. The comment said this was to keep the spot cache from depending on scheduling. The guarantee is stronger than that and survives: a frame's spots are a pure function of that frame, so whatever order the workers finish in, the cache ends up holding exactly what the serial loop put there. The feed order and the point at which the consecutive scheme stops accumulating are untouched, so the schemes see the same frames as before. The phase was never CPU-bound, which is why this works: its one thread was spin-waiting in the driver on the same six full-image passes, about 5.5 ms a frame. One engine per worker for the whole pass, not per batch. An engine on a 16M-pixel detector costs ~0.13 s to build - a first attempt rebuilt them per batch and came out SLOWER than the serial loop, four batches of eight engines against ~1.5 s of frame work. Two smaller things, in the same per-image path: The strong-pixel flagging kernel reads four pixels at a time and issues one atomicOr per four instead of per pixel. Same values, confirmed end to end. The beam-stop pre-scan sized its per-worker accumulators for the worker count rather than for the workers that use them. Two int64 arrays per pixel per shard, eight shards: 2.3 GB allocated and cleared on a 16M-pixel detector for a path the GPU never reads. They are allocated on use now, which is most of what made that phase look like eight workers with two of them running. Measured, three A/B pairs with the order alternated, whole run: 16M-pixel 3600-frame 42.8 s -> 38.8 s, 16M-pixel 1800-frame 38.7 s -> 36.1 s, both at two GPUs; a 4M-pixel control 23.3 s -> 23.1 s. First-pass indexing itself -44 % and -36 %. The two smaller datasets merge byte-identically over six runs each; the largest is not bit-reproducible in the unmodified binary either, and stayed inside that spread. Co-Authored-By: Claude Opus 5 --- docs/CHANGELOG.md | 3 +- image_analysis/beam_stop/ShadowFinder.cpp | 20 ++- .../spot_finding/AdaptiveSpotFinderGPU.cu | 33 +++- rugnux/Rugnux.cpp | 144 +++++++++++++++--- 4 files changed, 167 insertions(+), 33 deletions(-) diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index aad3bf52..c4cdca5c 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -3,7 +3,8 @@ ### Unreleased * **rugnux: scaling and merging are much faster on crystals that integrate far beyond the resolution they merge at.** Observations outside the scaling resolution range are dropped as they are ingested instead of being scaled, combined and post-refined first, and the geometry post-refinement fits on a bounded sample of them. On a large-cell dataset merging at less than half the resolution its detector reaches, a run drops from 68 s to 37 s; merged statistics are unchanged. * rugnux: the detector-frame modulation correction is fitted on a grid spanning the detector rather than the reflections that happen to be present, so whether it is applied no longer depends on how far integration reached. On a crystal where it was being refused, merged R_meas improves by 4 percentage points. -* rugnux: the beam-stop pre-scan reads and accumulates its frames on several threads instead of one. On a 16M-pixel rotation dataset it drops from 7.9 s to 4.8 s; the shadow it finds is unchanged. +* rugnux: the first-pass rotation indexing finds its spots on every worker rather than one, which is worth most on large detectors - on a 16M-pixel dataset that phase drops by 44% and the whole run by 9%. The lattice it picks is unchanged. +* rugnux: the beam-stop pre-scan reads and accumulates its frames on several threads instead of one. On a 16M-pixel rotation dataset it drops from 7.9 s to 4.8 s; the shadow it finds is unchanged. Its per-worker accumulators are now allocated only where they are used, saving 2.3 GB of host memory on a 16M-pixel detector. ### 1.0.0-rc.161 This is an UNSTABLE release. It includes many experimental features, as well as many AI generated fixes. We recommend using rc.152 for production use. diff --git a/image_analysis/beam_stop/ShadowFinder.cpp b/image_analysis/beam_stop/ShadowFinder.cpp index 269ea2e2..c856aca6 100644 --- a/image_analysis/beam_stop/ShadowFinder.cpp +++ b/image_analysis/beam_stop/ShadowFinder.cpp @@ -302,16 +302,12 @@ ShadowFinder::ShadowFinder(const DiffractionExperiment &experiment, const PixelM #endif } +// A shard's accumulators are allocated when a frame is first added to it, not here: with a GPU they +// are never used at all, and on a 16 Mpx detector eight of them are 2.9 GB to allocate and clear - +// which measured 0.8 s of the pre-scan, all of it wasted. void ShadowFinder::SetShardCount(size_t n) { - const size_t npixels = static_cast(width) * height; shards.clear(); shards.resize(std::max(1, n)); - for (auto &p : shards) { - p.max_value.assign(npixels, 0); - p.sum_value.assign(npixels, 0); - p.valid_count.assign(npixels, 0); - p.frames = 0; - } } ShadowFinder::Projection ShadowFinder::Reduce() const { @@ -376,7 +372,8 @@ ShadowFinder::Projection ShadowFinder::Reduce() const { } } #endif - for (const auto &p : shards) + for (const auto &p : shards) { + if (p.frames == 0) continue; // never used, and its accumulators were never allocated for (size_t i = lo; i < hi; i++) { if (p.valid_count[i] == 0) continue; @@ -385,6 +382,7 @@ ShadowFinder::Projection ShadowFinder::Reduce() const { out.sum_value[i] += p.sum_value[i]; out.valid_count[i] += p.valid_count[i]; } + } })); } for (auto &f : futures) f.get(); @@ -441,6 +439,12 @@ void ShadowFinder::AddImage(const DataMessage &data, std::vector &buffe #endif Projection &p = shards[shard]; + if (p.max_value.empty()) { + const size_t npixels = static_cast(width) * height; + p.max_value.assign(npixels, 0); + p.sum_value.assign(npixels, 0); + p.valid_count.assign(npixels, 0); + } const auto ptr = data.image.GetUncompressedPtr(buffer); switch (data.image.GetMode()) { case CompressedImageMode::Int8: Add(reinterpret_cast(ptr), p); break; diff --git a/image_analysis/spot_finding/AdaptiveSpotFinderGPU.cu b/image_analysis/spot_finding/AdaptiveSpotFinderGPU.cu index 8e1d144a..6ba4a8f2 100644 --- a/image_analysis/spot_finding/AdaptiveSpotFinderGPU.cu +++ b/image_analysis/spot_finding/AdaptiveSpotFinderGPU.cu @@ -230,7 +230,38 @@ __global__ void flag_strong(const int32_t *__restrict__ image, const float *__restrict__ thr, uint32_t *__restrict__ strong, size_t npix, int nbins) { - for (size_t idx = blockIdx.x * blockDim.x + threadIdx.x; idx < npix; idx += blockDim.x * gridDim.x) { + // Four pixels per thread, read as one 16-byte and one 8-byte transaction instead of four of + // each, exactly as the ring reduction above reads them - and flagged with a single atomicOr, + // because four consecutive pixels always fall in the same word of the bit buffer. The last + // npix % 4 pixels are done one at a time below, so nothing is read past the end. + const size_t stride = static_cast(blockDim.x) * gridDim.x; + const size_t nquad = npix / 4; + + for (size_t q = blockIdx.x * blockDim.x + threadIdx.x; q < nquad; q += stride) { + const int4 v4 = reinterpret_cast(image)[q]; + const ushort4 b4 = reinterpret_cast(pixel_to_bin)[q]; + const int32_t vq[4] = {v4.x, v4.y, v4.z, v4.w}; + const uint16_t bq[4] = {b4.x, b4.y, b4.z, b4.w}; + + uint32_t bits = 0; + #pragma unroll + for (int k = 0; k < 4; k++) { + const int32_t v = vq[k]; + if (v == INT32_MAX) { + bits |= 1u << k; + } else if (v != INT32_MIN) { + const int b = bq[k]; + if (b < nbins && static_cast(v) >= thr[b]) + bits |= 1u << k; + } + } + if (bits) { + const size_t idx = 4 * q; + atomicOr(&strong[idx / 32], bits << (idx % 32)); + } + } + + for (size_t idx = 4 * nquad + blockIdx.x * blockDim.x + threadIdx.x; idx < npix; idx += stride) { const int32_t v = image[idx]; bool s = false; if (v == INT32_MAX) { diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index 343d56ba..0605fd30 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -1177,20 +1177,16 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b }; // Spots for a first-pass image ordinal, cached (a frame is reused across schemes and the - // validation set; re-finding is expensive on the --redo-rotation-spots path). Accessed only - // from single-threaded sections (the scheme feed and the validation loop), so no locking. + // validation set; re-finding is expensive on the --redo-rotation-spots path). The workers below + // hand their spots back and this thread puts them in, so the map itself is still only ever + // touched from here and needs no locking. std::map> spot_cache; - // One analysis engine for the whole first pass. Constructing it allocates a CUDA stream and a - // full set of GPU buffers (preprocessing, spot finding, azimuthal integration, Bragg - // integration) - far more work than analysing a frame - so building it per image would pay for - // hundreds of them, serially. This section is single-threaded, so one engine is enough. - auto analysis = std::make_unique(experiment_, mapping, pixel_mask_, *indexer, - /*enable_fused_adaptive_gpu=*/true); - auto profile = std::make_unique(mapping); - auto get_spots = [&](int ordinal) -> const std::vector & { - auto it = spot_cache.find(ordinal); - if (it != spot_cache.end()) - return it->second; + // The spots of one frame, found with the caller's own analysis engine. A frame's spots are a + // pure function of that frame and the settings - the engine carries nothing from one image to + // the next, which is why the main image loop can give one to each of its workers - so it does + // not matter which engine, or which thread, finds them. + const auto find_spots = [&](MXAnalysisWithoutFPGA &analysis, AzimuthalIntegrationProfile &profile, + int ordinal) { const int image_idx = start_image + ordinal * config_.stride; std::vector spots; try { @@ -1204,13 +1200,92 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b m.image = img->image; if (dataset->efficiency.size() > image_idx) m.image_collection_efficiency = dataset->efficiency[image_idx]; - analysis->Analyze(m, *profile, first_pass); + analysis.Analyze(m, profile, first_pass); spots = std::move(m.spots); } } catch (const std::exception &e) { logger.Warning("First-pass spot read failed for image {}: {}", image_idx, e.what()); } - return spot_cache.emplace(ordinal, std::move(spots)).first->second; + return spots; + }; + + // Find the spots of every ordinal in the list that is not cached yet, on several workers, and + // put them in the cache. + // + // The cache ends up holding exactly what a serial loop would have put there: the results are + // collected into a vector indexed by position and inserted afterwards, so neither the spots of + // a frame nor the set of frames in the cache depends on how the workers interleaved. + // + // Worker t keeps its analysis engine for the whole first pass. Building one is a CUDA stream + // and a full set of device buffers - measured at 0.13 s against 4 ms for a frame - and the + // "wedge" scheme only discovers how far into the sweep it has to read as the frames arrive, + // so this is called several times and rebuilding them each time cost a quarter of the pass. + // The engine belongs to a card, not to a thread, so a later call only has to make that card + // current again; workers take a card by index so an engine always meets the one it was + // built on. + const int spot_workers = get_gpu_count() > 0 + ? std::min(config_.nthreads, std::max(8, 4 * get_gpu_count())) + : config_.nthreads; + struct SpotEngine { + std::unique_ptr analysis; + std::unique_ptr profile; + }; + std::vector engines(std::max(spot_workers, 1)); + const auto use_gpu_of_worker = [](size_t t) { + if (const int32_t gpus = get_gpu_count(); gpus > 0) + set_gpu(static_cast(t) % gpus); + }; + int prefetch_batches = 0; + double prefetch_time_s = 0.0; + const auto prefetch_spots = [&](const std::vector &ordinals) { + std::vector wanted; + for (const int ordinal : ordinals) + if (!spot_cache.contains(ordinal)) + wanted.push_back(ordinal); + if (wanted.empty()) + return; + const auto prefetch_start = std::chrono::steady_clock::now(); + ++prefetch_batches; + + std::vector> found(wanted.size()); + std::atomic next{0}; + std::vector> futures; + const size_t nworkers = std::min(engines.size(), wanted.size()); + futures.reserve(nworkers); + for (size_t t = 0; t < nworkers; t++) + futures.emplace_back(std::async(std::launch::async, [&, t] { + use_gpu_of_worker(t); // must precede engine construction + SpotEngine &e = engines[t]; + if (!e.analysis) { + e.analysis = std::make_unique( + experiment_, mapping, pixel_mask_, *indexer, + /*enable_fused_adaptive_gpu=*/true); + e.profile = std::make_unique(mapping); + } + for (size_t i = next.fetch_add(1); i < wanted.size(); i = next.fetch_add(1)) + found[i] = find_spots(*e.analysis, *e.profile, wanted[i]); + })); + for (auto &f : futures) + f.get(); + + for (size_t i = 0; i < wanted.size(); i++) + spot_cache.emplace(wanted[i], std::move(found[i])); + prefetch_time_s += std::chrono::duration( + std::chrono::steady_clock::now() - prefetch_start).count(); + }; + + // Release the engines on workers of their own, so a stream and its device buffers are freed + // with the card they were made on current. + const auto release_spot_engines = [&] { + std::vector> futures; + futures.reserve(engines.size()); + for (size_t t = 0; t < engines.size(); t++) + futures.emplace_back(std::async(std::launch::async, [&, t] { + use_gpu_of_worker(t); + engines[t] = SpotEngine{}; + })); + for (auto &f : futures) + f.get(); }; // How many validation frames a candidate global lattice actually indexes per image, scored on @@ -1227,10 +1302,8 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b idx.ForceRotationIndexerResult(r); // Fill the spot cache first. It is a plain map filled on demand, and the loop below runs the // frames in parallel, where a lookup racing an insert is not something the map survives. - // Spot finding stays serial and in frame order, so what ends up in the cache does not depend - // on scheduling; after the first candidate this costs nothing. - for (const int ordinal : validation) - get_spots(ordinal); + // After the first candidate this costs nothing. + prefetch_spots(validation); // Each frame's verdict is its own - the score is just how many of them index - so they run // together. Refining one frame's geometry is a few Ceres solves and no GPU, and it is the // same call the main image loop already makes from every one of its workers on this same @@ -1251,14 +1324,23 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b }; // Feed one first-pass scheme (a set of image ordinals) into its own rotation indexer, ready to - // be indexed. Spots are pulled from the single-threaded cache here; the FFT + refinement runs - // separately (RunIndexing) so the two schemes' indexing can overlap. + // be indexed. Spots are pulled from the cache here; the FFT + refinement runs separately + // (RunIndexing) so the two schemes' indexing can overlap. + // + // The feed itself stays in frame order and stops where it always did, but it keeps the cache a + // chunk of frames ahead of itself so those frames' spots are found on several workers. A chunk + // rather than the whole list because the "wedge" scheme stops on the accumulation cap, and how + // far into the sweep that is only becomes known as the frames arrive. + constexpr int SPOT_PREFETCH_CHUNK = 64; auto feed_scheme = [&](IndexerThreadPool &pool, const std::vector &ordinals) { auto ri = std::make_unique(experiment_, pool); - for (const int ordinal : ordinals) { + for (size_t i = 0; i < ordinals.size(); i++) { if (cancelled_ || ri->AccumulationFull()) break; - ri->ProcessImage(ordinal, get_spots(ordinal), rot_angle(ordinal)); + if (!spot_cache.contains(ordinals[i])) + prefetch_spots({ordinals.begin() + i, + ordinals.begin() + std::min(ordinals.size(), i + SPOT_PREFETCH_CHUNK)}); + ri->ProcessImage(ordinals[i], spot_cache.at(ordinals[i]), rot_angle(ordinals[i])); } return ri; }; @@ -1390,6 +1472,19 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b return bp; }; + // Find the spots of everything the first pass is certain to want in one go - both schemes' + // opening frames and the validation set. A batch of workers costs an analysis engine each, + // which is far more than a frame is to analyse, so one batch for the lot beats one per scheme. + { + std::vector upfront = spread_ordinals; + upfront.insert(upfront.end(), validation.begin(), validation.end()); + for (int i = 0; i < std::min(SPOT_PREFETCH_CHUNK, images_to_process); i++) + upfront.push_back(i); + std::sort(upfront.begin(), upfront.end()); + upfront.erase(std::unique(upfront.begin(), upfront.end()), upfront.end()); + prefetch_spots(upfront); + } + FirstPass best = pick_best(*indexer_pool, *indexer); // Long-axis rescue. When the de-novo cell indexes few validation frames, a long, finely-spaced @@ -1446,6 +1541,7 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b best = constrained; } } + release_spot_engines(); // no more frames to read; give the cards their memory back if (!cancelled_) { // The two ways the first pass can end without a usable lattice. They are one situation from @@ -1482,6 +1578,8 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b c.a, c.b, c.c, c.alpha, c.beta, c.gamma, next_steps)); } indexer->ForceRotationIndexerResult(*best.result); + logger.Info("First-pass spot finding: {} frames in {} batches, {:.2f} s", + spot_cache.size(), prefetch_batches, prefetch_time_s); logger.Info("Two-pass rotation indexing found lattice (scheme '{}': {}/{} validation frames)", best.name, best.score, static_cast(validation.size()));