[11/n][ET-VK] Enable vTensor creation from external image

Pull Request resolved: #5993

Nearly all metadata is initialized to null/dummy values, except those absolutely needed in the pipeline: (1) image extents, (2) logical limits.
ghstack-source-id: 246919998
@exported-using-ghexport

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

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,
@@ -281,6 +303,21 @@ vTensorStorage::vTensorStorage(
       last_access_{},
       has_copies_{false} {}
 
+vTensorStorage::vTensorStorage(
+    Context* const context,
+    const vkapi::VulkanImage& image)
+    : context_(context),
+      storage_type_{storage_type(image)},
+      image_extents_(
+          {image.extents().width,
+           image.extents().height,
+           image.extents().depth}),
+      buffer_length_{0},
+      buffer_offset_{0},
+      image_(image),
+      buffer_(vkapi::VulkanBuffer()),
+      last_access_{} {}
+
 vTensorStorage::vTensorStorage(
     vTensorStorage& other,
     const int64_t buffer_offset)
@@ -446,6 +483,33 @@ vTensor::vTensor(
 }
 
 // NOLINTNEXTLINE
+vTensor::vTensor(
+    Context* context,
+    const vkapi::VulkanImage& image,
+    const utils::GPUMemoryLayout memory_layout)
+    : dtype_(vkapi::element_scalartype(image.format())),
+      // Calculate tensor metadata
+      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_(utils::multiply_integers(sizes_)),
+      padded_sizes_(calculate_padded_sizes(sizes_, packed_dim_)),
+      unsqueezed_strides_(),
+      padded_numel_(),
+      logical_limits_(),
+      // Utility Uniform Buffers that can be passed to shaders as arguments
+      sizes_uniform_(),
+      strides_uniform_(),
+      numel_uniform_(),
+      axis_map_uniform_(),
+      logical_limits_uniform_(),
+      // Construct Tensor storage
+      storage_(context, image) {
+  set_logical_limits(storage_.image_extents_);
+}
+
 vTensor::vTensor(vTensor& other)
     : dtype_(other.dtype_),
       // Copy tensor size metadata

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

@@ -95,6 +95,8 @@ class vTensorStorage final {
       const vkapi::ScalarType dtype,
       const bool allocate_memory = true);
 
+  vTensorStorage(Context* const context, const vkapi::VulkanImage& image);
+
  protected:
   /*
    * This allows for creation of tensors that use the same underlying storage
@@ -185,6 +187,11 @@ class vTensor final {
 
   vTensor(const vTensor& other) = delete;
 
+  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
    * same buffer resource of another vTensor, with the same sizes and strides

backends/vulkan/runtime/graph/ComputeGraph.cpp

@@ -296,6 +296,13 @@ ValueRef ComputeGraph::add_tensor(
       sizes, dtype, suggested_memory_layout(sizes), shared_object_idx);
 }
 
+ValueRef ComputeGraph::add_tensor(const vkapi::VulkanImage& image) {
+  ValueRef idx(static_cast<int>(values_.size()));
+  check_no_active_value_ptrs();
+  values_.emplace_back(api::vTensor(context(), image));
+  return idx;
+}
+
 ValueRef ComputeGraph::add_tensor_view(const ValueRef vref) {
   const vTensorPtr t = get_tensor(vref);
   ValueRef idx(static_cast<int>(values_.size()));

backends/vulkan/runtime/graph/ComputeGraph.h

@@ -466,6 +466,11 @@ class ComputeGraph final {
       const vkapi::ScalarType dtype,
       const int64_t shared_object_idx = -1);
 
+  /*
+   * Add a `api::vTensor` value to the graph with the specified image.
+   */
+  ValueRef add_tensor(const vkapi::VulkanImage& image);
+
   /*
    * Add a `api::vTensor` value to the graph with the properties of `vref`.
    */

backends/vulkan/runtime/vk_api/Types.h

@@ -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/memory/Image.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;
   }