vbatts/llama.cpp
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Vulkan memory management

Browse source
This commit is contained in:
0cc4m 2023-06-11 19:26:52 +02:00
parent 88d4ec05a8
commit a4004d4fa8
1 changed files with 46 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

71
ggml-vulkan.cpp
View file

@ -1,6 +1,7 @@
#include "ggml-vulkan.h" #include "ggml-vulkan.h"
#include <vulkan/vulkan.hpp> #include <vulkan/vulkan.hpp>
#define VMA_IMPLEMENTATION
#include "external/vk_mem_alloc.h" #include "external/vk_mem_alloc.h"
#include <atomic> #include <atomic>
@ -106,13 +107,14 @@ struct scoped_spin_lock {
struct vk_buffer { struct vk_buffer {
vk::Buffer buffer; vk::Buffer buffer;
vk::DeviceMemory memory;
size_t size = 0; size_t size = 0;
}; };
static vk_buffer g_vk_buffer_pool[MAX_VK_BUFFERS]; static vk_buffer g_vk_buffer_pool[MAX_VK_BUFFERS];
static std::atomic_flag g_vk_pool_lock = ATOMIC_FLAG_INIT; static std::atomic_flag g_vk_pool_lock = ATOMIC_FLAG_INIT;
static vk::Buffer ggml_vk_pool_malloc(size_t size, size_t * actual_size) { static void ggml_vk_pool_malloc(size_t size, vk_buffer* buf) {
scoped_spin_lock lock(g_vk_pool_lock); scoped_spin_lock lock(g_vk_pool_lock);
int best_i = -1; int best_i = -1;
@ -121,33 +123,33 @@ static vk::Buffer ggml_vk_pool_malloc(size_t size, size_t * actual_size) {
size_t worst_size = 0; //largest unused buffer seen so far size_t worst_size = 0; //largest unused buffer seen so far
for (int i = 0; i < MAX_VK_BUFFERS; ++i) { for (int i = 0; i < MAX_VK_BUFFERS; ++i) {
vk_buffer &b = g_vk_buffer_pool[i]; vk_buffer &b = g_vk_buffer_pool[i];
if (b.size > 0 && b.size >= size && b.size < best_size) if (b.size > 0 && b.size >= size && b.size < best_size) {
{
best_i = i; best_i = i;
best_size = b.size; best_size = b.size;
} }
if (b.size > 0 && b.size > worst_size) if (b.size > 0 && b.size > worst_size) {
{
worst_i = i; worst_i = i;
worst_size = b.size; worst_size = b.size;
} }
} }
if(best_i!=-1) //found the smallest buffer that fits our needs if(best_i != -1) {
{ //found the smallest buffer that fits our needs
vk_buffer& b = g_vk_buffer_pool[best_i]; vk_buffer& b = g_vk_buffer_pool[best_i];
vk::Buffer buffer = b.buffer; buf->buffer = b.buffer;
*actual_size = b.size; buf->memory = b.memory;
buf->size = b.size;
b.size = 0; b.size = 0;
return buffer; return;
} }
if(worst_i!=-1) //no buffer that fits our needs, resize largest one to save memory if(worst_i != -1) {
{ //no buffer that fits our needs, resize largest one to save memory
vk_buffer& b = g_vk_buffer_pool[worst_i]; vk_buffer& b = g_vk_buffer_pool[worst_i];
vk::Buffer buffer = b.buffer; b.size = 0;
b.size = 0; vk_device.freeMemory(b.memory);
// vkReleaseMemObject(buffer); vk_device.destroyBuffer(b.buffer);
} }
vk::Buffer buffer; buf = new vk_buffer;
buf->size = size;
vk::BufferCreateInfo buffer_create_info{ vk::BufferCreateInfo buffer_create_info{
vk::BufferCreateFlags(), vk::BufferCreateFlags(),
@ -172,27 +174,46 @@ static vk::Buffer ggml_vk_pool_malloc(size_t size, size_t * actual_size) {
VmaAllocation buffer_allocation; VmaAllocation buffer_allocation;
vmaCreateBuffer(allocator, vmaCreateBuffer(allocator,
&static_cast<VkBufferCreateInfo>(buffer_create_info), (VkBufferCreateInfo*)&buffer_create_info,
&allocation_info, &allocation_info,
&static_cast<VkBuffer>(buffer), (VkBuffer*)&buf->buffer,
&buffer_allocation, &buffer_allocation,
nullptr); nullptr);
*actual_size = size; vk::MemoryRequirements buffer_memory_requirements = vk_device.getBufferMemoryRequirements(buf->buffer);
return buffer; vk::PhysicalDeviceMemoryProperties memory_properties = vk_physical_device.getMemoryProperties();
uint32_t memory_type_index = uint32_t(~0);
for (uint32_t current_memory_type_index = 0; current_memory_type_index < memory_properties.memoryTypeCount; current_memory_type_index++) {
vk::MemoryType memory_type = memory_properties.memoryTypes[current_memory_type_index];
if ((vk::MemoryPropertyFlagBits::eHostVisible & memory_type.propertyFlags) &&
(vk::MemoryPropertyFlagBits::eHostCoherent & memory_type.propertyFlags))
{
memory_type_index = current_memory_type_index;
break;
}
}
vk::MemoryAllocateInfo buffer_memory_allocate_info(buffer_memory_requirements.size, memory_type_index);
buf->memory = vk_device.allocateMemory(buffer_memory_allocate_info);
} }
static void ggml_vk_pool_free(vk::Buffer buffer, size_t size) { static void ggml_vk_pool_free(vk_buffer* buffer) {
scoped_spin_lock lock(g_vk_pool_lock); scoped_spin_lock lock(g_vk_pool_lock);
for (int i = 0; i < MAX_VK_BUFFERS; ++i) { for (int i = 0; i < MAX_VK_BUFFERS; ++i) {
vk_buffer& b = g_vk_buffer_pool[i]; vk_buffer& b = g_vk_buffer_pool[i];
if (b.size == 0) { if (b.size == 0) {
b.buffer = buffer; b.buffer = buffer->buffer;
b.size = size; b.memory = buffer->memory;
b.size = buffer->size;
return; return;
} }
} }
fprintf(stderr, "WARNING: vk buffer pool full, increase MAX_VK_BUFFERS\n"); fprintf(stderr, "WARNING: vk buffer pool full, increase MAX_VK_BUFFERS\n");
vkReleaseMemObject(mem); buffer->size = 0;
vk_device.freeMemory(buffer->memory);
vk_device.destroyBuffer(buffer->buffer);
} }