Reworked vulkan descriptor set binding to set dynamic buffer offsets at draw time

6 files changed, 134 insertions, 57 deletions

diff --git a/Graphics/GraphicsEngineVulkan/include/DeviceContextVkImpl.h b/Graphics/GraphicsEngineVulkan/include/DeviceContextVkImpl.h
index 4eb5920b..020c3270 100644
--- a/Graphics/GraphicsEngineVulkan/include/DeviceContextVkImpl.h
+++ b/Graphics/GraphicsEngineVulkan/include/DeviceContextVkImpl.h
@@ -36,6 +36,7 @@
 #include "VulkanDynamicHeap.h"
 #include "ResourceReleaseQueue.h"
 #include "DescriptorPoolManager.h"
+#include "PipelineLayout.h"
 
 #ifdef _DEBUG
 #   define VERIFY_CONTEXT_BINDINGS
@@ -149,8 +150,6 @@ public:
 
     VulkanDynamicAllocation AllocateDynamicSpace(Uint32 SizeInBytes);
 
-    std::vector<uint32_t>& GetDynamicBufferOffsets(){return m_DynamicBufferOffsets;}
-
 private:
     void CommitRenderPassAndFramebuffer(class PipelineStateVkImpl *pPipelineStateVk);
     void CommitVkVertexBuffers();
@@ -181,7 +180,6 @@ private:
 #if 0
     GenerateMipsHelper m_MipsGenerator;
 #endif
-    class ShaderResourceCacheVk *m_pCommittedResourceCache = nullptr;
 
     FixedBlockMemoryAllocator m_CmdListAllocator;
 
@@ -204,7 +202,7 @@ private:
     // be submitted next
     Atomics::AtomicInt64 m_NextCmdBuffNumber;
 
-    std::vector<uint32_t> m_DynamicBufferOffsets;
+    PipelineLayout::DescriptorSetBindInfo m_DesrSetBindInfo;
     VulkanDynamicHeap m_DynamicHeap;
 };
 
diff --git a/Graphics/GraphicsEngineVulkan/include/PipelineLayout.h b/Graphics/GraphicsEngineVulkan/include/PipelineLayout.h
index 556fab20..a0810ab7 100644
--- a/Graphics/GraphicsEngineVulkan/include/PipelineLayout.h
+++ b/Graphics/GraphicsEngineVulkan/include/PipelineLayout.h
@@ -81,10 +81,46 @@ public:
     void AllocateDynamicDescriptorSet(DeviceContextVkImpl*    pCtxVkImpl,
                                       ShaderResourceCacheVk&  ResourceCache)const;
 
-    // Binds Vulkan descriptor sets to the command buffer in the cmd context
-    void BindDescriptorSets(DeviceContextVkImpl*    pCtxVkImpl,
-                            bool                    IsCompute,
-                            ShaderResourceCacheVk&  ResourceCache)const;
+    struct DescriptorSetBindInfo
+    {
+        std::vector<VkDescriptorSet> vkSets;
+        std::vector<uint32_t>        DynamicOffsets;
+        ShaderResourceCacheVk*       pResourceCache = nullptr;
+        VkPipelineBindPoint          BindPoint      = VK_PIPELINE_BIND_POINT_MAX_ENUM;
+
+#ifdef _DEBUG
+        const PipelineLayout *pDbgPipelineLayout = nullptr;
+#endif
+        DescriptorSetBindInfo()
+        {
+            vkSets.reserve(2);
+            DynamicOffsets.reserve(64);
+        }
+
+        void Reset()
+        {
+            vkSets.clear();
+            DynamicOffsets.clear();
+            pResourceCache = nullptr;
+            BindPoint = VK_PIPELINE_BIND_POINT_MAX_ENUM;
+#ifdef _DEBUG
+            pDbgPipelineLayout = nullptr;
+#endif
+        }
+    };
+
+    // Prepares Vulkan descriptor sets for binding. Actual binding 
+    // may not be possible until draw command time because dynamic offsets are 
+    // set by the same Vulkan command. If there are no dynamic descriptors, this
+    // function also binds descriptor sets rightaway.
+    void PrepareDescriptorSets(DeviceContextVkImpl*    pCtxVkImpl,
+                               bool                    IsCompute,
+                               ShaderResourceCacheVk&  ResourceCache,
+                               DescriptorSetBindInfo&  BindInfo)const;
+
+    // Computes dynamic offsets and binds descriptor sets
+    void BindDescriptorSetsWithDynamicOffsets(DeviceContextVkImpl*    pCtxVkImpl,
+                                              DescriptorSetBindInfo&  BindInfo)const;
 
 private:
 
