Vulkan: Support binding texel buffers in UtilityShaderDraw

3 files changed, 62 insertions, 18 deletions

diff --git a/Source/Core/VideoBackends/Vulkan/ShaderCompiler.cpp b/Source/Core/VideoBackends/Vulkan/ShaderCompiler.cpp
index ded63ce3ed..2265a34364 100644
--- a/Source/Core/VideoBackends/Vulkan/ShaderCompiler.cpp
+++ b/Source/Core/VideoBackends/Vulkan/ShaderCompiler.cpp
@@ -52,6 +52,7 @@ static const char SHADER_HEADER[] = R"(
   #define UBO_BINDING(packing, x) layout(packing, set = 0, binding = (x - 1))
   #define SAMPLER_BINDING(x) layout(set = 1, binding = x)
   #define SSBO_BINDING(x) layout(set = 2, binding = x)
+  #define TEXEL_BUFFER_BINDING(x) layout(set = 2, binding = x)
   #define VARYING_LOCATION(x) layout(location = x)
   #define FORCE_EARLY_Z layout(early_fragment_tests) in
 
diff --git a/Source/Core/VideoBackends/Vulkan/Util.cpp b/Source/Core/VideoBackends/Vulkan/Util.cpp
index f2b23a25ee..f85321bac5 100644
--- a/Source/Core/VideoBackends/Vulkan/Util.cpp
+++ b/Source/Core/VideoBackends/Vulkan/Util.cpp
@@ -364,6 +364,15 @@ void UtilityShaderDraw::SetPSSampler(size_t index, VkImageView view, VkSampler s
   m_ps_samplers[index].imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
 }
 
+void UtilityShaderDraw::SetPSTexelBuffer(VkBufferView view)
+{
+  // Should only be used with the texture conversion pipeline layout.
+  _assert_(m_pipeline_info.pipeline_layout ==
+           g_object_cache->GetPipelineLayout(PIPELINE_LAYOUT_TEXTURE_CONVERSION));
+
+  m_ps_texel_buffer = view;
+}
+
 void UtilityShaderDraw::SetRasterizationState(const RasterizationState& state)
 {
   m_pipeline_info.rasterization_state.bits = state.bits;
@@ -600,29 +609,59 @@ void UtilityShaderDraw::BindDescriptors()
   vkUpdateDescriptorSets(g_vulkan_context->GetDevice(), num_set_writes, set_writes.data(), 0,
                          nullptr);
 
-  // Bind only the sets we updated
-  if (bind_descriptor_sets[0] != VK_NULL_HANDLE && bind_descriptor_sets[1] == VK_NULL_HANDLE)
+  if (m_ps_texel_buffer != VK_NULL_HANDLE)
   {
-    // UBO only
-    vkCmdBindDescriptorSets(m_command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS,
-                            m_pipeline_info.pipeline_layout,
-                            DESCRIPTOR_SET_BIND_POINT_UNIFORM_BUFFERS, 1, &bind_descriptor_sets[0],
-                            NUM_UBO_DESCRIPTOR_SET_BINDINGS, m_ubo_offsets.data());
+    // TODO: Handle case where this fails.
+    // This'll only be when we do over say, 1024 allocations per frame, which shouldn't happen.
+    // TODO: Execute the command buffer, reset render passes and then try again.
+    VkDescriptorSet set = g_command_buffer_mgr->AllocateDescriptorSet(
+        g_object_cache->GetDescriptorSetLayout(DESCRIPTOR_SET_LAYOUT_TEXEL_BUFFERS));
+    if (set == VK_NULL_HANDLE)
+    {
+      PanicAlert("Failed to allocate texel buffer descriptor set for utility draw");
+      return;
+    }
+
+    VkWriteDescriptorSet set_write = {VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
+                                      nullptr,
+                                      set,
+                                      0,
+                                      0,
+                                      1,
+                                      VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER,
+                                      nullptr,
+                                      nullptr,
+                                      &m_ps_texel_buffer};
+    vkUpdateDescriptorSets(g_vulkan_context->GetDevice(), 1, &set_write, 0, nullptr);
+    bind_descriptor_sets[DESCRIPTOR_SET_BIND_POINT_STORAGE_OR_TEXEL_BUFFER] = set;
+  }
+
+  // Fast path when there are no gaps in the set bindings
+  u32 bind_point_index;
+  for (bind_point_index = 0; bind_point_index < NUM_DESCRIPTOR_SET_BIND_POINTS; bind_point_index++)
+  {
+    if (bind_descriptor_sets[bind_point_index] == VK_NULL_HANDLE)
+      break;
   }
-  else if (bind_descriptor_sets[0] == VK_NULL_HANDLE && bind_descriptor_sets[1] != VK_NULL_HANDLE)
+  if (bind_point_index > 0)
   {
-    // Samplers only
-    vkCmdBindDescriptorSets(
-        m_command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS, m_pipeline_info.pipeline_layout,
-        DESCRIPTOR_SET_BIND_POINT_PIXEL_SHADER_SAMPLERS, 1, &bind_descriptor_sets[1], 0, nullptr);
+    // Bind the contiguous sets, any others after any gaps will be handled below
+    vkCmdBindDescriptorSets(m_command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS,
+                            m_pipeline_info.pipeline_layout, 0, bind_point_index,
+                            &bind_descriptor_sets[0], NUM_UBO_DESCRIPTOR_SET_BINDINGS,
+                            m_ubo_offsets.data());
   }
-  else if (bind_descriptor_sets[0] != VK_NULL_HANDLE && bind_descriptor_sets[1] != VK_NULL_HANDLE)
+
+  // Handle any remaining sets
+  for (u32 i = bind_point_index; i < NUM_DESCRIPTOR_SET_BIND_POINTS; i++)
   {
-    // Both
-    vkCmdBindDescriptorSets(
-        m_command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS, m_pipeline_info.pipeline_layout,
-        DESCRIPTOR_SET_BIND_POINT_UNIFORM_BUFFERS, 2, bind_descriptor_sets.data(),
-        NUM_UBO_DESCRIPTOR_SET_BINDINGS, m_ubo_offsets.data());
+    if (bind_descriptor_sets[i] == VK_NULL_HANDLE)
+      continue;
+
+    // No need to worry about dynamic offsets here, since #0 will always be bound above.
+    vkCmdBindDescriptorSets(m_command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS,
+                            m_pipeline_info.pipeline_layout, i, 1, &bind_descriptor_sets[i], 0,
+                            nullptr);
   }
 }
 
diff --git a/Source/Core/VideoBackends/Vulkan/Util.h b/Source/Core/VideoBackends/Vulkan/Util.h
index 6900cce08b..bb701716c2 100644
--- a/Source/Core/VideoBackends/Vulkan/Util.h
+++ b/Source/Core/VideoBackends/Vulkan/Util.h
@@ -137,6 +137,8 @@ public:
 
   void SetPSSampler(size_t index, VkImageView view, VkSampler sampler);
 
+  void SetPSTexelBuffer(VkBufferView view);
+
   void SetRasterizationState(const RasterizationState& state);
   void SetDepthStencilState(const DepthStencilState& state);
   void SetBlendState(const BlendState& state);
@@ -182,6 +184,8 @@ private:
 
   std::array<VkDescriptorImageInfo, NUM_PIXEL_SHADER_SAMPLERS> m_ps_samplers = {};
 
+  VkBufferView m_ps_texel_buffer = VK_NULL_HANDLE;
+
   PipelineInfo m_pipeline_info = {};
 };