diff --git a/backends/vulkan/runtime/api/containers/Tensor.cpp b/backends/vulkan/runtime/api/containers/Tensor.cpp index c9a31b92da5..b3dff832339 100644 --- a/backends/vulkan/runtime/api/containers/Tensor.cpp +++ b/backends/vulkan/runtime/api/containers/Tensor.cpp @@ -682,13 +682,9 @@ vTensorStorage::vTensorStorage( image_extents_(other.image_extents_), buffer_length_{other.buffer_length_}, buffer_offset_{buffer_offset}, - image_(), + image_(other.image_), buffer_(other.buffer_, buffer_offset), - last_access_{other.last_access_} { - if (other.storage_type_ != utils::kBuffer) { - VK_THROW("Tensors with texture storage cannot be copied!"); - } -} + last_access_{other.last_access_} {} vTensorStorage::~vTensorStorage() { flush(); @@ -758,14 +754,10 @@ void vTensorStorage::transition( } bool vTensorStorage::is_copy_of(const vTensorStorage& other) const { - if (storage_type_ != other.storage_type_) { - return false; - } - // Copies are only enabled for buffer storage at the moment - if (storage_type_ != utils::kBuffer) { - return false; + if (storage_type_ == utils::kBuffer) { + return buffer_.is_copy_of(other.buffer_); } - return buffer_.is_copy_of(other.buffer_); + return image_.is_copy_of(other.image_); } void vTensorStorage::discard_and_reallocate( diff --git a/backends/vulkan/runtime/vk_api/memory/Allocation.h b/backends/vulkan/runtime/vk_api/memory/Allocation.h index 44e8277a35c..1b9c6c0dadf 100644 --- a/backends/vulkan/runtime/vk_api/memory/Allocation.h +++ b/backends/vulkan/runtime/vk_api/memory/Allocation.h @@ -80,6 +80,7 @@ struct Allocation final { } friend class VulkanBuffer; + friend class VulkanImage; }; } // namespace vkapi diff --git a/backends/vulkan/runtime/vk_api/memory/Buffer.cpp b/backends/vulkan/runtime/vk_api/memory/Buffer.cpp index 5a78dab764d..93a87f01ae1 100644 --- a/backends/vulkan/runtime/vk_api/memory/Buffer.cpp +++ b/backends/vulkan/runtime/vk_api/memory/Buffer.cpp @@ -83,7 +83,7 @@ VulkanBuffer::VulkanBuffer( : buffer_properties_(other.buffer_properties_), allocator_(other.allocator_), memory_(other.memory_), - owns_memory_(other.owns_memory_), + owns_memory_(false), is_copy_(true), handle_(other.handle_) { // TODO: set the offset and range appropriately diff --git a/backends/vulkan/runtime/vk_api/memory/Image.cpp b/backends/vulkan/runtime/vk_api/memory/Image.cpp index 42352cfb7e7..34c5d0e2c29 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), + is_copy_(false), handles_{ VK_NULL_HANDLE, VK_NULL_HANDLE, @@ -120,6 +121,7 @@ VulkanImage::VulkanImage( allocator_(vma_allocator), memory_{}, owns_memory_{allocate_memory}, + is_copy_(false), handles_{ VK_NULL_HANDLE, VK_NULL_HANDLE, @@ -175,6 +177,17 @@ VulkanImage::VulkanImage( } } +VulkanImage::VulkanImage(const VulkanImage& other) noexcept + : image_properties_(other.image_properties_), + view_properties_(other.view_properties_), + sampler_properties_(other.sampler_properties_), + allocator_(other.allocator_), + memory_(other.memory_), + owns_memory_{false}, + is_copy_(true), + handles_(other.handles_), + layout_(other.layout_) {} + VulkanImage::VulkanImage(VulkanImage&& other) noexcept : image_properties_(other.image_properties_), view_properties_(other.view_properties_), @@ -182,6 +195,7 @@ VulkanImage::VulkanImage(VulkanImage&& other) noexcept allocator_(other.allocator_), memory_(std::move(other.memory_)), owns_memory_(other.owns_memory_), + is_copy_(other.is_copy_), handles_(other.handles_), layout_(other.layout_) { other.handles_.image = VK_NULL_HANDLE; @@ -201,6 +215,7 @@ VulkanImage& VulkanImage::operator=(VulkanImage&& other) noexcept { allocator_ = other.allocator_; memory_ = std::move(other.memory_); owns_memory_ = other.owns_memory_; + is_copy_ = other.is_copy_; handles_ = other.handles_; layout_ = other.layout_; @@ -212,6 +227,13 @@ VulkanImage& VulkanImage::operator=(VulkanImage&& other) noexcept { } VulkanImage::~VulkanImage() { + // 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. + if (is_copy_) { + return; + } + if (VK_NULL_HANDLE != handles_.image_view) { vkDestroyImageView(this->device(), handles_.image_view, nullptr); } diff --git a/backends/vulkan/runtime/vk_api/memory/Image.h b/backends/vulkan/runtime/vk_api/memory/Image.h index 1e78f84a5c5..1f93bb68a53 100644 --- a/backends/vulkan/runtime/vk_api/memory/Image.h +++ b/backends/vulkan/runtime/vk_api/memory/Image.h @@ -22,6 +22,12 @@ #include namespace vkcompute { + +// Forward declare vTensor classes such that they can be set as friend classes +namespace api { +class vTensorStorage; +} // namespace api + namespace vkapi { class ImageSampler final { @@ -96,7 +102,23 @@ class VulkanImage final { VkSampler, const bool allocate_memory = true); - VulkanImage(const VulkanImage&) = delete; + protected: + /* + * The Copy constructor allows for creation of a class instance that are + * "aliases" of another class instance. The resulting class instance will not + * have ownership of the underlying VkImage. + * + * This behaviour is analogous to creating a copy of a pointer, thus it is + * unsafe, as the original class instance may be destroyed before the copy. + * These constructors are therefore marked protected so that they may be used + * only in situations where the lifetime of the original class instance is + * guaranteed to exceed, or at least be the same as, the lifetime of the + * copied class instance. + */ + VulkanImage(const VulkanImage& other) noexcept; + + public: + // To discourage creating copies, the assignment operator is still deleted. VulkanImage& operator=(const VulkanImage&) = delete; VulkanImage(VulkanImage&&) noexcept; @@ -123,6 +145,9 @@ class VulkanImage final { Allocation memory_; // Indicates whether the underlying memory is owned by this resource bool owns_memory_; + // Indicates whether this VulkanImage was copied from another VulkanImage, + // thus it does not have ownership of the underlying VKBuffer + bool is_copy_; Handles handles_; // Layout VkImageLayout layout_; @@ -193,10 +218,18 @@ class VulkanImage final { return owns_memory_; } + inline bool is_copy() const { + return is_copy_; + } + inline operator bool() const { return (handles_.image != VK_NULL_HANDLE); } + inline bool is_copy_of(const VulkanImage& other) const { + return (handles_.image == other.handles_.image) && is_copy_; + } + 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)); @@ -207,6 +240,8 @@ class VulkanImage final { } VkMemoryRequirements get_memory_requirements() const; + + friend class api::vTensorStorage; }; struct ImageMemoryBarrier final { diff --git a/backends/vulkan/test/vulkan_compute_api_test.cpp b/backends/vulkan/test/vulkan_compute_api_test.cpp index dd65095b329..59b3ee42dc8 100644 --- a/backends/vulkan/test/vulkan_compute_api_test.cpp +++ b/backends/vulkan/test/vulkan_compute_api_test.cpp @@ -604,26 +604,30 @@ TEST_F(VulkanComputeAPITest, texture_add_sanity_check) { } } -TEST_F(VulkanComputeAPITest, tensor_copy_test) { - std::vector sizes = {9, 9}; - std::vector strides = - get_reference_strides(sizes, utils::kWidthPacked); - std::vector dim_order = {0, 1}; +TEST_F(VulkanComputeAPITest, tensor_alias_test) { + for (utils::StorageType storage_type : {utils::kTexture3D, utils::kBuffer}) { + std::vector sizes = {9, 9}; - vTensor original = CREATE_FLOAT_BUFFER(sizes, /*allocate_memory=*/true); - vTensor copy = vTensor(original, sizes, dim_order); - EXPECT_TRUE(get_vma_allocation_count() == 1); - EXPECT_TRUE(copy.is_view_of(original)); + const size_t alloc_count_before = get_vma_allocation_count(); - // Fill original tensor with some data - fill_vtensor(original, 2.5f, true); + vTensor original = vTensor(context(), sizes, vkapi::kFloat, storage_type); - std::vector data_out(copy.staging_buffer_numel()); - // Extract the copy tensor; should contain the data of the original tensor - extract_vtensor(copy, data_out); + vTensor copy = vTensor(original); - for (size_t i = 0; i < data_out.size(); ++i) { - CHECK_VALUE(data_out, i, 2.5f + i); + // Two tensors but only one additional allocation. + EXPECT_TRUE(get_vma_allocation_count() == alloc_count_before + 1); + EXPECT_TRUE(copy.is_view_of(original)); + + // Fill original tensor with some data + fill_vtensor(original, 2.5f, true); + + std::vector data_out(copy.staging_buffer_numel()); + // Extract the copy tensor; should contain the data of the original tensor + extract_vtensor(copy, data_out); + + for (size_t i = 0; i < original.numel(); ++i) { + CHECK_VALUE(data_out, i, 2.5f + i); + } } }