From ccf346930751d1ab85888f3ad331e34d26a38b11 Mon Sep 17 00:00:00 2001
From: Wojtek Figat <wojtek@figat.pl>
Date: Tue, 10 Oct 2023 14:17:34 +0200
Subject: [PATCH] Minor Vulkan tweaks

---
 .../GraphicsDevice/Vulkan/GPUContextVulkan.cpp       | 10 ++++------
 .../GraphicsDevice/Vulkan/GPUTimerQueryVulkan.cpp    | 12 ++++++------
 .../GraphicsDevice/Vulkan/GPUTimerQueryVulkan.h      |  2 +-
 3 files changed, 11 insertions(+), 13 deletions(-)

diff --git a/Source/Engine/GraphicsDevice/Vulkan/GPUContextVulkan.cpp b/Source/Engine/GraphicsDevice/Vulkan/GPUContextVulkan.cpp
index d758ec999..79e882f6d 100644
--- a/Source/Engine/GraphicsDevice/Vulkan/GPUContextVulkan.cpp
+++ b/Source/Engine/GraphicsDevice/Vulkan/GPUContextVulkan.cpp
@@ -1210,16 +1210,16 @@ void GPUContextVulkan::ResolveMultisample(GPUTexture* sourceMultisampleTexture,
 
 void GPUContextVulkan::DrawInstanced(uint32 verticesCount, uint32 instanceCount, int32 startInstance, int32 startVertex)
 {
-    const auto cmdBuffer = _cmdBufferManager->GetCmdBuffer();
     OnDrawCall();
+    const auto cmdBuffer = _cmdBufferManager->GetCmdBuffer();
     vkCmdDraw(cmdBuffer->GetHandle(), verticesCount, instanceCount, startVertex, startInstance);
     RENDER_STAT_DRAW_CALL(verticesCount * instanceCount, verticesCount * instanceCount / 3);
 }
 
 void GPUContextVulkan::DrawIndexedInstanced(uint32 indicesCount, uint32 instanceCount, int32 startInstance, int32 startVertex, int32 startIndex)
 {
-    const auto cmdBuffer = _cmdBufferManager->GetCmdBuffer();
     OnDrawCall();
+    const auto cmdBuffer = _cmdBufferManager->GetCmdBuffer();
     vkCmdDrawIndexed(cmdBuffer->GetHandle(), indicesCount, instanceCount, startIndex, startVertex, startInstance);
     RENDER_STAT_DRAW_CALL(0, indicesCount / 3 * instanceCount);
 }
@@ -1227,10 +1227,9 @@ void GPUContextVulkan::DrawIndexedInstanced(uint32 indicesCount, uint32 instance
 void GPUContextVulkan::DrawInstancedIndirect(GPUBuffer* bufferForArgs, uint32 offsetForArgs)
 {
     ASSERT(bufferForArgs && EnumHasAnyFlags(bufferForArgs->GetFlags(), GPUBufferFlags::Argument));
-
+    OnDrawCall();
     auto bufferForArgsVK = (GPUBufferVulkan*)bufferForArgs;
     const auto cmdBuffer = _cmdBufferManager->GetCmdBuffer();
-    OnDrawCall();
     vkCmdDrawIndirect(cmdBuffer->GetHandle(), bufferForArgsVK->GetHandle(), (VkDeviceSize)offsetForArgs, 1, sizeof(VkDrawIndirectCommand));
     RENDER_STAT_DRAW_CALL(0, 0);
 }
@@ -1238,10 +1237,9 @@ void GPUContextVulkan::DrawInstancedIndirect(GPUBuffer* bufferForArgs, uint32 of
 void GPUContextVulkan::DrawIndexedInstancedIndirect(GPUBuffer* bufferForArgs, uint32 offsetForArgs)
 {
     ASSERT(bufferForArgs && EnumHasAnyFlags(bufferForArgs->GetFlags(), GPUBufferFlags::Argument));
-
+    OnDrawCall();
     auto bufferForArgsVK = (GPUBufferVulkan*)bufferForArgs;
     const auto cmdBuffer = _cmdBufferManager->GetCmdBuffer();
-    OnDrawCall();
     vkCmdDrawIndexedIndirect(cmdBuffer->GetHandle(), bufferForArgsVK->GetHandle(), (VkDeviceSize)offsetForArgs, 1, sizeof(VkDrawIndexedIndirectCommand));
     RENDER_STAT_DRAW_CALL(0, 0);
 }
diff --git a/Source/Engine/GraphicsDevice/Vulkan/GPUTimerQueryVulkan.cpp b/Source/Engine/GraphicsDevice/Vulkan/GPUTimerQueryVulkan.cpp
index 9c58b77b0..cce86799a 100644
--- a/Source/Engine/GraphicsDevice/Vulkan/GPUTimerQueryVulkan.cpp
+++ b/Source/Engine/GraphicsDevice/Vulkan/GPUTimerQueryVulkan.cpp
@@ -16,7 +16,7 @@ void GPUTimerQueryVulkan::Interrupt(CmdBufferVulkan* cmdBuffer)
     if (!_interrupted)
     {
         _interrupted = true;
-        WriteTimestamp(cmdBuffer, _queries[_queryIndex].End);
+        WriteTimestamp(cmdBuffer, _queries[_queryIndex].End, VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT);
     }
 }
 
@@ -28,7 +28,7 @@ void GPUTimerQueryVulkan::Resume(CmdBufferVulkan* cmdBuffer)
     e.End.Pool = nullptr;
 
     _interrupted = false;
-    WriteTimestamp(cmdBuffer, e.Begin);
+    WriteTimestamp(cmdBuffer, e.Begin, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT);
 
     _queries.Add(e);
     _queryIndex++;
@@ -56,13 +56,13 @@ bool GPUTimerQueryVulkan::GetResult(Query& query)
     return false;
 }
 
-void GPUTimerQueryVulkan::WriteTimestamp(CmdBufferVulkan* cmdBuffer, Query& query) const
+void GPUTimerQueryVulkan::WriteTimestamp(CmdBufferVulkan* cmdBuffer, Query& query, VkPipelineStageFlagBits stage) const
 {
     auto pool = _device->FindAvailableTimestampQueryPool();
     uint32 index;
     pool->AcquireQuery(index);
 
-    vkCmdWriteTimestamp(cmdBuffer->GetHandle(), VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT, pool->GetHandle(), index);
+    vkCmdWriteTimestamp(cmdBuffer->GetHandle(), stage, pool->GetHandle(), index);
     pool->MarkQueryAsStarted(index);
 
     query.Pool = pool;
@@ -168,7 +168,7 @@ void GPUTimerQueryVulkan::Begin()
 
         _queryIndex = 0;
         _interrupted = false;
-        WriteTimestamp(cmdBuffer, e.Begin);
+        WriteTimestamp(cmdBuffer, e.Begin, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT);
         context->GetCmdBufferManager()->OnQueryBegin(this);
 
         ASSERT(_queries.IsEmpty());
@@ -193,7 +193,7 @@ void GPUTimerQueryVulkan::End()
 
         if (!_interrupted)
         {
-            WriteTimestamp(cmdBuffer, _queries[_queryIndex].End);
+            WriteTimestamp(cmdBuffer, _queries[_queryIndex].End, VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT);
         }
         context->GetCmdBufferManager()->OnQueryEnd(this);
     }
diff --git a/Source/Engine/GraphicsDevice/Vulkan/GPUTimerQueryVulkan.h b/Source/Engine/GraphicsDevice/Vulkan/GPUTimerQueryVulkan.h
index b081b946d..ddf461b13 100644
--- a/Source/Engine/GraphicsDevice/Vulkan/GPUTimerQueryVulkan.h
+++ b/Source/Engine/GraphicsDevice/Vulkan/GPUTimerQueryVulkan.h
@@ -59,7 +59,7 @@ public:
 private:
 
     bool GetResult(Query& query);
-    void WriteTimestamp(CmdBufferVulkan* cmdBuffer, Query& query) const;
+    void WriteTimestamp(CmdBufferVulkan* cmdBuffer, Query& query, VkPipelineStageFlagBits stage) const;
     bool TryGetResult();
     bool UseQueries();