Merge branch 'regr-gpu-oom' into rc175

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