Enable externally allocated VulkanImages (#5992)

Summary:
Pull Request resolved: #5992

We now include handle to `VkDevice` in `VulkanImage` since our allocations don't go through VMA.

**Internal:**
Following image creation of `myvulkanahb`.
https://www.internalfb.com/code/aosp-vendor-meta-tools-camera-tools/[wrist-sw5100-14.0%3Af8460689765eee9332f461491a1edcfc6ad8fb92]/myvulkanahb/myvulkanahb.cpp?lines=22%2C94-96
https://www.internalfb.com/code/aosp-vendor-meta-tools-camera-tools/[wrist-sw5100-14.0%3Af8460689765eee9332f461491a1edcfc6ad8fb92]/myvulkanahb/vulkan/ImageProcessorAhbRgbaInRgbaOut.cpp?lines=26-36%2C40-47
ghstack-source-id: 247129375
exported-using-ghexport

Reviewed By: SS-JIA

Differential Revision: D63298759

fbshipit-source-id: 0583a92524ff9d10e6313a631601bf0ca7908d34

Author: Jorge Pineda

backends/vulkan/runtime/api/containers/Tensor.cpp

@@ -221,6 +221,7 @@ vkapi::VulkanImage allocate_image(
   VkSampler sampler = adapter_ptr->sampler_cache().retrieve(sampler_props);
 
   return adapter_ptr->vma().create_image(
+      context_ptr->device(),
       vkapi::create_extent3d(image_extents),
       image_format,
       image_type,

backends/vulkan/runtime/vk_api/Descriptor.cpp

@@ -116,8 +116,12 @@ DescriptorSet& DescriptorSet::bind(
 DescriptorSet& DescriptorSet::bind(
     const uint32_t idx,
     const VulkanImage& image) {
+  // If the image does not have an allocator attached, then it is externally
+  // allocated; assume it is already bound to memory. Otherwise, it must be
+  // bound to a VmaAllocation to be used.
   VK_CHECK_COND(
-      image.has_memory(), "Image must be bound to memory for it to be usable");
+      image.vma_allocator() == VK_NULL_HANDLE || image.has_memory(),
+      "Image must be bound to memory for it to be usable");
 
   VkImageLayout binding_layout = image.layout();
   if (shader_layout_signature_[idx] == VK_DESCRIPTOR_TYPE_STORAGE_IMAGE) {

backends/vulkan/runtime/vk_api/memory/Allocator.cpp

@@ -95,6 +95,7 @@ Allocation Allocator::create_allocation(
 }
 
 VulkanImage Allocator::create_image(
+    const VkDevice device,
     const VkExtent3D& extents,
     const VkFormat image_format,
     const VkImageType image_type,
@@ -127,13 +128,14 @@ VulkanImage Allocator::create_image(
   const VkImageLayout initial_layout = VK_IMAGE_LAYOUT_UNDEFINED;
 
   return VulkanImage(
+      device,
       allocator_,
       alloc_create_info,
       image_props,
       view_props,
       sampler_props,
-      initial_layout,
       sampler,
+      initial_layout,
       allocate_memory);
 }
 

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

@@ -55,6 +55,7 @@ class Allocator final {
       const VmaAllocationCreateInfo& create_info);
 
   VulkanImage create_image(
+      const VkDevice,
       const VkExtent3D&,
       const VkFormat,
       const VkImageType,

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

@@ -92,7 +92,8 @@ void swap(ImageSampler& lhs, ImageSampler& rhs) noexcept {
 //
 
 VulkanImage::VulkanImage()
-    : image_properties_{},
+    : device_{VK_NULL_HANDLE},
+      image_properties_{},
       view_properties_{},
       sampler_properties_{},
       allocator_(VK_NULL_HANDLE),
@@ -108,15 +109,17 @@ VulkanImage::VulkanImage()
       layout_{} {}
 
 VulkanImage::VulkanImage(
+    VkDevice device,
     VmaAllocator vma_allocator,
     const VmaAllocationCreateInfo& allocation_create_info,
     const ImageProperties& image_props,
     const ViewProperties& view_props,
     const SamplerProperties& sampler_props,
-    const VkImageLayout layout,
     VkSampler sampler,
+    const VkImageLayout layout,
     const bool allocate_memory)
-    : image_properties_(image_props),
+    : device_{device},
+      image_properties_(image_props),
       view_properties_(view_props),
       sampler_properties_(sampler_props),
       allocator_(vma_allocator),
@@ -178,8 +181,31 @@ VulkanImage::VulkanImage(
   }
 }
 
+VulkanImage::VulkanImage(
+    VkDevice device,
+    const ImageProperties& image_props,
+    VkImage image,
+    VkImageView image_view,
+    VkSampler sampler,
+    const VkImageLayout layout)
+    : device_{device},
+      image_properties_{image_props},
+      view_properties_{},
+      sampler_properties_{},
+      allocator_(VK_NULL_HANDLE),
+      memory_{},
+      owns_memory_(false),
+      is_copy_(false),
+      handles_{
+          image,
+          image_view,
+          sampler,
+      },
+      layout_{layout} {}
+
 VulkanImage::VulkanImage(const VulkanImage& other) noexcept
-    : image_properties_(other.image_properties_),
+    : device_(other.device_),
+      image_properties_(other.image_properties_),
       view_properties_(other.view_properties_),
       sampler_properties_(other.sampler_properties_),
       allocator_(other.allocator_),
@@ -191,7 +217,8 @@ VulkanImage::VulkanImage(const VulkanImage& other) noexcept
       layout_(other.layout_) {}
 
 VulkanImage::VulkanImage(VulkanImage&& other) noexcept
-    : image_properties_(other.image_properties_),
+    : device_(other.device_),
+      image_properties_(other.image_properties_),
       view_properties_(other.view_properties_),
       sampler_properties_(other.sampler_properties_),
       allocator_(other.allocator_),
@@ -212,6 +239,7 @@ VulkanImage& VulkanImage::operator=(VulkanImage&& other) noexcept {
   VkImageView tmp_image_view = handles_.image_view;
   bool tmp_owns_memory = owns_memory_;
 
+  device_ = other.device_;
   image_properties_ = other.image_properties_;
   view_properties_ = other.view_properties_;
   sampler_properties_ = other.sampler_properties_;

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

@@ -93,15 +93,24 @@ class VulkanImage final {
   explicit VulkanImage();
 
   explicit VulkanImage(
+      VkDevice,
       const VmaAllocator,
       const VmaAllocationCreateInfo&,
       const ImageProperties&,
       const ViewProperties&,
       const SamplerProperties&,
-      const VkImageLayout layout,
       VkSampler,
+      const VkImageLayout,
       const bool allocate_memory = true);
 
+  explicit VulkanImage(
+      VkDevice,
+      const ImageProperties&,
+      VkImage,
+      VkImageView,
+      VkSampler,
+      const VkImageLayout);
+
  protected:
   /*
    * The Copy constructor allows for creation of a class instance that are
@@ -136,6 +145,7 @@ class VulkanImage final {
   friend struct ImageMemoryBarrier;
 
  private:
+  VkDevice device_;
   ImageProperties image_properties_;
   ViewProperties view_properties_;
   SamplerProperties sampler_properties_;
@@ -159,9 +169,7 @@ class VulkanImage final {
   void create_image_view();
 
   inline VkDevice device() const {
-    VmaAllocatorInfo allocator_info{};
-    vmaGetAllocatorInfo(allocator_, &allocator_info);
-    return allocator_info.device;
+    return device_;
   }
 
   inline VmaAllocator vma_allocator() const {