Vulkan: expose allocator lifecycle

xFile3160 · xFile3160 · commit f2afbd8e1ee8 · 2026-04-28T12:25:09.000-07:00
When an application provides its own Vulkan context, Halide still needs its runtime allocator for shader modules, staging buffers, and internal Vulkan resources.

Add public acquire/release helpers for the opaque Vulkan allocator. The acquire helper creates or validates an allocator for the supplied instance/device pair. The release helper tears down Halide-owned allocator state without destroying the caller's Vulkan instance, device, or queue.

Key the Vulkan compilation cache by allocator instead of VkDevice so cache cleanup follows the allocator lifetime.
diff --git a/src/runtime/HalideRuntimeVulkan.h b/src/runtime/HalideRuntimeVulkan.h
@@ -105,6 +105,41 @@ extern int halide_vulkan_release_context(void *user_context,
                                          VkDevice device,
                                          VkQueue queue,
                                          VkDebugUtilsMessengerEXT messenger);
+
+/** Ensure a Halide Vulkan memory allocator exists for an externally-managed
+ * Vulkan context. Intended for embedders that override
+ * halide_vulkan_acquire_context()/halide_vulkan_release_context().
+ *
+ * The embedder should store the returned allocator with the same object that
+ * owns the external context, return it from later acquire-context calls for
+ * that context, and release it when that external context is torn down.
+ *
+ * This call refreshes Halide's Vulkan dispatch tables for the supplied
+ * instance/device. If `*allocator` is null, a new allocator bound to
+ * `device`/`physical_device` is created and stored back. If `*allocator` is
+ * non-null, it must already be bound to the supplied device.
+ */
+extern int halide_vulkan_acquire_memory_allocator(void *user_context,
+                                                  struct halide_vulkan_memory_allocator **allocator,
+                                                  VkInstance instance,
+                                                  VkDevice device,
+                                                  VkPhysicalDevice physical_device);
+
+/** Destroy a Halide Vulkan memory allocator created for an externally-managed
+ * Vulkan context after the embedder has ensured no in-flight Halide work is
+ * using it. This only releases Halide-owned allocator and shader-module state;
+ * it does not destroy the Vulkan instance, device, queue, or any
+ * embedder-owned debug messenger.
+ *
+ * This call refreshes Halide's Vulkan dispatch tables for the supplied
+ * instance/device. The supplied device and physical_device must match the
+ * allocator's context.
+ */
+extern int halide_vulkan_release_memory_allocator(void *user_context,
+                                                  struct halide_vulkan_memory_allocator *allocator,
+                                                  VkInstance instance,
+                                                  VkDevice device,
+                                                  VkPhysicalDevice physical_device);
 // --
 
 // Override the default allocation callbacks (default uses Vulkan runtime implementation)
