From e0bcaccbe50a773c8addec49bbfc8b7ba953c63b Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Tue, 13 Oct 2026 15:05:10 +0200 Subject: [PATCH] vulkan: allow multiple devices to access gpu-visible host memory (#29741) * vulkan: allow multiple devices to access gpu-visible host memory * clean up erroneous newlines * add Windows support * handle nullptr * fix enable condition --- ggml/src/ggml-vulkan/ggml-vulkan-buffers.cpp | 156 ++++++++++++++----- ggml/src/ggml-vulkan/ggml-vulkan-common.h | 4 +- ggml/src/ggml-vulkan/ggml-vulkan-types.h | 37 ++--- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 27 +++- 4 files changed, 153 insertions(+), 71 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-buffers.cpp b/ggml/src/ggml-vulkan/ggml-vulkan-buffers.cpp index 351951638..eabd8974a 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-buffers.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan-buffers.cpp @@ -1,5 +1,27 @@ #include "ggml-vulkan-common.h" +#include +#include +#include +#include +#include + +static void * portable_aligned_malloc(size_t align, size_t size) { +#ifdef _WIN32 + return _aligned_malloc(size, align); +#else + return std::aligned_alloc(align, size); +#endif +} + +static void portable_aligned_free(void * ptr) { +#ifdef _WIN32 + _aligned_free(ptr); +#else + std::free(ptr); +#endif +} + 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, @@ -243,67 +265,125 @@ void ggml_vk_destroy_buffer(vk_buffer& buf) { buf.reset(); } -void * ggml_vk_host_malloc(vk_device& device, size_t size) { - VK_LOG_MEMORY("ggml_vk_host_malloc(" << size << ")"); - vk_buffer buf = ggml_vk_create_buffer(device, size, - {vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached, - vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent}); +// 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; +} - 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); +void * ggml_vk_host_malloc(size_t size) { + VK_LOG_MEMORY("ggml_vk_host_malloc(" << size << ")"); + + ggml_vk_instance_init(); + + if (vk_instance.multi_device) { + const size_t align = std::max(vk_instance.host_import_alignment, TENSOR_ALIGNMENT); + const size_t alloc_size = (size + align - 1) / align * align; + + void * ptr = portable_aligned_malloc(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; + try { + buf = ggml_vk_create_buffer(dev0, size, + {vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached, + vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent}); + } catch (vk::SystemError&) { return nullptr; } - std::lock_guard guard(device->pinned_memory_mutex); - device->pinned_memory.push_back(std::make_tuple(buf->ptr, size, buf)); + 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; + } + 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) { + portable_aligned_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 8922025ed..66b7557be 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-common.h +++ b/ggml/src/ggml-vulkan/ggml-vulkan-common.h @@ -45,8 +45,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 4158ff779..6ad178328 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-types.h +++ b/ggml/src/ggml-vulkan/ggml-vulkan-types.h @@ -48,43 +48,24 @@ typedef struct VkPhysicalDeviceCooperativeMatrixDecodeVectorFeaturesNV { #endif #include - #include - #include - #include - #include - #include - #include - #include - #include - #include - #include - #include - #include - #include - #include - #include - #include - #include - #include - #include #if defined(_MSC_VER) @@ -110,9 +91,7 @@ typedef struct VkPhysicalDeviceCooperativeMatrixDecodeVectorFeaturesNV { #endif #include "ggml-impl.h" - #include "ggml-backend-impl.h" - #include "ggml-vulkan-shaders.hpp" #if !defined(VK_KHR_shader_bfloat16) @@ -680,7 +659,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 @@ -1019,8 +997,6 @@ struct vk_device_struct { std::vector all_pipelines; - std::vector> pinned_memory; - vk::Fence fence; vk_buffer sync_staging; @@ -1333,6 +1309,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; @@ -1347,6 +1329,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 5435171f1..e8589dc0e 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -5411,22 +5411,37 @@ void ggml_vk_instance_init() { } GGML_LOG_DEBUG("ggml_vulkan: Found %zu Vulkan devices:\n", vk_instance.device_indices.size()); + bool all_external_memory_host = vk_instance.device_indices.size() > 1; for (size_t i = 0; i < vk_instance.device_indices.size(); i++) { vk::PhysicalDevice vkdev = devices[vk_instance.device_indices[i]]; 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); + } else { + all_external_memory_host = false; + } + ggml_vk_print_gpu_info(i); } + + vk_instance.multi_device = all_external_memory_host; } void ggml_vk_init(ggml_backend_vk_context * ctx, size_t idx) { @@ -13118,18 +13133,16 @@ 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) { VK_LOG_MEMORY("ggml_backend_vk_host_buffer_type_alloc_buffer(" << size << ")"); size += 32; // Behave like the CPU buffer type - void * ptr = nullptr; - try { - ptr = ggml_vk_host_malloc(vk_instance.devices[0], size); - } catch (vk::SystemError& e) { - GGML_LOG_WARN("ggml_vulkan: Failed to allocate pinned memory (%s)\n", e.what()); + void * ptr = ggml_vk_host_malloc(size); + if (ptr == nullptr) { + GGML_LOG_WARN("ggml_vulkan: Failed to allocate pinned memory\n"); // fallback to cpu buffer return ggml_backend_buft_alloc_buffer(ggml_backend_cpu_buffer_type(), size); }