gfx: Add `resolveQuery` command. (#2125)

Co-authored-by: Yong He <yhe@nvidia.com>

8 files changed, 128 insertions, 2 deletions

diff --git a/tools/gfx/cuda/render-cuda.cpp b/tools/gfx/cuda/render-cuda.cpp
index bacaa42ee..638414358 100644
--- a/tools/gfx/cuda/render-cuda.cpp
+++ b/tools/gfx/cuda/render-cuda.cpp
@@ -1132,6 +1132,21 @@ public:
                 SLANG_UNIMPLEMENTED_X("resolveResource");
             }
 
+            virtual SLANG_NO_THROW void SLANG_MCALL resolveQuery(
+                IQueryPool* queryPool,
+                uint32_t index,
+                uint32_t count,
+                IBufferResource* buffer,
+                uint64_t offset) override
+            {
+                SLANG_UNUSED(queryPool);
+                SLANG_UNUSED(index);
+                SLANG_UNUSED(count);
+                SLANG_UNUSED(buffer);
+                SLANG_UNUSED(offset);
+                SLANG_UNIMPLEMENTED_X("resolveQuery");
+            }
+
             virtual SLANG_NO_THROW void SLANG_MCALL copyTextureToBuffer(
                 IBufferResource* dst,
                 size_t dstOffset,
diff --git a/tools/gfx/d3d12/render-d3d12.cpp b/tools/gfx/d3d12/render-d3d12.cpp
index 0539b8111..36ec173d7 100644
--- a/tools/gfx/d3d12/render-d3d12.cpp
+++ b/tools/gfx/d3d12/render-d3d12.cpp
@@ -563,8 +563,7 @@ public:
     /// Used for query types that does not correspond to a D3D query,
     /// such as ray-tracing acceleration structure post-build info.
     class PlainBufferProxyQueryPoolImpl
-        : public IQueryPool
-        , public ComObject
+        : public QueryPoolBase
     {
     public:
         SLANG_COM_OBJECT_IUNKNOWN_ALL
@@ -4596,6 +4595,62 @@ public:
                 }
             }
 
