Support dispatchComputeIndirect and draw(Indexed)Indirect with count buffer for Vulkan (#5929)

* Support gfx ComputeCommandEncoder::dispatchComputeIndirect for Vulkan

* Support count buffer for Vulkan in gfx RenderCommandEncoder::drawIndirect and RenderCommandEncoder::drawIndexedIndirect

* Fix an unintended change

* Fix format issue

---------

Co-authored-by: Yong He <yonghe@outlook.com>

2 files changed, 58 insertions, 22 deletions

diff --git a/tools/gfx/vulkan/vk-api.h b/tools/gfx/vulkan/vk-api.h
index 8b86bc281..9d2535f09 100644
--- a/tools/gfx/vulkan/vk-api.h
+++ b/tools/gfx/vulkan/vk-api.h
@@ -86,10 +86,13 @@ namespace gfx
     x(vkCmdFillBuffer) \
     x(vkCmdBindDescriptorSets) \
     x(vkCmdDispatch) \
+    x(vkCmdDispatchIndirect) \
     x(vkCmdDraw) \
     x(vkCmdDrawIndexed) \
     x(vkCmdDrawIndirect) \
+    x(vkCmdDrawIndirectCount) \
     x(vkCmdDrawIndexedIndirect) \
+    x(vkCmdDrawIndexedIndirectCount) \
     x(vkCmdSetScissor) \
     x(vkCmdSetViewport) \
     x(vkCmdBindVertexBuffers) \
diff --git a/tools/gfx/vulkan/vk-command-encoder.cpp b/tools/gfx/vulkan/vk-command-encoder.cpp
index 42af3bee9..b4f4cc06b 100644
--- a/tools/gfx/vulkan/vk-command-encoder.cpp
+++ b/tools/gfx/vulkan/vk-command-encoder.cpp
@@ -1175,19 +1175,31 @@ Result RenderCommandEncoder::drawIndirect(
     IBufferResource* countBuffer,
     Offset countOffset)
 {
-    // Vulkan does not support sourcing the count from a buffer.
-    if (countBuffer)
-        return SLANG_FAIL;
-
     SLANG_RETURN_ON_FAIL(prepareDraw());
     auto& api = *m_api;
     auto argBufferImpl = static_cast<BufferResourceImpl*>(argBuffer);
-    api.vkCmdDrawIndirect(
-        m_vkCommandBuffer,
-        argBufferImpl->m_buffer.m_buffer,
-        argOffset,
-        maxDrawCount,
-        sizeof(VkDrawIndirectCommand));
+
+    if (countBuffer)
+    {
+        auto countBufferImpl = static_cast<BufferResourceImpl*>(countBuffer);
+        api.vkCmdDrawIndirectCount(
+            m_vkCommandBuffer,
+            argBufferImpl->m_buffer.m_buffer,
+            argOffset,
+            countBufferImpl->m_buffer.m_buffer,
+            countOffset,
+            maxDrawCount,
+            sizeof(VkDrawIndirectCommand));
+    }
+    else
+    {
+        api.vkCmdDrawIndirect(
+            m_vkCommandBuffer,
+            argBufferImpl->m_buffer.m_buffer,
+            argOffset,
+            maxDrawCount,
+            sizeof(VkDrawIndirectCommand));
+    }
     return SLANG_OK;
 }
 
@@ -1198,20 +1210,31 @@ Result RenderCommandEncoder::drawIndexedIndirect(
     IBufferResource* countBuffer,
     Offset countOffset)
 {
-    // Vulkan does not support sourcing the count from a buffer.
-    if (countBuffer)
-        return SLANG_FAIL;
-
     SLANG_RETURN_ON_FAIL(prepareDraw());
-
     auto& api = *m_api;
     auto argBufferImpl = static_cast<BufferResourceImpl*>(argBuffer);
-    api.vkCmdDrawIndexedIndirect(
-        m_vkCommandBuffer,
-        argBufferImpl->m_buffer.m_buffer,
-        argOffset,
-        maxDrawCount,
-        sizeof(VkDrawIndexedIndirectCommand));
+
+    if (countBuffer)
+    {
+        auto countBufferImpl = static_cast<BufferResourceImpl*>(countBuffer);
+        api.vkCmdDrawIndexedIndirectCount(
+            m_vkCommandBuffer,
+            argBufferImpl->m_buffer.m_buffer,
+            argOffset,
+            countBufferImpl->m_buffer.m_buffer,
+            countOffset,
+            maxDrawCount,
+            sizeof(VkDrawIndexedIndirectCommand));
+    }
+    else
+    {
+        api.vkCmdDrawIndexedIndirect(
+            m_vkCommandBuffer,
+            argBufferImpl->m_buffer.m_buffer,
+            argOffset,
+            maxDrawCount,
+            sizeof(VkDrawIndexedIndirectCommand));
+    }
     return SLANG_OK;
 }
 
@@ -1311,7 +1334,17 @@ Result ComputeCommandEncoder::dispatchCompute(int x, int y, int z)
 
 Result ComputeCommandEncoder::dispatchComputeIndirect(IBufferResource* argBuffer, Offset offset)
 {
-    SLANG_UNIMPLEMENTED_X("dispatchComputeIndirect");
+    auto pipeline = static_cast<PipelineStateImpl*>(m_currentPipeline.Ptr());
+    if (!pipeline)
+    {
+        return SLANG_FAIL;
+    }
+
+    // Also create descriptor sets based on the given pipeline layout
+    SLANG_RETURN_ON_FAIL(bindRenderState(VK_PIPELINE_BIND_POINT_COMPUTE));
+    auto argBufferImpl = static_cast<BufferResourceImpl*>(argBuffer);
+    m_api->vkCmdDispatchIndirect(m_vkCommandBuffer, argBufferImpl->m_buffer.m_buffer, offset);
+    return SLANG_OK;
 }
 
 void RayTracingCommandEncoder::_memoryBarrier(