Enable vTensor creation from external image (#5993)

Summary:
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: 247078909
exported-using-ghexport

Reviewed By: SS-JIA

Differential Revision: D63843819

fbshipit-source-id: 71a52800fc5b79953ab98a8b549b12b4760db757

Author: Jorge Pineda

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_(utils::multiply_integers(padded_sizes_)),
+      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

@@ -180,6 +180,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;
   }

backends/vulkan/test/vulkan_compute_api_test.cpp

@@ -23,6 +23,8 @@
 
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/QPackUtils.h>
 
+#include <executorch/backends/vulkan/runtime/vk_api/VkUtils.h>
+
 #include <executorch/backends/vulkan/test/utils/test_utils.h>
 
 using namespace vkcompute;
@@ -1094,6 +1096,49 @@ TEST_F(VulkanComputeAPITest, print_object_sizes) {
   EXPECT_TRUE(sizeof(ExecuteNode) < 500);
 }
 
+TEST_F(VulkanComputeAPITest, test_tensor_creation_from_vulkan_image) {
+  const auto w = 16;
+  const auto h = 12;
+  const auto d = 1;
+  const utils::uvec3 image_extents = {w, h, d};
+
+  vkapi::Adapter* adapter_ptr = context()->adapter_ptr();
+
+  vkapi::ImageSampler::Properties sampler_props{
+      VK_FILTER_NEAREST,
+      VK_SAMPLER_MIPMAP_MODE_NEAREST,
+      VK_SAMPLER_ADDRESS_MODE_REPEAT,
+      VK_BORDER_COLOR_FLOAT_TRANSPARENT_BLACK,
+  };
+
+  VkFormat image_format = VK_FORMAT_R32G32B32A32_SFLOAT;
+  VkImageType image_type = VK_IMAGE_TYPE_3D;
+  VkImageViewType image_view_type = VK_IMAGE_VIEW_TYPE_3D;
+
+  VkSampler sampler = adapter_ptr->sampler_cache().retrieve(sampler_props);
+
+  auto image = adapter_ptr->vma().create_image(
+      context()->device(),
+      vkapi::create_extent3d(image_extents),
+      image_format,
+      image_type,
+      image_view_type,
+      sampler_props,
+      sampler,
+      /*allow_transfer = */ true,
+      /*allocate_memory = */ true);
+
+  auto tensor = vTensor(context(), image);
+
+  const auto exp_sizes = std::vector<int64_t>{w, h, d * 4};
+  EXPECT_TRUE(tensor.sizes() == exp_sizes);
+  EXPECT_TRUE(tensor.packed_dim() == 2);
+
+  const auto exp_numel = w * h * d * 4;
+  EXPECT_TRUE(tensor.numel() == exp_numel);
+  EXPECT_TRUE(tensor.padded_numel() == exp_numel);
+}
+
 TEST(VulkanComputeGraphTest, test_values_scalars) {
   GraphConfig config;
   ComputeGraph graph(config);