diff --git a/Graphics/GraphicsEngineVulkan/include/PipelineStateVkImpl.h b/Graphics/GraphicsEngineVulkan/include/PipelineStateVkImpl.h
index 0b0393dd..b355809d 100644
--- a/Graphics/GraphicsEngineVulkan/include/PipelineStateVkImpl.h
+++ b/Graphics/GraphicsEngineVulkan/include/PipelineStateVkImpl.h
@@ -64,11 +64,18 @@ public:
 
     virtual VkPipeline GetVkPipeline()const override final { return m_Pipeline; }
 
-    ShaderResourceCacheVk* CommitAndTransitionShaderResources(IShaderResourceBinding*   pShaderResourceBinding, 
-                                                              DeviceContextVkImpl*      pCtxVkImpl,
-                                                              bool                      CommitResources,
-                                                              bool                      TransitionResources)const;
-    
+    void CommitAndTransitionShaderResources(IShaderResourceBinding*                 pShaderResourceBinding, 
+                                            DeviceContextVkImpl*                    pCtxVkImpl,
+                                            bool                                    CommitResources,
+                                            bool                                    TransitionResources,
+                                            PipelineLayout::DescriptorSetBindInfo*  pDescrSetBindInfo)const;
+
+    void BindDescriptorSetsWithDynamicOffsets(DeviceContextVkImpl*                   pCtxVkImpl,
+                                              PipelineLayout::DescriptorSetBindInfo& BindInfo)
+    {
+        m_PipelineLayout.BindDescriptorSetsWithDynamicOffsets(pCtxVkImpl, BindInfo);
+    }
+
     const PipelineLayout& GetPipelineLayout()const{return m_PipelineLayout;}
     
     const ShaderResourceLayoutVk& GetShaderResLayout(Uint32 ShaderInd)const
diff --git a/Graphics/GraphicsEngineVulkan/src/DeviceContextVkImpl.cpp b/Graphics/GraphicsEngineVulkan/src/DeviceContextVkImpl.cpp
index 1f9abad0..9062d962 100644
--- a/Graphics/GraphicsEngineVulkan/src/DeviceContextVkImpl.cpp
+++ b/Graphics/GraphicsEngineVulkan/src/DeviceContextVkImpl.cpp
@@ -91,7 +91,6 @@ namespace Diligent
             bIsDeferred ? Attribs.DeferredCtxDynamicHeapPageSize : Attribs.ImmediateCtxDynamicHeapPageSize
         }
     {
-        m_DynamicBufferOffsets.reserve(64);
     }
 
     DeviceContextVkImpl::~DeviceContextVkImpl()
@@ -213,10 +212,8 @@ namespace Diligent
                 CommitScissorRects();
             }
         }
-        
-#if 0
-        m_pCommittedResourceCache = nullptr;
-#endif
+
+        m_DesrSetBindInfo.Reset();
     }
 
     void DeviceContextVkImpl::TransitionShaderResources(IPipelineState *pPipelineState, IShaderResourceBinding *pShaderResourceBinding)
@@ -224,7 +221,7 @@ namespace Diligent
         VERIFY_EXPR(pPipelineState != nullptr);
 
         auto *pPipelineStateVk = ValidatedCast<PipelineStateVkImpl>(pPipelineState);
-        pPipelineStateVk->CommitAndTransitionShaderResources(pShaderResourceBinding, this, false, true);
+        pPipelineStateVk->CommitAndTransitionShaderResources(pShaderResourceBinding, this, false, true, nullptr);
     }
 
     void DeviceContextVkImpl::CommitShaderResources(IShaderResourceBinding *pShaderResourceBinding, Uint32 Flags)
@@ -233,7 +230,7 @@ namespace Diligent
             return;
 
         auto *pPipelineStateVk = m_pPipelineState.RawPtr<PipelineStateVkImpl>();
-        m_pCommittedResourceCache = pPipelineStateVk->CommitAndTransitionShaderResources(pShaderResourceBinding, this, true, (Flags & COMMIT_SHADER_RESOURCES_FLAG_TRANSITION_RESOURCES)!=0);
+        pPipelineStateVk->CommitAndTransitionShaderResources(pShaderResourceBinding, this, true, (Flags & COMMIT_SHADER_RESOURCES_FLAG_TRANSITION_RESOURCES) != 0, &m_DesrSetBindInfo);
     }
 
     void DeviceContextVkImpl::SetStencilRef(Uint32 StencilRef)
