Merge branch 'gpu-oom' into rc175: merge helpers wait for GPU work beside before failing; MR score-engine cache bounded per card; scaling engine released before MR on the predicted-model path

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>

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 <typename T>
 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)));

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)

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)

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;

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)

image_analysis/structure_refinement/MolecularReplacement.cpp

@@ -7,6 +7,7 @@
 #include <array>
 #include <atomic>
 #include <chrono>
+#include <condition_variable>
 #include <cmath>
 #include <future>
 #include <functional>
@@ -277,6 +278,13 @@ PlacementScore r_factors(const gemmi::AsuData<std::complex<float>> &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<std::complex<float>> 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<Engine> 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<ModelStructureFactorsGPU> sf;
+        size_t points = 0, bytes = 0, quarter = 0;
+        try {
+            sf = std::make_unique<ModelStructureFactorsGPU>(RigidBodyGPUEngine::CurrentDevice(), st.models[0],
+                                                            st.cell, sg, d_min);
+            points = sf->Reflections();
+            bytes = sf->DeviceBytes() + ModelScaleGPU::DeviceBytes(points) + 2 * points * sizeof(std::complex<float>);
+            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<ModelStructureFactorsGPU>(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<float>);
-                if (sf->Supported() && bytes <= sf->DeviceTotalMemory() / 4) {
-                    sf->Reserve();
-                    const int device = RigidBodyGPUEngine::CurrentDevice();
-                    set_gpu(sf->Device());
-                    e.engine = std::make_unique<Engine>();
-                    e.engine->stream = std::make_unique<CudaStream>();
-                    e.engine->scale = std::make_unique<ModelScaleGPU>(*e.engine->stream, points);
-                    e.engine->fcalc = CudaDevicePtr<std::complex<float>>(points);
-                    e.engine->fmask = CudaDevicePtr<std::complex<float>>(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> engine;
+        try {
+            sf->Reserve();
+            const int device = RigidBodyGPUEngine::CurrentDevice();
+            set_gpu(card);
+            engine = std::make_shared<Engine>();
+            engine->stream = std::make_unique<CudaStream>();
+            engine->scale = std::make_unique<ModelScaleGPU>(*engine->stream, points);
+            engine->fcalc = CudaDevicePtr<std::complex<float>>(points);
+            engine->fmask = CudaDevicePtr<std::complex<float>>(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> engine;
+        std::shared_ptr<Engine> 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<Engine> Lease(Entry &e) {
+        if (!e.engine)
+            return nullptr;
+        return std::shared_ptr<Engine>(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<std::tuple<std::string, size_t, double>, Entry> engines_;
+    std::map<int, size_t> 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<std::complex<float>> fmask, fmodel;
 #ifdef JFJOCH_USE_CUDA
-    ScoreEngines::Engine *gpu = engines != nullptr ? engines->get(st, sg, d_min) : nullptr;
+    const std::shared_ptr<ScoreEngines::Engine> gpu = engines != nullptr ? engines->get(st, sg, d_min) : nullptr;
     if (gpu)
         gpu->sf->Compute(st.models[0], fmodel, fmask);
     else

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<MRMerge> 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,

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<CudaDevicePtr<uint8_t>> 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<GPUWorkBeside> work(std::in_place);
+        fill();
+        CHECK_THROWS_AS(CudaDevicePtr<uint8_t>(BYTES, CudaAlloc::Synchronous), JFJochException);
+        std::thread ends([&] {
+            std::this_thread::sleep_for(std::chrono::milliseconds(500));
+            held.clear();
+            work.reset();
+        });
+        CHECK_NOTHROW(CudaDevicePtr<uint8_t>(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