diff --git a/src/runtime/gpu_context_common.h b/src/runtime/gpu_context_common.h
@@ -127,7 +127,7 @@ class GPUCompilationCache {
         }
 
         for (int i = 0; i < (1 << log2_compilations_size); i++) {
-            if (compilations[i].kernel_id > kInvalidId &&
+            if (compilations[i].kernel_id > kDeletedId &&
                 (all || (compilations[i].context == context)) &&
                 compilations[i].use_count == 0) {
                 debug(user_context) << "Releasing cached compilation: " << compilations[i].module_state
diff --git a/src/runtime/runtime_api.cpp b/src/runtime/runtime_api.cpp
@@ -213,8 +213,10 @@ extern "C" __attribute__((used)) void *halide_runtime_api_functions[] = {
     (void *)&halide_d3d12compute_release_context,
     (void *)&halide_d3d12compute_run,
     (void *)&halide_vulkan_acquire_context,
+    (void *)&halide_vulkan_acquire_memory_allocator,
     (void *)&halide_vulkan_device_interface,
     (void *)&halide_vulkan_initialize_kernels,
+    (void *)&halide_vulkan_release_memory_allocator,
     (void *)&halide_vulkan_release_context,
     (void *)&halide_vulkan_run,
     (void *)&halide_webgpu_device_interface,
diff --git a/src/runtime/vulkan.cpp b/src/runtime/vulkan.cpp
@@ -13,6 +13,42 @@ using namespace Halide::Runtime::Internal::Vulkan;
 
 // --------------------------------------------------------------------------
 
+namespace Halide {
+namespace Runtime {
+namespace Internal {
+namespace Vulkan {
+
+ALWAYS_INLINE int vk_load_external_context_functions(void *user_context, VkInstance instance, VkDevice device) {
+    if (vkGetInstanceProcAddr == nullptr) {
+        vk_load_vulkan_loader_functions(user_context);
+        if (vkGetInstanceProcAddr == nullptr) {
+            error(user_context) << "Vulkan: Failed to resolve loader functions for external context!\n";
+            return halide_error_code_symbol_not_found;
+        }
+    }
+
+    vk_load_vulkan_instance_functions(user_context, instance);
+    if (vkGetPhysicalDeviceProperties == nullptr || vkGetDeviceProcAddr == nullptr) {
+        error(user_context) << "Vulkan: Failed to resolve instance functions for external context!\n";
+        return halide_error_code_symbol_not_found;
+    }
+
+    vk_load_vulkan_device_functions(user_context, device);
+    if (vkCreateBuffer == nullptr || vkAllocateMemory == nullptr) {
+        error(user_context) << "Vulkan: Failed to resolve device functions for external context!\n";
+        return halide_error_code_symbol_not_found;
+    }
+
+    return halide_error_code_success;
+}
+
+}  // namespace Vulkan
+}  // namespace Internal
+}  // namespace Runtime
+}  // namespace Halide
+
+// --------------------------------------------------------------------------
+
 extern "C" {
 
 // --------------------------------------------------------------------------
@@ -79,6 +115,82 @@ WEAK int halide_vulkan_release_context(void *user_context, VkInstance instance,
     return halide_error_code_success;
 }
 
+WEAK int halide_vulkan_acquire_memory_allocator(void *user_context,
+                                                halide_vulkan_memory_allocator **allocator,
+                                                VkInstance instance,
+                                                VkDevice device,
+                                                VkPhysicalDevice physical_device) {
+    if (allocator == nullptr) {
+        error(user_context) << "Vulkan: allocator output pointer is null!\n";
+        return halide_error_code_buffer_argument_is_null;
+    }
+    if (instance == VK_NULL_HANDLE || device == VK_NULL_HANDLE || physical_device == VK_NULL_HANDLE) {
+        error(user_context) << "Vulkan: invalid external context handles for allocator acquisition!\n";
+        return halide_error_code_device_interface_no_device;
+    }
+
+    int error_code = vk_load_external_context_functions(user_context, instance, device);
+    if (error_code != halide_error_code_success) {
+        return error_code;
+    }
+
+    VulkanMemoryAllocator *runtime_allocator =
+        reinterpret_cast<VulkanMemoryAllocator *>(*allocator);
+    if (runtime_allocator != nullptr) {
+        if (runtime_allocator->current_device() != device ||
+            runtime_allocator->current_physical_device() != physical_device) {
+            error(user_context) << "Vulkan: external allocator does not match supplied device handles!\n";
+            return halide_error_code_internal_error;
+        }
+        return halide_error_code_success;
+    }
+
+    const VkAllocationCallbacks *alloc_callbacks =
+        halide_vulkan_get_allocation_callbacks(user_context);
+    runtime_allocator =
+        vk_create_memory_allocator(user_context, device, physical_device, alloc_callbacks);
+    if (runtime_allocator == nullptr) {
+        error(user_context) << "Vulkan: Failed to create memory allocator for external context!\n";
+        return halide_error_code_out_of_memory;
+    }
+
+    *allocator = reinterpret_cast<halide_vulkan_memory_allocator *>(runtime_allocator);
+    return halide_error_code_success;
+}
+
+WEAK int halide_vulkan_release_memory_allocator(void *user_context,
+                                                halide_vulkan_memory_allocator *allocator,
+                                                VkInstance instance,
+                                                VkDevice device,
+                                                VkPhysicalDevice physical_device) {
+    VulkanMemoryAllocator *runtime_allocator =
+        reinterpret_cast<VulkanMemoryAllocator *>(allocator);
+    if (runtime_allocator == nullptr) {
+        return halide_error_code_success;
+    }
+    if (instance == VK_NULL_HANDLE || device == VK_NULL_HANDLE || physical_device == VK_NULL_HANDLE) {
+        error(user_context) << "Vulkan: invalid external context handles for allocator release!\n";
+        return halide_error_code_device_interface_no_device;
+    }
+    if (runtime_allocator->current_device() != device ||
+        runtime_allocator->current_physical_device() != physical_device) {
+        error(user_context) << "Vulkan: external allocator does not match supplied device handles during release!\n";
+        return halide_error_code_internal_error;
+    }
+
+    int error_code = vk_load_external_context_functions(user_context, instance, device);
+    if (error_code != halide_error_code_success) {
+        return error_code;
+    }
+    if (vkDestroyShaderModule == nullptr || vkFreeMemory == nullptr) {
+        error(user_context) << "Vulkan: Failed to resolve device functions for external allocator release!\n";
+        return halide_error_code_symbol_not_found;
+    }
+
+    vk_destroy_shader_modules(user_context, runtime_allocator);
+    return vk_destroy_memory_allocator(user_context, runtime_allocator);
+}
+
 WEAK bool halide_vulkan_is_initialized() {
     halide_mutex_lock(&thread_lock);
     bool is_initialized = (cached_instance != nullptr) && (cached_device != nullptr);
@@ -159,7 +271,7 @@ WEAK int halide_vulkan_initialize_kernels(void *user_context, void **state_ptr,
     debug(user_context) << "halide_vulkan_initialize_kernels got compilation_cache mutex.\n";
 
     VulkanCompilationCacheEntry *cache_entry = nullptr;
-    if (!compilation_cache.kernel_state_setup(user_context, state_ptr, ctx.device, cache_entry,
+    if (!compilation_cache.kernel_state_setup(user_context, state_ptr, ctx.allocator, cache_entry,
                                               Halide::Runtime::Internal::Vulkan::vk_compile_kernel_module,
                                               user_context, ctx.allocator, src, size)) {
         error(user_context) << "Vulkan: Failed to setup compilation cache!\n";
@@ -185,7 +297,7 @@ WEAK void halide_vulkan_finalize_kernels(void *user_context, void *state_ptr) {
 
     VulkanContext ctx(user_context);
     if (ctx.error == halide_error_code_success) {
-        compilation_cache.release_hold(user_context, ctx.device, state_ptr);
+        compilation_cache.release_hold(user_context, ctx.allocator, state_ptr);
     }
 
 #ifdef DEBUG_RUNTIME
@@ -1151,7 +1263,7 @@ WEAK int halide_vulkan_run(void *user_context,
 
     // 1. Get the shader module cache entry
     VulkanCompilationCacheEntry *cache_entry = nullptr;
-    bool found = compilation_cache.lookup(ctx.device, state_ptr, cache_entry);
+    bool found = compilation_cache.lookup(ctx.allocator, state_ptr, cache_entry);
     if (!found || (cache_entry == nullptr)) {
         error(user_context) << "Vulkan: Failed to locate shader module! Unable to proceed!\n";
         return halide_error_code_internal_error;
diff --git a/src/runtime/vulkan_resources.h b/src/runtime/vulkan_resources.h
@@ -72,7 +72,7 @@ struct VulkanCompilationCacheEntry {
     uint32_t module_count = 0;
 };
 
-WEAK Halide::Internal::GPUCompilationCache<VkDevice, VulkanCompilationCacheEntry *> compilation_cache;
+WEAK Halide::Internal::GPUCompilationCache<VulkanMemoryAllocator *, VulkanCompilationCacheEntry *> compilation_cache;
 
 // --------------------------------------------------------------------------
 
@@ -1665,29 +1665,20 @@ void vk_destroy_compiled_shader_module(VulkanCompiledShaderModule *shader_module
         return;
     }
 
-    if (shader_module->descriptor_set_layouts) {
-        for (uint32_t n = 0; n < shader_module->shader_count; n++) {
-            debug(user_context) << "  destroying descriptor set layout [" << n << "] " << shader_module->descriptor_set_layouts[n] << "\n";
-            vk_destroy_descriptor_set_layout(user_context, allocator, shader_module->descriptor_set_layouts[n]);
-            shader_module->descriptor_set_layouts[n] = VK_NULL_HANDLE;
-        }
-        debug(user_context) << "  destroying descriptor set layout " << (void *)shader_module->descriptor_set_layouts << "\n";
-        vk_host_free(user_context, shader_module->descriptor_set_layouts, allocator->callbacks());
-        shader_module->descriptor_set_layouts = nullptr;
-    }
-    if (shader_module->pipeline_layout) {
-        debug(user_context) << "  destroying pipeline layout " << (void *)shader_module->pipeline_layout << "\n";
-        vk_destroy_pipeline_layout(user_context, allocator, shader_module->pipeline_layout);
-        shader_module->pipeline_layout = VK_NULL_HANDLE;
-    }
-
     if (shader_module->shader_bindings) {
 #ifdef DEBUG_RUNTIME
         debug(user_context)
             << "  destroying shader bindings ("
             << "shader_module: " << shader_module << ", "
             << "shader_bindings: " << shader_module->shader_bindings << ")\n";
 #endif
+        for (uint32_t n = 0; n < shader_module->shader_count; n++) {
+            if (shader_module->shader_bindings[n].compute_pipeline) {
+                debug(user_context) << "  destroying shader binding compute pipeline [" << n << "]\n";
+                vk_destroy_compute_pipeline(user_context, allocator, shader_module->shader_bindings[n].compute_pipeline);
+                shader_module->shader_bindings[n].compute_pipeline = VK_NULL_HANDLE;
+            }
+        }
         for (uint32_t n = 0; n < shader_module->shader_count; n++) {
             debug(user_context) << "  destroying shader binding [" << n << "] ";
             if (shader_module->shader_bindings[n].entry_point_name) {
@@ -1717,15 +1708,25 @@ void vk_destroy_compiled_shader_module(VulkanCompiledShaderModule *shader_module
                 vk_host_free(user_context, shader_module->shader_bindings[n].shared_memory_allocations, allocator->callbacks());
                 shader_module->shader_bindings[n].shared_memory_allocations = nullptr;
             }
-            if (shader_module->shader_bindings[n].compute_pipeline) {
-                debug(user_context) << "  destroying shader binding compute pipeline [" << n << "]\n";
-                vk_destroy_compute_pipeline(user_context, allocator, shader_module->shader_bindings[n].compute_pipeline);
-                shader_module->shader_bindings[n].compute_pipeline = VK_NULL_HANDLE;
-            }
         }
         vk_host_free(user_context, shader_module->shader_bindings, allocator->callbacks());
         shader_module->shader_bindings = nullptr;
     }
+    if (shader_module->pipeline_layout) {
+        debug(user_context) << "  destroying pipeline layout " << (void *)shader_module->pipeline_layout << "\n";
+        vk_destroy_pipeline_layout(user_context, allocator, shader_module->pipeline_layout);
+        shader_module->pipeline_layout = VK_NULL_HANDLE;
+    }
+    if (shader_module->descriptor_set_layouts) {
+        for (uint32_t n = 0; n < shader_module->shader_count; n++) {
+            debug(user_context) << "  destroying descriptor set layout [" << n << "] " << shader_module->descriptor_set_layouts[n] << "\n";
+            vk_destroy_descriptor_set_layout(user_context, allocator, shader_module->descriptor_set_layouts[n]);
+            shader_module->descriptor_set_layouts[n] = VK_NULL_HANDLE;
+        }
+        debug(user_context) << "  destroying descriptor set layout " << (void *)shader_module->descriptor_set_layouts << "\n";
+        vk_host_free(user_context, shader_module->descriptor_set_layouts, allocator->callbacks());
+        shader_module->descriptor_set_layouts = nullptr;
+    }
     if (shader_module->shader_module) {
         debug(user_context) << "   destroying shader module " << (void *)shader_module->shader_module << "\n";
         vkDestroyShaderModule(allocator->current_device(), shader_module->shader_module, allocator->callbacks());
@@ -1778,7 +1779,7 @@ int vk_destroy_shader_modules(void *user_context, VulkanMemoryAllocator *allocat
     uint64_t t_before = halide_current_time_ns(user_context);
 #endif
     if (allocator != nullptr) {
-        compilation_cache.delete_context(user_context, allocator->current_device(), vk_destroy_compilation_cache_entry);
+        compilation_cache.delete_context(user_context, allocator, vk_destroy_compilation_cache_entry);
     }
 
 #ifdef DEBUG_RUNTIME

Original file line number	Diff line number	Diff line change
`@@ -127,7 +127,7 @@ class GPUCompilationCache {`
`127`	`127`	`}`
`128`	`128`
`129`	`129`	`for (int i = 0; i < (1 << log2_compilations_size); i++) {`
`130`		`- if (compilations[i].kernel_id > kInvalidId &&`
	`130`	`+ if (compilations[i].kernel_id > kDeletedId &&`
`131`	`131`	`(all \|\| (compilations[i].context == context)) &&`
`132`	`132`	`compilations[i].use_count == 0) {`
`133`	`133`	`debug(user_context) << "Releasing cached compilation: " << compilations[i].module_state`