diff --git a/image_analysis/indexing/CUDAMemHelpers.h b/image_analysis/indexing/CUDAMemHelpers.h index adbd98a20..002a83f31 100644 --- a/image_analysis/indexing/CUDAMemHelpers.h +++ b/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");