Vulkan: Use guardedalloc for driver allocations.

Vulkan has a pluggable memory allocation feature, which allows internal
driver allocations to be done by the client application provided
allocator. Vulkan uses this for more client application allocations
done inside the driver, but can also do it for more internal oriented
allocations.

VK_ALLOCATION_CALLBACKS initializes allocation callbacks for host allocations.
The macro creates a local static variable with the name vk_allocation_callbacks
that can be passed to vulkan API functions that expect
const VkAllocationCallbacks *pAllocator.

When WITH_VULKAN_GUARDEDALLOC=Off the memory allocation implemented
in the vulkan device driver is used for both internal and application
oriented memory operations.

For now this would help during the development of Vulkan backend to
detect hidden memory leaks that are hidden inside the driver part
of the stack. In a later stage we need to measure the overhead and
if this should become the default behavior.

Pull Request #104434

CMakeLists.txt

@@ -625,8 +625,10 @@ mark_as_advanced(
 
 # Vulkan
 option(WITH_VULKAN_BACKEND "Enable Vulkan as graphics backend (only for development)" OFF)
+option(WITH_VULKAN_GUARDEDALLOC "Use guardedalloc for host allocations done inside Vulkan (development option)" OFF)
 mark_as_advanced(
   WITH_VULKAN_BACKEND
+  WITH_VULKAN_GUARDEDALLOC
 )
 
 # Metal

source/blender/gpu/CMakeLists.txt

@@ -195,6 +195,7 @@ set(VULKAN_SRC
   vulkan/vk_fence.cc
   vulkan/vk_framebuffer.cc
   vulkan/vk_index_buffer.cc
+  vulkan/vk_memory.cc
   vulkan/vk_pixel_buffer.cc
   vulkan/vk_query.cc
   vulkan/vk_shader.cc
@@ -211,6 +212,7 @@ set(VULKAN_SRC
   vulkan/vk_fence.hh
   vulkan/vk_framebuffer.hh
   vulkan/vk_index_buffer.hh
+  vulkan/vk_memory.hh
   vulkan/vk_pixel_buffer.hh
   vulkan/vk_query.hh
   vulkan/vk_shader.hh
@@ -303,6 +305,10 @@ if(WITH_VULKAN_BACKEND)
   add_definitions(-DWITH_VULKAN_BACKEND)
 endif()
 
+if(WITH_VULKAN_GUARDEDALLOC)
+  add_definitions(-DWITH_VULKAN_GUARDEDALLOC)
+endif()
+
 set(MSL_SRC
   shaders/metal/mtl_shader_defines.msl
   shaders/metal/mtl_shader_common.msl

source/blender/gpu/vulkan/vk_memory.cc

@@ -0,0 +1,45 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later
+ * Copyright 2023 Blender Foundation. All rights reserved. */
+
+/** \file
+ * \ingroup gpu
+ */
+
+#include "vk_memory.hh"
+
+#include "MEM_guardedalloc.h"
+
+namespace blender::gpu {
+
+#ifdef WITH_VULKAN_GUARDEDALLOC
+
+void *vk_memory_allocation(void *user_data,
+                           size_t size,
+                           size_t alignment,
+                           VkSystemAllocationScope /*scope*/)
+{
+  const char *name = static_cast<const char *>(const_cast<const void *>(user_data));
+  if (alignment) {
+    return MEM_mallocN_aligned(size, alignment, name);
+  }
+  return MEM_mallocN(size, name);
+}
+
+void *vk_memory_reallocation(void *user_data,
+                             void *original,
+                             size_t size,
+                             size_t /*alignment*/,
+                             VkSystemAllocationScope /*scope*/)
+{
+  const char *name = static_cast<const char *>(const_cast<const void *>(user_data));
+  return MEM_reallocN_id(original, size, name);
+}
+
+void vk_memory_free(void * /*user_data*/, void *memory)
+{
+  MEM_freeN(memory);
+}
+
+#endif
+
+}  // namespace blender::gpu

source/blender/gpu/vulkan/vk_memory.hh

@@ -0,0 +1,64 @@
+/* SPDX-License-Identifier: GPL-2.0-or-later
+ * Copyright 2023 Blender Foundation. All rights reserved. */
+
+/** \file
+ * \ingroup gpu
+ */
+
+#pragma once
+
+#ifdef __APPLE__
+#  include <MoltenVK/vk_mvk_moltenvk.h>
+#else
+#  include <vulkan/vulkan.h>
+#endif
+
+namespace blender::gpu {
+
+/**
+ * `VK_ALLOCATION_CALLBACKS` initializes allocation callbacks for host allocations.
+ * The macro creates a local static variable with the name `vk_allocation_callbacks`
+ * that can be passed to vulkan API functions that expect
+ * `const VkAllocationCallbacks *pAllocator`.
+ *
+ * When Blender is compiled with `WITH_VULKAN_GUARDEDALLOC` this will use
+ * `MEM_guardedalloc` for host allocations that the driver does on behalf
+ * of blender. More internal allocations are still being allocated via the
+ * implemention inside the vulkan device driver.
+ *
+ * When `WITH_VULKAN_GUARDEDALLOC=Off` the memory allocation implemented
+ * in the vulkan device driver is used for both internal and application
+ * focussed memory operations.
+ */
+
+#ifdef WITH_VULKAN_GUARDEDALLOC
+void *vk_memory_allocation(void *user_data,
+                           size_t size,
+                           size_t alignment,
+                           VkSystemAllocationScope scope);
+void *vk_memory_reallocation(
+    void *user_data, void *original, size_t size, size_t alignment, VkSystemAllocationScope scope);
+void vk_memory_free(void *user_data, void *memory);
+
+constexpr VkAllocationCallbacks vk_allocation_callbacks_init(const char *name)
+{
+  VkAllocationCallbacks callbacks = {};
+  callbacks.pUserData = const_cast<char *>(name);
+  callbacks.pfnAllocation = vk_memory_allocation;
+  callbacks.pfnReallocation = vk_memory_reallocation;
+  callbacks.pfnFree = vk_memory_free;
+  callbacks.pfnInternalAllocation = nullptr;
+  callbacks.pfnInternalFree = nullptr;
+  return callbacks;
+}
+
+#  define VK_ALLOCATION_CALLBACKS \
+    static constexpr const VkAllocationCallbacks vk_allocation_callbacks_ = \
+        vk_allocation_callbacks_init(__func__); \
+    static constexpr const VkAllocationCallbacks *vk_allocation_callbacks = &vk_allocation_callbacks_;
+#else
+#  define VK_ALLOCATION_CALLBACKS \
+    static constexpr const VkAllocationCallbacks *vk_allocation_callbacks = nullptr;
+#endif
+
+}  // namespace blender::gpu

source/blender/gpu/vulkan/vk_shader.cc

@@ -8,6 +8,7 @@
 #include "vk_shader.hh"
 
 #include "vk_backend.hh"
+#include "vk_memory.hh"
 #include "vk_shader_log.hh"
 
 #include "BLI_string_utils.h"
@@ -559,6 +560,8 @@ Vector<uint32_t> VKShader::compile_glsl_to_spirv(Span<const char *> sources,
 
 void VKShader::build_shader_module(Span<uint32_t> spirv_module, VkShaderModule *r_shader_module)
 {
+  VK_ALLOCATION_CALLBACKS;
+
   VkShaderModuleCreateInfo create_info = {};
   create_info.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
   create_info.codeSize = spirv_module.size() * sizeof(uint32_t);
@@ -567,7 +570,7 @@ void VKShader::build_shader_module(Span<uint32_t> spirv_module, VkShaderModule *
   VKContext &context = *static_cast<VKContext *>(VKContext::get());
 
   VkResult result = vkCreateShaderModule(
-      context.device_get(), &create_info, nullptr, r_shader_module);
+      context.device_get(), &create_info, vk_allocation_callbacks, r_shader_module);
   if (result != VK_SUCCESS) {
     compilation_failed_ = true;
     *r_shader_module = VK_NULL_HANDLE;
@@ -581,21 +584,23 @@ VKShader::VKShader(const char *name) : Shader(name)
 
 VKShader::~VKShader()
 {
+  VK_ALLOCATION_CALLBACKS
+
   VkDevice device = context_->device_get();
   if (vertex_module_ != VK_NULL_HANDLE) {
-    vkDestroyShaderModule(device, vertex_module_, nullptr);
+    vkDestroyShaderModule(device, vertex_module_, vk_allocation_callbacks);
     vertex_module_ = VK_NULL_HANDLE;
   }
   if (geometry_module_ != VK_NULL_HANDLE) {
-    vkDestroyShaderModule(device, geometry_module_, nullptr);
+    vkDestroyShaderModule(device, geometry_module_, vk_allocation_callbacks);
     geometry_module_ = VK_NULL_HANDLE;
   }
   if (fragment_module_ != VK_NULL_HANDLE) {
-    vkDestroyShaderModule(device, fragment_module_, nullptr);
+    vkDestroyShaderModule(device, fragment_module_, vk_allocation_callbacks);
     fragment_module_ = VK_NULL_HANDLE;
   }
   if (compute_module_ != VK_NULL_HANDLE) {
-    vkDestroyShaderModule(device, compute_module_, nullptr);
+    vkDestroyShaderModule(device, compute_module_, vk_allocation_callbacks);
     compute_module_ = VK_NULL_HANDLE;
   }
 }