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

@@ -497,7 +497,7 @@ vTensor::vTensor(
       numel_(utils::multiply_integers(sizes_)),
       padded_sizes_(calculate_padded_sizes(sizes_, packed_dim_)),
       unsqueezed_strides_(),
-      padded_numel_(),
+      padded_numel_(utils::multiply_integers(padded_sizes_)),
       logical_limits_(),
       // Utility Uniform Buffers that can be passed to shaders as arguments
       sizes_uniform_(),

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

@@ -134,8 +134,8 @@ VulkanImage Allocator::create_image(
       image_props,
       view_props,
       sampler_props,
-      initial_layout,
       sampler,
+      initial_layout,
       allocate_memory);
 }
 

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

@@ -115,8 +115,8 @@ VulkanImage::VulkanImage(
     const ImageProperties& image_props,
     const ViewProperties& view_props,
     const SamplerProperties& sampler_props,
-    const VkImageLayout layout,
     VkSampler sampler,
+    const VkImageLayout layout,
     const bool allocate_memory)
     : device_{device},
       image_properties_(image_props),
@@ -186,7 +186,8 @@ VulkanImage::VulkanImage(
     const ImageProperties& image_props,
     VkImage image,
     VkImageView image_view,
-    VkSampler sampler)
+    VkSampler sampler,
+    const VkImageLayout layout)
     : device_{device},
       image_properties_{image_props},
       view_properties_{},
@@ -200,7 +201,7 @@ VulkanImage::VulkanImage(
           image_view,
           sampler,
       },
-      layout_{VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL} {}
+      layout_{layout} {}
 
 VulkanImage::VulkanImage(const VulkanImage& other) noexcept
     : device_(other.device_),

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

@@ -99,16 +99,17 @@ class VulkanImage final {
       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);
+      VkSampler,
+      const VkImageLayout);
 
  protected:
   /*

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