@@ -423,11 +420,9 @@ namespace Diligent
         else
             CommitVkVertexBuffers();
 
+        if(!m_DesrSetBindInfo.DynamicOffsets.empty())
+            pPipelineStateVk->BindDescriptorSetsWithDynamicOffsets(this, m_DesrSetBindInfo);
 #if 0
-        if(m_pCommittedResourceCache != nullptr)
-        {
-            pPipelineStateVk->GetRootSignature().CommitRootViews(*m_pCommittedResourceCache, GraphCtx, false, m_ContextId);
-        }
 #ifdef _DEBUG
         else
         {
@@ -454,10 +449,10 @@ namespace Diligent
         {
             if( auto *pBufferVk = ValidatedCast<BufferVkImpl>(DrawAttribs.pIndirectDrawAttribs) )
             {
-//#ifdef _DEBUG
-//                if(pBufferVk->GetDesc().Usage == USAGE_DYNAMIC)
-//                    pBufferVk->DbgVerifyDynamicAllocation(m_ContextId);
-//#endif
+#ifdef _DEBUG
+                if(pBufferVk->GetDesc().Usage == USAGE_DYNAMIC)
+                    pBufferVk->DbgVerifyDynamicAllocation(m_ContextId);
+#endif
                 if(!pBufferVk->CheckAccessFlags(VK_ACCESS_INDIRECT_COMMAND_READ_BIT))
                     BufferMemoryBarrier(*pBufferVk, VK_ACCESS_INDIRECT_COMMAND_READ_BIT);
 
@@ -815,6 +810,7 @@ namespace Diligent
         ReleaseStaleContextResources(SubmittedCmdBuffNumber, SubmittedFenceValue, CompletedFenceValue);
 
         m_State = ContextState{};
+        m_DesrSetBindInfo.Reset();
         m_CommandBuffer.Reset();
         m_pPipelineState.Release(); 
     }
@@ -832,6 +828,7 @@ namespace Diligent
 
         TDeviceContextBase::InvalidateState();
         m_State = ContextState{};
+        m_DesrSetBindInfo.Reset();
         VERIFY(m_CommandBuffer.GetState().RenderPass == VK_NULL_HANDLE, "Invalidating context with unifinished render pass");
         m_CommandBuffer.Reset();
     }
@@ -1115,6 +1112,7 @@ namespace Diligent
 
         m_CommandBuffer.Reset();
         m_State = ContextState{};
+        m_DesrSetBindInfo.Reset();
         m_pPipelineState.Release();
 
         InvalidateState();
diff --git a/Graphics/GraphicsEngineVulkan/src/PipelineLayout.cpp b/Graphics/GraphicsEngineVulkan/src/PipelineLayout.cpp
index 6ec059fe..8c746f5d 100644
--- a/Graphics/GraphicsEngineVulkan/src/PipelineLayout.cpp
+++ b/Graphics/GraphicsEngineVulkan/src/PipelineLayout.cpp
@@ -425,12 +425,12 @@ void PipelineLayout::AllocateDynamicDescriptorSet(DeviceContextVkImpl*    pCtxVk
     }
 }
 
