Vulkan: export allocator for custom contexts

src/runtime/HalideRuntimeVulkan.h

@@ -105,6 +105,64 @@ extern int halide_vulkan_release_context(void *user_context,
                                          VkDevice device,
                                          VkQueue queue,
                                          VkDebugUtilsMessengerEXT messenger);
+
+typedef int (*halide_vulkan_acquire_context_t)(void *user_context,
+                                               struct halide_vulkan_memory_allocator **allocator,
+                                               VkInstance *instance,
+                                               VkDevice *device,
+                                               VkPhysicalDevice *physical_device,
+                                               VkQueue *queue,
+                                               uint32_t *queue_family_index,
+                                               VkDebugUtilsMessengerEXT *messenger,
+                                               bool create);
+typedef int (*halide_vulkan_release_context_t)(void *user_context,
+                                               VkInstance instance,
+                                               VkDevice device,
+                                               VkQueue queue,
+                                               VkDebugUtilsMessengerEXT messenger);
+
+/** Override the Vulkan context acquisition callback. Returns the previous
+ * handler. If unset, Halide uses its built-in Vulkan context management.
+ */
+extern halide_vulkan_acquire_context_t halide_set_vulkan_acquire_context(halide_vulkan_acquire_context_t handler);
+
+/** Override the Vulkan context release callback. Returns the previous handler. */
+extern halide_vulkan_release_context_t halide_set_vulkan_release_context(halide_vulkan_release_context_t handler);
+
+/** 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)

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

src/runtime/runtime_api.cpp

@@ -213,10 +213,14 @@ 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_set_vulkan_acquire_context,
+    (void *)&halide_set_vulkan_release_context,
     (void *)&halide_webgpu_device_interface,
     (void *)&halide_webgpu_initialize_kernels,
     (void *)&halide_webgpu_finalize_kernels,

src/runtime/vulkan.cpp

@@ -13,9 +13,34 @@ using namespace Halide::Runtime::Internal::Vulkan;
 
 // --------------------------------------------------------------------------
 
-extern "C" {
+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;
+}
 
 // The default implementation of halide_acquire_vulkan_context uses
 // the global pointers above, and serializes access with a spin lock.
@@ -29,15 +54,15 @@ extern "C" {
 //   call to halide_release_vulkan_context. halide_acquire_vulkan_context
 //   should block while a previous call (if any) has not yet been
 //   released via halide_release_vulkan_context.
-WEAK int halide_vulkan_acquire_context(void *user_context,
-                                       halide_vulkan_memory_allocator **allocator,
-                                       VkInstance *instance,
-                                       VkDevice *device,
-                                       VkPhysicalDevice *physical_device,
-                                       VkQueue *queue,
-                                       uint32_t *queue_family_index,
-                                       VkDebugUtilsMessengerEXT *messenger,
-                                       bool create) {
+WEAK int default_vulkan_acquire_context(void *user_context,
+                                        halide_vulkan_memory_allocator **allocator,
+                                        VkInstance *instance,
+                                        VkDevice *device,
+                                        VkPhysicalDevice *physical_device,
+                                        VkQueue *queue,
+                                        uint32_t *queue_family_index,
+                                        VkDebugUtilsMessengerEXT *messenger,
+                                        bool create) {
 #ifdef DEBUG_RUNTIME
     halide_start_clock(user_context);
 #endif
@@ -74,11 +99,133 @@ WEAK int halide_vulkan_acquire_context(void *user_context,
     return halide_error_code_success;
 }
 
-WEAK int halide_vulkan_release_context(void *user_context, VkInstance instance, VkDevice device, VkQueue queue, VkDebugUtilsMessengerEXT messenger) {
+WEAK int default_vulkan_release_context(void *user_context, VkInstance instance, VkDevice device, VkQueue queue, VkDebugUtilsMessengerEXT messenger) {
     halide_mutex_unlock(&thread_lock);
     return halide_error_code_success;
 }
 
+WEAK halide_vulkan_acquire_context_t vulkan_acquire_context_handler =
+    default_vulkan_acquire_context;
+WEAK halide_vulkan_release_context_t vulkan_release_context_handler =
+    default_vulkan_release_context;
+
+}  // namespace Vulkan
+}  // namespace Internal
+}  // namespace Runtime
+}  // namespace Halide
+
+// --------------------------------------------------------------------------
+
+extern "C" {
+
+// --------------------------------------------------------------------------
+
+WEAK int halide_vulkan_acquire_context(void *user_context,
+                                       halide_vulkan_memory_allocator **allocator,
+                                       VkInstance *instance,
+                                       VkDevice *device,
+                                       VkPhysicalDevice *physical_device,
+                                       VkQueue *queue,
+                                       uint32_t *queue_family_index,
+                                       VkDebugUtilsMessengerEXT *messenger,
+                                       bool create) {
+    return vulkan_acquire_context_handler(user_context, allocator, instance, device,
+                                          physical_device, queue, queue_family_index,
+                                          messenger, create);
+}
+
+WEAK int halide_vulkan_release_context(void *user_context, VkInstance instance, VkDevice device, VkQueue queue, VkDebugUtilsMessengerEXT messenger) {
+    return vulkan_release_context_handler(user_context, instance, device, queue, messenger);
+}
+
+WEAK halide_vulkan_acquire_context_t halide_set_vulkan_acquire_context(halide_vulkan_acquire_context_t handler) {
+    halide_vulkan_acquire_context_t result = vulkan_acquire_context_handler;
+    vulkan_acquire_context_handler = handler ? handler : default_vulkan_acquire_context;
+    return result;
+}
+
+WEAK halide_vulkan_release_context_t halide_set_vulkan_release_context(halide_vulkan_release_context_t handler) {
+    halide_vulkan_release_context_t result = vulkan_release_context_handler;
+    vulkan_release_context_handler = handler ? handler : default_vulkan_release_context;
+    return result;
+}
+
+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 +306,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 +332,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 +1298,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;