[ET-VK] Only use LINEAR tiling if it's available

## Context

Recently, we switched to using `VK_IMAGE_TILING_LINEAR` to minimize memory footprint.

However, according to the [Vulkan Spec](https://registry.khronos.org/vulkan/specs/1.3-extensions/man/html/VkImageCreateInfo.html) the `VK_IMAGE_TILING_LINEAR` tiling may only be available for 2D textures.

## Changes

* When constructing the `Adapter` class, check if it's possible to create a 3D texture with LINEAR tiling
* Add a way to query preferred tiling from `Context`
* Construct VkImage with the preferred tiling.

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

[ghstack-poisoned]

Author: SS-JIA

backends/vulkan/runtime/api/Context.cpp

@@ -39,7 +39,12 @@ Context::Context(size_t adapter_i, const ContextConfig& config)
       buffer_clearlist_mutex_{},
       buffers_to_clear_{},
       image_clearlist_mutex_{},
-      images_to_clear_{} {}
+      images_to_clear_{},
+      preferred_image_tiling_{VK_IMAGE_TILING_OPTIMAL} {
+  if (adapter_p_->linear_tiling_3d_enabled()) {
+    preferred_image_tiling_ = VK_IMAGE_TILING_LINEAR;
+  }
+}
 
 Context::~Context() {
   try {

backends/vulkan/runtime/api/Context.h

@@ -73,6 +73,8 @@ class Context final {
   std::vector<vkapi::VulkanBuffer> buffers_to_clear_;
   std::mutex image_clearlist_mutex_;
   std::vector<vkapi::VulkanImage> images_to_clear_;
+  // Misc
+  VkImageTiling preferred_image_tiling_;
 
  public:
   // Adapter access
@@ -123,6 +125,10 @@ class Context final {
     return querypool_;
   }
 
+  inline VkImageTiling preferred_image_tiling() {
+    return preferred_image_tiling_;
+  }
+
   /*
    * By default, the querypool attached to a Context instance is uninitialized.
    * This function triggers the querypool to be created via vkCreateQueryPool.

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

@@ -245,6 +245,7 @@ vkapi::VulkanImage allocate_image(
       vkapi::create_extent3d(image_extents),
       image_format,
       image_type,
+      context_ptr->preferred_image_tiling(),
       image_view_type,
       sampler_props,
       sampler,

backends/vulkan/runtime/vk_api/Adapter.cpp

@@ -156,7 +156,39 @@ Adapter::Adapter(
       pipeline_layout_cache_(device_.handle),
       compute_pipeline_cache_(device_.handle, cache_data_path),
       sampler_cache_(device_.handle),
-      vma_(instance_, physical_device_.handle, device_.handle) {}
+      vma_(instance_, physical_device_.handle, device_.handle),
+      linear_tiling_3d_enabled_{true} {
+  // Test creating a 3D image with linear tiling to see if it is supported.
+  // According to the Vulkan spec, linear tiling may not be supported for 3D
+  // images.
+  VkExtent3D image_extents{1u, 1u, 1u};
+  const VkImageCreateInfo image_create_info{
+      VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO, // sType
+      nullptr, // pNext
+      0u, // flags
+      VK_IMAGE_TYPE_3D, // imageType
+      VK_FORMAT_R32G32B32A32_SFLOAT, // format
+      image_extents, // extents
+      1u, // mipLevels
+      1u, // arrayLayers
+      VK_SAMPLE_COUNT_1_BIT, // samples
+      VK_IMAGE_TILING_LINEAR, // tiling
+      VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_STORAGE_BIT, // usage
+      VK_SHARING_MODE_EXCLUSIVE, // sharingMode
+      0u, // queueFamilyIndexCount
+      nullptr, // pQueueFamilyIndices
+      VK_IMAGE_LAYOUT_UNDEFINED, // initialLayout
+  };
+  VkImage image = VK_NULL_HANDLE;
+  VkResult res =
+      vkCreateImage(device_.handle, &image_create_info, nullptr, &image);
+  if (res == VK_ERROR_FEATURE_NOT_PRESENT) {
+    linear_tiling_3d_enabled_ = false;
+  } else if (res == VK_SUCCESS) {
+    vkDestroyImage(device_.handle, image, nullptr);
+  }
+  return;
+}
 
 Adapter::Queue Adapter::request_queue() {
   // Lock the mutex as multiple threads can request a queue at the same time

backends/vulkan/runtime/vk_api/Adapter.h

@@ -92,6 +92,8 @@ class Adapter final {
   // Memory Management
   SamplerCache sampler_cache_;
   Allocator vma_;
+  // Miscellaneous
+  bool linear_tiling_3d_enabled_;
 
  public:
   // Physical Device metadata
@@ -153,6 +155,10 @@ class Adapter final {
     return vma_;
   }
 
+  inline bool linear_tiling_3d_enabled() const {
+    return linear_tiling_3d_enabled_;
+  }
+
   // Physical Device Features
 
   inline bool supports_16bit_storage_buffers() {

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

@@ -99,6 +99,7 @@ VulkanImage Allocator::create_image(
     const VkExtent3D& extents,
     const VkFormat image_format,
     const VkImageType image_type,
+    const VkImageTiling image_tiling,
     const VkImageViewType image_view_type,
     const VulkanImage::SamplerProperties& sampler_props,
     VkSampler sampler,
@@ -117,6 +118,7 @@ VulkanImage Allocator::create_image(
       image_type,
       image_format,
       extents,
+      image_tiling,
       usage,
   };
 

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

@@ -59,6 +59,7 @@ class Allocator final {
       const VkExtent3D&,
       const VkFormat,
       const VkImageType,
+      const VkImageTiling,
       const VkImageViewType,
       const VulkanImage::SamplerProperties&,
       VkSampler,

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

@@ -156,7 +156,7 @@ VulkanImage::VulkanImage(
       1u, // mipLevels
       1u, // arrayLayers
       VK_SAMPLE_COUNT_1_BIT, // samples
-      VK_IMAGE_TILING_LINEAR, // tiling
+      image_properties_.image_tiling, // tiling
       image_properties_.image_usage, // usage
       VK_SHARING_MODE_EXCLUSIVE, // sharingMode
       0u, // queueFamilyIndexCount

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

@@ -74,6 +74,7 @@ class VulkanImage final {
     VkImageType image_type;
     VkFormat image_format;
     VkExtent3D image_extents;
+    VkImageTiling image_tiling;
     VkImageUsageFlags image_usage;
   };
 

backends/vulkan/test/vulkan_compute_api_test.cpp

@@ -1123,6 +1123,7 @@ TEST_F(VulkanComputeAPITest, test_tensor_creation_from_vulkan_image) {
       vkapi::create_extent3d(image_extents),
       image_format,
       image_type,
+      context()->preferred_image_tiling(),
       image_view_type,
       sampler_props,
       sampler,