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()));