vulkan: Find optimal memory type but with fallback

Some memory properties are nice to have, but not critical.
`eHostCached`, for instance, isn't essential, and yet we fail on devices
where this memory property isn't available.

    ggml_vulkan: No suitable memory type found: ErrorOutOfDeviceMemory

This change differentiates between those properties that are critical
and those that are just nice-to-have, and will fail only when critical
properties aren't available.

Fixes #5319.

Author: luciferous

ggml-vulkan.cpp

@@ -652,9 +652,21 @@ static void ggml_vk_queue_cleanup(vk_queue& q) {
     q.cmd_buffer_idx = 0;
 }
 
-static vk_buffer ggml_vk_create_buffer(size_t size, vk::MemoryPropertyFlags req_flags) {
+static int32_t find_properties(const vk::PhysicalDeviceMemoryProperties* mem_props, vk::MemoryRequirements* mem_req, vk::MemoryPropertyFlags flags) {
+    for (uint32_t i = 0; i < mem_props->memoryTypeCount; ++i) {
+        vk::MemoryType memory_type = mem_props->memoryTypes[i];
+        if ((mem_req->memoryTypeBits & ((uint64_t)1 << i)) &&
+            (flags & memory_type.propertyFlags) == flags &&
+            mem_props->memoryHeaps[memory_type.heapIndex].size >= mem_req->size) {
+            return static_cast<int32_t>(i);
+        }
+    }
+    return -1;
+}
+
+static vk_buffer ggml_vk_create_buffer(size_t size, vk::MemoryPropertyFlags req_flags, vk::MemoryPropertyFlags desired_flags = vk::MemoryPropertyFlags(0)) {
 #ifdef GGML_VULKAN_DEBUG
-    std::cerr << "ggml_vk_create_buffer(" << size << ", " << to_string(req_flags) << ")" << std::endl;
+    std::cerr << "ggml_vk_create_buffer(" << size << ", " << to_string(req_flags) << ", " << to_string(desired_flags) << ")" << std::endl;
 #endif
     GGML_ASSERT(size > 0);
 
@@ -676,17 +688,15 @@ static vk_buffer ggml_vk_create_buffer(size_t size, vk::MemoryPropertyFlags req_
 
     vk::PhysicalDeviceMemoryProperties mem_props = vk_device.physical_device.getMemoryProperties();
 
-    uint32_t memory_type_index = UINT32_MAX;
-
-    for (uint32_t i = 0; i < mem_props.memoryTypeCount; ++i) {
-        vk::MemoryType memory_type = mem_props.memoryTypes[i];
-        if ((mem_req.memoryTypeBits & ((uint64_t)1 << i)) && (req_flags & memory_type.propertyFlags) == req_flags && mem_props.memoryHeaps[memory_type.heapIndex].size >= mem_req.size) {
-            memory_type_index = i;
-            break;
-        }
+    uint32_t memory_type_index = -1;
+    if (desired_flags) {
+      memory_type_index = find_properties(&mem_props, &mem_req, req_flags | desired_flags);
+    }
+    if (memory_type_index == -1) {
+        memory_type_index = find_properties(&mem_props, &mem_req, req_flags);
     }
 
-    if (memory_type_index >= mem_props.memoryTypeCount) {
+    if (memory_type_index == -1) {
         throw vk::OutOfDeviceMemoryError("No suitable memory type found");
     }
 
@@ -712,9 +722,9 @@ static vk_buffer ggml_vk_create_buffer(size_t size, vk::MemoryPropertyFlags req_
     return buf;
 }
 
-static vk_buffer ggml_vk_create_buffer_check(size_t size, vk::MemoryPropertyFlags req_flags) {
+static vk_buffer ggml_vk_create_buffer_check(size_t size, vk::MemoryPropertyFlags req_flags, vk::MemoryPropertyFlags desired_flags = vk::MemoryPropertyFlags(0)) {
     try {
-        return ggml_vk_create_buffer(size, req_flags);
+        return ggml_vk_create_buffer(size, req_flags, desired_flags);
     } catch (const vk::SystemError& e) {
         std::cerr << "ggml_vulkan: Memory allocation of size " << size << " failed." << std::endl;
         std::cerr << "ggml_vulkan: " << e.what() << std::endl;
@@ -729,7 +739,10 @@ static vk_buffer ggml_vk_create_buffer_device(size_t size) {
     } catch (const vk::SystemError& e) {
         if (vk_device.uma) {
             // Fall back to host memory type
-            buf = ggml_vk_create_buffer_check(size, vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent);
+            buf = ggml_vk_create_buffer_check(
+                size,
+                /* required */ vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent,
+                /* optional */ vk::MemoryPropertyFlagBits::eHostCached);
         } else {
             std::cerr << "ggml_vulkan: Device memory allocation of size " << size << " failed." << std::endl;
             std::cerr << "ggml_vulkan: " << e.what() << std::endl;
@@ -1261,7 +1274,10 @@ static void * ggml_vk_host_malloc(size_t size) {
 #ifdef GGML_VULKAN_DEBUG
     std::cerr << "ggml_vk_host_malloc(" << size << ")" << std::endl;
 #endif
-    vk_buffer buf = ggml_vk_create_buffer(size, vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached);
+    vk_buffer buf = ggml_vk_create_buffer(
+        size,
+        /* required */ vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent,
+        /* optional */ vk::MemoryPropertyFlagBits::eHostCached);
 
     if(!(buf.memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible)) {
         fprintf(stderr, "WARNING: failed to allocate %.2f MB of pinned memory\n",
@@ -1408,7 +1424,10 @@ static void deferred_memcpy(void * dst, const void * src, size_t size, std::vect
 static void ensure_sync_staging_buffer(size_t size) {
     if (vk_sync_staging.size < size) {
         ggml_vk_destroy_buffer(vk_sync_staging);
-        vk_sync_staging = ggml_vk_create_buffer_check(size, vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached);
+        vk_sync_staging = ggml_vk_create_buffer_check(
+            size,
+            /* required */ vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent,
+            /* optional */ vk::MemoryPropertyFlagBits::eHostCached);
     }
 }
 
@@ -3812,7 +3831,10 @@ void ggml_vk_preallocate_buffers() {
     std::cerr << "qx_size: " << vk_prealloc_size_qx << " qy_size: " << vk_prealloc_size_qy << " x_size: " << vk_prealloc_size_x << " y_size: " << vk_prealloc_size_y << " split_k_size: " << vk_prealloc_size_split_k << std::endl;
 #endif
 #if defined(GGML_VULKAN_RUN_TESTS)
-    vk_staging = ggml_vk_create_buffer_check(100ul * 1024ul * 1024ul, vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached);
+    vk_staging = ggml_vk_create_buffer_check(
+        100ul * 1024ul * 1024ul,
+        /* required */ vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent,
+        /* optional */ vk::MemoryPropertyFlagBits::eHostCached);
     ggml_vk_test_transfer(8192 * 1000, false);
     ggml_vk_test_transfer(8192 * 1000, true);
 
@@ -3904,7 +3926,10 @@ void ggml_vk_preallocate_buffers() {
         if (vk_staging.size > 0) {
             ggml_vk_destroy_buffer(vk_staging);
         }
-        vk_staging = ggml_vk_create_buffer_check(vk_staging_size, vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostCached);
+        vk_staging = ggml_vk_create_buffer_check(
+            vk_staging_size,
+            /* required */ vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent,
+            /* optional */ vk::MemoryPropertyFlagBits::eHostCached);
     }
 }