From 8efc40afac580477e60059a37f0306a9612dcaac Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Wed, 2 Oct 2024 14:54:43 -0700 Subject: [PATCH] [ET-VK] Fix tensor views with tensors that use deferred memory allocation ## Context Fix one (hopefully) final kink with tensor views. ### The issue When creating a tensor view of a tensor that uses deferred memory allocation, at the time the view is created, it will copy the handle of the original tensor's `VkImage`. However, since the `VkImage` is not yet bounded to memory, it will have have a `VkImageView` associated with it, and the tensor view will copy the `VkImageView` handle and `VK_NULL_HANDLE`. Later, the original tensor is expected to be bound to an allocation, at which point it will create an image view and bind the `VkImage` to memory. However, this does not update its views, which will still have a null `VkImageView`. If the tensor view is also bound to the allocation, then the same `VkImage` will be bound to the same memory multiple times which is against the Vulkan spec. ### The fix Allow binding of `VulkanImage` and `VulkanBuffer` copies to only call the binding function if it is not a copy; assume that the original tensor is responsible for binding to memory. In `ComputeGraph`, when adding a tensor view, add the created `Value` as a user of any `SharedObject` of which the original tensor is also a user. Differential Revision: [D63790718](https://our.internmc.facebook.com/intern/diff/D63790718/) [ghstack-poisoned] --- backends/vulkan/runtime/graph/ComputeGraph.cpp | 10 ++++++++++ .../runtime/graph/containers/SharedObject.cpp | 4 ++++ .../vulkan/runtime/graph/containers/SharedObject.h | 1 + backends/vulkan/runtime/vk_api/memory/Buffer.h | 4 +++- backends/vulkan/runtime/vk_api/memory/Image.cpp | 13 +++++++++---- backends/vulkan/runtime/vk_api/memory/Image.h | 12 +++++++++++- backends/vulkan/test/vulkan_compute_api_test.cpp | 10 +++++----- 7 files changed, 43 insertions(+), 11 deletions(-) diff --git a/backends/vulkan/runtime/graph/ComputeGraph.cpp b/backends/vulkan/runtime/graph/ComputeGraph.cpp index 967c892b0b4..011d62c0ea4 100644 --- a/backends/vulkan/runtime/graph/ComputeGraph.cpp +++ b/backends/vulkan/runtime/graph/ComputeGraph.cpp @@ -300,6 +300,11 @@ ValueRef ComputeGraph::add_tensor_view(const ValueRef vref) { const vTensorPtr t = get_tensor(vref); ValueRef idx(static_cast(values_.size())); values_.emplace_back(api::vTensor(*t)); + for (SharedObject& sobj : shared_objects_) { + if (sobj.has_user(vref)) { + sobj.add_user(this, idx); + } + } return idx; } @@ -311,6 +316,11 @@ ValueRef ComputeGraph::add_tensor_view( const vTensorPtr t = get_tensor(vref); ValueRef idx(static_cast(values_.size())); values_.emplace_back(api::vTensor(*t, sizes, strides, offset_numel)); + for (SharedObject& sobj : shared_objects_) { + if (sobj.has_user(vref)) { + sobj.add_user(this, idx); + } + } return idx; } diff --git a/backends/vulkan/runtime/graph/containers/SharedObject.cpp b/backends/vulkan/runtime/graph/containers/SharedObject.cpp index f2474da6673..10ddd6f2ca3 100644 --- a/backends/vulkan/runtime/graph/containers/SharedObject.cpp +++ b/backends/vulkan/runtime/graph/containers/SharedObject.cpp @@ -12,6 +12,10 @@ namespace vkcompute { +bool SharedObject::has_user(const ValueRef idx) const { + return std::find(users.begin(), users.end(), idx) != users.end(); +} + void SharedObject::add_user(ComputeGraph* const graph, const ValueRef idx) { vTensorPtr t = graph->get_tensor(idx); diff --git a/backends/vulkan/runtime/graph/containers/SharedObject.h b/backends/vulkan/runtime/graph/containers/SharedObject.h index bd77f6f39ba..f9b16e6c202 100644 --- a/backends/vulkan/runtime/graph/containers/SharedObject.h +++ b/backends/vulkan/runtime/graph/containers/SharedObject.h @@ -31,6 +31,7 @@ struct SharedObject { std::vector users; vkapi::Allocation allocation; + bool has_user(const ValueRef idx) const; void add_user(ComputeGraph* const graph, const ValueRef idx); void allocate(ComputeGraph* const graph); void bind_users(ComputeGraph* const graph); diff --git a/backends/vulkan/runtime/vk_api/memory/Buffer.h b/backends/vulkan/runtime/vk_api/memory/Buffer.h index 865ca7866f6..14722511f4f 100644 --- a/backends/vulkan/runtime/vk_api/memory/Buffer.h +++ b/backends/vulkan/runtime/vk_api/memory/Buffer.h @@ -161,7 +161,9 @@ class VulkanBuffer final { inline void bind_allocation(const Allocation& memory) { VK_CHECK_COND(!memory_, "Cannot bind an already bound allocation!"); - VK_CHECK(vmaBindBufferMemory(allocator_, memory.allocation, handle_)); + if (!is_copy_) { + VK_CHECK(vmaBindBufferMemory(allocator_, memory.allocation, handle_)); + } memory_.allocation = memory.allocation; } diff --git a/backends/vulkan/runtime/vk_api/memory/Image.cpp b/backends/vulkan/runtime/vk_api/memory/Image.cpp index 10fce5e07f2..604cab995f1 100644 --- a/backends/vulkan/runtime/vk_api/memory/Image.cpp +++ b/backends/vulkan/runtime/vk_api/memory/Image.cpp @@ -98,6 +98,7 @@ VulkanImage::VulkanImage() allocator_(VK_NULL_HANDLE), memory_{}, owns_memory_(false), + owns_view_(false), is_copy_(false), handles_{ VK_NULL_HANDLE, @@ -121,6 +122,7 @@ VulkanImage::VulkanImage( allocator_(vma_allocator), memory_{}, owns_memory_{allocate_memory}, + owns_view_{false}, is_copy_(false), handles_{ VK_NULL_HANDLE, @@ -168,6 +170,7 @@ VulkanImage::VulkanImage( &(memory_.allocation), nullptr)); // Only create the image view if the image has been bound to memory + owns_view_ = true; create_image_view(); } else { VK_CHECK(vkCreateImage( @@ -182,6 +185,7 @@ VulkanImage::VulkanImage(const VulkanImage& other) noexcept allocator_(other.allocator_), memory_(other.memory_), owns_memory_{false}, + owns_view_{false}, is_copy_(true), handles_(other.handles_), layout_(other.layout_) {} @@ -193,6 +197,7 @@ VulkanImage::VulkanImage(VulkanImage&& other) noexcept allocator_(other.allocator_), memory_(std::move(other.memory_)), owns_memory_(other.owns_memory_), + owns_view_(other.owns_view_), is_copy_(other.is_copy_), handles_(other.handles_), layout_(other.layout_) { @@ -225,6 +230,10 @@ VulkanImage& VulkanImage::operator=(VulkanImage&& other) noexcept { } VulkanImage::~VulkanImage() { + if (owns_view_ && handles_.image_view != VK_NULL_HANDLE) { + vkDestroyImageView(this->device(), handles_.image_view, nullptr); + } + // Do not destroy any resources if this class instance is a copy of another // class instance, since this means that this class instance does not have // ownership of the underlying resource. @@ -232,10 +241,6 @@ VulkanImage::~VulkanImage() { return; } - if (handles_.image_view != VK_NULL_HANDLE) { - vkDestroyImageView(this->device(), handles_.image_view, nullptr); - } - if (handles_.image != VK_NULL_HANDLE) { if (owns_memory_) { vmaDestroyImage(allocator_, handles_.image, memory_.allocation); diff --git a/backends/vulkan/runtime/vk_api/memory/Image.h b/backends/vulkan/runtime/vk_api/memory/Image.h index 447e980595f..7f10301412f 100644 --- a/backends/vulkan/runtime/vk_api/memory/Image.h +++ b/backends/vulkan/runtime/vk_api/memory/Image.h @@ -145,6 +145,9 @@ class VulkanImage final { Allocation memory_; // Indicates whether the underlying memory is owned by this resource bool owns_memory_; + // In some cases, a VulkanImage may be a copy of another VulkanImage but still + // own a unique view of the VkImage. + bool owns_view_; // Indicates whether this VulkanImage was copied from another VulkanImage, // thus it does not have ownership of the underlying VKBuffer bool is_copy_; @@ -228,10 +231,17 @@ class VulkanImage final { inline void bind_allocation(const Allocation& memory) { VK_CHECK_COND(!memory_, "Cannot bind an already bound allocation!"); - VK_CHECK(vmaBindImageMemory(allocator_, memory.allocation, handles_.image)); + // To prevent multiple instances of binding the same VkImage to a memory + // block, do not actually bind memory if this VulkanImage is a copy. Assume + // that the original VulkanImage is responsible for binding the image. + if (!is_copy_) { + VK_CHECK( + vmaBindImageMemory(allocator_, memory.allocation, handles_.image)); + } memory_.allocation = memory.allocation; // Only create the image view if the image has been bound to memory + owns_view_ = true; create_image_view(); } diff --git a/backends/vulkan/test/vulkan_compute_api_test.cpp b/backends/vulkan/test/vulkan_compute_api_test.cpp index fc92739669a..f0d95127a8c 100644 --- a/backends/vulkan/test/vulkan_compute_api_test.cpp +++ b/backends/vulkan/test/vulkan_compute_api_test.cpp @@ -3146,12 +3146,12 @@ void test_transpose_view_mm( mat2_t_small_size = {B, N - 1, K - 3}; } - // Build graph + // Build graph; use shared objects to test views of shared objects IOValueRef mat1 = - graph.add_input_tensor(mat1_size, vkapi::kFloat, utils::kWidthPacked); - IOValueRef mat2_transpose = - graph.add_input_tensor(mat2_t_size, vkapi::kFloat, utils::kWidthPacked); + graph.add_input_tensor(mat1_size, vkapi::kFloat, utils::kWidthPacked, 0); + IOValueRef mat2_transpose = graph.add_input_tensor( + mat2_t_size, vkapi::kFloat, utils::kWidthPacked, 1); ValueRef mat2 = graph.add_tensor_view(mat2_transpose.value); @@ -3167,7 +3167,7 @@ void test_transpose_view_mm( } IOValueRef out; - out.value = graph.add_tensor(out_size, vkapi::kFloat, utils::kWidthPacked); + out.value = graph.add_tensor(out_size, vkapi::kFloat, utils::kWidthPacked, 2); VK_GET_OP_FN("aten.transpose.int") (graph, {mat2_transpose.value, dim0, dim1, mat2});