From fe57f9261e51b07b2716b11d7c204a68328c4aba Mon Sep 17 00:00:00 2001
From: Sergey Kosarevsky <sk@linderdaum.com>
Date: Tue, 22 Oct 2024 01:15:38 -0700
Subject: [PATCH] Create image view with identity swizzle for storage images

---
 lvk/vulkan/VulkanClasses.cpp | 17 ++++++++++++++++-
 lvk/vulkan/VulkanClasses.h   |  1 +
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/lvk/vulkan/VulkanClasses.cpp b/lvk/vulkan/VulkanClasses.cpp
index 6336ca8297..f6c4331ed5 100644
--- a/lvk/vulkan/VulkanClasses.cpp
+++ b/lvk/vulkan/VulkanClasses.cpp
@@ -3851,6 +3851,16 @@ lvk::Holder<lvk::TextureHandle> lvk::VulkanContext::createTexture(const TextureD
   image.imageView_ = image.createImageView(
       vkDevice_, vkImageViewType, vkFormat, aspect, 0, VK_REMAINING_MIP_LEVELS, 0, numLayers, mapping, ycbcrInfo, debugNameImageView);
 
+  if (image.vkUsageFlags_ & VK_IMAGE_USAGE_STORAGE_BIT) {
+    if (desc.swizzle.r != Swizzle_Default || desc.swizzle.g != Swizzle_Default || desc.swizzle.b != Swizzle_Default ||
+        desc.swizzle.a != Swizzle_Default) {
+      // use identity swizzle for storage images
+      image.imageViewStorage_ = image.createImageView(
+          vkDevice_, vkImageViewType, vkFormat, aspect, 0, VK_REMAINING_MIP_LEVELS, 0, numLayers, {}, ycbcrInfo, debugNameImageView);
+      LVK_ASSERT(image.imageViewStorage_ != VK_NULL_HANDLE);
+    }
+  }
+
   if (!LVK_VERIFY(image.imageView_ != VK_NULL_HANDLE)) {
     Result::setResult(outResult, Result::Code::RuntimeError, "Cannot create VkImageView");
     return {};
@@ -4914,6 +4924,10 @@ void lvk::VulkanContext::destroy(lvk::TextureHandle handle) {
 
   deferredTask(std::packaged_task<void()>(
       [device = getVkDevice(), imageView = tex->imageView_]() { vkDestroyImageView(device, imageView, nullptr); }));
+  if (tex->imageViewStorage_) {
+    deferredTask(std::packaged_task<void()>(
+        [device = getVkDevice(), imageView = tex->imageViewStorage_]() { vkDestroyImageView(device, imageView, nullptr); }));
+  }
 
   for (size_t i = 0; i != LVK_MAX_MIP_LEVELS; i++) {
     for (size_t j = 0; j != LVK_ARRAY_NUM_ELEMENTS(tex->imageViewForFramebuffer_[0]); j++) {
@@ -6544,6 +6558,7 @@ void lvk::VulkanContext::checkAndUpdateDescriptorSets() {
   for (const auto& obj : texturesPool_.objects_) {
     const VulkanImage& img = obj.obj_;
     const VkImageView view = obj.obj_.imageView_;
+    const VkImageView storageView = obj.obj_.imageViewStorage_ ? obj.obj_.imageViewStorage_ : view;
     // multisampled images cannot be directly accessed from shaders
     const bool isTextureAvailable = (img.vkSamples_ & VK_SAMPLE_COUNT_1_BIT) == VK_SAMPLE_COUNT_1_BIT;
     const bool isYUVImage = isTextureAvailable && img.isSampledImage() && lvk::getNumImagePlanes(img.vkImageFormat_) > 1;
@@ -6557,7 +6572,7 @@ void lvk::VulkanContext::checkAndUpdateDescriptorSets() {
     LVK_ASSERT(infoSampledImages.back().imageView != VK_NULL_HANDLE);
     infoStorageImages.push_back(VkDescriptorImageInfo{
         .sampler = VK_NULL_HANDLE,
-        .imageView = isStorageImage ? view : dummyImageView,
+        .imageView = isStorageImage ? storageView : dummyImageView,
         .imageLayout = VK_IMAGE_LAYOUT_GENERAL,
     });
     if (hasYcbcrSamplers) {
diff --git a/lvk/vulkan/VulkanClasses.h b/lvk/vulkan/VulkanClasses.h
index 30aae71898..dedefb10e4 100644
--- a/lvk/vulkan/VulkanClasses.h
+++ b/lvk/vulkan/VulkanClasses.h
@@ -113,6 +113,7 @@ struct VulkanImage final {
   mutable VkImageLayout vkImageLayout_ = VK_IMAGE_LAYOUT_UNDEFINED;
   // precached image views - owned by this VulkanImage
   VkImageView imageView_ = VK_NULL_HANDLE; // default view with all mip-levels
+  VkImageView imageViewStorage_ = VK_NULL_HANDLE; // default view with identity swizzle (all mip-levels)
   VkImageView imageViewForFramebuffer_[LVK_MAX_MIP_LEVELS][6] = {}; // max 6 faces for cubemap rendering
 };