RotationScaleMerge: hand back the surface buffers before the merge, so the GPU Wilson test fits

The device correction-surface fit (gpu-tail2) keeps its per-term arrays - the
terms, their parity, the group CSR, the three subsets' orders and the per-term
weights, about 60 bytes a term - as engine members, and nothing freed them
after the last surface. They stayed on the card through the merge, where the
GPU Wilson outlier test (gpu-tail) then asks for ~96 bytes per full. On a
large rotation set (139.8 M partials, 13.35 M fulls, 16 GB card) the engine's
resident arrays left 1242 MiB free for the first Wilson test and 998 MiB for
the one after the automatic resolution cut, against 1222 MiB needed: the run
stopped with "Wilson outlier test ... run the CPU build" although nothing ran
beside it. The surfaces' buffers are released once the last surface is
fitted (the next pass's surfaces build their own), which gives back 905 MiB
there.

The up-front checks of the three dataset-sized tail steps (Wilson test, pair
sort, error-model fit) also read free memory as cudaMemGetInfo reports it,
which counts what the allocation pool holds with no buffer in it as taken.
They now hand that back and, only if it is still short, wait for GPU work
running beside the merge (GPUWorkBeside, the speculative probe) to end before
giving up - as the engine's own allocations already do. That changes when a
step runs, never what it computes; a set that does not fit still fails at
once with the same "run the CPU build" message.

p.mtz: the large set completes (and is identical with the Wilson test forced
onto the host); myob 9bdd6916 and 8a1a 309ea9a8 unchanged.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01SVmAWnzCmRKAXVUCdc4iNi

image_analysis/indexing/CUDAMemHelpers.h

@@ -10,6 +10,7 @@
 #include <stdexcept>
 #include <vector>
 #include "../common/JFJochException.h"
+#include "../common/CUDAWrapper.h"
 
 class CudaStream {
     cudaStream_t stream_ = nullptr;
@@ -218,6 +219,31 @@ inline void cuda_trim_allocation_pool() {
     cudaGetLastError();
 }
 
+// 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. cudaMemGetInfo counts what the
+// pool holds with no buffer in it (up to CUDA_MEM_POOL_RELEASE_THRESHOLD) as taken, so that is handed
+// back before the step gives up; and GPU work running beside it in this process (GPUWorkBeside) is
+// waited for, as its memory comes back when it ends. Only when the step runs depends on this, never
+// what it computes. `free_bytes` is the last reading, for the message of a step that does not fit.
+inline bool cuda_memory_available(size_t bytes, size_t &free_bytes) {
+    const auto read_free = [&] {
+        size_t total_bytes = 0;
+        if (cudaMemGetInfo(&free_bytes, &total_bytes) != cudaSuccess)
+            throw JFJochException(JFJochExceptionCategory::GPUCUDAError, "Failed to read free device memory");
+        return bytes <= free_bytes;
+    };
+    if (read_free())
+        return true;
+    cuda_trim_allocation_pool();
+    if (read_free())
+        return true;
+    // That work is a probe pass of seconds, so this long means it is stuck: the step fails as it would
+    // have without it.
+    wait_for_gpu_work_beside(std::chrono::minutes(10));
+    cuda_trim_allocation_pool();
+    return read_free();
+}
+
 // Which allocator a device buffer uses. Pooled is what every per-worker buffer wants. Synchronous is
 // 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

image_analysis/scale_merge/ErrorModelGPU.cu

@@ -91,9 +91,8 @@ ErrorModelFit FitErrorModelGPU(const std::vector<ErrorModelSample> &pool) {
 
     // The rank sort is the peak: the pool, two key and two index arrays and the sort scratch.
     constexpr size_t BYTES_PER_SAMPLE = 32 + 2 * 8 + 2 * 4 + 16;
-    size_t free_bytes = 0, total_bytes = 0;
-    cuda_err(cudaMemGetInfo(&free_bytes, &total_bytes));
-    if (BYTES_PER_SAMPLE * pool.size() > free_bytes)
+    size_t free_bytes = 0;
+    if (!cuda_memory_available(BYTES_PER_SAMPLE * pool.size(), free_bytes))
         throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
                               "Error model: " + std::to_string(pool.size()) + " samples need "
                               + std::to_string(BYTES_PER_SAMPLE * pool.size() >> 20) + " MiB of GPU memory, "

image_analysis/scale_merge/PairSortGPU.cu

@@ -46,9 +46,8 @@ void SortPairsGPU(std::vector<std::pair<double, double>> &v) {
     if (n < 2)
         return;
     constexpr size_t BYTES_PER_PAIR = 16 + 16 + 16 + 8;   // pairs in and out, keys, indices
-    size_t free_bytes = 0, total_bytes = 0;
-    cuda_err(cudaMemGetInfo(&free_bytes, &total_bytes));
-    if (BYTES_PER_PAIR * v.size() > free_bytes)
+    size_t free_bytes = 0;
+    if (!cuda_memory_available(BYTES_PER_PAIR * v.size(), free_bytes))
         throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
                               "Sorting " + std::to_string(v.size()) + " pairs needs "
                               + std::to_string(BYTES_PER_PAIR * v.size() >> 20) + " MiB of GPU memory, "

image_analysis/scale_merge/RotationScaleMerge.cpp

@@ -6800,6 +6800,8 @@ RotationScaleMerge::Result RotationScaleMerge::RunArm(bool scale_partials, bool
     // merge reads it. (The collapsed-scale drop was pushed where it happened.)
     if ((corrections || rejected_frames) && combined_on_gpu && scaled_fulls_on_gpu)
         upload_fulls_corr();
+    if (gpu_active_)
+        gpu_->SurfaceRelease();
 #endif
 
     surfaces_mark.Stop();

image_analysis/scale_merge/RotationScaleMergeGPU.cu

@@ -2238,6 +2238,19 @@ int RotationScaleMergeGPU::SurfaceBuildTerms(const int32_t *cell, const float *c
     return n_terms;
 }
 
+void RotationScaleMergeGPU::SurfaceRelease() {
+    DeviceGuard guard(impl_->device, impl_->available);
+    auto &d = *impl_;
+    d.s_slots = 0;
+    d.s_term = {}; d.s_parity = {}; d.s_gperm = {}; d.s_gstart = {};
+    for (int subset = 0; subset < 3; ++subset) {
+        d.s_perm[subset] = {}; d.s_seg_start[subset] = {};
+        d.s_n_blocks[subset] = d.s_n_sel[subset] = 0;
+    }
+    d.s_A = {}; d.s_sw = {}; d.s_swI = {}; d.s_tcross = {}; d.s_tref2 = {}; d.s_cross = {}; d.s_ref2 = {};
+    d.s_w_Is = {}; d.s_w_Iref = {}; d.s_Iref = {};
+}
+
 void RotationScaleMergeGPU::SurfaceShellWeights(int n_terms, int nshell, int32_t *group_shell, double *shw) {
     DeviceGuard guard(impl_->device, impl_->available);
     auto &d = *impl_;

image_analysis/scale_merge/RotationScaleMergeGPU.h

@@ -264,6 +264,11 @@ public:
     // cross = sum w Is Iref and ref2 = sum w Iref^2 (length ncell each).
     void SurfaceFitSums(int subset, double *cross, double *ref2);
 
+    // Hand back the surface fits' buffers - several bytes per term, sized by the fulls - once the last
+    // surface is fitted: the merge after the surfaces needs the card, and the next pass's surfaces
+    // build their own (SurfaceBuildTerms).
+    void SurfaceRelease();
+
 private:
     struct Impl;
     std::unique_ptr<Impl> impl_;

image_analysis/scale_merge/WilsonOutliersGPU.cu

@@ -229,9 +229,8 @@ WilsonOutlierResult WilsonOutliersGPU(const std::vector<WilsonObservation> &obs,
 
     // Observations (28 bytes) + ~13 more 4-8 byte arrays of the same length, and the sort scratch.
     constexpr size_t BYTES_PER_OBS = 96;
-    size_t free_bytes = 0, total_bytes = 0;
-    cuda_err(cudaMemGetInfo(&free_bytes, &total_bytes));
-    if (BYTES_PER_OBS * obs.size() > free_bytes)
+    size_t free_bytes = 0;
+    if (!cuda_memory_available(BYTES_PER_OBS * obs.size(), free_bytes))
         throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
                               "Wilson outlier test: " + std::to_string(obs.size()) + " observations need "
                               + std::to_string(BYTES_PER_OBS * obs.size() >> 20) + " MiB of GPU memory, "