Vulkan: Allow re-use of uniform buffers when doing per-stage uploads

This is safe now because we invalidate the pointers after submitting a
command buffer.

Source/Core/VideoBackends/Vulkan/StateTracker.cpp

@@ -357,20 +357,9 @@ bool StateTracker::CheckForShaderChanges(u32 gx_primitive_type, DSTALPHA_MODE ds
 
 void StateTracker::UpdateVertexShaderConstants()
 {
-  if (!VertexShaderManager::dirty)
+  if (!VertexShaderManager::dirty || !ReserveConstantStorage())
     return;
 
-  // Since the other stages uniform buffers' may be still be using the earlier data,
-  // we can't reuse the earlier part of the buffer without re-uploading everything.
-  if (!m_uniform_stream_buffer->ReserveMemory(m_uniform_buffer_reserve_size,
-                                              g_vulkan_context->GetUniformBufferAlignment(), false,
-                                              false, false))
-  {
-    // Re-upload all constants to a new portion of the buffer.
-    UploadAllConstants();
-    return;
-  }
-
   // Buffer allocation changed?
   if (m_uniform_stream_buffer->GetBuffer() !=
       m_bindings.uniform_buffer_bindings[UBO_DESCRIPTOR_SET_BINDING_VS].buffer)
@@ -394,17 +383,9 @@ void StateTracker::UpdateVertexShaderConstants()
 void StateTracker::UpdateGeometryShaderConstants()
 {
   // Skip updating geometry shader constants if it's not in use.
-  if (m_pipeline_state.gs == VK_NULL_HANDLE || !GeometryShaderManager::dirty)
+  if (m_pipeline_state.gs == VK_NULL_HANDLE || !GeometryShaderManager::dirty ||
-    return;
+      !ReserveConstantStorage())
-
-  // Since the other stages uniform buffers' may be still be using the earlier data,
-  // we can't reuse the earlier part of the buffer without re-uploading everything.
-  if (!m_uniform_stream_buffer->ReserveMemory(m_uniform_buffer_reserve_size,
-                                              g_vulkan_context->GetUniformBufferAlignment(), false,
-                                              false, false))
   {
-    // Re-upload all constants to a new portion of the buffer.
-    UploadAllConstants();
     return;
   }
 
@@ -430,20 +411,9 @@ void StateTracker::UpdateGeometryShaderConstants()
 
 void StateTracker::UpdatePixelShaderConstants()
 {
-  if (!PixelShaderManager::dirty)
+  if (!PixelShaderManager::dirty || !ReserveConstantStorage())
     return;
 
-  // Since the other stages uniform buffers' may be still be using the earlier data,
-  // we can't reuse the earlier part of the buffer without re-uploading everything.
-  if (!m_uniform_stream_buffer->ReserveMemory(m_uniform_buffer_reserve_size,
-                                              g_vulkan_context->GetUniformBufferAlignment(), false,
-                                              false, false))
-  {
-    // Re-upload all constants to a new portion of the buffer.
-    UploadAllConstants();
-    return;
-  }
-
   // Buffer allocation changed?
   if (m_uniform_stream_buffer->GetBuffer() !=
       m_bindings.uniform_buffer_bindings[UBO_DESCRIPTOR_SET_BINDING_PS].buffer)
@@ -464,6 +434,27 @@ void StateTracker::UpdatePixelShaderConstants()
   PixelShaderManager::dirty = false;
 }
 
+bool StateTracker::ReserveConstantStorage()
+{
+  // Since we invalidate all constants on command buffer execution, it doesn't matter if this
+  // causes the stream buffer to be resized.
+  if (m_uniform_stream_buffer->ReserveMemory(m_uniform_buffer_reserve_size,
+                                             g_vulkan_context->GetUniformBufferAlignment(), true,
+                                             true, false))
+  {
+    return true;
+  }
+
+  // The only places that call constant updates are safe to have state restored.
+  WARN_LOG(VIDEO, "Executing command buffer while waiting for space in uniform buffer");
+  Util::ExecuteCurrentCommandsAndRestoreState(false);
+
+  // Since we are on a new command buffer, all constants have been invalidated, and we need
+  // to reupload them. We may as well do this now, since we're issuing a draw anyway.
+  UploadAllConstants();
+  return false;
+}
+
 void StateTracker::UploadAllConstants()
 {
   // We are free to re-use parts of the buffer now since we're uploading all constants.
@@ -476,17 +467,12 @@ void StateTracker::UploadAllConstants()
   size_t allocation_size = geometry_constants_offset + sizeof(GeometryShaderConstants);
 
   // Allocate everything at once.
-  if (!m_uniform_stream_buffer->ReserveMemory(allocation_size, ub_alignment, true, true, false))
+  // We should only be here if the buffer was full and a command buffer was submitted anyway.
-  {
-    // The only places that call constant updates are safe to have state restored.
-    WARN_LOG(VIDEO, "Executing command buffer while waiting for space in uniform buffer");
-    Util::ExecuteCurrentCommandsAndRestoreState(false);
   if (!m_uniform_stream_buffer->ReserveMemory(allocation_size, ub_alignment, true, true, false))
   {
     PanicAlert("Failed to allocate space for constants in streaming buffer");
     return;
   }
-  }
 
   // Update bindings
   for (size_t i = 0; i < NUM_UBO_DESCRIPTOR_SET_BINDINGS; i++)

Source/Core/VideoBackends/Vulkan/StateTracker.h

@@ -174,6 +174,11 @@ private:
 
   bool UpdatePipeline();
   bool UpdateDescriptorSet();
+
+  // Allocates storage in the uniform buffer of the specified size. If this storage cannot be
+  // allocated immediately, the current command buffer will be submitted and all stage's
+  // constants will be re-uploaded. false will be returned in this case, otherwise true.
+  bool ReserveConstantStorage();
   void UploadAllConstants();
 
   // Which bindings/state has to be updated before the next draw.