vulkan: implement zero-copy buffer creation from host memory on UMA

Add support for importing host-allocated memory as Vulkan buffers using
VK_EXT_external_memory_host extension.  This enables zero-copy model
loading on UMA (Unified Memory Architecture) systems by allowing the
GPU to directly access mmap'd model weight data without requiring a
copy to device memory.

Signed-off-by: Giuseppe Scrivano <[email protected]>

Author: giuseppe

ggml/src/ggml-vulkan/ggml-vulkan.cpp

@@ -423,6 +423,7 @@ struct vk_device_struct {
     bool multi_add;
     bool shader_int64;
     bool buffer_device_address;
+    bool ext_external_memory_host;
 
     bool add_rms_fusion;
     uint32_t partials_binding_alignment;
@@ -680,6 +681,9 @@ struct vk_buffer_struct {
 
     vk_device device;
 
+    bool from_host_ptr = false;
+    size_t alignment_offset = 0;
+
     ~vk_buffer_struct() {
         if (size == 0) {
             return;
@@ -1500,13 +1504,6 @@ struct ggml_backend_vk_context {
 
 static void * const vk_ptr_base = (void *)(uintptr_t) 0x1000;  // NOLINT
 
-static uint64_t vk_tensor_offset(const ggml_tensor * tensor) {
-    if (tensor->view_src) {
-        return (uint8_t *) tensor->view_src->data - (uint8_t *) vk_ptr_base;
-    }
-    return (uint8_t *) tensor->data - (uint8_t *) vk_ptr_base;
-}
-
 struct ggml_backend_vk_buffer_context {
     vk_device_ref device;
     vk_buffer dev_buffer;
@@ -1523,6 +1520,16 @@ struct ggml_backend_vk_buffer_context {
     }
 };
 
+static uint64_t vk_tensor_offset(const ggml_tensor * tensor) {
+    ggml_backend_vk_buffer_context * buf_ctx = static_cast<ggml_backend_vk_buffer_context *>(tensor->buffer->context);
+    vk_buffer buf = buf_ctx->dev_buffer;
+
+    void * base_addr = buf->from_host_ptr ? buf->ptr : vk_ptr_base;
+    void * tensor_data = tensor->view_src ? tensor->view_src->data : tensor->data;
+
+    return (uint8_t *)tensor_data - (uint8_t *)base_addr;
+}
+
 #ifdef GGML_VULKAN_MEMORY_DEBUG
 static std::mutex log_mutex;
 
@@ -2180,6 +2187,76 @@ static vk_buffer ggml_vk_create_buffer_device(vk_device& device, size_t size) {
     return buf;
 }
 
+static vk_buffer ggml_vk_create_buffer_from_host_ptr(vk_device& device, void * ptr, size_t size) {
+    if (!device->ext_external_memory_host) {
+        throw std::runtime_error("VK_EXT_external_memory_host extension not available");
+    }
+
+    const size_t page_size = device->physical_device.getProperties().limits.minMemoryMapAlignment;
+    uintptr_t ptr_addr = reinterpret_cast<uintptr_t>(ptr);
+    uintptr_t page_aligned_base = ptr_addr & ~(page_size - 1);
+    void* aligned_ptr = reinterpret_cast<void*>(page_aligned_base);
+    size_t offset = ptr_addr - page_aligned_base;
+    size_t aligned_size = (size + offset + page_size - 1) & ~(page_size - 1);
+
+    vk::BufferUsageFlags usage_flags = vk::BufferUsageFlagBits::eStorageBuffer |
+                                       vk::BufferUsageFlagBits::eTransferSrc |
+                                       vk::BufferUsageFlagBits::eShaderDeviceAddress;
+
+    vk_buffer buf = std::make_shared<vk_buffer_struct>();
+
+    vk::BufferCreateInfo buffer_create_info{{}, aligned_size, usage_flags, vk::SharingMode::eExclusive};
+    buf->buffer = device->device.createBuffer(buffer_create_info);
+
+    vk::MemoryRequirements mem_req = device->device.getBufferMemoryRequirements(buf->buffer);
+    vk::MemoryPropertyFlags req_flags = vk::MemoryPropertyFlagBits::eHostVisible |
+                                        vk::MemoryPropertyFlagBits::eHostCoherent;
+
+    vk::MemoryRequirements modified_req = mem_req;
+
+    vk::PhysicalDeviceMemoryProperties mem_props = device->physical_device.getMemoryProperties();
+    uint32_t memory_type_index = find_properties(&mem_props, &modified_req, req_flags);
+    if (memory_type_index == UINT32_MAX) {
+        device->device.destroyBuffer(buf->buffer);
+        throw vk::OutOfDeviceMemoryError("No compatible memory type found");
+    }
+
+    VkImportMemoryHostPointerInfoEXT import_info = {
+        VK_STRUCTURE_TYPE_IMPORT_MEMORY_HOST_POINTER_INFO_EXT,
+        nullptr,
+        VK_EXTERNAL_MEMORY_HANDLE_TYPE_HOST_ALLOCATION_BIT_EXT,
+        aligned_ptr
+    };
+
+    VkMemoryAllocateInfo alloc_info = {
+        VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO,
+        &import_info,
+        aligned_size,
+        memory_type_index
+    };
+
+    buf->device_memory = device->device.allocateMemory(alloc_info);
+    device->device.bindBufferMemory(buf->buffer, buf->device_memory, 0);
+
+    buf->ptr = aligned_ptr;
+    buf->size = aligned_size;
+    buf->alignment_offset = offset;
+    buf->from_host_ptr = true;
+    buf->device = device;
+    buf->memory_property_flags = req_flags;
+
+    if (device->buffer_device_address) {
+        const vk::BufferDeviceAddressInfo addressInfo(buf->buffer);
+        buf->bda_addr = device->device.getBufferAddress(addressInfo);
+    }
+
+#ifdef GGML_VULKAN_MEMORY_DEBUG
+    device->memory_logger->log_allocation(buf, size);
+#endif
+
+    return buf;
+}
+
 static void ggml_vk_destroy_buffer(vk_buffer& buf) {
     if (buf == nullptr) {
         return;
@@ -3819,6 +3896,8 @@ static vk_device ggml_vk_get_device(size_t idx) {
                 pipeline_robustness = true;
             } else if (strcmp("VK_EXT_subgroup_size_control", properties.extensionName) == 0) {
                 device->subgroup_size_control = true;
+            } else if (strcmp("VK_EXT_external_memory_host", properties.extensionName) == 0) {
+                device->ext_external_memory_host = true;
 #if defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT)
             } else if (strcmp("VK_KHR_cooperative_matrix", properties.extensionName) == 0 &&
                        !getenv("GGML_VK_DISABLE_COOPMAT")) {
@@ -4223,6 +4302,10 @@ static vk_device ggml_vk_get_device(size_t idx) {
             device_extensions.push_back("VK_KHR_shader_float16_int8");
         }
 
+        if (device->ext_external_memory_host) {
+            device_extensions.push_back("VK_EXT_external_memory_host");
+        }
+
 #if defined(VK_KHR_cooperative_matrix)
         if (device->coopmat_support) {
             // Query supported shapes
@@ -11835,9 +11918,13 @@ static void ggml_backend_vk_buffer_free_buffer(ggml_backend_buffer_t buffer) {
 }
 
 static void * ggml_backend_vk_buffer_get_base(ggml_backend_buffer_t buffer) {
-    return vk_ptr_base;
+    ggml_backend_vk_buffer_context * buf_ctx = (ggml_backend_vk_buffer_context *)buffer->context;
+    vk_buffer buf = buf_ctx->dev_buffer;
 
-    UNUSED(buffer);
+    if (buf->from_host_ptr) {
+        return (uint8_t*)buf->ptr + buf->alignment_offset;
+    }
+    return vk_ptr_base;
 }
 
 static enum ggml_status ggml_backend_vk_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) {
@@ -12876,6 +12963,25 @@ static ggml_backend_buffer_type_t ggml_backend_vk_device_get_host_buffer_type(gg
     return ggml_backend_vk_host_buffer_type();
 }
 
+static ggml_backend_buffer_t ggml_backend_vk_device_buffer_from_host_ptr(ggml_backend_dev_t dev, void * ptr, size_t size, size_t max_tensor_size) {
+    VK_LOG_MEMORY("ggml_backend_vk_device_buffer_from_host_ptr(" << size << ")");
+
+    ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context;
+    vk_device device = ggml_vk_get_device(ctx->device);
+
+    if (!device->uma) {
+        GGML_ABORT("ggml_backend_vk_device_buffer_from_host_ptr works only with UMA devices");
+    }
+
+    vk_buffer dev_buffer = ggml_vk_create_buffer_from_host_ptr(device, ptr, size);
+
+    ggml_backend_vk_buffer_context * bufctx = new ggml_backend_vk_buffer_context(device, std::move(dev_buffer), ctx->name);
+    ggml_backend_buffer_type_t buft = ggml_backend_vk_device_get_buffer_type(dev);
+
+    UNUSED(max_tensor_size);
+    return ggml_backend_buffer_init(buft, ggml_backend_vk_buffer_interface, bufctx, size);
+}
+
 static enum ggml_backend_dev_type ggml_backend_vk_device_get_type(ggml_backend_dev_t dev) {
     ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context;
 
@@ -12884,6 +12990,7 @@ static enum ggml_backend_dev_type ggml_backend_vk_device_get_type(ggml_backend_d
 
 static void ggml_backend_vk_device_get_props(ggml_backend_dev_t dev, struct ggml_backend_dev_props * props) {
     ggml_backend_vk_device_context * ctx = (ggml_backend_vk_device_context *)dev->context;
+    vk_device device = ggml_vk_get_device(ctx->device);
 
     props->name        = ggml_backend_vk_device_get_name(dev);
     props->description = ggml_backend_vk_device_get_description(dev);
@@ -12893,7 +13000,7 @@ static void ggml_backend_vk_device_get_props(ggml_backend_dev_t dev, struct ggml
     props->caps = {
         /* .async                 = */ false,
         /* .host_buffer           = */ true,
-        /* .buffer_from_host_ptr  = */ false,
+        /* .buffer_from_host_ptr  = */ device->uma,
         /* .events                = */ false,
     };
 }
@@ -13338,7 +13445,7 @@ static const struct ggml_backend_device_i ggml_backend_vk_device_i = {
     /* .init_backend         = */ ggml_backend_vk_device_init,
     /* .get_buffer_type      = */ ggml_backend_vk_device_get_buffer_type,
     /* .get_host_buffer_type = */ ggml_backend_vk_device_get_host_buffer_type,
-    /* .buffer_from_host_ptr = */ NULL,
+    /* .buffer_from_host_ptr = */ ggml_backend_vk_device_buffer_from_host_ptr,
     /* .supports_op          = */ ggml_backend_vk_device_supports_op,
     /* .supports_buft        = */ ggml_backend_vk_device_supports_buft,
     /* .offload_op           = */ ggml_backend_vk_device_offload_op,

src/llama-mmap.cpp

@@ -285,7 +285,7 @@ struct llama_mmap::impl {
         }
         if (prefetch) { flags |= MAP_POPULATE; }
 #endif
-        addr = mmap(NULL, file->size(), PROT_READ, flags, fd, 0);
+        addr = mmap(NULL, file->size(), PROT_READ | PROT_WRITE, flags, fd, 0);
         if (addr == MAP_FAILED) {
             throw std::runtime_error(format("mmap failed: %s", strerror(errno)));
         }

src/llama-model-loader.cpp

@@ -503,7 +503,7 @@ llama_model_loader::llama_model_loader(
     get_key(llm_kv(LLM_KV_GENERAL_ARCHITECTURE), arch_name, false);
     llm_kv = LLM_KV(llm_arch_from_string(arch_name));
 
-    files.emplace_back(new llama_file(fname.c_str(), "rb"));
+    files.emplace_back(new llama_file(fname.c_str(), "r+b"));
     contexts.emplace_back(ctx);
 
     // Save tensors data offset of the main file.
@@ -571,7 +571,7 @@ llama_model_loader::llama_model_loader(
                 }
             }
 
-            files.emplace_back(new llama_file(fname_split, "rb"));
+            files.emplace_back(new llama_file(fname_split, "r+b"));
             contexts.emplace_back(ctx);
 
             // Save tensors data offset info of the shard.