diff --git a/common/CUDAWrapper.cpp b/common/CUDAWrapper.cpp index 5eebbd99e..5d3f45ecc 100644 --- a/common/CUDAWrapper.cpp +++ b/common/CUDAWrapper.cpp @@ -28,27 +28,50 @@ std::string get_gpu_description() { } namespace { - std::mutex gpu_work_beside_mutex; - std::condition_variable gpu_work_beside_done; - int gpu_work_beside = 0; + // How many of one kind of GPU work are running, for a thread to wait until none is. + struct GPUWorkCount { + std::mutex m; + std::condition_variable none; + int n = 0; + void Add() { + std::lock_guard lock(m); + ++n; + } + void Remove() { + { + std::lock_guard lock(m); + --n; + } + none.notify_all(); + } + }; + GPUWorkCount gpu_work_beside, gpu_image_passes; } GPUWorkBeside::GPUWorkBeside() { - std::lock_guard lock(gpu_work_beside_mutex); - ++gpu_work_beside; + gpu_work_beside.Add(); } GPUWorkBeside::~GPUWorkBeside() { - { - std::lock_guard lock(gpu_work_beside_mutex); - --gpu_work_beside; - } - gpu_work_beside_done.notify_all(); + gpu_work_beside.Remove(); } bool wait_for_gpu_work_beside(std::chrono::seconds timeout) { - std::unique_lock lock(gpu_work_beside_mutex); - return gpu_work_beside_done.wait_for(lock, timeout, [] { return gpu_work_beside == 0; }); + std::unique_lock lock(gpu_work_beside.m); + return gpu_work_beside.none.wait_for(lock, timeout, [] { return gpu_work_beside.n == 0; }); +} + +GPUImagePass::GPUImagePass() { + gpu_image_passes.Add(); +} + +GPUImagePass::~GPUImagePass() { + gpu_image_passes.Remove(); +} + +void wait_for_gpu_image_passes() { + std::unique_lock lock(gpu_image_passes.m); + gpu_image_passes.none.wait(lock, [] { return gpu_image_passes.n == 0; }); } #ifndef JFJOCH_USE_CUDA diff --git a/common/CUDAWrapper.h b/common/CUDAWrapper.h index 7941dbd60..a61937f5a 100644 --- a/common/CUDAWrapper.h +++ b/common/CUDAWrapper.h @@ -100,3 +100,18 @@ public: // Wait until no GPUWorkBeside is held, for at most `timeout`. True once none is - at once if none was. bool wait_for_gpu_work_beside(std::chrono::seconds timeout); + +// The image pass of a run's main line - not of a probe holding a GPUWorkBeside: the stretch that holds +// the most of the card, and that never waits for a probe, so a probe that does not fit beside it can +// wait for it to end. Hold one for as long as the image pass runs. Build-independent - it only counts. +class GPUImagePass { +public: + GPUImagePass(); + ~GPUImagePass(); + GPUImagePass(const GPUImagePass &) = delete; + GPUImagePass &operator=(const GPUImagePass &) = 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 4afe38a97..ec9e750e3 100644 --- a/image_analysis/indexing/IndexerThreadPool.cpp +++ b/image_analysis/indexing/IndexerThreadPool.cpp @@ -24,12 +24,40 @@ void WarmUpCuFFT() { #endif } +#ifdef JFJOCH_USE_CUDA +// An FFT indexer is a cuFFT plan and ~0.6-0.9 GB of device buffers, and several pipeline copies of one +// process run on the card at the same time: probe passes (GPUWorkBeside) beside the main line, for a +// few seconds each. So one that does not fit hands back what the allocation pool holds unused and waits +// for the work beside it to end, then is built again - as the merge's allocations do. The main line +// waits for the probes; a probe waits for the main line's image pass (GPUImagePass), never for a probe, +// itself included. Neither of those waits for the other, so nothing waits in a circle. Only when the +// indexer is built depends on this, never what it computes; one that does not fit the card fails as +// before. +static std::unique_ptr MakeFFTIndexerGPU(const IndexingSettings &settings, bool beside) { + try { + return std::make_unique(settings); + } catch (const JFJochException &) { + cuda_throw_if_context_lost(); // not memory: waiting would only hide the real error + cuda_clear_error(); + } + if (beside) + wait_for_gpu_image_passes(); + else + // That work is a probe pass of seconds, so this long means it is stuck: the indexer fails as + // it would have without it. + wait_for_gpu_work_beside(std::chrono::minutes(10)); + cuda_trim_allocation_pool(); + return std::make_unique(settings); +} +#endif + // The indexer for one RESOLVED algorithm, or nullptr if this build/host cannot serve it. -static std::unique_ptr MakeIndexer(IndexingAlgorithmEnum algorithm, const IndexingSettings &settings) { +static std::unique_ptr MakeIndexer(IndexingAlgorithmEnum algorithm, const IndexingSettings &settings, + bool beside) { #ifdef JFJOCH_USE_CUDA if (get_gpu_count() > 0) { if (algorithm == IndexingAlgorithmEnum::FFT) - return std::make_unique(settings); + return MakeFFTIndexerGPU(settings, beside); if (algorithm == IndexingAlgorithmEnum::FFBIDX) return std::make_unique(); } @@ -41,8 +69,9 @@ static std::unique_ptr MakeIndexer(IndexingAlgorithmEnum algorithm, con return nullptr; } -IndexerThread::IndexerThread(const IndexingSettings &settings, int threadid, IndexerConstruction construction) - : settings_(settings), construction_(construction) { +IndexerThread::IndexerThread(const IndexingSettings &settings, int threadid, IndexerConstruction construction, + bool beside) + : settings_(settings), construction_(construction), beside_(beside) { std::unique_lock lock(m); state = TaskState::STARTING; worker_thread = std::thread(&IndexerThread::Worker, this, threadid); @@ -74,14 +103,14 @@ void IndexerThread::Worker(int threadid) { const auto requested = settings_.GetAlgorithm(); if (requested == IndexingAlgorithmEnum::Auto || requested == IndexingAlgorithmEnum::FFT) - fft_indexer = MakeIndexer(IndexingAlgorithmEnum::FFT, settings_); + fft_indexer = MakeIndexer(IndexingAlgorithmEnum::FFT, settings_, beside_); if (requested == IndexingAlgorithmEnum::Auto || requested == IndexingAlgorithmEnum::FFBIDX) - ffbidx_indexer = MakeIndexer(IndexingAlgorithmEnum::FFBIDX, settings_); + ffbidx_indexer = MakeIndexer(IndexingAlgorithmEnum::FFBIDX, settings_, beside_); if ((requested == IndexingAlgorithmEnum::Auto && get_gpu_count() == 0) || requested == IndexingAlgorithmEnum::FFTW) - fftw_indexer = MakeIndexer(IndexingAlgorithmEnum::FFTW, settings_); + fftw_indexer = MakeIndexer(IndexingAlgorithmEnum::FFTW, settings_, beside_); } catch (const std::exception &e) { spdlog::error("Failed to initialize indexer: {}", e.what()); { @@ -134,7 +163,7 @@ void IndexerThread::Worker(int threadid) { // A preconstructing worker already holds it; an OnFirstUse worker builds it here, // on the first frame that resolves to this algorithm. if (slot && !*slot) - *slot = MakeIndexer(algorithm, settings_); + *slot = MakeIndexer(algorithm, settings_, beside_); if (!slot || !*slot) { // Algorithm is already resolved here (never Auto/None - see @@ -208,14 +237,15 @@ IndexerThread::~IndexerThread() { Finalize(); } -IndexerThreadPool::IndexerThreadPool(const IndexingSettings &settings, IndexerConstruction construction) +IndexerThreadPool::IndexerThreadPool(const IndexingSettings &settings, IndexerConstruction construction, + bool beside) : worker_busy(settings.GetIndexingThreads(), 0), worker_free_count(settings.GetIndexingThreads()), viable_cell_min_spots(settings.GetViableCellMinSpots()), blocking(settings.GetBlockingBehavior()), settings_(settings) { for (size_t i = 0; i < settings.GetIndexingThreads(); ++i) - tasks.emplace_back(std::make_unique(std::cref(settings), i, construction)); + tasks.emplace_back(std::make_unique(std::cref(settings), i, construction, beside)); } int IndexerThreadPool::GetFreeWorker() { diff --git a/image_analysis/indexing/IndexerThreadPool.h b/image_analysis/indexing/IndexerThreadPool.h index 35f4beccb..709f984a3 100644 --- a/image_analysis/indexing/IndexerThreadPool.h +++ b/image_analysis/indexing/IndexerThreadPool.h @@ -52,6 +52,7 @@ class IndexerThread { // IndexerThreadPool(experiment.GetIndexingSettings()), which returns by value. const IndexingSettings settings_; const IndexerConstruction construction_; + const bool beside_; bool stop = false; enum class TaskState {STARTING, IDLE, READY, RUNNING, COMPLETED, ERROR} state = TaskState::STARTING; @@ -65,7 +66,7 @@ class IndexerThread { void Worker(int threadid); public: - IndexerThread(const IndexingSettings& settings, int threadid, IndexerConstruction construction); + IndexerThread(const IndexingSettings& settings, int threadid, IndexerConstruction construction, bool beside); ~IndexerThread(); std::unique_ptr Run(const DiffractionExperiment &experiment, const std::vector &recip, bool severity_only = false); @@ -83,8 +84,12 @@ class IndexerThreadPool { const IndexingSettings settings_; int GetFreeWorker(); public: + // beside: the pool serves a probe pass that holds GPUWorkBeside, so an FFT indexer that does not fit + // the card waits for the main line's image pass to end, not for the probes - it is one of them (see + // MakeFFTIndexerGPU). IndexerThreadPool(const IndexingSettings& settings, - IndexerConstruction construction = IndexerConstruction::Preconstruct); + IndexerConstruction construction = IndexerConstruction::Preconstruct, + bool beside = false); // The settings the pool's indexers were built with. [[nodiscard]] const IndexingSettings &Settings() const { return settings_; } // severity_only skips indexing and produces just the spindle severity - see Indexer::Run. diff --git a/rugnux/Rugnux.cpp b/rugnux/Rugnux.cpp index f4117de35..f39cb79e5 100644 --- a/rugnux/Rugnux.cpp +++ b/rugnux/Rugnux.cpp @@ -440,6 +440,7 @@ void Rugnux::StartNearTieRecheck(std::shared_ptr run, bool geometry_prep run->speculative_probe_.reset(); run->speculate_geometry_probe_ = false; run->rotation_scale_check_.reset(); + run->gpu_work_beside_ = false; // it runs after the run, on its own std::lock_guard lock(near_tie_rechecks_->m); near_tie_rechecks_->pending.push_back({std::move(run), geometry_prepass, first_pass_outcome_}); } @@ -510,6 +511,7 @@ void Rugnux::StartSpeculativeGeometryProbe(std::shared_ptr run, const st run->postrefine_probe_ = false; run->postrefine_probe_only_ = false; run->rerun_if_starved_ = false; + run->gpu_work_beside_ = true; // The walk's other probe, at the geometry in hand, on a copy of its own: where this pass's own // first pass has not left its memo the walk would index there too, after the pass. auto hand = std::make_shared(*run); @@ -880,7 +882,8 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b RaiseFFTBoundForKnownCell(indexing_settings, experiment_); } indexing_bound_A = indexing_settings.GetFFT_MaxUnitCell_A(); - indexer_pool = std::make_unique(indexing_settings, IndexerConstruction::OnFirstUse); + indexer_pool = std::make_unique(indexing_settings, IndexerConstruction::OnFirstUse, + gpu_work_beside_); indexer = std::make_unique(experiment_, indexer_pool.get()); // With no per-image file to write, nothing reads the message's reflection list - except the // geometry pre-pass, which measures the spots against it (SpotFootprint.h). @@ -993,7 +996,11 @@ ProcessResult Rugnux::RunPipeline(RugnuxObserver *observer, bool write_output, b } TimingMark image_pass_mark(logger, "image pass"); + std::optional gpu_image_pass; // what a probe that does not fit beside it waits for + if (!gpu_work_beside_) + gpu_image_pass.emplace(); ImagePass(p); + gpu_image_pass.reset(); image_pass_mark.Stop(); // Scaling and merging (full analysis only). diff --git a/rugnux/Rugnux.h b/rugnux/Rugnux.h index efe8986cc..fc77ec497 100644 --- a/rugnux/Rugnux.h +++ b/rugnux/Rugnux.h @@ -1063,6 +1063,10 @@ class Rugnux { }; std::shared_ptr speculative_probe_; bool speculate_geometry_probe_ = false; + // Whether this run is a probe pass running beside the main line under a GPUWorkBeside - the + // speculative probe, the rotation-scale walk's indexing probes. Its indexer pools then wait for the + // main line's image pass, never for that work, being that work themselves (IndexerThreadPool). + bool gpu_work_beside_ = false; // run: the copy of this run the probe runs on, taken by the caller on the thread that owns the // run - the post-refinement that decides whether to start it runs beside the merge, which does // not leave the run's members alone. diff --git a/rugnux/RugnuxFirstPass.cpp b/rugnux/RugnuxFirstPass.cpp index e1bd20d49..9afa3cacd 100644 --- a/rugnux/RugnuxFirstPass.cpp +++ b/rugnux/RugnuxFirstPass.cpp @@ -1552,7 +1552,7 @@ void Rugnux::FirstPassRun::StartShortAxisSpeculation() { sp.spots_per_image = first_pass_spots_per_image; sp.experiment = std::make_unique(rx.experiment_); sp.pool = std::make_unique(ShortAxisSettings(*sp.experiment), - IndexerConstruction::OnFirstUse); + IndexerConstruction::OnFirstUse, rx.gpu_work_beside_); for (const auto &[name, ordinals] : schemes) sp.sris.push_back(FeedScheme(*sp.pool, ordinals, sp.experiment.get())); for (auto &ri : sp.sris) { diff --git a/rugnux/RugnuxFirstPassRescues.cpp b/rugnux/RugnuxFirstPassRescues.cpp index 765007839..c48a3b68e 100644 --- a/rugnux/RugnuxFirstPassRescues.cpp +++ b/rugnux/RugnuxFirstPassRescues.cpp @@ -184,7 +184,7 @@ auto Rugnux::FirstPassRun::LongAxisRescue(const FirstPass &standing) -> std::opt coarse_settings.FFT_MaxUnitCell_A(RESCUE_MAX_UNIT_CELL_A); coarse_settings.IndexingThreads(2); coarse_settings.RefineThreads(std::clamp(rx.config_.nthreads / 2, 1, MAX_REFINE_THREADS)); - IndexerThreadPool coarse_pool(coarse_settings, IndexerConstruction::OnFirstUse); + IndexerThreadPool coarse_pool(coarse_settings, IndexerConstruction::OnFirstUse, rx.gpu_work_beside_); // Coarse first pass: keep the recovered cell with the LONGEST axis directly. Its full- // resolution per-frame validation would be low (the coarse cell is metrically right but // imprecise), so do NOT score it here - it is only the reference for the constrained re-index. @@ -213,7 +213,7 @@ auto Rugnux::FirstPassRun::LongAxisRescue(const FirstPass &standing) -> std::opt RaiseFFTBoundForKnownCell(wide_settings, rx.experiment_); FirstPass constrained; if (wide_settings.GetFFT_MaxUnitCell_A() > indexing_bound_A) { - IndexerThreadPool wide_pool(wide_settings, IndexerConstruction::OnFirstUse); + IndexerThreadPool wide_pool(wide_settings, IndexerConstruction::OnFirstUse, rx.gpu_work_beside_); constrained = PickBest(wide_pool, *indexer); } else { constrained = PickBest(*indexer_pool, *indexer); @@ -808,7 +808,7 @@ auto Rugnux::FirstPassRun::ShortAxisHypothesis(FirstPass best) -> FirstPass { } else { short_spec.reset(); short_pool = std::make_unique(ShortAxisSettings(rx.experiment_), - IndexerConstruction::OnFirstUse); + IndexerConstruction::OnFirstUse, rx.gpu_work_beside_); sris = FeedSchemes(*short_pool); RunIndexing(sris); } diff --git a/rugnux/RugnuxPasses.cpp b/rugnux/RugnuxPasses.cpp index 3c1af7b85..4464e65ce 100644 --- a/rugnux/RugnuxPasses.cpp +++ b/rugnux/RugnuxPasses.cpp @@ -669,6 +669,13 @@ RotationScaleWalk Rugnux::WalkRotationScaleBeside(const ProcessResult &pass1, Ro // to take, which the validation frames of the whole sweep then judge - so it is fitted on // the frames pass 1 fitted the first one on. r->prepass_frames_only_ = !index_only && !decide_late_; + // An indexing probe runs beside the refined passes, on the same card, and gives its memory + // back in seconds: what of theirs does not fit beside it waits for it to end. A refit + // runs only once they have ended, so it is nobody's work beside. + std::optional beside; + if (index_only) + beside.emplace(); + r->gpu_work_beside_ = index_only; ProcessResult res; try { res = r->RunPipeline(nullptr, /*write_output=*/false, /*geometry_prepass=*/false); diff --git a/rugnux/RugnuxStills.cpp b/rugnux/RugnuxStills.cpp index 2c938876d..b710ffcbb 100644 --- a/rugnux/RugnuxStills.cpp +++ b/rugnux/RugnuxStills.cpp @@ -117,7 +117,7 @@ void Rugnux::RefineStillsGeometry(int start_image, int end_image, int images_to_ // First-pass engines, built from the current (nominal) geometry: an index-only pass (no // integration) to obtain each frame's spots + assigned HKL + orientation. AzimuthalIntegrationMapping mapping(experiment_, pixel_mask_); - IndexerThreadPool pool(experiment_.GetIndexingSettings(), IndexerConstruction::OnFirstUse); + IndexerThreadPool pool(experiment_.GetIndexingSettings(), IndexerConstruction::OnFirstUse, gpu_work_beside_); IndexAndRefine indexer(experiment_, &pool, /*retain_outcomes=*/false); auto pass_settings = config_.spot_finding; diff --git a/tests/CUDAMemHelpersTest.cpp b/tests/CUDAMemHelpersTest.cpp index 0e8bb2311..8c08d74f0 100644 --- a/tests/CUDAMemHelpersTest.cpp +++ b/tests/CUDAMemHelpersTest.cpp @@ -15,6 +15,7 @@ #include #include "../image_analysis/indexing/CUDAMemHelpers.h" +#include "../image_analysis/indexing/IndexerThreadPool.h" // The broker starts fresh worker threads for every data collection. A stream per thread ever started // is a device-memory leak that ends a long-running broker, so threads that follow one another have to @@ -127,6 +128,65 @@ TEST_CASE("CudaPinnedPool_ReusedByTheNextEngine", "[CUDAMemHelpers]") { CHECK(cuda_release_pinned_pool() == 0); } +// An FFT indexer that finds the card full waits for the GPU work beside it to give its memory back, and +// is built then: the main line's for the probes, a probe's for the main line's image pass - never for a +// probe, itself included. With nothing to wait for it fails at once, as one that does not fit always has. +TEST_CASE("IndexerThreadPool_FFTIndexerWaitsForGPUWorkBeside", "[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}}; + + // The whole card, held by the work beside. + std::optional>> held; + const auto fill = [&held] { + held.emplace(); + for (size_t chunk = size_t(256) << 20; chunk >= (size_t(1) << 20); chunk /= 2) + while (true) { + try { + held->emplace_back(chunk, CudaAlloc::Synchronous); + } catch (const JFJochException &) { + break; + } + } + }; + IndexerThreadPool probe(settings, IndexerConstruction::OnFirstUse, /*beside=*/true); + IndexerThreadPool pool(settings, IndexerConstruction::OnFirstUse); + + fill(); + { + GPUWorkBeside work; + CHECK(probe.Run(x, recip).error != nullptr); + } + { + std::optional image_pass(std::in_place); + std::thread ends([&] { + std::this_thread::sleep_for(std::chrono::milliseconds(500)); + held.reset(); + image_pass.reset(); + }); + GPUWorkBeside work; + CHECK(probe.Run(x, recip).error == nullptr); + ends.join(); + } + + fill(); + { + std::optional work(std::in_place); + std::thread ends([&] { + std::this_thread::sleep_for(std::chrono::milliseconds(500)); + held.reset(); + work.reset(); + }); + CHECK(pool.Run(x, recip).error == nullptr); + ends.join(); + } +} + #endif // A merge that finds the card full waits for the GPU work running beside it: until the last one