Vulkan: export allocator for custom contexts

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)

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

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<halide_vulkan_memory_allocator *>(allocator));
+            if (result != halide_error_code_success) {
+                error = static_cast<halide_error_code_t>(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(

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);