Enable tensor aliases with texture storage (#5347)

Summary:
Pull Request resolved: #5347

## Context

Allow aliasing for `VulkanImage` objects, and by extension texture backed Tensors.

Essentially just a re-application/extension of what was done in #4769.

bypass-github-export-checks

ghstack-source-id: 242452078
exported-using-ghexport

Reviewed By: jorgep31415

Differential Revision: D62606929

fbshipit-source-id: 446b07d3176d99679750674e6ae82e29e88bc27d

Author: SS-JIA

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(

backends/vulkan/runtime/vk_api/memory/Allocation.h

@@ -80,6 +80,7 @@ struct Allocation final {
   }
 
   friend class VulkanBuffer;
+  friend class VulkanImage;
 };
 
 } // namespace vkapi

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

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,13 +177,25 @@ 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_),
       sampler_properties_(other.sampler_properties_),
       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);
   }

backends/vulkan/runtime/vk_api/memory/Image.h

@@ -22,6 +22,12 @@
 #include <unordered_map>
 
 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 {

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<int64_t> sizes = {9, 9};
-  std::vector<int64_t> strides =
-      get_reference_strides(sizes, utils::kWidthPacked);
-  std::vector<int64_t> dim_order = {0, 1};
+TEST_F(VulkanComputeAPITest, tensor_alias_test) {
+  for (utils::StorageType storage_type : {utils::kTexture3D, utils::kBuffer}) {
+    std::vector<int64_t> 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<float> 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<float> 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);
+    }
   }
 }