-void PipelineLayout::BindDescriptorSets(DeviceContextVkImpl*    pCtxVkImpl,
-                                        bool                    IsCompute,
-                                        ShaderResourceCacheVk&  ResourceCache)const
+void PipelineLayout::PrepareDescriptorSets(DeviceContextVkImpl*    pCtxVkImpl,
+                                           bool                    IsCompute,
+                                           ShaderResourceCacheVk&  ResourceCache,
+                                           DescriptorSetBindInfo&  BindInfo)const
 {
-    uint32_t SetCount = 0;
-    std::array<VkDescriptorSet, 2> vkSets = {};
+    BindInfo.vkSets.clear();
 
     VERIFY(m_LayoutMgr.GetDescriptorSet(SHADER_VARIABLE_TYPE_STATIC).SetIndex == m_LayoutMgr.GetDescriptorSet(SHADER_VARIABLE_TYPE_MUTABLE).SetIndex, 
            "Static and mutable variables are expected to share the same descriptor set");
@@ -440,36 +440,74 @@ void PipelineLayout::BindDescriptorSets(DeviceContextVkImpl*    pCtxVkImpl,
         const auto &Set = m_LayoutMgr.GetDescriptorSet(VarType);
         if(Set.SetIndex >= 0)
         {
-            SetCount = std::max(SetCount, static_cast<uint32_t>(Set.SetIndex+1));
-            VERIFY_EXPR(vkSets[Set.SetIndex] == VK_NULL_HANDLE);
-            vkSets[Set.SetIndex] = ResourceCache.GetDescriptorSet(Set.SetIndex).GetVkDescriptorSet();
-            VERIFY(vkSets[Set.SetIndex] != VK_NULL_HANDLE, "Descriptor set must not be null");
+            if(Set.SetIndex >= BindInfo.vkSets.size())
+                BindInfo.vkSets.resize(Set.SetIndex + 1);
+            VERIFY_EXPR(BindInfo.vkSets[Set.SetIndex] == VK_NULL_HANDLE);
+            BindInfo.vkSets[Set.SetIndex] = ResourceCache.GetDescriptorSet(Set.SetIndex).GetVkDescriptorSet();
+            VERIFY(BindInfo.vkSets[Set.SetIndex] != VK_NULL_HANDLE, "Descriptor set must not be null");
         }
         TotalDynamicDescriptors += Set.NumDynamicDescriptors;
     }
 
 #ifdef _DEBUG
-    for (uint32_t i = 0; i < SetCount; ++i)
-        VERIFY(vkSets[i] != VK_NULL_HANDLE, "Descriptor set must not be null");
+    for (const auto& set : BindInfo.vkSets)
+        VERIFY(set != VK_NULL_HANDLE, "Descriptor set must not be null");
 #endif
 
-    auto& DynamicOffsets = pCtxVkImpl->GetDynamicBufferOffsets();
-    DynamicOffsets.resize(TotalDynamicDescriptors);
-    ResourceCache.GetDynamicBufferOffsets(pCtxVkImpl->GetContextId(), DynamicOffsets);
+    BindInfo.DynamicOffsets.resize(TotalDynamicDescriptors);
+    BindInfo.BindPoint = IsCompute ? VK_PIPELINE_BIND_POINT_COMPUTE : VK_PIPELINE_BIND_POINT_GRAPHICS;
+    BindInfo.pResourceCache = &ResourceCache;
+#ifdef _DEBUG
+    BindInfo.pDbgPipelineLayout = this;
+#endif
+
+    if(TotalDynamicDescriptors == 0)
+    {
+        // There are no dynamic descriptors, so we can bind descriptor sets right now
+        auto& CmdBuffer = pCtxVkImpl->GetCommandBuffer();
+        CmdBuffer.BindDescriptorSets(BindInfo.BindPoint,
+                                     m_LayoutMgr.GetVkPipelineLayout(), 
+                                     0, // First set
+                                     static_cast<uint32_t>(BindInfo.vkSets.size()),
+                                     !BindInfo.vkSets.empty() ? BindInfo.vkSets.data() : nullptr,
+                                     0, 
+                                     nullptr);
+    }
+}
+
+void PipelineLayout::BindDescriptorSetsWithDynamicOffsets(DeviceContextVkImpl*    pCtxVkImpl,
+                                                          DescriptorSetBindInfo&  BindInfo)const
+{
+    VERIFY(BindInfo.pDbgPipelineLayout != nullptr, "Pipeline layout is not initialized, which most likely means that CommitShaderResources() has never been called");
+    VERIFY(BindInfo.pDbgPipelineLayout->IsSameAs(*this), "Inconsistent pipeline layout");
+    VERIFY(!BindInfo.DynamicOffsets.empty(), "This function should only be called for pipelines that contain dynamic descriptors");
+
+    VERIFY_EXPR(BindInfo.pResourceCache != nullptr);
+#ifdef _DEBUG
+    Uint32 TotalDynamicDescriptors = 0;
+    for (SHADER_VARIABLE_TYPE VarType = SHADER_VARIABLE_TYPE_MUTABLE; VarType <= SHADER_VARIABLE_TYPE_DYNAMIC; VarType = static_cast<SHADER_VARIABLE_TYPE>(VarType + 1))
+    {
+        const auto &Set = m_LayoutMgr.GetDescriptorSet(VarType);
+        TotalDynamicDescriptors += Set.NumDynamicDescriptors;
+    }
+    VERIFY(BindInfo.DynamicOffsets.size() == TotalDynamicDescriptors, "Incosistent dynamic buffer size");
+#endif
+
+    BindInfo.pResourceCache->GetDynamicBufferOffsets(pCtxVkImpl->GetContextId(), BindInfo.DynamicOffsets);
 
     auto& CmdBuffer = pCtxVkImpl->GetCommandBuffer();
     // vkCmdBindDescriptorSets causes the sets numbered [firstSet .. firstSet+descriptorSetCount-1] to use the 
     // bindings stored in pDescriptorSets[0 .. descriptorSetCount-1] for subsequent rendering commands 
     // (either compute or graphics, according to the pipelineBindPoint). Any bindings that were previously 
     // applied via these sets are no longer valid (13.2.5)
-    CmdBuffer.BindDescriptorSets(IsCompute ? VK_PIPELINE_BIND_POINT_COMPUTE : VK_PIPELINE_BIND_POINT_GRAPHICS,
+    CmdBuffer.BindDescriptorSets(BindInfo.BindPoint,
                                  m_LayoutMgr.GetVkPipelineLayout(), 
                                  0, // First set
-                                 SetCount,
-                                 vkSets.data(),
+                                 static_cast<uint32_t>(BindInfo.vkSets.size()),
+                                 !BindInfo.vkSets.empty() ? BindInfo.vkSets.data() : nullptr,
                                  // dynamicOffsetCount must equal the total number of dynamic descriptors in the sets being bound (13.2.5)
-                                 static_cast<uint32_t>(DynamicOffsets.size()),
-                                 DynamicOffsets.data());
+                                 static_cast<uint32_t>(BindInfo.DynamicOffsets.size()),
+                                 !BindInfo.DynamicOffsets.empty() ? BindInfo.DynamicOffsets.data() : nullptr);
 }
 
 }
