Merge branch 'gpu-oom-8qaw' into rc173 (CUDA: trim the stream-ordered pool and retry when cudaMalloc fails)

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

image_analysis/indexing/CUDAMemHelpers.h

@@ -204,6 +204,20 @@ inline cudaStream_t cuda_allocation_stream() {
     }
 }
 
+// Return the current device's pool memory that no buffer holds to the driver. Frees ordered on a
+// stream count only once the stream has reached them, hence the device-wide wait first.
+inline void cuda_trim_allocation_pool() {
+    int device = 0;
+    cudaMemPool_t pool = nullptr;
+    if (cudaGetDevice(&device) != cudaSuccess || cudaDeviceGetDefaultMemPool(&pool, device) != cudaSuccess) {
+        cudaGetLastError();
+        return;
+    }
+    cudaDeviceSynchronize();
+    cudaMemPoolTrimTo(pool, 0);
+    cudaGetLastError();
+}
+
 // 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
@@ -245,6 +259,14 @@ public:
             cudaGetLastError();
             stream_ = nullptr;
         }
+        if (cudaMalloc(&ptr_, count * sizeof(T)) == cudaSuccess)
+            return;
+        // cudaMalloc cannot take memory the pool is holding: what pooled buffers handed back stays
+        // reserved by the pool, up to gigabytes after the image loop's engines are gone, until the
+        // pool is trimmed. Hand it back to the driver and try once more; the failed attempt must not
+        // stay behind as the last error (see above).
+        cudaGetLastError();
+        cuda_trim_allocation_pool();
         if (cudaMalloc(&ptr_, count * sizeof(T)) != cudaSuccess)
             throw JFJochException(JFJochExceptionCategory::GPUCUDAError,
                                   "Failed to allocate device memory");