From ef7f13a1eba9ebdbb0cdaac75ea3e89f3fe95fc7 Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Wed, 30 Sep 2026 13:54:25 +0200 Subject: [PATCH] vulkan: allow multiple devices to access gpu-visible host memory --- ggml/src/ggml-vulkan/ggml-vulkan-buffers.cpp | 129 ++++++++++++++----- ggml/src/ggml-vulkan/ggml-vulkan-common.h | 4 +- ggml/src/ggml-vulkan/ggml-vulkan-types.h | 16 ++- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 18 ++- 4 files changed, 124 insertions(+), 43 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-buffers.cpp b/ggml/src/ggml-vulkan/ggml-vulkan-buffers.cpp index 4d4c849513..c24bd99743 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-buffers.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan-buffers.cpp @@ -1,5 +1,11 @@ #include "ggml-vulkan-common.h" +#include +#include +#include +#include +#include + ggml_backend_buffer_type_i ggml_backend_vk_buffer_type_interface = { /* .get_name = */ ggml_backend_vk_buffer_type_name, /* .alloc_buffer = */ ggml_backend_vk_buffer_type_alloc_buffer, @@ -241,67 +247,120 @@ void ggml_vk_destroy_buffer(vk_buffer& buf) { buf.reset(); } -void * ggml_vk_host_malloc(vk_device& device, size_t size) { +// caller must hold host_allocations_mutex +static size_t ggml_vk_host_allocation_find(const void * ptr) { + for (size_t i = 0; i < vk_instance.host_allocations.size(); i++) { + const uint8_t * base = (const uint8_t *) vk_instance.host_allocations[i].base; + if (ptr >= base && ptr < base + vk_instance.host_allocations[i].size) { + return i; + } + } + return SIZE_MAX; +} + +void * ggml_vk_host_malloc(size_t size) { VK_LOG_MEMORY("ggml_vk_host_malloc(" << size << ")"); - vk_buffer buf = ggml_vk_create_buffer(device, size, + + ggml_vk_instance_init(); + + if (vk_instance.multi_device) { + const size_t align = std::max(vk_instance.host_import_alignment, sizeof(void *)); + const size_t alloc_size = (size + align - 1) / align * align; + + void * ptr = std::aligned_alloc(align, alloc_size); + if (ptr == nullptr) { + return nullptr; + } + vk_host_allocation entry; + entry.base = ptr; + entry.size = alloc_size; + + std::unique_lock guard(vk_instance.host_allocations_mutex); + vk_instance.host_allocations.push_back(std::move(entry)); + return ptr; + } + + vk_device dev0 = ggml_vk_get_device(0); + vk_buffer buf = ggml_vk_create_buffer(dev0, size, {vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached, vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent}); - if(!(buf->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible)) { - fprintf(stderr, "WARNING: failed to allocate %.2f MB of pinned memory\n", - size/1024.0/1024.0); - device->device.freeMemory(buf->device_memory); - device->device.destroyBuffer(buf->buffer); + if (!(buf->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible)) { + fprintf(stderr, "WARNING: failed to allocate %.2f MB of pinned memory\n", size / 1024.0 / 1024.0); + ggml_vk_destroy_buffer(buf); return nullptr; } - std::lock_guard guard(device->pinned_memory_mutex); - device->pinned_memory.push_back(std::make_tuple(buf->ptr, size, buf)); + vk_host_allocation entry; + entry.base = buf->ptr; + entry.size = size; + entry.views[dev0->idx] = buf; + std::unique_lock guard(vk_instance.host_allocations_mutex); + vk_instance.host_allocations.push_back(std::move(entry)); return buf->ptr; } -void ggml_vk_host_free(vk_device& device, void* ptr) { +void ggml_vk_host_free(void * ptr) { if (ptr == nullptr) { return; } VK_LOG_MEMORY("ggml_vk_host_free(" << ptr << ")"); - std::lock_guard guard(device->pinned_memory_mutex); - vk_buffer buf; - size_t index; - for (size_t i = 0; i < device->pinned_memory.size(); i++) { - const uint8_t* addr = (const uint8_t*) std::get<0>(device->pinned_memory[i]); - const uint8_t* endr = addr + std::get<1>(device->pinned_memory[i]); - if (ptr >= addr && ptr < endr) { - buf = std::get<2>(device->pinned_memory[i]); - index = i; - break; + vk_host_allocation entry; + { + std::unique_lock guard(vk_instance.host_allocations_mutex); + const size_t index = ggml_vk_host_allocation_find(ptr); + if (index == SIZE_MAX) { + fprintf(stderr, "WARNING: failed to free pinned memory: memory not in map\n"); + return; } - } - if (buf == nullptr) { - fprintf(stderr, "WARNING: failed to free pinned memory: memory not in map\n"); - return; + entry = std::move(vk_instance.host_allocations[index]); + vk_instance.host_allocations.erase(vk_instance.host_allocations.begin() + index); } - ggml_vk_destroy_buffer(buf); - - device->pinned_memory.erase(device->pinned_memory.begin() + index); + for (auto & view : entry.views) { + ggml_vk_destroy_buffer(view); + } + if (vk_instance.multi_device) { + std::free(entry.base); + } } void ggml_vk_host_get(const vk_device& device, const void * ptr, vk_buffer& buf, size_t& buf_offset) { - std::shared_lock guard(device->pinned_memory_mutex); buf = nullptr; buf_offset = 0; - for (size_t i = 0; i < device->pinned_memory.size(); i++) { - const uint8_t* addr = (const uint8_t*) std::get<0>(device->pinned_memory[i]); - const uint8_t* endr = addr + std::get<1>(device->pinned_memory[i]); - if (ptr >= addr && ptr < endr) { - buf = std::get<2>(device->pinned_memory[i]); - buf_offset = ((const uint8_t *)ptr) - addr; - break; + + std::unique_lock guard(vk_instance.host_allocations_mutex); + + const size_t index = ggml_vk_host_allocation_find(ptr); + if (index == SIZE_MAX) { + return; + } + vk_host_allocation & entry = vk_instance.host_allocations[index]; + + if (entry.views[device->idx] == nullptr) { + if (!vk_instance.multi_device) { + return; + } + if (!device->external_memory_host) { + static std::atomic warned_mask{0}; + const uint32_t bit = 1u << device->idx; + if (!(warned_mask.fetch_or(bit) & bit)) { + GGML_LOG_WARN("ggml_vulkan: %s does not support VK_EXT_external_memory_host; host memory " + "will be staged, not pinned, on this device\n", device->name.c_str()); + } + return; + } + // const only propagates from the const context; the pointee is mutable + entry.views[device->idx] = ggml_vk_buffer_from_host_ptr(const_cast(device), entry.base, entry.size); + if (entry.views[device->idx] == nullptr) { + return; } } + + buf = entry.views[device->idx]; + buf_offset = (const uint8_t *) ptr - (const uint8_t *) entry.base; } void ggml_vk_ensure_sync_staging_buffer(vk_device& device, size_t size) { diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-common.h b/ggml/src/ggml-vulkan/ggml-vulkan-common.h index 125e7608a3..b217ccc6f4 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-common.h +++ b/ggml/src/ggml-vulkan/ggml-vulkan-common.h @@ -41,8 +41,8 @@ bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, const vk_ vk_buffer ggml_vk_create_buffer_check(vk_device& device, size_t size, vk::MemoryPropertyFlags req_flags, vk::MemoryPropertyFlags fallback_flags = vk::MemoryPropertyFlags(0)); vk_buffer ggml_vk_create_buffer_device(vk_device& device, size_t size); void ggml_vk_destroy_buffer(vk_buffer& buf); -void * ggml_vk_host_malloc(vk_device& device, size_t size); -void ggml_vk_host_free(vk_device& device, void* ptr); +void * ggml_vk_host_malloc(size_t size); +void ggml_vk_host_free(void * ptr); void ggml_vk_host_get(const vk_device& device, const void * ptr, vk_buffer& buf, size_t& buf_offset); void ggml_vk_ensure_sync_staging_buffer(vk_device& device, size_t size); void ggml_vk_ensure_sync_staging_buffer(ggml_backend_vk_context * ctx, size_t size); diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-types.h b/ggml/src/ggml-vulkan/ggml-vulkan-types.h index ee50f51a38..edc1560c07 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-types.h +++ b/ggml/src/ggml-vulkan/ggml-vulkan-types.h @@ -678,7 +678,6 @@ class vk_memory_logger; struct vk_device_struct { std::recursive_mutex mutex; std::mutex queue_submit_mutex; - mutable std::shared_mutex pinned_memory_mutex; // Guards compile_pending, all_pipelines, and the dynamic pipeline maps // (flash_attn, fa_mask_opt, solve_tri, conv2d, etc). The actual compile @@ -1017,8 +1016,6 @@ struct vk_device_struct { std::vector all_pipelines; - std::vector> pinned_memory; - vk::Fence fence; vk_buffer sync_staging; @@ -1331,6 +1328,12 @@ struct ggml_backend_vk_buffer_context { }; +struct vk_host_allocation { + void * base = nullptr; + size_t size = 0; + vk_buffer views[GGML_VK_MAX_DEVICES]; // indexed by device->idx; null = not (yet) mapped here +}; + struct vk_instance_t { vk::Instance instance; @@ -1345,6 +1348,13 @@ struct vk_instance_t { std::vector device_indices; std::vector device_supports_membudget; vk_device devices[GGML_VK_MAX_DEVICES]; + + std::vector host_allocations; + std::shared_mutex host_allocations_mutex; + + size_t host_import_alignment = 0; + + bool multi_device = false; }; typedef void (*ggml_vk_func_t)(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 38ec1f9e82..2022db911e 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -5388,6 +5388,8 @@ void ggml_vk_instance_init() { GGML_LOG_INFO("ggml_vulkan: No devices found.\n"); return; } + + vk_instance.multi_device = vk_instance.device_indices.size() > 1; } GGML_LOG_DEBUG("ggml_vulkan: Found %zu Vulkan devices:\n", vk_instance.device_indices.size()); @@ -5396,15 +5398,25 @@ void ggml_vk_instance_init() { std::vector extensionprops = vkdev.enumerateDeviceExtensionProperties(); bool membudget_supported = false; + bool external_memory_host_supported = false; for (const auto & ext : extensionprops) { if (strcmp(VK_EXT_MEMORY_BUDGET_EXTENSION_NAME, ext.extensionName) == 0) { membudget_supported = true; - break; + } else if (strcmp("VK_EXT_external_memory_host", ext.extensionName) == 0) { + external_memory_host_supported = true; } } vk_instance.device_supports_membudget.push_back(membudget_supported); + if (external_memory_host_supported) { + vk::PhysicalDeviceProperties2 props2; + vk::PhysicalDeviceExternalMemoryHostPropertiesEXT ext_props; + props2.pNext = &ext_props; + vkdev.getProperties2(&props2); + vk_instance.host_import_alignment = std::max(vk_instance.host_import_alignment, ext_props.minImportedHostPointerAlignment); + } + ggml_vk_print_gpu_info(i); } } @@ -13054,7 +13066,7 @@ static const char * ggml_backend_vk_host_buffer_type_name(ggml_backend_buffer_ty static void ggml_backend_vk_host_buffer_free_buffer(ggml_backend_buffer_t buffer) { VK_LOG_MEMORY("ggml_backend_vk_host_buffer_free_buffer()"); - ggml_vk_host_free(vk_instance.devices[0], buffer->context); + ggml_vk_host_free(buffer->context); } static ggml_backend_buffer_t ggml_backend_vk_host_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { @@ -13063,7 +13075,7 @@ static ggml_backend_buffer_t ggml_backend_vk_host_buffer_type_alloc_buffer(ggml_ size += 32; // Behave like the CPU buffer type void * ptr = nullptr; try { - ptr = ggml_vk_host_malloc(vk_instance.devices[0], size); + ptr = ggml_vk_host_malloc(size); } catch (vk::SystemError& e) { GGML_LOG_WARN("ggml_vulkan: Failed to allocate pinned memory (%s)\n", e.what()); // fallback to cpu buffer