From 07659374bd693324d62cd11be7c9c3aa73a26249 Mon Sep 17 00:00:00 2001
From: Michael Vigovsky <upliner@gmail.com>
Date: Mon, 21 Oct 2024 19:17:51 +0300
Subject: [PATCH] Always check allocate_memory_* return values

---
 src/allocator.cpp | 48 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 48 insertions(+)

diff --git a/src/allocator.cpp b/src/allocator.cpp
index ee070aa0b24..9307d0af22b 100644
--- a/src/allocator.cpp
+++ b/src/allocator.cpp
@@ -760,6 +760,12 @@ VkBufferMemory* VkBlobAllocator::fastMalloc(size_t size)
     }
 
     block->memory = allocate_memory(memoryRequirements.size, buffer_memory_type_index);
+    if (!block->memory)
+    {
+        vkDestroyBuffer(vkdev->vkdevice(), block->buffer, 0);
+        delete block;
+        return 0;
+    }
 
     // ignore memoryRequirements.alignment as we always bind at zero offset
     vkBindBufferMemory(vkdev->vkdevice(), block->buffer, block->memory, 0);
@@ -1032,6 +1038,12 @@ VkImageMemory* VkBlobAllocator::fastMalloc(int w, int h, int c, size_t elemsize,
 
     // bind at memory offset
     ptr->memory = allocate_memory(new_block_size, image_memory_type_index);
+    if (!ptr->memory)
+    {
+        vkDestroyImage(vkdev->vkdevice(), ptr->image, 0);
+        delete ptr;
+        return 0;
+    }
     ptr->bind_offset = 0;
     ptr->bind_capacity = aligned_size;
 
@@ -1341,6 +1353,12 @@ VkBufferMemory* VkWeightAllocator::fastMalloc(size_t size)
             }
 
             block->memory = allocate_dedicated_memory(memoryRequirements2.memoryRequirements.size, buffer_memory_type_index, 0, block->buffer);
+            if (!block->memory)
+            {
+                vkDestroyBuffer(vkdev->vkdevice(), block->buffer, 0);
+                delete block;
+                return 0;
+            }
 
             // ignore memoryRequirements2.memoryRequirements.alignment as we always bind at zero offset
             vkBindBufferMemory(vkdev->vkdevice(), block->buffer, block->memory, 0);
@@ -1400,6 +1418,12 @@ VkBufferMemory* VkWeightAllocator::fastMalloc(size_t size)
     }
 
     block->memory = allocate_memory(memoryRequirements.size, buffer_memory_type_index);
+    if (!block->memory)
+    {
+        vkDestroyBuffer(vkdev->vkdevice(), block->buffer, 0);
+        delete block;
+        return 0;
+    }
 
     // ignore memoryRequirements.alignment as we always bind at zero offset
     vkBindBufferMemory(vkdev->vkdevice(), block->buffer, block->memory, 0);
@@ -1547,6 +1571,12 @@ VkImageMemory* VkWeightAllocator::fastMalloc(int w, int h, int c, size_t elemsiz
 
             // bind memory
             ptr->memory = allocate_dedicated_memory(memoryRequirements2.memoryRequirements.size, image_memory_type_index, ptr->image, 0);
+            if (!ptr->memory)
+            {
+                vkDestroyImage(vkdev->vkdevice(), ptr->image, 0);
+                delete ptr;
+                return 0;
+            }
             ptr->bind_offset = 0;
             ptr->bind_capacity = memoryRequirements2.memoryRequirements.size;
 
@@ -1654,6 +1684,12 @@ VkImageMemory* VkWeightAllocator::fastMalloc(int w, int h, int c, size_t elemsiz
 
     // bind at memory offset
     ptr->memory = allocate_memory(new_block_size, image_memory_type_index);
+    if (!ptr->memory)
+    {
+        vkDestroyImage(vkdev->vkdevice(), ptr->image, 0);
+        delete ptr;
+        return 0;
+    }
     ptr->bind_offset = 0;
     ptr->bind_capacity = aligned_size;
 
@@ -1788,6 +1824,12 @@ VkBufferMemory* VkStagingAllocator::fastMalloc(size_t size)
     }
 
     ptr->memory = allocate_memory(memoryRequirements.size, buffer_memory_type_index);
+    if (!ptr->memory)
+    {
+        vkDestroyBuffer(vkdev->vkdevice(), ptr->buffer, 0);
+        delete ptr;
+        return 0;
+    }
 
     // ignore memoryRequirements.alignment as we always bind at zero offset
     vkBindBufferMemory(vkdev->vkdevice(), ptr->buffer, ptr->memory, 0);
@@ -1897,6 +1939,12 @@ VkBufferMemory* VkWeightStagingAllocator::fastMalloc(size_t size)
     }
 
     ptr->memory = allocate_memory(memoryRequirements.size, buffer_memory_type_index);
+    if (!ptr->memory)
+    {
+        vkDestroyBuffer(vkdev->vkdevice(), ptr->buffer, 0);
+        delete ptr;
+        return 0;
+    }
 
     // ignore memoryRequirements.alignment as we always bind at zero offset
     vkBindBufferMemory(vkdev->vkdevice(), ptr->buffer, ptr->memory, 0);