From 1c3a2e202cc9e3ddda38f4e4928177e7a6a09325 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 15 Oct 2026 11:39:19 +0200 Subject: [PATCH 1/3] GPU memory over the run: FFT indexers given back after the last indexing; speculative probe takes turns with the merge on a small card - IndexerThreadPool::ReleaseIndexers(): each worker drops the indexers it has built (once its current frame is done) and the allocation pool is trimmed, so their memory goes back to the card, not only to the pool. A later frame builds them again, as on first use. A rotation pass calls it once its first pass is done and again when the leftover-lattice census - the last of its indexing - ends. Measured on the RTX 5080: the 1.24 GB (pre-pass) / 0.62 GB (later passes) the pass's pool held through the image pass and the scaling is gone; 4nwv's pre-pass image pass peaks 5.97 instead of 7.18 GB of device memory. - GPUMerge: the main line holds one through ScaleMergeAndSymmetry. The speculative geometry probe (GPUWorkBeside(after_merges=true)) waits for it on a card below the take-turns size (12 GB), where the probe's FFT indexers and image workers beside the merge's arrays would overfill the card - on a Windows (WDDM) card that does not fail an allocation but pages it and crawls. The merge waits for nothing, so nothing the merge waits on (the rotation-scale walk's probes) can wait in a circle with it. No change on a 16 GB card. Output unchanged: p.hkl byte-identical on cytc_x10sa, myob_x10sa, 4nwv and 8qaw on this card; also identical with the card's total read as 8.55 GB (10 / 9 image workers instead of 16) and with 4 image workers, so the worker count does not change the result. Co-Authored-By: Claude Opus 5.5 --- common/CUDAWrapper.cpp | 17 +++- common/CUDAWrapper.h | 19 ++++- image_analysis/indexing/IndexerThreadPool.cpp | 42 +++++++++- image_analysis/indexing/IndexerThreadPool.h | 8 ++ rugnux/Rugnux.cpp | 13 ++- rugnux/RugnuxFirstPassAccept.cpp | 2 + tests/CUDAMemHelpersTest.cpp | 83 +++++++++++++++++++ 7 files changed, 174 insertions(+), 10 deletions(-) diff --git a/common/CUDAWrapper.cpp b/common/CUDAWrapper.cpp index 6b5eb1272..3a71878f4 100644 --- a/common/CUDAWrapper.cpp +++ b/common/CUDAWrapper.cpp @@ -33,7 +33,7 @@ namespace { // thread to wait until none of one kind is. std::mutex gpu_work_m; std::condition_variable gpu_work_changed; - int gpu_work_beside_n = 0, gpu_image_passes_n = 0; + int gpu_work_beside_n = 0, gpu_image_passes_n = 0, gpu_merges_n = 0; // A card too small for a probe pass beside the main line's image pass, or beside another probe pass: // measured on a 16 Mpx detector, two of them together held 8.6 GB. On such a card a probe pass waits, @@ -56,10 +56,12 @@ namespace { } } -GPUWorkBeside::GPUWorkBeside() { +GPUWorkBeside::GPUWorkBeside(bool after_merges) { std::unique_lock lock(gpu_work_m); if (take_turns()) - gpu_work_changed.wait(lock, [] { return gpu_image_passes_n == 0 && gpu_work_beside_n == 0; }); + gpu_work_changed.wait(lock, [after_merges] { + return gpu_image_passes_n == 0 && gpu_work_beside_n == 0 && (!after_merges || gpu_merges_n == 0); + }); ++gpu_work_beside_n; } @@ -83,6 +85,15 @@ GPUImagePass::~GPUImagePass() { remove(gpu_image_passes_n); } +GPUMerge::GPUMerge() { + std::lock_guard lock(gpu_work_m); + ++gpu_merges_n; +} + +GPUMerge::~GPUMerge() { + remove(gpu_merges_n); +} + void wait_for_gpu_image_passes() { std::unique_lock lock(gpu_work_m); gpu_work_changed.wait(lock, [] { return gpu_image_passes_n == 0; }); diff --git a/common/CUDAWrapper.h b/common/CUDAWrapper.h index b3a4cbf1f..f772cb81f 100644 --- a/common/CUDAWrapper.h +++ b/common/CUDAWrapper.h @@ -115,10 +115,11 @@ void cuda_throw_if_context_lost(); // GPU work that runs beside the main line of a process and gives all of its device memory back when it // ends: a probe pass run on a copy of the run while the run itself scales and merges. Hold one for as // long as that work runs. On a small card it waits first for the main line's image passes (GPUImagePass) -// and for other work beside to end. Build-independent - it only counts. +// and for other work beside to end - and, with after_merges, for the main line's merges (GPUMerge). +// Build-independent - it only counts. class GPUWorkBeside { public: - GPUWorkBeside(); + explicit GPUWorkBeside(bool after_merges = false); ~GPUWorkBeside(); GPUWorkBeside(const GPUWorkBeside &) = delete; GPUWorkBeside &operator=(const GPUWorkBeside &) = delete; @@ -139,6 +140,20 @@ public: GPUImagePass &operator=(const GPUImagePass &) = delete; }; +// The scaling, merging and model validation of a run's main line - not of a probe. A probe that the main +// line never waits for (the speculative geometry probe) waits for it on a small card: the merge holds as +// much of the card as the image pass, and on a card that pages what does not fit rather than failing the +// allocation, the two together would not fail but crawl. It waits for nothing itself, so nothing that the +// main line waits for inside its merge - the rotation-scale walk's probes - can wait in a circle with it. +// Build-independent - it only counts. +class GPUMerge { +public: + GPUMerge(); + ~GPUMerge(); + GPUMerge(const GPUMerge &) = delete; + GPUMerge &operator=(const GPUMerge &) = delete; +}; + // Wait until no GPUImagePass is held - at once if none was. With no timeout: an image pass waits for no // probe, so it ends. void wait_for_gpu_image_passes(); diff --git a/image_analysis/indexing/IndexerThreadPool.cpp b/image_analysis/indexing/IndexerThreadPool.cpp index 40845cd77..820020e15 100644 --- a/image_analysis/indexing/IndexerThreadPool.cpp +++ b/image_analysis/indexing/IndexerThreadPool.cpp @@ -147,14 +147,28 @@ void IndexerThread::Worker(int threadid) { while (true) { std::unique_ptr input; + bool drop = false; // Look for task + handle stop { std::unique_lock lock(m); - c_start.wait(lock, [this] { return stop || state == TaskState::READY; }); + c_start.wait(lock, [this] { return stop || drop_indexers || state == TaskState::READY; }); if (stop && (state != TaskState::READY)) return; - state = TaskState::RUNNING; - input = std::move(task_input); + if (state == TaskState::READY) { + state = TaskState::RUNNING; + input = std::move(task_input); + } else + drop = true; + } + if (drop) { + fft_indexer.reset(); + ffbidx_indexer.reset(); + fftw_indexer.reset(); + { + std::unique_lock lock(m); + drop_indexers = false; + } + c_done.notify_all(); } if (input) { std::unique_ptr tmp_result; @@ -210,6 +224,15 @@ void IndexerThread::Worker(int threadid) { } } +void IndexerThread::Release() { + std::unique_lock lock(m); + if (stop) + return; + drop_indexers = true; + c_start.notify_all(); + c_done.wait(lock, [this] { return !drop_indexers; }); +} + void IndexerThread::Finalize() { { std::unique_lock lock(m); @@ -257,6 +280,19 @@ IndexerThreadPool::IndexerThreadPool(const IndexingSettings &settings, IndexerCo tasks.emplace_back(std::make_unique(std::cref(settings), i, construction, beside)); } +void IndexerThreadPool::ReleaseIndexers() { + for (auto &t : tasks) + t->Release(); +#ifdef JFJOCH_USE_CUDA + // Their buffers went back to the allocation pool, which would keep them reserved for the next + // buffer of whoever asks; hand them to the card. + for (int32_t d = 0; d < get_gpu_count(); d++) { + ScopedGpuPin pin(d); + cuda_trim_allocation_pool(); + } +#endif +} + int IndexerThreadPool::GetFreeWorker() { std::unique_lock lock(m); diff --git a/image_analysis/indexing/IndexerThreadPool.h b/image_analysis/indexing/IndexerThreadPool.h index 7415ea7e6..ea011f9cb 100644 --- a/image_analysis/indexing/IndexerThreadPool.h +++ b/image_analysis/indexing/IndexerThreadPool.h @@ -68,6 +68,7 @@ class IndexerThread { const bool beside_; bool stop = false; + bool drop_indexers = false; // see Release() enum class TaskState {STARTING, IDLE, READY, RUNNING, COMPLETED, ERROR} state = TaskState::STARTING; std::mutex m; std::condition_variable c_running; @@ -83,6 +84,9 @@ public: ~IndexerThread(); std::unique_ptr Run(const DiffractionExperiment &experiment, const std::vector &recip, bool severity_only = false); + // The thread drops the indexers it has built - once the frame it is on, if any, is done - and builds + // them again on the next frame that needs one. Returns when they are gone. + void Release(); void Finalize(); }; @@ -108,6 +112,10 @@ public: // severity_only skips indexing and produces just the spindle severity - see Indexer::Run. IndexerResult Run(const DiffractionExperiment& experiment, const std::vector& recip, bool severity_only = false); + // Give back to the card the device memory of the indexers built so far - an FFT indexer holds + // 0.6-1.2 GB of it - for a stretch with nothing to index. A frame that comes after all builds its + // indexer again, as on first use, and the indexer it gets computes what the dropped one would. + void ReleaseIndexers(); }; diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index 53f284e53..33483f7a0 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -556,9 +556,10 @@ void Rugnux::StartSpeculativeGeometryProbe(std::shared_ptr run, const st sp->hand = hand; // A probe that fails leaves no memo, and the walk's own probe then runs it again and meets the // failure itself. It runs beside the merge of the pass that started it, on the same card: a merge - // that does not fit beside it waits for it to end (GPUWorkBeside). + // that does not fit beside it waits for it to end (GPUWorkBeside) - and on a small card it waits for + // that merge to end instead (GPUMerge). sp->done = std::async(std::launch::async, [run, hand] { - GPUWorkBeside beside; + GPUWorkBeside beside(/*after_merges=*/true); for (const auto &r : {run, hand}) { try { r->RunPipeline(nullptr, /*write_output=*/false, /*geometry_prepass=*/false); @@ -1025,6 +1026,10 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b if (probe_done) return result; } + // A rotation pass indexes nothing after its first pass but the leftover-lattice census, which gives + // them back again when it ends; the image pass and the merge get the card without them. + if (indexer_pool && experiment_.IsRotationIndexing()) + indexer_pool->ReleaseIndexers(); // The rotation-scale walk beside this pass starts once its frames are indexed (see RunAllPasses). if (rotation_scale_check_) rotation_scale_check_->Set(rotation_scale_check_->output_indexed); @@ -1061,7 +1066,11 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b if (full && !cancelled_ && (result.indexing_rate.value_or(0.0f) > 0 || result.rotation_lattice_found) && (config_.run_scaling || !config_.reference_data.empty())) { TimingMark mark(logger, "scale, merge and symmetry"); + std::optional gpu_merge; // what the speculative probe waits for on a small card + if (!gpu_work_beside_) + gpu_merge.emplace(); const bool finished = ScaleMergeAndSymmetry(p); + gpu_merge.reset(); if (p.postrefine_ahead.valid()) { p.postrefine_ahead.get(); p.postrefine_log.ReplayInto(logger); diff --git a/rugnux/RugnuxFirstPassAccept.cpp b/rugnux/RugnuxFirstPassAccept.cpp index 1df384d0b..f0eb897c6 100644 --- a/rugnux/RugnuxFirstPassAccept.cpp +++ b/rugnux/RugnuxFirstPassAccept.cpp @@ -395,6 +395,8 @@ void Rugnux::FirstPassRun::StartLeftoverCensus(const FirstPass &best) { census.primary_spots_pct, n_all, census.primary_intensity_pct, census.frames, census.lattices.size(), std::chrono::duration(std::chrono::steady_clock::now() - census_start).count()); + // The last of the pass's indexing (see RunPipeline). + indexer_pool->ReleaseIndexers(); return census; }); } diff --git a/tests/CUDAMemHelpersTest.cpp b/tests/CUDAMemHelpersTest.cpp index ae944a2cd..33b59cb0f 100644 --- a/tests/CUDAMemHelpersTest.cpp +++ b/tests/CUDAMemHelpersTest.cpp @@ -4,6 +4,7 @@ #include #include "../common/CUDAWrapper.h" +#include #include #include @@ -329,6 +330,65 @@ TEST_CASE("GetGpuFreeMemory_IntegratedCountsThePageCache", "[CUDAMemHelpers]") { CHECK(double(free_bytes) == Catch::Approx(double(cuda_free)).margin(SLACK)); } + +// The pool's indexers go back to the card - not only to the allocation pool - when it is asked, and the +// next frame builds them again. +TEST_CASE("IndexerThreadPool_ReleaseIndexersGivesTheCardBack", "[CUDAMemHelpers]") { + if (get_gpu_count() == 0) + SKIP("No CUDA GPU present"); + + IndexingSettings settings; + settings.Algorithm(IndexingAlgorithmEnum::FFT).IndexingThreads(1); + DiffractionExperiment x; + x.ImportIndexingSettings(settings); + const std::vector recip{{0.1, 0, 0}, {0, 0.1, 0}, {0, 0, 0.1}}; + + // What the allocation pools of all cards hand out: an indexer worker takes its card round-robin. + const auto pooled = [] { + uint64_t sum = 0; + for (int d = 0; d < get_gpu_count(); d++) { + cudaMemPool_t pool = nullptr; + uint64_t used = 0; + if (cudaDeviceGetDefaultMemPool(&pool, d) == cudaSuccess + && cudaMemPoolGetAttribute(pool, cudaMemPoolAttrUsedMemCurrent, &used) == cudaSuccess) + sum += used; + } + return sum; + }; + const auto reserved = [] { + uint64_t sum = 0; + for (int d = 0; d < get_gpu_count(); d++) { + cudaMemPool_t pool = nullptr; + uint64_t r = 0; + if (cudaDeviceGetDefaultMemPool(&pool, d) == cudaSuccess + && cudaMemPoolGetAttribute(pool, cudaMemPoolAttrReservedMemCurrent, &r) == cudaSuccess) + sum += r; + } + return sum; + }; + const uint64_t before = pooled(); + const uint64_t reserved_before = reserved(); + { + IndexerThreadPool pool(settings, IndexerConstruction::OnFirstUse); + const IndexerResult first = pool.Run(x, recip); + REQUIRE(first.error == nullptr); + REQUIRE(pooled() > before + (uint64_t(100) << 20)); + + pool.ReleaseIndexers(); + CHECK(pooled() <= before + (uint64_t(16) << 20)); + CHECK(reserved() <= reserved_before + (uint64_t(64) << 20)); + + const IndexerResult again = pool.Run(x, recip); + CHECK(again.error == nullptr); + CHECK(again.executed == first.executed); + CHECK(again.lattice.size() == first.lattice.size()); + } + for (int d = 0; d < get_gpu_count(); d++) { + ScopedGpuPin pin(d); + cuda_trim_allocation_pool(); + } +} + #endif // A merge that finds the card full waits for the GPU work running beside it: until the last one @@ -347,3 +407,26 @@ TEST_CASE("GPUWorkBeside_WaitEndsWithTheWork", "[CUDAMemHelpers]") { CHECK(wait_for_gpu_work_beside(std::chrono::seconds(60))); ends.join(); } + +// The speculative geometry probe waits for the main line's merge on a card too small for both, and on a +// larger one runs beside it; the rest of the work beside never waits for a merge. +TEST_CASE("GPUWorkBeside_AfterMergesTakesTurnsOnASmallCard", "[CUDAMemHelpers]") { + // As take_turns() in CUDAWrapper.cpp. + const bool small = get_gpu_count() > 0 && get_gpu_total_memory(0) < 12'000'000'000ull; + + std::optional merge(std::in_place); + { GPUWorkBeside other; } // does not wait + + std::atomic started = false; + std::thread probe([&started] { + GPUWorkBeside beside(/*after_merges=*/true); + started = true; + }); + if (small) { + std::this_thread::sleep_for(std::chrono::milliseconds(200)); + CHECK_FALSE(started); + } + merge.reset(); + probe.join(); + CHECK(started); +} From fd1fb3e25e144d898e0f59ba4b26a75bc094b256 Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 15 Oct 2026 12:05:48 +0200 Subject: [PATCH 2/3] Rotation scaling on the GPU: a partial's frame from the frame CSR, its two flag bytes in one, output unchanged - The partials lie in frame order, so a partial's frame is the last frame starting at or before it in the frame CSR (FrameOf, a binary search) and is no longer stored: UpdateCorr, SmoothCorr (which takes the fulls' stored frames through the same FrameOf), FilterFrame, the combine, and the ingest's flux, delta_phi and partiality kernels. 4 B less per observation. - on_ice and clipped/overloaded share one flags byte (on ice at bit 0, the host's OBS_CLIPPED / OBS_OVERLOADED one bit up); a full's on_ice byte is 0/1, so the kernels shared by the two read either alike. 1 B less. Measured on an RTX 5080, largest engine of each set, round 2 -> round 3, MB: set obs/hkl ingest peak (B/refl) post-decision peak 4nwv 32.7 M/15.4 M 2336 (72) -> 2172 (67) 3625 -> 3461 8qaw 122.7 M/ 8.6 M 7421 (61) -> 6808 (56) 8156 -> 7543 9v16 54.7 M/ 5.7 M 3419 (63) -> 3146 (58) 3774 -> 3500 7qij 78.1 M/ 2.2 M 4707 (60) -> 4317 (55) 4707 -> 4317 8sa8 50.0 M/ 5.0 M 3134 (63) -> 2884 (58) 3603 -> 3353 Kernel time on 8qaw (nsys, all calls): UpdateCorr 364 -> 365 ms, SmoothCorr 143 -> 112 ms, combine 477 -> 476 ms, ingest partiality/delta_phi/flux 113 -> 109 ms - the search costs nothing measurable, the saved reads pay for it. GPUScalingPeakBytes refitted: 56 + 99 -> 51 B/obs + 98 B/hkl, within 10%. DeviceBytes: 58 -> 53 B per observation, ingest 60 -> 55 B per reflection. p.hkl byte-identical to the rc175 baseline on 15 battery sets (4nwv, 8qaw, 9v16, 7qij, 8sa8, 11if, 5reo, 6oel, 9jzo, 9min, metformin, lyso_x06da_ref, thau_bl1a_3p8keV, insu_I_x06da_5keV_2, myob_x06da); compare shows no change. [RotationScale] and ErrorModel_DeviceFitIsTheHostFit pass. Co-Authored-By: Claude Opus 5.5 --- .../scale_merge/RotationScaleMerge.cpp | 21 ++- .../scale_merge/RotationScaleMergeGPU.cu | 151 +++++++++--------- .../scale_merge/RotationScaleMergeGPU.h | 8 +- 3 files changed, 89 insertions(+), 91 deletions(-) diff --git a/image_analysis/scale_merge/RotationScaleMerge.cpp b/image_analysis/scale_merge/RotationScaleMerge.cpp index 7ab124f79..7cfc12586 100644 --- a/image_analysis/scale_merge/RotationScaleMerge.cpp +++ b/image_analysis/scale_merge/RotationScaleMerge.cpp @@ -752,10 +752,10 @@ size_t RotationScaleMerge::GPUScalingPeakBytes() const { // The device scaling's peak from this decision on - the observations taken from the reflections // (IngestObservations) included, the sort before it not - fitted to the engine's device buffers // measured at their peak on fourteen engines of five large data sets (23-123 M observations, - // 1.6-15 M distinct hkl): within 9% on each, the furthest under a P1 merge of a large cell. Unlike + // 1.6-15 M distinct hkl): within 10% on each, the furthest under a P1 merge of a large cell. Unlike // DeviceBytes, an estimate and not an upper bound - it decides where the GPU clearly cannot take the // scaling. - constexpr size_t PEAK_BYTES_PER_OBS = 56, PEAK_BYTES_PER_HKL = 99; + constexpr size_t PEAK_BYTES_PER_OBS = 51, PEAK_BYTES_PER_HKL = 98; return PEAK_BYTES_PER_OBS * static_cast(n_partials_obs) + PEAK_BYTES_PER_HKL * rawrun_start.size(); } @@ -763,18 +763,18 @@ size_t RotationScaleMerge::DeviceBytes() const { // What each costs at the peak: an observation's arrays, and its group CSR until the combine; a full's // arrays with its error-model sample, or with the surface terms; an hkl's raw run and its ASU group's // sums. Fitted as an upper bound to the engine's device buffers measured at their peak on fourteen - // engines of five large data sets (23-123 M observations): up to 14% over. The events are counted + // engines of five large data sets (23-123 M observations): up to 15% over. The events are counted // on the device path's ingest; elsewhere every observation may be an event of its own. - constexpr size_t BYTES_PER_OBS = 58, BYTES_PER_EVENT = 30, BYTES_PER_HKL = 86; + constexpr size_t BYTES_PER_OBS = 53, BYTES_PER_EVENT = 30, BYTES_PER_HKL = 86; const auto events = static_cast(resident_ingest ? rocking_events_at_ingest : n_partials_obs); const size_t bytes = BYTES_PER_OBS * static_cast(n_partials_obs) + BYTES_PER_EVENT * events + BYTES_PER_HKL * rawrun_start.size(); // The resident ingest's own peak, over every reflection read and every run before the resolution // cut, measured on the same engines: the observations taken from the reflections hold the - // observation arrays, the place the cut gave each reflection and the runs it kept, 60 to 72 bytes + // observation arrays, the place the cut gave each reflection and the runs it kept, 55 to 67 bytes // per reflection; the sort and the cut before them 32 per reflection and 56 per run. The two 64 MB // staging buffers beside either. - constexpr size_t INGEST_BYTES_PER_REFLECTION = 60, INGEST_BYTES_PER_RUN = 42; + constexpr size_t INGEST_BYTES_PER_REFLECTION = 55, INGEST_BYTES_PER_RUN = 42; constexpr size_t INGEST_STAGING_BYTES = size_t(160) << 20; const size_t ingest = INGEST_BYTES_PER_REFLECTION * ingest_reflections + INGEST_BYTES_PER_RUN * ingest_runs + INGEST_STAGING_BYTES; @@ -826,8 +826,7 @@ void RotationScaleMerge::IngestDevice(const RotationScaleMerge *observations_fro constexpr int NFLOAT = 10; std::unique_ptr fbuf[NFLOAT]; for (auto &b : fbuf) b.reset(new float[slice]); - std::unique_ptr frm(new int32_t[slice]); - std::unique_ptr onice(new uint8_t[slice]), clip(new uint8_t[slice]); + std::unique_ptr flags(new uint8_t[slice]); for (int at = 0; at < n; at += slice) { const int cnt = std::min(n - at, slice); ParallelChunks(cnt, nthreads, [&](int lo, int hi) { @@ -837,7 +836,7 @@ void RotationScaleMerge::IngestDevice(const RotationScaleMerge *observations_fro fbuf[3][i] = o.partiality; fbuf[4][i] = o.zeta; fbuf[5][i] = o.corr; fbuf[6][i] = o.bkg; fbuf[7][i] = o.var_bkg; fbuf[8][i] = o.image_number; fbuf[9][i] = o.d; - frm[i] = o.frame; onice[i] = o.on_ice; clip[i] = o.clipped; + flags[i] = static_cast((o.on_ice ? 1 : 0) | (o.clipped << 1)); // SetObsFlags } }); constexpr F fields[NFLOAT] = {F::I, F::Sigma, F::PrescalingCorr, F::Partiality, @@ -845,9 +844,7 @@ void RotationScaleMerge::IngestDevice(const RotationScaleMerge *observations_fro F::D}; for (int k = 0; k < NFLOAT; ++k) gpu_->SetObsField(fields[k], at, cnt, fbuf[k].get()); - gpu_->SetObsFrame(at, cnt, frm.get()); - gpu_->SetObsOnIce(at, cnt, onice.get()); - gpu_->SetObsClipped(at, cnt, clip.get()); + gpu_->SetObsFlags(at, cnt, flags.get()); } auto pos = std::make_shared(); pos->px.resize(n); pos->py.resize(n); diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu index 612072e94..e14faab8b 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.cu +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.cu @@ -60,6 +60,32 @@ namespace { // (the group reduction above all) then fetch one sector for it instead of one per field. struct alignas(16) ObsRecord { float I, sigma, partiality, prescaling_corr; }; // floats 0, 1, 2, 3 + // A partial's flags, one byte: on an ice ring, then RotationScaleMerge::OBS_CLIPPED and + // OBS_OVERLOADED one bit up. A full's on_ice byte is 0 or OBS_ON_ICE_FLAG, so the kernels shared by + // the two read either the same way. + constexpr uint8_t OBS_ON_ICE_FLAG = 1; + constexpr uint8_t OBS_CLIPPED_FLAG = 2; + constexpr uint8_t OBS_OVERLOADED_FLAG = 4; + + // An observation's frame. The fulls carry theirs (frame); the partials lie in frame order, so theirs + // is the last frame starting at or before them in the frame CSR (start, n_frames entries) - an empty + // frame starts where the next one does, so it is never the last. + struct FrameOf { + const int32_t *frame; + const int32_t *start; + int n_frames; + __device__ __forceinline__ int operator()(int i) const { + if (frame) return frame[i]; + int lo = 0, hi = n_frames - 1; + while (lo < hi) { + const int mid = (lo + hi + 1) / 2; + if (start[mid] <= i) lo = mid; + else hi = mid - 1; + } + return lo; + } + }; + // The observation's inverse-variance weight in the scaling reference (0 = not in it), and its // scaled intensity. The same filter as the host's ReferenceWeight. __device__ __forceinline__ double ReferenceWeight(float I, float sigma, float partiality, float corr, @@ -108,7 +134,7 @@ namespace { const uint8_t *__restrict__ on_ice, const double *__restrict__ group_mean, float &coeff) { const int g = group[i]; - if (!((g >= 0) && !on_ice[i] && isfinite(zeta[i]) && zeta[i] > 0.0f && o.partiality >= min_partiality)) + if (!((g >= 0) && !(on_ice[i] & OBS_ON_ICE_FLAG) && isfinite(zeta[i]) && zeta[i] > 0.0f && o.partiality >= min_partiality)) return false; const double mean = group_mean[g]; if (!isfinite(mean)) return false; @@ -178,7 +204,7 @@ namespace { const int32_t *__restrict__ frame_start, const int32_t *__restrict__ frame_count, const ObsRecord *__restrict__ obs, const float *__restrict__ corr, - const uint8_t *__restrict__ on_ice, const int32_t *__restrict__ group, + const uint8_t *__restrict__ flags, const int32_t *__restrict__ group, const double *__restrict__ group_mean, double *__restrict__ cc_out, int64_t *__restrict__ cc_n_out) { const int f = blockIdx.x; @@ -188,7 +214,7 @@ namespace { double sx = 0, sy = 0, sx2 = 0, sy2 = 0, sxy = 0; long nl = 0; for (int i = lo + threadIdx.x; i < hi; i += blockDim.x) { - if (on_ice[i]) continue; + if (flags[i] & OBS_ON_ICE_FLAG) continue; const int g = group[i]; if (g < 0) continue; const ObsRecord o = obs[i]; @@ -220,13 +246,14 @@ namespace { // SmoothG corr adjust: corr[i] *= ratio[frame[i]] for frames flagged apply, in double then stored // as float - matching CPU SmoothG's `corr = float(corr * (g/g_smooth))`. Grid-stride, resident corr. - __global__ void SmoothCorrKernel(int n_obs, const int32_t *__restrict__ frame, + __global__ void SmoothCorrKernel(int n_obs, FrameOf frame_of, const uint8_t *__restrict__ apply, const double *__restrict__ ratio, float *__restrict__ corr) { for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) { - if (!apply[frame[i]]) continue; + const int f = frame_of(i); + if (!apply[f]) continue; const float c = corr[i]; - if (isfinite(c)) corr[i] = float(double(c) * ratio[frame[i]]); + if (isfinite(c)) corr[i] = float(double(c) * ratio[f]); } } @@ -247,19 +274,19 @@ namespace { } // Zero corr on the rejected frames (--min-image-cc), grid-stride over the resident corr. - __global__ void FilterFrameKernel(int n_obs, const int32_t *__restrict__ frame, + __global__ void FilterFrameKernel(int n_obs, FrameOf frame_of, const uint8_t *__restrict__ reject, float *__restrict__ corr) { for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) - if (reject[frame[i]]) corr[i] = 0.0f; + if (reject[frame_of(i)]) corr[i] = 0.0f; } // corr = prescaling_corr / (partiality * G[frame]) for fitted frames; unchanged otherwise (grid-stride). - __global__ void UpdateCorrKernel(int n_obs, const int32_t *__restrict__ frame, + __global__ void UpdateCorrKernel(int n_obs, FrameOf frame_of, const ObsRecord *__restrict__ obs, const double *__restrict__ g, const uint8_t *__restrict__ scaled, float *__restrict__ corr) { for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n_obs; i += gridDim.x * blockDim.x) { - const int f = frame[i]; + const int f = frame_of(i); if (!scaled[f]) continue; const ObsRecord o = obs[i]; const double denom = double(o.partiality) * g[f]; @@ -281,9 +308,6 @@ namespace { return isfinite(obs[i].I) && isfinite(obs[i].sigma) && obs[i].sigma > 0.0f; } - // RotationScaleMerge::OBS_CLIPPED / OBS_OVERLOADED, the flag bits a partial's `clipped` byte carries. - constexpr uint8_t OBS_CLIPPED_FLAG = 1; - constexpr uint8_t OBS_OVERLOADED_FLAG = 2; // All device pointers + scalars the combine kernels need, passed by value. struct CombineParams { @@ -295,8 +319,8 @@ namespace { const ObsRecord *__restrict__ obs; const float *__restrict__ corr, *__restrict__ bkg, *__restrict__ var_bkg, *__restrict__ image_number, *__restrict__ d; - const int32_t *__restrict__ frame; - const uint8_t *__restrict__ on_ice, *__restrict__ clipped; + FrameOf frame_of; + const uint8_t *__restrict__ flags; const int32_t *__restrict__ perm, *__restrict__ rr_start, *__restrict__ rr_count, *__restrict__ rr_h, *__restrict__ rr_k, *__restrict__ rr_l, *__restrict__ rr_group; @@ -361,13 +385,13 @@ namespace { double sum_w = 0.0, sum_wI = 0.0, sum_partiality = 0.0; float d = NAN; const int first = p.perm[ev_start]; - int peak_outcome = p.frame[first]; + int peak_outcome = p.frame_of(first); float peak_frame = p.image_number[first]; int peak_obs = first; float peak_partiality = -1.0f; - const bool on_ice = p.on_ice[first]; + const bool on_ice = p.flags[first] & OBS_ON_ICE_FLAG; uint8_t flags = 0; - for (int m = ev_start; m <= ev_end; ++m) flags |= p.clipped[p.perm[m]]; + for (int m = ev_start; m <= ev_end; ++m) flags |= p.flags[p.perm[m]]; // An overloaded event is not a measurement of the reflection; see the host combine. if (flags & OBS_OVERLOADED_FLAG) { ++n_overloaded; continue; } const bool clipped = flags & OBS_CLIPPED_FLAG; @@ -381,7 +405,7 @@ namespace { sum_partiality += p.obs[i].partiality; if (p.obs[i].partiality > peak_partiality) { peak_partiality = p.obs[i].partiality; - peak_outcome = p.frame[i]; + peak_outcome = p.frame_of(i); peak_frame = p.image_number[i]; peak_obs = i; } @@ -1115,10 +1139,10 @@ namespace { std::string("RotationScaleMergeGPU: ") + what + ": " + cudaGetErrorString(e)); } - // Device bytes per observation in the arrays SetPartialsLayout allocates: twelve floats (the record's - // I, sigma, partiality and prescaling_corr; zeta, corr, bkg, var_bkg, image_number, d, px, py), frame - // (4 bytes), on_ice and clipped (1 each). The merge's other buffers come on top of these. - constexpr size_t OBS_BYTES = 12 * 4 + 4 + 2 * 1; + // Device bytes per observation in the arrays SetPartialsLayout allocates: ten floats (the record's + // I, sigma, partiality and prescaling_corr; zeta, corr, bkg, var_bkg, image_number, d) and the flags + // byte. The merge's other buffers come on top of these. + constexpr size_t OBS_BYTES = 10 * 4 + 1; // Every buffer here is sized by the data set, and the scaling does not fall back to the CPU // mid-run: the two paths differ in the last bits, so the result would depend on the card. A @@ -1230,8 +1254,9 @@ struct RotationScaleMergeGPU::Impl { // arrays that replace them. void ReleaseFulls(); CudaDevicePtr zeta; - CudaDevicePtr on_ice, clipped; - CudaDevicePtr frame; + CudaDevicePtr flags; // OBS_ON_ICE_FLAG | OBS_CLIPPED_FLAG | OBS_OVERLOADED_FLAG + // A partial's frame, from the frame CSR (the partials lie in frame order). + FrameOf PartialFrame() const { return FrameOf{nullptr, frame_start.get(), n_frames}; } CudaDevicePtr corr; // mutable, resident across iterations CudaDevicePtr frame_start, frame_count; // per space group @@ -1378,9 +1403,7 @@ void RotationScaleMergeGPU::Impl::AllocObservations(int n_in) { zeta = Alloc(n); corr = Alloc(n); bkg = Alloc(n); var_bkg = Alloc(n); image_number = Alloc(n); d_obs = Alloc(n); - frame = Alloc(n); - on_ice = Alloc(n); - clipped = Alloc(n); + flags = Alloc(n); } void RotationScaleMergeGPU::Impl::AllocFrames(const int32_t *fstart, const int32_t *fcount) { @@ -1471,7 +1494,7 @@ void RotationScaleMergeGPU::CopyObservationsFrom(const RotationScaleMergeGPU &fr copy(d.obs, f.obs, n); copy(d.zeta, f.zeta, n); copy(d.corr, f.corr, n); copy(d.bkg, f.bkg, n); copy(d.var_bkg, f.var_bkg, n); copy(d.image_number, f.image_number, n); copy(d.d_obs, f.d_obs, n); - copy(d.frame, f.frame, n); copy(d.on_ice, f.on_ice, n); copy(d.clipped, f.clipped, n); + copy(d.flags, f.flags, n); const size_t nr = std::max(1, d.n_runs); d.perm = d.Alloc(std::max(1, d.n_perm)); d.rr_start = d.Alloc(nr); d.rr_count = d.Alloc(nr); @@ -1482,19 +1505,9 @@ void RotationScaleMergeGPU::CopyObservationsFrom(const RotationScaleMergeGPU &fr CudaCheck(cudaStreamSynchronize(d.s()), "copy observations"); } -void RotationScaleMergeGPU::SetObsFrame(int offset, int count, const int32_t *frame) { +void RotationScaleMergeGPU::SetObsFlags(int offset, int count, const uint8_t *flags) { DeviceGuard guard(impl_->device, impl_->available); - UploadChunk(impl_->frame, offset, count, frame, impl_->s()); -} - -void RotationScaleMergeGPU::SetObsOnIce(int offset, int count, const uint8_t *on_ice) { - DeviceGuard guard(impl_->device, impl_->available); - UploadChunk(impl_->on_ice, offset, count, on_ice, impl_->s()); -} - -void RotationScaleMergeGPU::SetObsClipped(int offset, int count, const uint8_t *clipped) { - DeviceGuard guard(impl_->device, impl_->available); - UploadChunk(impl_->clipped, offset, count, clipped, impl_->s()); + UploadChunk(impl_->flags, offset, count, flags, impl_->s()); } void RotationScaleMergeGPU::BuildGroups(int n_groups, const int32_t *rawrun_group, int n_sorted, @@ -1566,10 +1579,10 @@ void RotationScaleMergeGPU::ScalePartials(int iters, double min_partiality, bool d.obs.get(), d.corr.get(), d.group_mean.get()); CudaCheck(cudaGetLastError(), "ReduceGroupMeansKernel launch"); FitPerFrameGKernel<<s()>>>(d.n_frames, d.frame_start.get(), d.frame_count.get(), - nullptr, min_partiality, d.group.get(), d.obs.get(), d.zeta.get(), d.on_ice.get(), d.group_mean.get(), + nullptr, min_partiality, d.group.get(), d.obs.get(), d.zeta.get(), d.flags.get(), d.group_mean.get(), long(MIN_REFLECTIONS), d.g.get(), d.scaled.get(), d.info.get()); CudaCheck(cudaGetLastError(), "FitPerFrameGKernel launch"); - UpdateCorrKernel<<s()>>>(d.n_obs, d.frame.get(), d.obs.get(), + UpdateCorrKernel<<s()>>>(d.n_obs, d.PartialFrame(), d.obs.get(), d.g.get(), d.scaled.get(), d.corr.get()); } CudaCheck(cudaGetLastError(), "kernel launch"); @@ -1892,7 +1905,7 @@ void RotationScaleMergeGPU::SmoothCorr(const uint8_t *apply, const double *ratio d.Upload(d.smooth_apply, apply, d.n_frames); d.Upload(d.smooth_ratio, ratio, d.n_frames); const int blocks = std::min(65535, (d.n_obs + BLK - 1) / BLK); - SmoothCorrKernel<<s()>>>(d.n_obs, d.frame.get(), d.smooth_apply.get(), + SmoothCorrKernel<<s()>>>(d.n_obs, d.PartialFrame(), d.smooth_apply.get(), d.smooth_ratio.get(), d.corr.get()); CudaCheck(cudaGetLastError(), "smooth corr launch"); CudaCheck(cudaStreamSynchronize(impl_->s()), "smooth corr sync"); @@ -1905,7 +1918,7 @@ void RotationScaleMergeGPU::SmoothFullsCorr(const uint8_t *apply, const double * d.Upload(d.smooth_apply, apply, d.n_frames); d.Upload(d.smooth_ratio, ratio, d.n_frames); const int blocks = std::min(65535, (d.n_fulls + BLK - 1) / BLK); - SmoothCorrKernel<<s()>>>(d.n_fulls, d.f_frame.get(), d.smooth_apply.get(), + SmoothCorrKernel<<s()>>>(d.n_fulls, FrameOf{d.f_frame.get(), nullptr, 0}, d.smooth_apply.get(), d.smooth_ratio.get(), d.f_corr.get()); CudaCheck(cudaGetLastError(), "smooth fulls corr launch"); CudaCheck(cudaStreamSynchronize(impl_->s()), "smooth fulls corr sync"); @@ -1931,7 +1944,7 @@ void RotationScaleMergeGPU::FilterCorrByFrame(const uint8_t *reject) { auto &d = *impl_; d.Upload(d.filter_reject, reject, d.n_frames); const int blocks = std::min(65535, (d.n_obs + BLK - 1) / BLK); - FilterFrameKernel<<s()>>>(d.n_obs, d.frame.get(), d.filter_reject.get(), d.corr.get()); + FilterFrameKernel<<s()>>>(d.n_obs, d.PartialFrame(), d.filter_reject.get(), d.corr.get()); CudaCheck(cudaGetLastError(), "frame filter launch"); CudaCheck(cudaStreamSynchronize(impl_->s()), "frame filter sync"); } @@ -1947,7 +1960,7 @@ void RotationScaleMergeGPU::ComputePartialCC(double min_partiality, double *cc_o d.obs.get(), d.corr.get(), d.group_mean.get()); CudaCheck(cudaGetLastError(), "ReduceGroupMeansKernel launch"); PerFrameCCKernel<<s()>>>(d.n_frames, min_partiality, - d.frame_start.get(), d.frame_count.get(), d.obs.get(), d.corr.get(), d.on_ice.get(), d.group.get(), + d.frame_start.get(), d.frame_count.get(), d.obs.get(), d.corr.get(), d.flags.get(), d.group.get(), d.group_mean.get(), d.cc.get(), d.cc_n.get()); CudaCheck(cudaGetLastError(), "partial CC launch"); CudaCheck(cudaStreamSynchronize(impl_->s()), "partial CC sync"); @@ -1994,7 +2007,7 @@ int RotationScaleMergeGPU::Combine(const int32_t *rawrun_group, double min_parti p.max_frame_gap = max_frame_gap; p.obs = d.obs.get(); p.corr = d.corr.get(); p.bkg = d.bkg.get(); p.var_bkg = d.var_bkg.get(); p.image_number = d.image_number.get(); p.d = d.d_obs.get(); - p.frame = d.frame.get(); p.on_ice = d.on_ice.get(); p.clipped = d.clipped.get(); + p.frame_of = d.PartialFrame(); p.flags = d.flags.get(); p.perm = d.perm.get(); p.rr_start = d.rr_start.get(); p.rr_count = d.rr_count.get(); p.rr_h = d.rr_h.get(); p.rr_k = d.rr_k.get(); p.rr_l = d.rr_l.get(); p.rr_group = rr_group.get(); p.rr_nevents = rr_nevents.get(); @@ -2497,7 +2510,7 @@ namespace { struct IngestArrays { ObsRecord *obs; float *zeta, *corr, *bkg, *var_bkg, *image_number, *d, *dphi; - uint8_t *on_ice, *clipped; + uint8_t *flags; const int32_t *new_idx; // the cut's place of each reflection (an exclusive prefix over the kept) int n, n_keep; // reflections and observations }; @@ -2524,15 +2537,10 @@ namespace { a.image_number[i] = r.image_number; a.d[i] = r.d; a.dphi[i] = r.delta_phi_deg; - a.on_ice[i] = r.on_ice_ring ? 1 : 0; - a.clipped[i] = (r.clipped ? 1 : 0) | (r.overloaded ? 2 : 0); // OBS_CLIPPED | OBS_OVERLOADED + a.flags[i] = (r.on_ice_ring ? OBS_ON_ICE_FLAG : 0) | (r.clipped ? OBS_CLIPPED_FLAG : 0) + | (r.overloaded ? OBS_OVERLOADED_FLAG : 0); } - __global__ void IngestFrameKernel(int n_frames, const int32_t *__restrict__ start, int32_t *__restrict__ frame) { - for (int o = blockIdx.x; o < n_frames; o += gridDim.x) - for (int i = start[o] + threadIdx.x; i < start[o + 1]; i += blockDim.x) - frame[i] = o; - } // One thread per frame, its reflections in order: the host's sum to the bit. __global__ void IngestMeanBkgKernel(int n_frames, const int32_t *__restrict__ start, @@ -2548,10 +2556,10 @@ namespace { } } - __global__ void IngestFluxKernel(int n, const int32_t *__restrict__ frame, const double *__restrict__ flux, + __global__ void IngestFluxKernel(int n, FrameOf frame_of, const double *__restrict__ flux, ObsRecord *__restrict__ obs) { for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) - obs[i].prescaling_corr = __double2float_rn(__ddiv_rn(double(obs[i].prescaling_corr), flux[frame[i]])); + obs[i].prescaling_corr = __double2float_rn(__ddiv_rn(double(obs[i].prescaling_corr), flux[frame_of(i)])); } __global__ void IngestImageKeyKernel(int n, const float *__restrict__ image_number, @@ -2677,7 +2685,7 @@ namespace { __global__ void IngestDeltaPhiKernel(int n_runs, const int32_t *__restrict__ start, const int32_t *__restrict__ count, const int32_t *__restrict__ perm, const int32_t *__restrict__ rr_h, const int32_t *__restrict__ rr_k, - const int32_t *__restrict__ rr_l, const int32_t *__restrict__ frame, + const int32_t *__restrict__ rr_l, FrameOf frame_of, DeltaPhiGeometry g, float *__restrict__ dphi, unsigned long long *__restrict__ n_resolved) { constexpr float PI_F = 3.14159265358979323846f; @@ -2686,7 +2694,7 @@ namespace { const float h = float(rr_h[r]), k = float(rr_k[r]), l = float(rr_l[r]); for (int q = start[r]; q < start[r] + count[r]; ++q) { const int i = perm[q]; - const int o = frame[i]; + const int o = frame_of(i); if (!g.ok[o]) continue; const float *rc = g.recip + 9 * o; const Vec3 a{rc[0], rc[1], rc[2]}, b{rc[3], rc[4], rc[5]}, c{rc[6], rc[7], rc[8]}; @@ -2795,13 +2803,13 @@ namespace { const float *mos; // per frame double wedge, bandwidth_sigma, half_wavelength; }; - __global__ void IngestPartialityKernel(int n, const int32_t *__restrict__ frame, const float *__restrict__ zeta, + __global__ void IngestPartialityKernel(int n, FrameOf frame_of, const float *__restrict__ zeta, const float *__restrict__ d, const float *__restrict__ dphi, PartialityParams q, ObsRecord *__restrict__ obs, float *__restrict__ corr) { constexpr double PI_D = 3.14159265358979323846; for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += gridDim.x * blockDim.x) { const float z = zeta[i]; - float mos = q.mos[frame[i]]; + float mos = q.mos[frame_of(i)]; if (q.bandwidth_sigma > 0.0 && d[i] > 0.0f) { const double sin_theta = __ddiv_rn(q.half_wavelength, double(d[i])); const double sigma_bw = __ddiv_rn(__dmul_rn(__dmul_rn(__ddiv_rn(180.0, PI_D), q.bandwidth_sigma), sin_theta), @@ -3116,8 +3124,8 @@ void RotationScaleMergeGPU::IngestObservations(const Reflection *const *refl, co d.AllocObservations(n_keep); d.in_dphi = d.Alloc(std::max(1, n_keep)); const IngestArrays a{d.obs.get(), d.zeta.get(), d.corr.get(), d.bkg.get(), d.var_bkg.get(), d.image_number.get(), - d.d_obs.get(), d.in_dphi.get(), d.on_ice.get(), - d.clipped.get(), d.in_new_idx.get(), d.in_src_start[d.n_frames], n_keep}; + d.d_obs.get(), d.in_dphi.get(), d.flags.get(), d.in_new_idx.get(), + d.in_src_start[d.n_frames], n_keep}; d.StageReflections(refl, count, nthreads, [&](int cnt, int base, const Reflection *src) { IngestExtractKernel<<<(cnt + BLK - 1) / BLK, BLK, 0, d.s()>>>(cnt, base, src, a); CudaCheck(cudaGetLastError(), "IngestExtractKernel launch"); @@ -3138,17 +3146,10 @@ void RotationScaleMergeGPU::IngestObservations(const Reflection *const *refl, co } }); - // The kept observations are in frame order, as the reflections were. - std::vector start(d.n_frames + 1, n_keep); - for (int o = 0; o < d.n_frames; ++o) start[o] = d.in_layout.frame_start[o]; - CudaDevicePtr d_start; - d.Upload(d_start, start.data(), d.n_frames + 1); - IngestFrameKernel<<>>(d.n_frames, d_start.get(), - d.frame.get()); - CudaCheck(cudaGetLastError(), "IngestFrameKernel launch"); + // The kept observations are in frame order, as the reflections were (IngestSort laid out the CSR). CudaDevicePtr flux; d.Upload(flux, frame_flux, d.n_frames); - IngestFluxKernel<<>>(n_keep, d.frame.get(), flux.get(), d.obs.get()); + IngestFluxKernel<<>>(n_keep, d.PartialFrame(), flux.get(), d.obs.get()); CudaCheck(cudaGetLastError(), "IngestFluxKernel launch"); CudaCheck(cudaStreamSynchronize(d.s()), "flux"); } @@ -3166,7 +3167,7 @@ int64_t RotationScaleMergeGPU::IngestDeltaPhi(const DeltaPhiFrame &g) { const DeltaPhiGeometry dg{recip.get(), ok.get(), v(g.S0), v(g.m1), v(g.m2), v(g.m3), g.m2_S0, g.m3_S0}; IngestDeltaPhiKernel<<>>(d.n_runs, d.rr_start.get(), d.rr_count.get(), d.perm.get(), d.rr_h.get(), d.rr_k.get(), d.rr_l.get(), - d.frame.get(), dg, d.in_dphi.get(), n_resolved.get()); + d.PartialFrame(), dg, d.in_dphi.get(), n_resolved.get()); CudaCheck(cudaGetLastError(), "IngestDeltaPhiKernel launch"); unsigned long long out = 0; CopyAndWait(&out, n_resolved.get(), sizeof(out), cudaMemcpyDeviceToHost, d.s(), "resolved count"); @@ -3209,7 +3210,7 @@ void RotationScaleMergeGPU::IngestPartiality(const float *mosaicity, double wedg CudaDevicePtr mos; d.Upload(mos, mosaicity, d.n_frames); const PartialityParams q{mos.get(), wedge_deg, bandwidth_sigma, half_wavelength_A}; - IngestPartialityKernel<<>>(d.n_obs, d.frame.get(), d.zeta.get(), d.d_obs.get(), + IngestPartialityKernel<<>>(d.n_obs, d.PartialFrame(), d.zeta.get(), d.d_obs.get(), d.in_dphi.get(), q, d.obs.get(), d.corr.get()); CudaCheck(cudaGetLastError(), "IngestPartialityKernel launch"); CudaCheck(cudaStreamSynchronize(d.s()), "partiality"); diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.h b/image_analysis/scale_merge/RotationScaleMergeGPU.h index 054439b13..b479fe174 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.h +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.h @@ -39,7 +39,7 @@ public: // One slice ([offset, offset+count) of n_obs) of one immutable per-observation field, so the // host can stage a bounded chunk at a time out of its AoS record: staged whole and all at once, - // the fourteen arrays are more than half the entire observation payload held live again. + // the arrays are more than half the entire observation payload held live again. // Corr0 seeds the resident, mutable corr (SetCorr refreshes it). Bkg/VarBkg/ImageNumber/D // are the combine inputs (image-local background, fractional frame position for event // contiguity and resolution; the predicted detector position carried to the full for the @@ -47,9 +47,9 @@ public: enum class ObsField { I, Sigma, PrescalingCorr, Partiality, Zeta, Corr0, Bkg, VarBkg, ImageNumber, D }; void SetObsField(ObsField f, int offset, int count, const float *v); - void SetObsFrame(int offset, int count, const int32_t *frame); - void SetObsOnIce(int offset, int count, const uint8_t *on_ice); - void SetObsClipped(int offset, int count, const uint8_t *clipped); + // The flags byte: on an ice ring (bit 0), then the OBS_CLIPPED / OBS_OVERLOADED bits of + // RotationScaleMerge one bit up. The frame is not uploaded: the CSR above gives it. + void SetObsFlags(int offset, int count, const uint8_t *flags); // Every per-observation field the uploads above set, the frame CSR and the raw-hkl runs, copied from // `from` (which has them all, on its own device or this one) instead of being staged again from the From 2091c0524a19ede5278eff25ee5ed78a9ea51a8c Mon Sep 17 00:00:00 2001 From: Filip Leonarski Date: Thu, 15 Oct 2026 12:52:14 +0200 Subject: [PATCH 3/3] GPU memory: a probe pass hands its memory to the card when it ends The speculative geometry probe and the rotation-scale walk's indexing probes gave their indexers' and image workers' buffers back to the allocation pool only, which kept them reserved under the main line's merge and analyses: 5.5 GB on 8qaw and 5.3 GB on 4nwv on the RTX 5080, until the end of the run. The pool is trimmed (every card: cuda_trim_allocation_pools) when such a probe ends, as IndexerThreadPool::ReleaseIndexers now does through the same helper. Measured (RTX 5080): device memory held at the end of the final merge 8.7 -> 3.4 GB (4nwv), 14.5 -> 7.7 GB (8qaw). Output unchanged: p.hkl byte-identical on cytc_x10sa, myob_x10sa, 4nwv and 8qaw, on this card and with its total read as 8.55 GB. Co-Authored-By: Claude Opus 5.5 --- image_analysis/indexing/CUDAMemHelpers.h | 8 ++++++++ image_analysis/indexing/IndexerThreadPool.cpp | 5 +---- rugnux/Rugnux.cpp | 5 +++++ rugnux/RugnuxPasses.cpp | 5 +++++ 4 files changed, 19 insertions(+), 4 deletions(-) diff --git a/image_analysis/indexing/CUDAMemHelpers.h b/image_analysis/indexing/CUDAMemHelpers.h index c5341a484..f0b2128bf 100644 --- a/image_analysis/indexing/CUDAMemHelpers.h +++ b/image_analysis/indexing/CUDAMemHelpers.h @@ -204,6 +204,14 @@ inline void cuda_trim_allocation_pool() { cudaGetLastError(); } +// The same on every card: for work that has ended and whose threads were pinned round-robin (pin_gpu). +inline void cuda_trim_allocation_pools() { + for (int32_t d = 0; d < get_gpu_count(); d++) { + ScopedGpuPin pin(d); + cuda_trim_allocation_pool(); + } +} + // Whether `bytes` more device memory can be had on the current device, for a step whose buffers are // sized by the data set and which ends the run when they do not fit. The free memory counts what the // pool holds with no buffer in it (up to CUDA_MEM_POOL_RELEASE_THRESHOLD) as taken, so that is handed diff --git a/image_analysis/indexing/IndexerThreadPool.cpp b/image_analysis/indexing/IndexerThreadPool.cpp index 820020e15..2707482ad 100644 --- a/image_analysis/indexing/IndexerThreadPool.cpp +++ b/image_analysis/indexing/IndexerThreadPool.cpp @@ -286,10 +286,7 @@ void IndexerThreadPool::ReleaseIndexers() { #ifdef JFJOCH_USE_CUDA // Their buffers went back to the allocation pool, which would keep them reserved for the next // buffer of whoever asks; hand them to the card. - for (int32_t d = 0; d < get_gpu_count(); d++) { - ScopedGpuPin pin(d); - cuda_trim_allocation_pool(); - } + cuda_trim_allocation_pools(); #endif } diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index 33483f7a0..d52b589ba 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -569,6 +569,11 @@ void Rugnux::StartSpeculativeGeometryProbe(std::shared_ptr run, const st Logger("Rugnux").Warning("Speculative geometry probe failed: {}", e.what()); } } +#ifdef JFJOCH_USE_CUDA + // What its indexers and image workers held went back to the allocation pool; hand it to the card, + // or it stays reserved under the merge and the validation beside which the probe ran. + cuda_trim_allocation_pools(); +#endif }); speculative_probe_ = std::move(sp); } diff --git a/rugnux/RugnuxPasses.cpp b/rugnux/RugnuxPasses.cpp index b620c9379..01ae4bc01 100644 --- a/rugnux/RugnuxPasses.cpp +++ b/rugnux/RugnuxPasses.cpp @@ -756,6 +756,11 @@ RotationScaleWalk Rugnux::WalkRotationScaleBeside(const ProcessResult &pass1, Ro // empty result reads as. logger.Info("Rotation scale {:.5f}: the probe pass did not complete ({})", k, e.what()); } +#ifdef JFJOCH_USE_CUDA + // As the speculative probe (StartSpeculativeGeometryProbe): what it held goes back to the card. + if (index_only) + cuda_trim_allocation_pools(); +#endif return res; }; const auto probe = [&](float k, bool index_only) {