From fa20ad1071a9d6ccc8a09cede5c44bed78ca9b46 Mon Sep 17 00:00:00 2001 From: cesaremercurio Date: Mon, 17 Nov 2025 13:44:33 -0800 Subject: [PATCH] Vulkan: export allocator for custom contexts Add documented hooks that let an override of halide_vulkan_acquire_context return its own VkInstance/device yet still bootstrap Halide's allocator. The weak runtime now exposes halide_vulkan_export_memory_allocator/halide_vulkan_memory_allocator_release plus a vk_release_memory_allocator helper that only tears down Halide-owned caches. VulkanContext will lazily allocate+export Halide's allocator when an override returns nullptr. Example integration and validation lives at https://github.com/xFile3160/Halide/commit/9041848759bc78e3cc712771b24c3cee35dc63ad. --- src/runtime/HalideRuntimeVulkan.h | 17 ++++++++++ src/runtime/vulkan.cpp | 26 +++++++++++++++ src/runtime/vulkan_context.h | 53 ++++++++++++++++++++++++++++++- src/runtime/vulkan_internal.h | 6 ++++ 4 files changed, 101 insertions(+), 1 deletion(-) diff --git a/src/runtime/HalideRuntimeVulkan.h b/src/runtime/HalideRuntimeVulkan.h index e150b7c6d00b..e3bba9e83ad6 100644 --- a/src/runtime/HalideRuntimeVulkan.h +++ b/src/runtime/HalideRuntimeVulkan.h @@ -105,6 +105,23 @@ extern int halide_vulkan_release_context(void *user_context, VkDevice device, VkQueue queue, VkDebugUtilsMessengerEXT messenger); + +// - halide_vulkan_export_memory_allocator +// exports the internally allocated memory allocator in case the user wants to just set +// up their own context but use Halide's memory allocator. Must have overridden halide_vulkan_acquire_context +// and halide_vulkan_release_context. Must override also halide_vulkan_export_memory_allocator and guard access +// with the same locking used by the custom acquire/release implementations. This allows the allocator to be +// saved for future halide_vulkan_acquire_context calls that Halide will automatically issue to retrieve +// the custom context. +extern int halide_vulkan_export_memory_allocator(void *user_context, + struct halide_vulkan_memory_allocator *allocator); +// - halide_vulkan_memory_allocator_release +// releases the internally allocated memory allocator, important for proper memory cleanup. Must have overridden halide_vulkan_acquire_context +// and halide_vulkan_release_context, and must coordinate with the same locking as the custom implementations. +extern int halide_vulkan_memory_allocator_release(void *user_context, + struct halide_vulkan_memory_allocator *allocator, + VkInstance instance, + VkDebugUtilsMessengerEXT messenger); // -- // Override the default allocation callbacks (default uses Vulkan runtime implementation) diff --git a/src/runtime/vulkan.cpp b/src/runtime/vulkan.cpp index 36fc3047ddfb..b4c87f1bee25 100644 --- a/src/runtime/vulkan.cpp +++ b/src/runtime/vulkan.cpp @@ -86,6 +86,17 @@ WEAK bool halide_vulkan_is_initialized() { return is_initialized; } +WEAK int halide_vulkan_export_memory_allocator(void *user_context, halide_vulkan_memory_allocator *allocator) { + halide_mutex_lock(&thread_lock); + halide_error_code_t status = halide_error_code_success; + if (allocator == nullptr) { + error(user_context) << "Vulkan: Memory allocator is null!\n"; + status = halide_error_code_buffer_argument_is_null; + } + halide_mutex_unlock(&thread_lock); + return status; +} + WEAK int halide_vulkan_device_free(void *user_context, halide_buffer_t *halide_buffer) { debug(user_context) << "halide_vulkan_device_free (user_context: " << user_context @@ -258,6 +269,21 @@ WEAK int halide_vulkan_device_release(void *user_context) { return destroy_status; } +WEAK int halide_vulkan_memory_allocator_release(void *user_context, + struct halide_vulkan_memory_allocator *allocator, + VkInstance instance, + VkDebugUtilsMessengerEXT messenger) { + debug(user_context) << "halide_vulkan_memory_allocator_release (user_context: " << user_context << ")\n"; + // Destroy the context if we created it + if (allocator == nullptr) { + error(user_context) << "Vulkan: Memory allocator is null!\n"; + return halide_error_code_buffer_argument_is_null; + } + + return vk_release_memory_allocator(user_context, (VulkanMemoryAllocator *)allocator, + instance, messenger); +} + WEAK int halide_vulkan_device_malloc(void *user_context, halide_buffer_t *buf) { debug(user_context) << "halide_vulkan_device_malloc (user_context: " << user_context diff --git a/src/runtime/vulkan_context.h b/src/runtime/vulkan_context.h index 9e82b9fb5b64..59db5ae5dbff 100644 --- a/src/runtime/vulkan_context.h +++ b/src/runtime/vulkan_context.h @@ -58,6 +58,43 @@ class VulkanContext { error = halide_error_code_device_interface_no_device; halide_error_no_device_interface(user_context); } + // If user overrode halide_vulkan_acquire_context and returned nullptr for allocator, + // create Halide's allocator for the provided device. User must override `halide_vulkan_export_memory_allocator` + // and make sure to propagate it back at the next call of `halide_vulkan_acquire_context` as he overrides it. + if (allocator == nullptr && + instance != VK_NULL_HANDLE && + device != VK_NULL_HANDLE && + physical_device != VK_NULL_HANDLE) { +#ifdef DEBUG_RUNTIME + // Initialize clock for debug timing - normally done in halide_vulkan_acquire_context + halide_start_clock(user_context); +#endif + // make sure halide vulkan is loaded BEFORE creating allocator + debug(user_context) << "VulkanContext: Loading Vulkan function pointers for context override...\n"; + + vk_load_vulkan_loader_functions(user_context); + if (vkGetInstanceProcAddr == nullptr) { + debug(user_context) << "VulkanContext: Failed to load vkGetInstanceProcAddr from loader!\n"; + } else { + debug(user_context) << "VulkanContext: vkGetInstanceProcAddr loaded successfully: " << (void *)vkGetInstanceProcAddr << "\n"; + vk_load_vulkan_instance_functions(user_context, instance); + vk_load_vulkan_device_functions(user_context, device); + } + + allocator = vk_create_memory_allocator(user_context, device, physical_device, + halide_vulkan_get_allocation_callbacks(user_context)); + if (allocator == nullptr) { + error = halide_error_code_out_of_memory; + debug(user_context) << "Vulkan: Failed to create memory allocator for device!\n"; + return; + } + int result = halide_vulkan_export_memory_allocator(user_context, reinterpret_cast(allocator)); + if (result != halide_error_code_success) { + error = static_cast(result); + debug(user_context) << "Vulkan: Failed to export memory allocator for device!\n"; + return; + } + } halide_debug_assert(user_context, allocator != nullptr); halide_debug_assert(user_context, instance != VK_NULL_HANDLE); halide_debug_assert(user_context, device != VK_NULL_HANDLE); @@ -546,8 +583,8 @@ int vk_destroy_context(void *user_context, VulkanMemoryAllocator *allocator, if (allocator != nullptr) { vk_destroy_shader_modules(user_context, allocator); - vk_destroy_memory_allocator(user_context, allocator); vk_destroy_debug_utils_messenger(user_context, instance, allocator, messenger); + vk_destroy_memory_allocator(user_context, allocator); } const VkAllocationCallbacks *alloc_callbacks = halide_vulkan_get_allocation_callbacks(user_context); @@ -560,6 +597,20 @@ int vk_destroy_context(void *user_context, VulkanMemoryAllocator *allocator, return halide_error_code_success; } +// Clean up only Halide's internal resources for external context (leaves device/instance alone) +int vk_release_memory_allocator(void *user_context, VulkanMemoryAllocator *allocator, + VkInstance instance, VkDebugUtilsMessengerEXT messenger) { + debug(user_context) << "vk_release_memory_allocator (user_context: " << user_context << ")\n"; + // Clean up only Halide's internal resources, not the device/instance we don't own + if (allocator != nullptr) { + vk_destroy_shader_modules(user_context, allocator); + vk_destroy_debug_utils_messenger(user_context, instance, allocator, messenger); + vk_destroy_memory_allocator(user_context, allocator); + } + + return halide_error_code_success; +} + // -------------------------------------------------------------------------- VKAPI_ATTR VkBool32 VKAPI_CALL vk_debug_utils_messenger_callback( diff --git a/src/runtime/vulkan_internal.h b/src/runtime/vulkan_internal.h index aeef545385cc..af56fb136cfc 100644 --- a/src/runtime/vulkan_internal.h +++ b/src/runtime/vulkan_internal.h @@ -66,6 +66,12 @@ int vk_destroy_context( VkPhysicalDevice physical_device, VkQueue queue); +int vk_release_memory_allocator( + void *user_context, + VulkanMemoryAllocator *allocator, + VkInstance instance, + VkDebugUtilsMessengerEXT messenger); + int vk_find_compute_capability(void *user_context, int *major, int *minor); int vk_create_instance(void *user_context, const StringTable &requested_layers, VkInstance *instance, const VkAllocationCallbacks *alloc_callbacks);