From 2ad1f8138771cef32b710f8c47d4c7beb3f4eab5 Mon Sep 17 00:00:00 2001
From: AdamYuan <y13916619121@126.com>
Date: Fri, 27 Dec 2024 02:11:28 +0800
Subject: 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>
---
 tools/gfx/vulkan/vk-command-encoder.cpp | 77 +++++++++++++++++++++++----------
 1 file changed, 55 insertions(+), 22 deletions(-)

(limited to 'tools/gfx/vulkan/vk-command-encoder.cpp')

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(
-- 
cgit v1.2.3