+            virtual SLANG_NO_THROW void SLANG_MCALL resolveQuery(
+                IQueryPool* queryPool,
+                uint32_t index,
+                uint32_t count,
+                IBufferResource* buffer,
+                uint64_t offset) override
+            {
+                auto queryBase = static_cast<QueryPoolBase*>(queryPool);
+                switch (queryBase->m_desc.type)
+                {
+                case QueryType::AccelerationStructureCompactedSize:
+                case QueryType::AccelerationStructureCurrentSize:
+                case QueryType::AccelerationStructureSerializedSize:
+                    {
+                        auto queryPoolImpl = static_cast<PlainBufferProxyQueryPoolImpl*>(queryPool);
+                        auto bufferImpl = static_cast<BufferResourceImpl*>(buffer);
+                        auto srcQueryBuffer =
+                            queryPoolImpl->m_bufferResource->m_resource.getResource();
+
+                        D3D12_RESOURCE_BARRIER barrier = {};
+                        barrier.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
+                        barrier.Transition.StateBefore = D3D12_RESOURCE_STATE_UNORDERED_ACCESS;
+                        barrier.Transition.StateAfter = D3D12_RESOURCE_STATE_COPY_SOURCE;
+                        barrier.Transition.pResource = srcQueryBuffer;
+                        m_commandBuffer->m_cmdList->ResourceBarrier(1, &barrier);
+
+                        m_commandBuffer->m_cmdList->CopyBufferRegion(
+                            bufferImpl->m_resource.getResource(),
+                            offset,
+                            srcQueryBuffer,
+                            index * sizeof(uint64_t),
+                            count * sizeof(uint64_t));
+
+                        barrier.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
+                        barrier.Transition.StateBefore = D3D12_RESOURCE_STATE_COPY_SOURCE;
+                        barrier.Transition.StateAfter = D3D12_RESOURCE_STATE_UNORDERED_ACCESS;
+                        barrier.Transition.pResource = srcQueryBuffer;
+                        m_commandBuffer->m_cmdList->ResourceBarrier(1, &barrier);
+                    }
+                    break;
+                default:
+                    {
+                        auto queryPoolImpl = static_cast<QueryPoolImpl*>(queryPool);
+                        auto bufferImpl = static_cast<BufferResourceImpl*>(buffer);
+                        m_commandBuffer->m_cmdList->ResolveQueryData(
+                            queryPoolImpl->m_queryHeap.get(),
+                            queryPoolImpl->m_queryType,
+                            index,
+                            count,
+                            bufferImpl->m_resource.getResource(),
+                            offset);
+                    }
+                    break;
+                }
+            }
+
             virtual SLANG_NO_THROW void SLANG_MCALL copyTextureToBuffer(
                 IBufferResource* dst,
                 size_t dstOffset,
@@ -7730,6 +7785,8 @@ Result D3D12Device::createComputePipelineState(const ComputePipelineStateDesc& i
 
 Result D3D12Device::QueryPoolImpl::init(const IQueryPool::Desc& desc, D3D12Device* device)
 {
+    m_desc = desc;
+
     // Translate query type.
     D3D12_QUERY_HEAP_DESC heapDesc = {};
     heapDesc.Count = (UINT)desc.count;
@@ -7816,6 +7873,7 @@ Result D3D12Device::PlainBufferProxyQueryPoolImpl::init(
     m_device = device;
     m_stride = stride;
     m_count = (uint32_t)desc.count;
+    m_desc = desc;
     return SLANG_OK;
 }
 
diff --git a/tools/gfx/debug-layer.cpp b/tools/gfx/debug-layer.cpp
index eed6afdb0..7cc714e0d 100644
--- a/tools/gfx/debug-layer.cpp
+++ b/tools/gfx/debug-layer.cpp
@@ -1362,6 +1362,13 @@ void DebugResourceCommandEncoder::resolveResource(
     baseObject->resolveResource(getInnerObj(source), sourceState, sourceRange, getInnerObj(dest), destState, destRange);
 }
 
+void DebugResourceCommandEncoder::resolveQuery(
+    IQueryPool* queryPool, uint32_t index, uint32_t count, IBufferResource* buffer, uint64_t offset)
+{
+    SLANG_GFX_API_FUNC;
+    baseObject->resolveQuery(getInnerObj(queryPool), index, count, buffer, offset);
+}
+
 void DebugResourceCommandEncoder::copyTextureToBuffer(
     IBufferResource* dst,
     size_t dstOffset,
diff --git a/tools/gfx/debug-layer.h b/tools/gfx/debug-layer.h
index cb894dfd4..567a399fe 100644
--- a/tools/gfx/debug-layer.h
+++ b/tools/gfx/debug-layer.h
@@ -490,6 +490,13 @@ public:
         ResourceState src,
         ResourceState dst) override;
 
+    virtual SLANG_NO_THROW void SLANG_MCALL resolveQuery(
+        IQueryPool* queryPool,
+        uint32_t index,
+        uint32_t count,
+        IBufferResource* buffer,
+        uint64_t offset) override;
+
 public:
     DebugCommandBuffer* commandBuffer;
     bool isOpen = false;
diff --git a/tools/gfx/immediate-renderer-base.cpp b/tools/gfx/immediate-renderer-base.cpp
index 5632eb600..c41debd1e 100644
--- a/tools/gfx/immediate-renderer-base.cpp
+++ b/tools/gfx/immediate-renderer-base.cpp
@@ -402,6 +402,21 @@ public:
             SLANG_UNIMPLEMENTED_X("resolveResource");
         }
 
+        virtual SLANG_NO_THROW void SLANG_MCALL resolveQuery(
+            IQueryPool* queryPool,
+            uint32_t index,
+            uint32_t count,
+            IBufferResource* buffer,
+            uint64_t offset) override
+        {
+            SLANG_UNUSED(queryPool);
+            SLANG_UNUSED(index);
+            SLANG_UNUSED(count);
+            SLANG_UNUSED(buffer);
+            SLANG_UNUSED(offset);
+            SLANG_UNIMPLEMENTED_X("resolveQuery");
+        }
+
         virtual SLANG_NO_THROW void SLANG_MCALL copyTextureToBuffer(
             IBufferResource* dst,
             size_t dstOffset,
diff --git a/tools/gfx/renderer-shared.h b/tools/gfx/renderer-shared.h
index 5ba6bbc50..b27134306 100644
--- a/tools/gfx/renderer-shared.h
+++ b/tools/gfx/renderer-shared.h
@@ -1039,6 +1039,8 @@ public:
     SLANG_COM_OBJECT_IUNKNOWN_ALL
     IQueryPool* getInterface(const Slang::Guid& guid);
     virtual SLANG_NO_THROW Result SLANG_MCALL reset() override { return SLANG_OK; }
+
+    IQueryPool::Desc m_desc;
 };
 
 enum class PipelineType
diff --git a/tools/gfx/vulkan/render-vk.cpp b/tools/gfx/vulkan/render-vk.cpp
index 80084ae7b..d3f0329d8 100644
--- a/tools/gfx/vulkan/render-vk.cpp
+++ b/tools/gfx/vulkan/render-vk.cpp
@@ -4876,6 +4876,27 @@ public:
                 }
             }
 
+            virtual SLANG_NO_THROW void SLANG_MCALL resolveQuery(
+                IQueryPool* queryPool,
+                uint32_t index,
+                uint32_t count,
+                IBufferResource* buffer,
+                uint64_t offset) override
+            {
+                auto& vkApi = m_commandBuffer->m_renderer->m_api;
+                auto poolImpl = static_cast<QueryPoolImpl*>(queryPool);
+                auto bufferImpl = static_cast<BufferResourceImpl*>(buffer);
+                vkApi.vkCmdCopyQueryPoolResults(
+                    m_commandBuffer->m_commandBuffer,
+                    poolImpl->m_pool,
+                    index,
+                    count,
+                    bufferImpl->m_buffer.m_buffer,
+                    offset,
+                    sizeof(uint64_t),
+                    VK_QUERY_RESULT_64_BIT | VK_QUERY_RESULT_WAIT_BIT);
+            }
+
             virtual SLANG_NO_THROW void SLANG_MCALL copyTextureToBuffer(
                 IBufferResource* dst,
                 size_t dstOffset,
diff --git a/tools/gfx/vulkan/vk-api.h b/tools/gfx/vulkan/vk-api.h
index d2c912161..372ecaeea 100644
--- a/tools/gfx/vulkan/vk-api.h
+++ b/tools/gfx/vulkan/vk-api.h
@@ -105,6 +105,7 @@ namespace gfx {
     x(vkCmdBeginQuery) \
     x(vkCmdEndQuery) \
     x(vkCmdResetQueryPool) \
+    x(vkCmdCopyQueryPoolResults) \
     \
     x(vkCreateFence) \
     x(vkDestroyFence) \