diff --git a/ggml-cuda.cu b/ggml-cuda.cu index ab3d37ea0..9d42efb0d 100644 --- a/ggml-cuda.cu +++ b/ggml-cuda.cu @@ -4293,45 +4293,53 @@ static void * ggml_cuda_pool_malloc(size_t size, size_t * actual_size) { scoped_spin_lock lock(g_cuda_pool_lock); int id; CUDA_CHECK(cudaGetDevice(&id)); - - int best_i = -1; - size_t best_size = std::numeric_limits::max(); //smallest unused buffer that fits our needs - int worst_i = -1; - size_t worst_size = 0; //largest unused buffer seen so far - +#ifdef DEBUG_CUDA_MALLOC + int nnz = 0; + size_t max_size = 0, tot_size = 0; +#endif + size_t best_diff = 1ull << 36; + int ibest = -1; for (int i = 0; i < MAX_CUDA_BUFFERS; ++i) { cuda_buffer& b = g_cuda_buffer_pool[id][i]; - if (b.size > 0 && b.size >= size && b.size < best_size) - { - best_i = i; - best_size = b.size; - } - if (b.size > 0 && b.size > worst_size) - { - worst_i = i; - worst_size = b.size; + if (b.ptr != nullptr) { +#ifdef DEBUG_CUDA_MALLOC + ++nnz; + tot_size += b.size; + if (b.size > max_size) max_size = b.size; +#endif + if (b.size >= size) { + size_t diff = b.size - size; + if (diff < best_diff) { + best_diff = diff; + ibest = i; + if (!best_diff) { + void * ptr = b.ptr; + *actual_size = b.size; + b.ptr = nullptr; + b.size = 0; + return ptr; + } + } + } } } - if(best_i!=-1) //found the smallest buffer that fits our needs - { - cuda_buffer& b = g_cuda_buffer_pool[id][best_i]; + if (ibest >= 0) { + cuda_buffer& b = g_cuda_buffer_pool[id][ibest]; void * ptr = b.ptr; *actual_size = b.size; b.ptr = nullptr; b.size = 0; return ptr; } - if(worst_i!=-1) //no buffer that fits our needs, resize largest one to save memory - { - cuda_buffer& b = g_cuda_buffer_pool[id][worst_i]; - b.size = 0; - void * ptr = b.ptr; - cudaFree(ptr); - b.ptr = ptr = nullptr; - } +#ifdef DEBUG_CUDA_MALLOC + fprintf(stderr, "%s: %d buffers, max_size = %u MB, tot_size = %u MB, requested %u MB\n", __func__, nnz, + (uint32_t)(max_size/1024/1024), (uint32_t)(tot_size/1024/1024), (uint32_t)(size/1024/1024)); +#endif void * ptr; - CUDA_CHECK(cudaMalloc((void **) &ptr, size)); - *actual_size = size; + size_t look_ahead_size = (size_t) (1.05 * size); + look_ahead_size = 256 * ((look_ahead_size + 255)/256); + CUDA_CHECK(cudaMalloc((void **) &ptr, look_ahead_size)); + *actual_size = look_ahead_size; return ptr; }