Fix tensor views with tensors that use deferred memory allocation (pytorch#5831)

Summary:
Pull Request resolved: pytorch#5831

## 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.
ghstack-source-id: 246029590
exported-using-ghexport

Reviewed By: jorgep31415

Differential Revision: D63790718

fbshipit-source-id: 1a93851ce77c3c685986dbc8a6300d9905ff60d2

Author: SS-JIA

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<int>(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<int>(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;
 }
 

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);
 

backends/vulkan/runtime/graph/containers/SharedObject.h

@@ -31,6 +31,7 @@ struct SharedObject {
   std::vector<ValueRef> 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);

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;
   }
 

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,17 +230,17 @@ 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.
   if (is_copy_) {
     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);

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();
   }
 

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});