Vulkan: Compute shader support

1 files changed, 123 insertions, 0 deletions

diff --git a/Source/Core/VideoBackends/Vulkan/Texture2D.cpp b/Source/Core/VideoBackends/Vulkan/Texture2D.cpp
index 9dda089b21..9b8111aa94 100644
--- a/Source/Core/VideoBackends/Vulkan/Texture2D.cpp
+++ b/Source/Core/VideoBackends/Vulkan/Texture2D.cpp
@@ -4,6 +4,7 @@
 
 #include <algorithm>
 
+#include "Common/Assert.h"
 #include "VideoBackends/Vulkan/CommandBufferManager.h"
 #include "VideoBackends/Vulkan/Texture2D.h"
 #include "VideoBackends/Vulkan/VulkanContext.h"
@@ -273,10 +274,132 @@ void Texture2D::TransitionToLayout(VkCommandBuffer command_buffer, VkImageLayout
     break;
   }
 
+  // If we were using a compute layout, the stages need to reflect that
+  switch (m_compute_layout)
+  {
+  case ComputeImageLayout::Undefined:
+    break;
+  case ComputeImageLayout::ReadOnly:
+    barrier.srcAccessMask = VK_ACCESS_SHADER_READ_BIT;
+    srcStageMask = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT;
+    break;
+  case ComputeImageLayout::WriteOnly:
+    barrier.srcAccessMask = VK_ACCESS_SHADER_WRITE_BIT;
+    srcStageMask = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT;
+    break;
+  case ComputeImageLayout::ReadWrite:
+    barrier.srcAccessMask = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT;
+    srcStageMask = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT;
+    break;
+  }
+  m_compute_layout = ComputeImageLayout::Undefined;
+
   vkCmdPipelineBarrier(command_buffer, srcStageMask, dstStageMask, 0, 0, nullptr, 0, nullptr, 1,
                        &barrier);
 
   m_layout = new_layout;
 }
 
+void Texture2D::TransitionToLayout(VkCommandBuffer command_buffer, ComputeImageLayout new_layout)
+{
+  _assert_(new_layout != ComputeImageLayout::Undefined);
+  if (m_compute_layout == new_layout)
+    return;
+
+  VkImageMemoryBarrier barrier = {
+      VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,  // VkStructureType            sType
+      nullptr,                                 // const void*                pNext
+      0,                                       // VkAccessFlags              srcAccessMask
+      0,                                       // VkAccessFlags              dstAccessMask
+      m_layout,                                // VkImageLayout              oldLayout
+      VK_IMAGE_LAYOUT_GENERAL,                 // VkImageLayout              newLayout
+      VK_QUEUE_FAMILY_IGNORED,                 // uint32_t                   srcQueueFamilyIndex
+      VK_QUEUE_FAMILY_IGNORED,                 // uint32_t                   dstQueueFamilyIndex
+      m_image,                                 // VkImage                    image
+      {static_cast<VkImageAspectFlags>(Util::IsDepthFormat(m_format) ? VK_IMAGE_ASPECT_DEPTH_BIT :
+                                                                       VK_IMAGE_ASPECT_COLOR_BIT),
+       0, m_levels, 0, m_layers}  // VkImageSubresourceRange    subresourceRange
+  };
+
+  VkPipelineStageFlags srcStageMask, dstStageMask;
+  switch (m_layout)
+  {
+  case VK_IMAGE_LAYOUT_UNDEFINED:
+    // Layout undefined therefore contents undefined, and we don't care what happens to it.
+    barrier.srcAccessMask = 0;
+    srcStageMask = VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT;
+    break;
+
+  case VK_IMAGE_LAYOUT_PREINITIALIZED:
+    // Image has been pre-initialized by the host, so ensure all writes have completed.
+    barrier.srcAccessMask = VK_ACCESS_HOST_WRITE_BIT;
+    srcStageMask = VK_PIPELINE_STAGE_HOST_BIT;
+    break;
+
+  case VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL:
+    // Image was being used as a color attachment, so ensure all writes have completed.
+    barrier.srcAccessMask =
+        VK_ACCESS_COLOR_ATTACHMENT_READ_BIT | VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT;
+    srcStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
+    break;
+
+  case VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL:
+    // Image was being used as a depthstencil attachment, so ensure all writes have completed.
+    barrier.srcAccessMask =
+        VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_READ_BIT | VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT;
+    srcStageMask = VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT;
+    break;
+
+  case VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL:
+    // Image was being used as a shader resource, make sure all reads have finished.
+    barrier.srcAccessMask = VK_ACCESS_SHADER_READ_BIT;
+    srcStageMask = VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT;
+    break;
+
+  case VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL:
+    // Image was being used as a copy source, ensure all reads have finished.
+    barrier.srcAccessMask = VK_ACCESS_TRANSFER_READ_BIT;
+    srcStageMask = VK_PIPELINE_STAGE_TRANSFER_BIT;
+    break;
+
+  case VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL:
+    // Image was being used as a copy destination, ensure all writes have finished.
+    barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
+    srcStageMask = VK_PIPELINE_STAGE_TRANSFER_BIT;
+    break;
+
+  default:
+    srcStageMask = VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT;
+    break;
+  }
+
+  switch (new_layout)
+  {
+  case ComputeImageLayout::ReadOnly:
+    barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;
+    barrier.newLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
+    dstStageMask = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT;
+    break;
+  case ComputeImageLayout::WriteOnly:
+    barrier.dstAccessMask = VK_ACCESS_SHADER_WRITE_BIT;
+    barrier.newLayout = VK_IMAGE_LAYOUT_GENERAL;
+    dstStageMask = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT;
+    break;
+  case ComputeImageLayout::ReadWrite:
+    barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT | VK_ACCESS_SHADER_WRITE_BIT;
+    barrier.newLayout = VK_IMAGE_LAYOUT_GENERAL;
+    dstStageMask = VK_PIPELINE_STAGE_COMPUTE_SHADER_BIT;
+    break;
+  default:
+    dstStageMask = 0;
+    break;
+  }
+
+  m_layout = barrier.newLayout;
+  m_compute_layout = new_layout;
+
+  vkCmdPipelineBarrier(command_buffer, srcStageMask, dstStageMask, 0, 0, nullptr, 0, nullptr, 1,
+                       &barrier);
+}
+
 }  // namespace Vulkan