diff --git a/image_analysis/indexing/CUDAMemHelpers.h b/image_analysis/indexing/CUDAMemHelpers.h index 43ae2ca5d..0da8b4719 100644 --- a/image_analysis/indexing/CUDAMemHelpers.h +++ b/image_analysis/indexing/CUDAMemHelpers.h @@ -293,7 +293,11 @@ inline bool cuda_memory_available(size_t bytes, size_t &free_bytes) { // for a buffer that is freed by a thread other than the one whose work read it: cudaFreeAsync orders // the free on the freeing side's allocation stream, which says nothing about kernels queued // elsewhere, while cudaFree synchronises the whole device and therefore cannot be early. -enum class CudaAlloc { Pooled, Synchronous }; +// SynchronousAfterWorkBeside is Synchronous for a step of the main line - the merge's helpers - whose +// buffers may be asked for while a probe pass runs beside it on the same card (GPUWorkBeside): one that +// does not fit waits for that work to end and is asked for once more, as the merge's own arrays are +// (RotationScaleMergeGPU). Never for the work beside itself, which would wait for itself. +enum class CudaAlloc { Pooled, Synchronous, SynchronousAfterWorkBeside }; template class CudaDevicePtr { @@ -338,9 +342,17 @@ public: // stay behind as the last error (see above). cudaGetLastError(); cuda_trim_allocation_pool(); + cudaError_t err = cudaMalloc(&ptr_, count * sizeof(T)); + if (err == cudaErrorMemoryAllocation && alloc == CudaAlloc::SynchronousAfterWorkBeside) { + cudaGetLastError(); + // That work is a probe pass of seconds, so this long means it is stuck. + wait_for_gpu_work_beside(std::chrono::minutes(10)); + cuda_trim_allocation_pool(); + err = cudaMalloc(&ptr_, count * sizeof(T)); + } // With the reason: a context lost to another kernel's fault fails here too, and must not read // as a card that is merely full. - if (const cudaError_t err = cudaMalloc(&ptr_, count * sizeof(T)); err != cudaSuccess) + if (err != cudaSuccess) throw JFJochException(JFJochExceptionCategory::GPUCUDAError, GPUMemoryFailure("Allocating GPU memory", count * sizeof(T), err == cudaErrorMemoryAllocation, cudaGetErrorString(err))); diff --git a/image_analysis/scale_merge/ErrorModelGPU.cu b/image_analysis/scale_merge/ErrorModelGPU.cu index 3443d8e09..cd8c5a4c4 100644 --- a/image_analysis/scale_merge/ErrorModelGPU.cu +++ b/image_analysis/scale_merge/ErrorModelGPU.cu @@ -16,7 +16,7 @@ namespace { constexpr int THREADS = 256; constexpr int CHUNKS = 64; // blocks per bin in the sum kernel; fixed, so the tree is too constexpr int NB = ERROR_MODEL_BINS; - constexpr CudaAlloc ALLOC = CudaAlloc::Synchronous; + constexpr CudaAlloc ALLOC = CudaAlloc::SynchronousAfterWorkBeside; void cuda_err(cudaError_t val) { if (val != cudaSuccess) diff --git a/image_analysis/scale_merge/PairSortGPU.cu b/image_analysis/scale_merge/PairSortGPU.cu index 693821158..50ef96811 100644 --- a/image_analysis/scale_merge/PairSortGPU.cu +++ b/image_analysis/scale_merge/PairSortGPU.cu @@ -13,7 +13,7 @@ namespace { constexpr int THREADS = 256; - constexpr CudaAlloc ALLOC = CudaAlloc::Synchronous; + constexpr CudaAlloc ALLOC = CudaAlloc::SynchronousAfterWorkBeside; void cuda_err(cudaError_t val) { if (val != cudaSuccess) diff --git a/image_analysis/scale_merge/RotationScaleMergeGPU.cu b/image_analysis/scale_merge/RotationScaleMergeGPU.cu index 62805d660..4315e4b35 100644 --- a/image_analysis/scale_merge/RotationScaleMergeGPU.cu +++ b/image_analysis/scale_merge/RotationScaleMergeGPU.cu @@ -36,7 +36,7 @@ namespace { // free while unwinding from a failed call; nor can compute-sanitizer --track-stream-ordered-races // see it, and it reported every reassigned merge buffer as a use-after-free. cudaFree // synchronises the device first. - constexpr CudaAlloc ALLOC = CudaAlloc::Synchronous; + constexpr CudaAlloc ALLOC = CudaAlloc::SynchronousAfterWorkBeside; __device__ __forceinline__ double SafeInvD(double x, double fallback) { return (isfinite(x) && x != 0.0) ? 1.0 / x : fallback; diff --git a/image_analysis/scale_merge/WilsonOutliersGPU.cu b/image_analysis/scale_merge/WilsonOutliersGPU.cu index aeb84cc10..b877d07a7 100644 --- a/image_analysis/scale_merge/WilsonOutliersGPU.cu +++ b/image_analysis/scale_merge/WilsonOutliersGPU.cu @@ -25,7 +25,7 @@ namespace { constexpr int THREADS = 256; constexpr size_t OBS_PER_SHELL = 2000; // WilsonOutliers.cpp - constexpr CudaAlloc ALLOC = CudaAlloc::Synchronous; + constexpr CudaAlloc ALLOC = CudaAlloc::SynchronousAfterWorkBeside; void cuda_err(cudaError_t val) { if (val != cudaSuccess) diff --git a/image_analysis/structure_refinement/MolecularReplacement.cpp b/image_analysis/structure_refinement/MolecularReplacement.cpp index 6d6763811..5cb3a53fc 100644 --- a/image_analysis/structure_refinement/MolecularReplacement.cpp +++ b/image_analysis/structure_refinement/MolecularReplacement.cpp @@ -7,6 +7,7 @@ #include #include #include +#include #include #include #include @@ -277,6 +278,13 @@ PlacementScore r_factors(const gemmi::AsuData> &fmodel, cons // where the first placement of that kind is scored, each with the scale fit on the GPU beside it - or none, // where there is no GPU, the cell is too small for its gridding or the two would take over a quarter of the // card. Decided on the card's total memory, so the same input always takes the same path. +// +// Together they keep within that quarter too. A search asks for one per group, copy count and resolution, +// and kept to the end of the search they summed to well over the card on a large cell (seven of 1.4 GB +// beside the rigid-body engines). So an engine that does not fit beside the ones kept takes the place of +// the least recently used of those no score is using, and waits for a score to finish where every one is +// in use; one taken out is made again where it is next asked for. Every engine computes the same numbers +// whenever it is made, so this changes when the card holds them, not what a score gives. #ifdef JFJOCH_USE_CUDA class ScoreEngines { public: @@ -290,54 +298,123 @@ public: CudaDevicePtr> fcalc, fmask; }; - Engine *get(const gemmi::Structure &st, const gemmi::SpaceGroup &sg, double d_min) { + // The engine for this kind of placement, kept on the card for as long as the returned pointer lives. + std::shared_ptr get(const gemmi::Structure &st, const gemmi::SpaceGroup &sg, double d_min) { if (get_gpu_count() == 0) return nullptr; size_t atoms = 0; for (const gemmi::Chain &ch : st.models[0].chains) for (const gemmi::Residue &r : ch.residues) atoms += r.atoms.size(); - Entry *entry; - { - std::lock_guard lock(m_); - entry = &engines_[{sg.xhm(), atoms, d_min}]; + std::unique_lock lock(m_); + Entry &e = engines_[{sg.xhm(), atoms, d_min}]; + e.last_use = ++clock_; + // Made by the first to ask for it, without holding up those asking for other ones. + cv_.wait(lock, [&] { return !e.making; }); + if (e.engine || e.refused) + return Lease(e); + e.making = true; + lock.unlock(); + + std::unique_ptr sf; + size_t points = 0, bytes = 0, quarter = 0; + try { + sf = std::make_unique(RigidBodyGPUEngine::CurrentDevice(), st.models[0], + st.cell, sg, d_min); + points = sf->Reflections(); + bytes = sf->DeviceBytes() + ModelScaleGPU::DeviceBytes(points) + 2 * points * sizeof(std::complex); + quarter = sf->DeviceTotalMemory() / 4; + } catch (const JFJochException &) { + cuda_clear_error(); + sf.reset(); } - // Each made once, by the first to ask for it, without holding up the others. - Entry &e = *entry; - std::call_once(e.made, [&] { - try { - auto sf = std::make_unique(RigidBodyGPUEngine::CurrentDevice(), st.models[0], - st.cell, sg, d_min); - const size_t points = sf->Reflections(); - const size_t bytes = sf->DeviceBytes() + ModelScaleGPU::DeviceBytes(points) - + 2 * points * sizeof(std::complex); - if (sf->Supported() && bytes <= sf->DeviceTotalMemory() / 4) { - sf->Reserve(); - const int device = RigidBodyGPUEngine::CurrentDevice(); - set_gpu(sf->Device()); - e.engine = std::make_unique(); - e.engine->stream = std::make_unique(); - e.engine->scale = std::make_unique(*e.engine->stream, points); - e.engine->fcalc = CudaDevicePtr>(points); - e.engine->fmask = CudaDevicePtr>(points); - e.engine->sf = std::move(sf); - set_gpu(device); - } - } catch (const JFJochException &) { - e.engine.reset(); - cuda_clear_error(); - } - }); - return e.engine.get(); + lock.lock(); + if (!sf || !sf->Supported() || bytes > quarter) { + e.refused = true; + e.making = false; + cv_.notify_all(); + return nullptr; + } + const int card = sf->Device(); + cv_.wait(lock, [&] { return MakeRoom(card, bytes, quarter); }); + resident_[card] += bytes; + lock.unlock(); + + std::shared_ptr engine; + try { + sf->Reserve(); + const int device = RigidBodyGPUEngine::CurrentDevice(); + set_gpu(card); + engine = std::make_shared(); + engine->stream = std::make_unique(); + engine->scale = std::make_unique(*engine->stream, points); + engine->fcalc = CudaDevicePtr>(points); + engine->fmask = CudaDevicePtr>(points); + engine->sf = std::move(sf); + set_gpu(device); + } catch (const JFJochException &) { + engine.reset(); + cuda_clear_error(); + } + lock.lock(); + e.making = false; + if (engine) { + e.engine = std::move(engine); + e.card = card; + e.bytes = bytes; + } else { + resident_[card] -= bytes; + e.refused = true; + } + cv_.notify_all(); + return Lease(e); } private: struct Entry { - std::once_flag made; - std::unique_ptr engine; + std::shared_ptr engine; + bool making = false, refused = false; + int card = 0; + size_t bytes = 0; + uint64_t last_use = 0; }; + + // The entry's engine for a score: the copy it holds is what marks the engine in use, and handing it back + // is what may let an engine waiting for room in. + std::shared_ptr Lease(Entry &e) { + if (!e.engine) + return nullptr; + return std::shared_ptr(e.engine.get(), [this, held = e.engine](Engine *) mutable { + { + std::lock_guard lock(m_); + held.reset(); + } + cv_.notify_all(); + }); + } + + // Whether `bytes` more fit the card's quarter beside the engines kept there, once the least recently + // used of those no score is using have been taken out. Under m_. + bool MakeRoom(int card, size_t bytes, size_t quarter) { + while (resident_[card] + bytes > quarter) { + Entry *oldest = nullptr; + for (auto &[key, entry] : engines_) + if (entry.engine && entry.card == card && entry.engine.use_count() == 1 + && (oldest == nullptr || entry.last_use < oldest->last_use)) + oldest = &entry; + if (oldest == nullptr) + return false; + resident_[card] -= oldest->bytes; + oldest->engine.reset(); + } + return true; + } + std::map, Entry> engines_; + std::map resident_; + uint64_t clock_ = 0; std::mutex m_; + std::condition_variable cv_; }; // The scale of `scaling`'s points fitted on the engine's card: FitModelScale's grid over the solvent box, or @@ -464,7 +541,7 @@ PlacementScore score_placement(const gemmi::Structure &st, const gemmi::SpaceGro const gemmi::Op *law = nullptr, const gemmi::Fractional *tncs = nullptr) { gemmi::AsuData> fmask, fmodel; #ifdef JFJOCH_USE_CUDA - ScoreEngines::Engine *gpu = engines != nullptr ? engines->get(st, sg, d_min) : nullptr; + const std::shared_ptr gpu = engines != nullptr ? engines->get(st, sg, d_min) : nullptr; if (gpu) gpu->sf->Compute(st.models[0], fmodel, fmask); else diff --git a/rugnux/RugnuxScaleMerge.cpp b/rugnux/RugnuxScaleMerge.cpp index 5fb3df3d8..3e49e3c99 100644 --- a/rugnux/RugnuxScaleMerge.cpp +++ b/rugnux/RugnuxScaleMerge.cpp @@ -4212,6 +4212,10 @@ bool Rugnux::ScaleMergeAndSymmetry(PipelineLocals &p) { phase("Molecular replacement"); Step("Molecular replacement"); TimingMark mark(logger, "model prior"); + // As before the validation above, which a predicted model or --mr does not run: the engine's card + // memory goes back before the search sizes its rigid-body engines from what the card has free. + if (!ledger_ahead.valid()) + rsm.reset(); std::vector merges{{"MERGE", &sm.merged, *result.consensus_cell, experiment_.GetSpaceGroupOrP1()}}; if (model_twin) merges.push_back({"TWIN_SUBGROUP_MERGE", &model_twin->merged, model_twin->cell, model_twin->space_group, diff --git a/tests/CUDAMemHelpersTest.cpp b/tests/CUDAMemHelpersTest.cpp index 2d0047319..c349d9d88 100644 --- a/tests/CUDAMemHelpersTest.cpp +++ b/tests/CUDAMemHelpersTest.cpp @@ -206,6 +206,39 @@ TEST_CASE("IndexerThreadPool_FFTIndexerWaitsForGPUWorkBeside", "[CUDAMemHelpers] } } +// A merge helper's buffer that finds the card full while a probe pass holds the rest of it waits for +// that work to end and its memory to come back (CudaAlloc::SynchronousAfterWorkBeside); a plain +// synchronous one fails at once. +TEST_CASE("CudaDevicePtr_AfterWorkBesideWaitsForItsMemory", "[CUDAMemHelpers]") { + if (get_gpu_count() == 0) + SKIP("no GPU"); + std::vector> held; + const auto fill = [&] { + 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; + } + } + }; + constexpr size_t BYTES = size_t(64) << 20; + { + std::optional work(std::in_place); + fill(); + CHECK_THROWS_AS(CudaDevicePtr(BYTES, CudaAlloc::Synchronous), JFJochException); + std::thread ends([&] { + std::this_thread::sleep_for(std::chrono::milliseconds(500)); + held.clear(); + work.reset(); + }); + CHECK_NOTHROW(CudaDevicePtr(BYTES, CudaAlloc::SynchronousAfterWorkBeside)); + ends.join(); + } + cuda_trim_allocation_pool(); +} + #endif // A merge that finds the card full waits for the GPU work running beside it: until the last one