diff --git a/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp b/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp
index 72d1b1c2..9f36dcb1 100644
--- a/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp
+++ b/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp
@@ -477,13 +477,14 @@ bool PipelineStateVkImpl::IsCompatibleWith(const IPipelineState *pPSO)const
 }
 
 
-ShaderResourceCacheVk* PipelineStateVkImpl::CommitAndTransitionShaderResources(IShaderResourceBinding*   pShaderResourceBinding, 
-                                                                               DeviceContextVkImpl*      pCtxVkImpl,
-                                                                               bool                      CommitResources,
-                                                                               bool                      TransitionResources)const
+void PipelineStateVkImpl::CommitAndTransitionShaderResources(IShaderResourceBinding*                pShaderResourceBinding, 
+                                                             DeviceContextVkImpl*                   pCtxVkImpl,
+                                                             bool                                   CommitResources,
+                                                             bool                                   TransitionResources,
+                                                            PipelineLayout::DescriptorSetBindInfo*  pDescrSetBindInfo)const
 {
     if(!m_HasStaticResources && !m_HasNonStaticResources)
-        return nullptr;
+        return;
 
 #ifdef VERIFY_SHADER_BINDINGS
     if (pShaderResourceBinding == nullptr && m_HasNonStaticResources)
@@ -502,7 +503,7 @@ ShaderResourceCacheVk* PipelineStateVkImpl::CommitAndTransitionShaderResources(I
         if ( IsIncompatibleWith(pRefPSO) )
         {
             LOG_ERROR_MESSAGE("Shader resource binding is incompatible with the pipeline state \"", m_Desc.Name, "\". Operation will be ignored.");
-            return nullptr;
+            return;
         }
     }
 #endif
@@ -550,16 +551,15 @@ ShaderResourceCacheVk* PipelineStateVkImpl::CommitAndTransitionShaderResources(I
         {
             m_ShaderResourceLayouts[s].CommitDynamicResources(ResourceCache);
         }
-        // Bind descriptor sets
-        m_PipelineLayout.BindDescriptorSets(pCtxVkImpl, m_Desc.IsComputePipeline, ResourceCache);
+        // Prepare descriptor sets, and also bind them if there are no dynamic descriptors
+        VERIFY_EXPR(pDescrSetBindInfo != nullptr);
+        m_PipelineLayout.PrepareDescriptorSets(pCtxVkImpl, m_Desc.IsComputePipeline, ResourceCache, *pDescrSetBindInfo);
     }
     else
     {
         VERIFY(TransitionResources, "Resources should be transitioned or committed or both");
         ResourceCache.TransitionResources<false>(pCtxVkImpl);
     }
-
-    return &ResourceCache;
 }
 
 }