From a4dac01978b11de00b6ebce5fa226f916bd7cfd5 Mon Sep 17 00:00:00 2001 From: Cesare Mercurio Date: Sun, 17 May 2026 23:44:52 -0700 Subject: [PATCH] Vulkan: wrap external buffers as regions --- src/runtime/HalideRuntimeVulkan.h | 15 ++ src/runtime/internal/memory_resources.h | 8 + src/runtime/runtime_api.cpp | 4 + src/runtime/vulkan.cpp | 186 ++++++++++++++++-- src/runtime/vulkan_memory.h | 118 +++++++---- src/runtime/vulkan_resources.h | 3 +- .../generator/gpu_object_lifetime_aottest.cpp | 37 ++++ 7 files changed, 320 insertions(+), 51 deletions(-) diff --git a/src/runtime/HalideRuntimeVulkan.h b/src/runtime/HalideRuntimeVulkan.h index e150b7c6d00b..31ed6795e374 100644 --- a/src/runtime/HalideRuntimeVulkan.h +++ b/src/runtime/HalideRuntimeVulkan.h @@ -105,6 +105,21 @@ extern int halide_vulkan_release_context(void *user_context, VkDevice device, VkQueue queue, VkDebugUtilsMessengerEXT messenger); + +/** Wrap an externally-owned VkBuffer with a byte offset. */ +extern int halide_vulkan_wrap_vk_buffer(void *user_context, + struct halide_buffer_t *buf, + uint64_t vk_buffer, + uint64_t offset); + +/** Detach a wrapped VkBuffer without destroying it. */ +extern int halide_vulkan_detach_vk_buffer(void *user_context, halide_buffer_t *buf); + +/** Return the underlying VkBuffer for a halide_buffer_t. */ +extern uintptr_t halide_vulkan_get_vk_buffer(void *user_context, halide_buffer_t *buf); + +/** Return the total byte offset for a wrapped or cropped Vulkan buffer. */ +extern uint64_t halide_vulkan_get_vk_crop_offset(void *user_context, halide_buffer_t *buf); // -- // Override the default allocation callbacks (default uses Vulkan runtime implementation) diff --git a/src/runtime/internal/memory_resources.h b/src/runtime/internal/memory_resources.h index ea5f5420263b..84022a5b64e9 100644 --- a/src/runtime/internal/memory_resources.h +++ b/src/runtime/internal/memory_resources.h @@ -78,6 +78,12 @@ struct RegionIndexing { int32_t offset = 0; //< indexing offset from start of region (used to adjust indices in compute shader to avoid alignment constraints for arbitrary crops) }; +enum class MemoryOwnership { + Owned, //< region metadata managed by allocator; handle points to native resource + CropAlias, //< alias metadata for a cropped view; owner points to root region + Wrapped //< metadata supplied by caller for an external native resource +}; + // Client-facing struct for exchanging memory region allocation requests struct MemoryRegion { void *handle = nullptr; //< client data storing native handle (managed by alloc_block_region/free_block_region) or a pointer to region owning allocation @@ -85,6 +91,8 @@ struct MemoryRegion { RegionIndexing indexing; //< indexing adjustments for controlling access bool dedicated = false; //< flag indicating whether allocation is one dedicated resource (or split/shared into other resources) bool is_owner = true; //< flag indicating whether allocation is owned by this region, in which case handle is a native handle. Otherwise handle points to owning region of allocation. + MemoryOwnership ownership = MemoryOwnership::Owned; + MemoryRegion *owner = nullptr; MemoryProperties properties; //< properties for the allocated region }; diff --git a/src/runtime/runtime_api.cpp b/src/runtime/runtime_api.cpp index 734af982bf91..941697b8ebe4 100644 --- a/src/runtime/runtime_api.cpp +++ b/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_detach_vk_buffer, (void *)&halide_vulkan_device_interface, + (void *)&halide_vulkan_get_vk_buffer, + (void *)&halide_vulkan_get_vk_crop_offset, (void *)&halide_vulkan_initialize_kernels, (void *)&halide_vulkan_release_context, (void *)&halide_vulkan_run, + (void *)&halide_vulkan_wrap_vk_buffer, (void *)&halide_webgpu_device_interface, (void *)&halide_webgpu_initialize_kernels, (void *)&halide_webgpu_finalize_kernels, diff --git a/src/runtime/vulkan.cpp b/src/runtime/vulkan.cpp index f44b60b90e94..60961aac2060 100644 --- a/src/runtime/vulkan.cpp +++ b/src/runtime/vulkan.cpp @@ -10,6 +10,86 @@ #include "vulkan_resources.h" using namespace Halide::Runtime::Internal::Vulkan; +using Halide::Runtime::Internal::MemoryOwnership; + +// -------------------------------------------------------------------------- + +namespace Halide { +namespace Runtime { +namespace Internal { +namespace Vulkan { + +ALWAYS_INLINE uint64_t vk_external_buffer_offset_bytes(void *user_context, VulkanMemoryAllocator *allocator, MemoryRegion *region) { + halide_debug_assert(user_context, region != nullptr); + MemoryRegion *owner = allocator->owner_of(user_context, region); + return (owner != nullptr && owner->ownership == MemoryOwnership::Wrapped) ? owner->allocation.offset : 0; +} + +ALWAYS_INLINE uint64_t vk_total_buffer_offset_bytes(void *user_context, VulkanMemoryAllocator *allocator, MemoryRegion *region, halide_type_t type) { + return vk_external_buffer_offset_bytes(user_context, allocator, region) + (region->indexing.offset * type.bytes()); +} + +ALWAYS_INLINE void vk_destroy_wrapped_region(void *user_context, + VulkanMemoryAllocator *allocator, + MemoryRegion *region) { + if (region == nullptr) { + return; + } + + MemoryRegion *owner = allocator->owner_of(user_context, region); + if (owner == nullptr) { + return; + } + + if (region != owner && region->ownership == MemoryOwnership::CropAlias) { + vk_host_free(user_context, region, allocator->callbacks()); + return; + } + + if (owner->ownership != MemoryOwnership::Wrapped) { + return; + } + + vk_host_free(user_context, owner->handle, allocator->callbacks()); + vk_host_free(user_context, owner, allocator->callbacks()); +} + +ALWAYS_INLINE MemoryRegion *vk_create_wrapped_buffer_region(void *user_context, + halide_buffer_t *buf, + VkBuffer vk_buffer, + uint64_t offset, + const VkAllocationCallbacks *callbacks) { + constexpr VkSystemAllocationScope alloc_scope = VK_SYSTEM_ALLOCATION_SCOPE_OBJECT; + VkBuffer *native_handle = reinterpret_cast( + vk_host_malloc(user_context, sizeof(VkBuffer), 0, alloc_scope, callbacks)); + if (native_handle == nullptr) { + error(user_context) << "Vulkan: Failed to allocate wrapped buffer handle metadata.\n"; + return nullptr; + } + + MemoryRegion *region = reinterpret_cast( + vk_host_malloc(user_context, sizeof(MemoryRegion), 0, alloc_scope, callbacks)); + if (region == nullptr) { + vk_host_free(user_context, native_handle, callbacks); + error(user_context) << "Vulkan: Failed to allocate wrapped buffer region metadata.\n"; + return nullptr; + } + + *native_handle = vk_buffer; + memset(region, 0, sizeof(MemoryRegion)); + region->handle = native_handle; + region->allocation.offset = offset; + region->allocation.size = buf->size_in_bytes(); + region->is_owner = true; + region->ownership = MemoryOwnership::Wrapped; + region->owner = nullptr; + return region; +} + +} // namespace Vulkan +} // namespace Internal +} // namespace Runtime +} // namespace Halide // -------------------------------------------------------------------------- @@ -110,12 +190,20 @@ WEAK int halide_vulkan_device_free(void *user_context, halide_buffer_t *halide_b // get the allocated region for the device MemoryRegion *device_region = reinterpret_cast(halide_buffer->device); +#ifdef DEBUG_RUNTIME + const uint64_t device_region_size = device_region->allocation.size; +#endif MemoryRegion *memory_region = ctx.allocator->owner_of(user_context, device_region); if (ctx.allocator && memory_region && memory_region->handle) { - if (halide_can_reuse_device_allocations(user_context)) { - ctx.allocator->release(user_context, memory_region); + if (memory_region->ownership == MemoryOwnership::Wrapped) { + debug(user_context) << "Vulkan: Releasing wrapped external buffer metadata only.\n"; + vk_destroy_wrapped_region(user_context, ctx.allocator, device_region); } else { - ctx.allocator->reclaim(user_context, memory_region); + if (halide_can_reuse_device_allocations(user_context)) { + ctx.allocator->release(user_context, memory_region); + } else { + ctx.allocator->reclaim(user_context, memory_region); + } } } halide_buffer->device = 0; @@ -126,7 +214,7 @@ WEAK int halide_vulkan_device_free(void *user_context, halide_buffer_t *halide_b debug(user_context) << "Vulkan: Released memory for device region (" << "user_context: " << user_context << ", " << "buffer: " << halide_buffer << ", " - << "size_in_bytes: " << (uint64_t)device_region->allocation.size << ")\n"; + << "size_in_bytes: " << device_region_size << ")\n"; uint64_t t_after = halide_current_time_ns(user_context); debug(user_context) << " Time: " << (t_after - t_before) / 1.0e6 << " ms\n"; @@ -272,15 +360,24 @@ WEAK int halide_vulkan_device_malloc(void *user_context, halide_buffer_t *buf) { size_t size = buf->size_in_bytes(); if (buf->device) { MemoryRegion *device_region = (MemoryRegion *)(buf->device); - if (device_region->allocation.size >= size) { + MemoryRegion *memory_region = ctx.allocator->owner_of(user_context, device_region); + if (memory_region != nullptr && memory_region->allocation.size >= size) { debug(user_context) << "Vulkan: Requested allocation for existing device memory ... using existing buffer!\n"; return halide_error_code_success; } else { + if (memory_region == nullptr) { + error(user_context) << "Vulkan: Failed to retrieve memory region for existing device buffer!\n"; + return halide_error_code_internal_error; + } + if (memory_region->ownership == MemoryOwnership::Wrapped) { + error(user_context) << "Vulkan: Wrapped external buffer is too small for requested allocation!\n"; + return halide_error_code_device_malloc_failed; + } debug(user_context) << "Vulkan: Requested allocation of different size ... reallocating buffer!\n"; if (halide_can_reuse_device_allocations(user_context)) { - ctx.allocator->release(user_context, device_region); + ctx.allocator->release(user_context, memory_region); } else { - ctx.allocator->reclaim(user_context, device_region); + ctx.allocator->reclaim(user_context, memory_region); } buf->device = 0; } @@ -487,7 +584,7 @@ WEAK int halide_vulkan_copy_to_device(void *user_context, halide_buffer_t *halid bool to_host = false; uint64_t src_offset = copy_helper.src_begin; - uint64_t dst_offset = copy_helper.dst_begin + (device_region->indexing.offset * halide_buffer->type.bytes()); + uint64_t dst_offset = copy_helper.dst_begin + vk_total_buffer_offset_bytes(user_context, ctx.allocator, device_region, halide_buffer->type); copy_helper.src = (uint64_t)(staging_buffer); copy_helper.dst = (uint64_t)(device_buffer); @@ -656,7 +753,7 @@ WEAK int halide_vulkan_copy_to_host(void *user_context, halide_buffer_t *halide_ bool from_host = false; bool to_host = true; uint64_t copy_dst = copy_helper.dst; - uint64_t src_offset = copy_helper.src_begin + (device_region->indexing.offset * halide_buffer->type.bytes()); + uint64_t src_offset = copy_helper.src_begin + vk_total_buffer_offset_bytes(user_context, ctx.allocator, device_region, halide_buffer->type); uint64_t dst_offset = copy_helper.dst_begin; copy_helper.src = (uint64_t)(device_buffer); @@ -937,8 +1034,8 @@ WEAK int halide_vulkan_buffer_copy(void *user_context, struct halide_buffer_t *s // define the src and dst config uint64_t copy_dst = copy_helper.dst; - uint64_t src_offset = copy_helper.src_begin + (src_buffer_region->indexing.offset * src->type.bytes()); - uint64_t dst_offset = copy_helper.dst_begin + (dst_buffer_region->indexing.offset * dst->type.bytes()); + uint64_t src_offset = copy_helper.src_begin + vk_total_buffer_offset_bytes(user_context, ctx.allocator, src_buffer_region, src->type); + uint64_t dst_offset = copy_helper.dst_begin + vk_total_buffer_offset_bytes(user_context, ctx.allocator, dst_buffer_region, dst->type); copy_helper.src = (uint64_t)(src_device_buffer); copy_helper.dst = (uint64_t)(dst_device_buffer); @@ -1345,13 +1442,33 @@ WEAK int halide_vulkan_device_and_host_free(void *user_context, struct halide_bu return halide_default_device_and_host_free(user_context, buf, &vulkan_device_interface); } -WEAK int halide_vulkan_wrap_vk_buffer(void *user_context, struct halide_buffer_t *buf, uint64_t vk_buffer) { +WEAK int halide_vulkan_wrap_vk_buffer(void *user_context, + struct halide_buffer_t *buf, + uint64_t vk_buffer, + uint64_t offset) { halide_debug_assert(user_context, buf->device == 0); if (buf->device != 0) { error(user_context) << "Vulkan: Unable to wrap buffer ... invalid device pointer!\n"; return halide_error_code_device_wrap_native_failed; } - buf->device = vk_buffer; + if (vk_buffer == 0) { + error(user_context) << "Vulkan: Unable to wrap buffer ... invalid VkBuffer handle!\n"; + return halide_error_code_device_wrap_native_failed; + } + + VulkanContext ctx(user_context); + if (ctx.error != halide_error_code_success) { + error(user_context) << "Vulkan: Failed to acquire context!\n"; + return ctx.error; + } + + MemoryRegion *region = vk_create_wrapped_buffer_region( + user_context, buf, reinterpret_cast(vk_buffer), offset, ctx.allocator->callbacks()); + if (region == nullptr) { + return halide_error_code_out_of_memory; + } + + buf->device = reinterpret_cast(region); buf->device_interface = &vulkan_device_interface; buf->device_interface->impl->use_module(); return halide_error_code_success; @@ -1365,6 +1482,19 @@ WEAK int halide_vulkan_detach_vk_buffer(void *user_context, halide_buffer_t *buf error(user_context) << "Vulkan: Unable to detach buffer ... invalid device interface!\n"; return halide_error_code_incompatible_device_interface; } + MemoryRegion *device_region = reinterpret_cast(buf->device); + VulkanContext ctx(user_context); + if (ctx.error != halide_error_code_success) { + error(user_context) << "Vulkan: Failed to acquire context!\n"; + return ctx.error; + } + + MemoryRegion *owner = ctx.allocator->owner_of(user_context, device_region); + if (owner == nullptr || owner->ownership != MemoryOwnership::Wrapped) { + error(user_context) << "Vulkan: Unable to detach buffer ... buffer is not externally wrapped!\n"; + return halide_error_code_device_detach_native_failed; + } + vk_destroy_wrapped_region(user_context, ctx.allocator, device_region); buf->device = 0; buf->device_interface->impl->release_module(); buf->device_interface = nullptr; @@ -1376,7 +1506,29 @@ WEAK uintptr_t halide_vulkan_get_vk_buffer(void *user_context, halide_buffer_t * return 0; } halide_debug_assert(user_context, buf->device_interface == &vulkan_device_interface); - return (uintptr_t)buf->device; + MemoryRegion *device_region = reinterpret_cast(buf->device); + VulkanContext ctx(user_context); + if (ctx.error != halide_error_code_success) { + return 0; + } + MemoryRegion *owner = ctx.allocator->owner_of(user_context, device_region); + if (owner == nullptr || owner->handle == nullptr) { + return 0; + } + return (uintptr_t)(*reinterpret_cast(owner->handle)); +} + +WEAK uint64_t halide_vulkan_get_vk_crop_offset(void *user_context, halide_buffer_t *buf) { + if (buf->device == 0) { + return 0; + } + halide_debug_assert(user_context, buf->device_interface == &vulkan_device_interface); + MemoryRegion *device_region = reinterpret_cast(buf->device); + VulkanContext ctx(user_context); + if (ctx.error != halide_error_code_success) { + return 0; + } + return vk_total_buffer_offset_bytes(user_context, ctx.allocator, device_region, buf->type); } WEAK const struct halide_device_interface_t *halide_vulkan_device_interface() { @@ -1445,6 +1597,10 @@ namespace Vulkan { // -------------------------------------------------------------------------- +int vk_wrap_native_vk_buffer(void *user_context, struct halide_buffer_t *buf, uint64_t vk_buffer) { + return halide_vulkan_wrap_vk_buffer(user_context, buf, vk_buffer, 0); +} + WEAK halide_device_interface_impl_t vulkan_device_interface_impl = { halide_use_jit_module, halide_release_jit_module, @@ -1460,7 +1616,7 @@ WEAK halide_device_interface_impl_t vulkan_device_interface_impl = { halide_vulkan_device_crop, halide_vulkan_device_slice, halide_vulkan_device_release_crop, - halide_vulkan_wrap_vk_buffer, + vk_wrap_native_vk_buffer, halide_vulkan_detach_vk_buffer, }; diff --git a/src/runtime/vulkan_memory.h b/src/runtime/vulkan_memory.h index b21a2e476d86..d5f740bdeacd 100644 --- a/src/runtime/vulkan_memory.h +++ b/src/runtime/vulkan_memory.h @@ -268,6 +268,10 @@ void *VulkanMemoryAllocator::map(void *user_context, MemoryRegion *region) { } MemoryRegion *owner = owner_of(user_context, region); + if (owner != nullptr && owner->ownership == MemoryOwnership::Wrapped) { + error(user_context) << "VulkanMemoryAllocator: Unable to map wrapped external region through allocator!\n"; + return nullptr; + } RegionAllocator *region_allocator = RegionAllocator::find_allocator(user_context, owner); if (region_allocator == nullptr) { error(user_context) << "VulkanMemoryAllocator: Unable to map region! Invalid region allocator handle!\n"; @@ -322,6 +326,10 @@ int VulkanMemoryAllocator::unmap(void *user_context, MemoryRegion *region) { } MemoryRegion *owner = owner_of(user_context, region); + if (owner != nullptr && owner->ownership == MemoryOwnership::Wrapped) { + error(user_context) << "VulkanMemoryAllocator: Unable to unmap wrapped external region through allocator!\n"; + return halide_error_code_internal_error; + } RegionAllocator *region_allocator = RegionAllocator::find_allocator(user_context, owner); if (region_allocator == nullptr) { error(user_context) << "VulkanMemoryAllocator: Unable to unmap region! Invalid region allocator handle!\n"; @@ -361,29 +369,42 @@ MemoryRegion *VulkanMemoryAllocator::create_crop(void *user_context, MemoryRegio } MemoryRegion *owner = owner_of(user_context, region); - RegionAllocator *region_allocator = RegionAllocator::find_allocator(user_context, owner); - if (region_allocator == nullptr) { - error(user_context) << "VulkanMemoryAllocator: Unable to unmap region! Invalid region allocator handle!\n"; - return nullptr; // NOTE: caller must handle nullptr + if (owner == nullptr) { + error(user_context) << "VulkanMemoryAllocator: Unable to crop region! Invalid owner region!\n"; + return nullptr; } - // increment usage count - int error_code = region_allocator->retain(this, owner); - if (error_code != halide_error_code_success) { - error(user_context) << "VulkanMemoryAllocator: Unable to crop region! Failed to retain memory region!\n"; - return nullptr; // NOTE: caller must handle nullptr + const bool allocator_owned = owner->ownership == MemoryOwnership::Owned; + if (allocator_owned) { + RegionAllocator *region_allocator = RegionAllocator::find_allocator(user_context, owner); + if (region_allocator == nullptr) { + error(user_context) << "VulkanMemoryAllocator: Unable to crop region! Invalid region allocator handle!\n"; + return nullptr; // NOTE: caller must handle nullptr + } + + // increment usage count + int error_code = region_allocator->retain(this, owner); + if (error_code != halide_error_code_success) { + error(user_context) << "VulkanMemoryAllocator: Unable to crop region! Failed to retain memory region!\n"; + return nullptr; // NOTE: caller must handle nullptr + } } // create a new region to return, and copy all the other region's properties - const BlockAllocator::MemoryAllocators &allocators = block_allocator->current_allocators(); - if (allocators.system.allocate == nullptr) { - error(user_context) << "VulkanMemoryAllocator: Unable to create crop! Missing system allocator interface!\n"; - return nullptr; + MemoryRegion *memory_region = nullptr; + if (allocator_owned) { + const BlockAllocator::MemoryAllocators &allocators = block_allocator->current_allocators(); + if (allocators.system.allocate == nullptr) { + error(user_context) << "VulkanMemoryAllocator: Unable to create crop! Missing system allocator interface!\n"; + return nullptr; + } + memory_region = reinterpret_cast( + allocators.system.allocate(user_context, sizeof(MemoryRegion))); + } else { + memory_region = reinterpret_cast( + vk_host_malloc(user_context, sizeof(MemoryRegion), 0, VK_SYSTEM_ALLOCATION_SCOPE_OBJECT, callbacks())); } - MemoryRegion *memory_region = reinterpret_cast( - allocators.system.allocate(user_context, sizeof(MemoryRegion))); - if (memory_region == nullptr) { error(user_context) << "VulkanMemoryAllocator: Failed to allocate memory region! Out of memory!\n"; return nullptr; @@ -393,6 +414,8 @@ MemoryRegion *VulkanMemoryAllocator::create_crop(void *user_context, MemoryRegio // point the handle to the owner of the allocated region, and update the indexing offset memory_region->is_owner = false; memory_region->handle = (void *)owner; + memory_region->ownership = MemoryOwnership::CropAlias; + memory_region->owner = owner; memory_region->indexing.offset += owner->indexing.offset + indexing.offset; return memory_region; } @@ -404,36 +427,46 @@ int VulkanMemoryAllocator::destroy_crop(void *user_context, MemoryRegion *region } MemoryRegion *owner = owner_of(user_context, region); - RegionAllocator *region_allocator = RegionAllocator::find_allocator(user_context, owner); - if (region_allocator == nullptr) { - error(user_context) << "VulkanMemoryAllocator: Unable to destroy crop region! Invalid region allocator handle!\n"; + if (owner == nullptr) { + error(user_context) << "VulkanMemoryAllocator: Unable to destroy crop region! Invalid owner region!\n"; return halide_error_code_internal_error; } - // decrement usage count - int error_code = region_allocator->release(this, owner); - if (error_code != halide_error_code_success) { - error(user_context) << "VulkanBlockAllocator: Unable to destroy crop region! Region allocator failed to release memory region!\n"; - return error_code; + if (owner->ownership == MemoryOwnership::Owned) { + RegionAllocator *region_allocator = RegionAllocator::find_allocator(user_context, owner); + if (region_allocator == nullptr) { + error(user_context) << "VulkanMemoryAllocator: Unable to destroy crop region! Invalid region allocator handle!\n"; + return halide_error_code_internal_error; + } + + // decrement usage count + int error_code = region_allocator->release(this, owner); + if (error_code != halide_error_code_success) { + error(user_context) << "VulkanBlockAllocator: Unable to destroy crop region! Region allocator failed to release memory region!\n"; + return error_code; + } } - // discard the copied region struct - const BlockAllocator::MemoryAllocators &allocators = block_allocator->current_allocators(); - if (allocators.system.deallocate == nullptr) { - error(user_context) << "VulkanBlockAllocator: Unable to destroy crop region! Missing system allocator interface!\n"; - return halide_error_code_internal_error; + if (owner->ownership == MemoryOwnership::Owned) { + // discard the copied region struct + const BlockAllocator::MemoryAllocators &allocators = block_allocator->current_allocators(); + if (allocators.system.deallocate == nullptr) { + error(user_context) << "VulkanBlockAllocator: Unable to destroy crop region! Missing system allocator interface!\n"; + return halide_error_code_internal_error; + } + allocators.system.deallocate(user_context, region); + } else { + vk_host_free(user_context, region, callbacks()); } - allocators.system.deallocate(user_context, region); return halide_error_code_success; } MemoryRegion *VulkanMemoryAllocator::owner_of(void *user_context, MemoryRegion *region) { - if (region->is_owner) { - return region; - } else { - // If this is a cropped region, use the handle to retrieve the owner of the allocation - return reinterpret_cast(region->handle); + MemoryRegion *current = region; + while (current != nullptr && current->ownership == MemoryOwnership::CropAlias) { + current = current->owner != nullptr ? current->owner : reinterpret_cast(current->handle); } + return current; } int VulkanMemoryAllocator::release(void *user_context, MemoryRegion *region) { @@ -452,6 +485,11 @@ int VulkanMemoryAllocator::release(void *user_context, MemoryRegion *region) { error(user_context) << "VulkanMemoryAllocator: Unable to release region! Invalid block allocator!\n"; return halide_error_code_generic_error; } + MemoryRegion *owner = owner_of(user_context, region); + if (owner != nullptr && owner->ownership == MemoryOwnership::Wrapped) { + error(user_context) << "VulkanMemoryAllocator: Unable to release wrapped external region through allocator!\n"; + return halide_error_code_internal_error; + } return block_allocator->release(this, region); } @@ -471,6 +509,11 @@ int VulkanMemoryAllocator::reclaim(void *user_context, MemoryRegion *region) { error(user_context) << "VulkanMemoryAllocator: Unable to reclaim region! Invalid block allocator!\n"; return halide_error_code_generic_error; } + MemoryRegion *owner = owner_of(user_context, region); + if (owner != nullptr && owner->ownership == MemoryOwnership::Wrapped) { + error(user_context) << "VulkanMemoryAllocator: Unable to reclaim wrapped external region through allocator!\n"; + return halide_error_code_internal_error; + } return block_allocator->reclaim(this, region); } @@ -490,6 +533,11 @@ int VulkanMemoryAllocator::retain(void *user_context, MemoryRegion *region) { error(user_context) << "VulkanMemoryAllocator: Unable to retain region! Invalid block allocator!\n"; return halide_error_code_generic_error; } + MemoryRegion *owner = owner_of(user_context, region); + if (owner != nullptr && owner->ownership == MemoryOwnership::Wrapped) { + error(user_context) << "VulkanMemoryAllocator: Unable to retain wrapped external region through allocator!\n"; + return halide_error_code_internal_error; + } return block_allocator->retain(this, region); } diff --git a/src/runtime/vulkan_resources.h b/src/runtime/vulkan_resources.h index d2ef2ee5ba6f..faacfbb39034 100644 --- a/src/runtime/vulkan_resources.h +++ b/src/runtime/vulkan_resources.h @@ -603,7 +603,8 @@ int vk_update_descriptor_set(void *user_context, return halide_error_code_internal_error; } - VkDeviceSize range_offset = 0; + VkDeviceSize range_offset = + owner->ownership == MemoryOwnership::Wrapped ? owner->allocation.offset : 0; VkDeviceSize range_size = device_region->allocation.size; VkDescriptorBufferInfo device_buffer_info = { *device_buffer, // the buffer diff --git a/test/generator/gpu_object_lifetime_aottest.cpp b/test/generator/gpu_object_lifetime_aottest.cpp index d0d0a7d544f3..81ede9b4748c 100644 --- a/test/generator/gpu_object_lifetime_aottest.cpp +++ b/test/generator/gpu_object_lifetime_aottest.cpp @@ -137,6 +137,11 @@ int main(int argc, char **argv) { native_handle = halide_metal_get_buffer(nullptr, output.raw_buffer()); can_rewrap = true; } +#elif defined(TEST_VULKAN) + if (output.raw_buffer()->device_interface == halide_vulkan_device_interface()) { + native_handle = halide_vulkan_get_vk_buffer(nullptr, output.raw_buffer()); + can_rewrap = true; + } #endif if (can_rewrap) { @@ -155,6 +160,38 @@ int main(int argc, char **argv) { if (i == 1) { wrap_test.device_detach_native(); } + +#if defined(TEST_VULKAN) + if (output.raw_buffer()->device_interface == halide_vulkan_device_interface()) { + Buffer offset_wrap(79); + int result = halide_vulkan_wrap_vk_buffer(nullptr, offset_wrap.raw_buffer(), native_handle, sizeof(int)); + if (result != 0) { + printf("Error! halide_vulkan_wrap_vk_buffer() returned: %d\n", result); + return 1; + } + uint64_t offset = halide_vulkan_get_vk_crop_offset(nullptr, offset_wrap.raw_buffer()); + if (offset != sizeof(int)) { + printf("Error! Vulkan wrapped buffer offset: %llu != %zu\n", + (unsigned long long)offset, sizeof(int)); + return 1; + } + + offset_wrap.set_device_dirty(); + offset_wrap.copy_to_host(); + for (int x = 0; x < offset_wrap.width(); x++) { + if (offset_wrap(x) != output(x + 1)) { + printf("Error! (Vulkan offset wrap test %d): %d != %d\n", i, offset_wrap(x), output(x + 1)); + return 1; + } + } + + result = halide_vulkan_detach_vk_buffer(nullptr, offset_wrap.raw_buffer()); + if (result != 0) { + printf("Error! halide_vulkan_detach_vk_buffer() returned: %d\n", result); + return 1; + } + } +#endif } }