Update on "[11/n][ET-VK] Introduce vTensor creation from external image"

Nearly all metadata is initialized to null/dummy values, except those absolutely needed in the pipeline: (1) image extents, (2) logical limits.

Differential Revision: [D63843819](https://our.internmc.facebook.com/intern/diff/D63843819/)

[ghstack-poisoned]

Author: jorgep31415

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

@@ -13,6 +13,16 @@
 namespace vkcompute {
 namespace api {
 
+std::vector<int64_t> calculate_sizes(
+    const vkapi::VulkanImage& image,
+    const utils::GPUMemoryLayout memory_layout) {
+  auto sizes = std::vector<int64_t>{
+      image.extents().width, image.extents().height, image.extents().depth};
+  const auto packed_dim = utils::to_packed_dim<int32_t>(memory_layout);
+  sizes.at(packed_dim) *= 4;
+  return sizes;
+}
+
 std::vector<int64_t> calculate_dim_order(
     const size_t ndim,
     const int32_t packed_dim) {
@@ -186,6 +196,18 @@ utils::uvec3 calculate_image_extents(
 // vTensorStorage
 //
 
+utils::StorageType storage_type(const vkapi::VulkanImage& image) {
+  const auto type = image.type();
+  switch (type) {
+    case VK_IMAGE_TYPE_3D:
+      return utils::kTexture3D;
+    case VK_IMAGE_TYPE_2D:
+      return utils::kTexture2D;
+    default:
+      VK_THROW("Unsupported image type", type);
+  }
+}
+
 vkapi::VulkanImage allocate_image(
     Context* const context_ptr,
     utils::uvec3& image_extents,
@@ -285,7 +307,7 @@ vTensorStorage::vTensorStorage(
     Context* const context,
     const vkapi::VulkanImage& image)
     : context_(context),
-      storage_type_{utils::kTexture2D},
+      storage_type_{storage_type(image)},
       image_extents_(
           {image.extents().width,
            image.extents().height,
@@ -461,16 +483,19 @@ vTensor::vTensor(
 }
 
 // NOLINTNEXTLINE
-vTensor::vTensor(Context* const context, const vkapi::VulkanImage& image)
-    : dtype_(),
+vTensor::vTensor(
+    Context* context,
+    const vkapi::VulkanImage& image,
+    const utils::GPUMemoryLayout memory_layout)
+    : dtype_(vkapi::element_scalartype(image.format())),
       // Calculate tensor metadata
-      sizes_(),
-      packed_dim_(),
+      sizes_(calculate_sizes(image, memory_layout)),
+      packed_dim_(utils::to_packed_dim<int32_t>(memory_layout)),
       dim_order_(),
       axis_map_(default_axis_map()),
       strides_(),
-      numel_(),
-      padded_sizes_(),
+      numel_(utils::multiply_integers(sizes_)),
+      padded_sizes_(calculate_padded_sizes(sizes_, packed_dim_)),
       unsqueezed_strides_(),
       padded_numel_(),
       logical_limits_(),

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

@@ -187,7 +187,10 @@ class vTensor final {
 
   vTensor(const vTensor& other) = delete;
 
-  explicit vTensor(Context* context, const vkapi::VulkanImage& image);
+  explicit vTensor(
+      Context* context,
+      const vkapi::VulkanImage& image,
+      const utils::GPUMemoryLayout memory_layout = utils::kChannelsPacked);
 
   /*
    * This constructor allows for the creation of a vTensor that references the

backends/vulkan/runtime/utils/AhbUtils.h

@@ -116,12 +116,12 @@ DescriptorSet& DescriptorSet::bind(
 DescriptorSet& DescriptorSet::bind(
     const uint32_t idx,
     const VulkanImage& image) {
-  // Check is only accurate for images allocated with VMA
-  if (image.vma_allocator() != VK_NULL_HANDLE) {
-    VK_CHECK_COND(
-        image.has_memory(),
-        "Image must be bound to memory for it to be usable");
-  }
+  // 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.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/Descriptor.cpp

@@ -84,6 +84,7 @@ inline ScalarType element_scalartype(const VkFormat vkformat) {
     case VK_FORMAT_R8G8B8A8_SINT:
       return kChar;
     case VK_FORMAT_R8G8B8A8_UINT:
+    case VK_FORMAT_R8G8B8A8_UNORM:
       return kByte;
     case VK_FORMAT_R32G32B32A32_SINT:
       return kInt;

backends/vulkan/runtime/vk_api/Types.h

@@ -179,6 +179,10 @@ class VulkanImage final {
     return memory_.allocation;
   }
 
+  inline VkImageType type() const {
+    return image_properties_.image_type;
+  }
+
   inline VkFormat format() const {
     return image_properties_.image_format;
   }