Implemented VK shader resource cache initialization

5 files changed, 109 insertions, 145 deletions

diff --git a/Graphics/GraphicsEngineVulkan/include/DescriptorPoolManager.h b/Graphics/GraphicsEngineVulkan/include/DescriptorPoolManager.h
index 80291531..907f369c 100644
--- a/Graphics/GraphicsEngineVulkan/include/DescriptorPoolManager.h
+++ b/Graphics/GraphicsEngineVulkan/include/DescriptorPoolManager.h
@@ -34,36 +34,61 @@
 namespace Diligent
 {
 
+class DescriptorPoolManager;
+
 class DescriptorPoolAllocation
 {
 public:
     DescriptorPoolAllocation(VkDescriptorSet _Set,
-                             VulkanUtilities::VulkanDescriptorPool& _ParentPool) :
+                             VulkanUtilities::VulkanDescriptorPool& _ParentPool,
+                             DescriptorPoolManager& _ParentPoolMgr) :
         Set(_Set),
-        ParentPool(_ParentPool)
+        ParentPool(&_ParentPool),
+        ParentPoolMgr(&_ParentPoolMgr)
     {}
+    DescriptorPoolAllocation(){}
+
     DescriptorPoolAllocation(const DescriptorPoolAllocation&) = delete;
     DescriptorPoolAllocation& operator = (const DescriptorPoolAllocation&) = delete;
 
     DescriptorPoolAllocation(DescriptorPoolAllocation&& rhs) : 
-        Set(rhs.Set),
-        ParentPool(rhs.ParentPool)
+        Set          (rhs.Set),
+        ParentPool   (rhs.ParentPool),
+        ParentPoolMgr(rhs.ParentPoolMgr)
     {
-        rhs.Set = VK_NULL_HANDLE;
+        rhs.Set           = VK_NULL_HANDLE;
+        rhs.ParentPool    = nullptr;
+        rhs.ParentPoolMgr = nullptr;
     }
-    DescriptorPoolAllocation& operator = (DescriptorPoolAllocation&&) = delete;
+    
+    DescriptorPoolAllocation& operator = (DescriptorPoolAllocation&& rhs)
+    {
+        Release();
+
+        Set           = rhs.Set;
+        ParentPool    = rhs.ParentPool;
+        ParentPoolMgr = rhs.ParentPoolMgr;
+
+        rhs.Set           = VK_NULL_HANDLE;
+        rhs.ParentPool    = nullptr;
+        rhs.ParentPoolMgr = nullptr;
+        
+        return *this;
+    }
+
+    void Release();
 
     ~DescriptorPoolAllocation()
     {
-        VERIFY(Set == VK_NULL_HANDLE, "Allocation must be explicitly disposed through DescriptorPoolManager::DisposeAllocation");
+        Release();
     }
     
     VkDescriptorSet GetVkDescriptorSet()const {return Set;}
 
 private:
-    friend class DescriptorPoolManager;
-    VkDescriptorSet Set;
-    VulkanUtilities::VulkanDescriptorPool& ParentPool;
+    VkDescriptorSet Set = VK_NULL_HANDLE;
+    VulkanUtilities::VulkanDescriptorPool* ParentPool = nullptr;
+    DescriptorPoolManager* ParentPoolMgr = nullptr;
 };
 
 class DescriptorPoolManager
@@ -99,11 +124,13 @@ public:
     DescriptorPoolManager& operator = (DescriptorPoolManager&&)      = delete;
 
     DescriptorPoolAllocation Allocate(VkDescriptorSetLayout SetLayout);
-    void FreeAllocation(DescriptorPoolAllocation&& Allocation);
     void DisposeAllocations(uint64_t FenceValue);
     void ReleaseStaleAllocations(uint64_t LastCompletedFence);
 
 private:
+    friend class DescriptorPoolAllocation;
+    void FreeAllocation(VkDescriptorSet Set, VulkanUtilities::VulkanDescriptorPool& Pool);
+
     void CreateNewPool();
 
     const std::vector<VkDescriptorPoolSize> m_PoolSizes;
@@ -112,8 +139,8 @@ private:
 
     std::mutex m_Mutex;
     std::shared_ptr<const VulkanUtilities::VulkanLogicalDevice> m_LogicalDevice;
-    std::deque<VulkanUtilities::VulkanDescriptorPool> m_DescriptorPools;
-    std::vector<DescriptorPoolAllocation> m_ReleasedAllocations;
+    std::deque< std::unique_ptr<VulkanUtilities::VulkanDescriptorPool> > m_DescriptorPools;
+    std::vector< std::pair<VkDescriptorSet, VulkanUtilities::VulkanDescriptorPool*> > m_ReleasedAllocations;
 };
 
 }
diff --git a/Graphics/GraphicsEngineVulkan/include/RenderDeviceVkImpl.h b/Graphics/GraphicsEngineVulkan/include/RenderDeviceVkImpl.h
index 483845c2..8caa4213 100644
--- a/Graphics/GraphicsEngineVulkan/include/RenderDeviceVkImpl.h
+++ b/Graphics/GraphicsEngineVulkan/include/RenderDeviceVkImpl.h
@@ -81,10 +81,6 @@ public:
 
     //virtual void CreateBufferFromD3DResource(IVkResource *pVkBuffer, const BufferDesc& BuffDesc, IBuffer **ppBuffer)override final;
 
-/*
-    DescriptorHeapAllocation AllocateDescriptor( Vk_DESCRIPTOR_HEAP_TYPE Type, UINT Count = 1 );
-    DescriptorHeapAllocation AllocateGPUDescriptors( Vk_DESCRIPTOR_HEAP_TYPE Type, UINT Count = 1 );
-*/
     Uint64 GetCompletedFenceValue();
 	virtual Uint64 GetNextFenceValue() override final
     {
@@ -113,6 +109,16 @@ public:
     DynamicUploadHeap* RequestUploadHeap();
     void ReleaseUploadHeap(DynamicUploadHeap* pUploadHeap);
     */
+
+    DescriptorPoolAllocation AllocateDescriptorSet(VkDescriptorSetLayout SetLayout)
+    {
+        return m_DescriptorPools[0].Allocate(SetLayout);
+    }
+    DescriptorPoolAllocation AllocateDynamicDescriptorSet(VkDescriptorSetLayout SetLayout, Uint32 CtxId)
+    {
+        return m_DescriptorPools[1 + CtxId].Allocate(SetLayout);
+    }
+
     
     std::shared_ptr<const VulkanUtilities::VulkanInstance> GetVulkanInstance()const{return m_VulkanInstance;}
     const VulkanUtilities::VulkanPhysicalDevice &GetPhysicalDevice(){return *m_PhysicalDevice;}
diff --git a/Graphics/GraphicsEngineVulkan/include/ShaderResourceCacheVk.h b/Graphics/GraphicsEngineVulkan/include/ShaderResourceCacheVk.h
index 912a3f14..ff6aee2d 100644
--- a/Graphics/GraphicsEngineVulkan/include/ShaderResourceCacheVk.h
+++ b/Graphics/GraphicsEngineVulkan/include/ShaderResourceCacheVk.h
@@ -56,23 +56,11 @@
 //
 // Dynamic resources are not VkDescriptorSet 
 
-#include "DescriptorHeap.h"
+#include "DescriptorPoolManager.h"
 
 namespace Diligent
 {
-#if 0
-enum class CachedResourceType : Int32
-{
-    Unknown = -1,
-    CBV = 0,
-    TexSRV,
-    BufSRV,
-    TexUAV,
-    BufUAV,
-    Sampler,
-    NumTypes
-};
-#endif
+
 class ShaderResourceCacheVk
 {
 public:
@@ -117,14 +105,20 @@ public:
 
         VkDescriptorSet GetVkDescriptorSet()const
         {
-            UNSUPPORTED("Not yet implemented");
-            return VK_NULL_HANDLE;
+            return m_DescriptorSetAllocation.GetVkDescriptorSet();
+        }
+
+        void AssignDescriptorSetAllocation(DescriptorPoolAllocation&& Allocation)
+        {
+            VERIFY(m_NumResources > 0, "Descriptor set is empty");
+            m_DescriptorSetAllocation = std::move(Allocation);
         }
 
         const Uint32 m_NumResources = 0;
     private:
         
         Resource* const m_pResources = nullptr;
+        DescriptorPoolAllocation m_DescriptorSetAllocation;
     };
 
     inline DescriptorSet& GetDescriptorSet(Uint32 Index)
@@ -134,101 +128,7 @@ public:
     }
 
     inline Uint32 GetNumDescriptorSets()const{return m_NumSets; }
-#if 0
-    void SetDescriptorHeapSpace(DescriptorHeapAllocation &&CbcSrvUavHeapSpace, DescriptorHeapAllocation &&SamplerHeapSpace)
-    {
-        VERIFY(m_SamplerHeapSpace.GetCpuHandle().ptr == 0 && m_CbvSrvUavHeapSpace.GetCpuHandle().ptr == 0, "Space has already been allocated in GPU descriptor heaps");
-#ifdef _DEBUG
-        Uint32 NumSamplerDescriptors = 0, NumSrvCbvUavDescriptors = 0;
-        for (Uint32 rt = 0; rt < m_NumTables; ++rt)
-        {
-            auto &Tbl = GetDescriptorSet(rt);
-            if(Tbl.m_TableStartOffset != InvalidDescriptorOffset)
-            {
-                if(Tbl.DbgGetHeapType() == Vk_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV)
-                {
-                    VERIFY(Tbl.m_TableStartOffset == NumSrvCbvUavDescriptors, "Descriptor space allocation is not continuous");
-                    NumSrvCbvUavDescriptors = std::max(NumSrvCbvUavDescriptors, Tbl.m_TableStartOffset + Tbl.GetSize());
-                }
-                else
-                {
-                    VERIFY(Tbl.m_TableStartOffset == NumSamplerDescriptors, "Descriptor space allocation is not continuous");
-                    NumSamplerDescriptors = std::max(NumSamplerDescriptors, Tbl.m_TableStartOffset + Tbl.GetSize());
-                }
-            }
-        }
-        VERIFY(NumSrvCbvUavDescriptors == CbcSrvUavHeapSpace.GetNumHandles() || NumSrvCbvUavDescriptors == 0 && CbcSrvUavHeapSpace.GetCpuHandle(0).ptr == 0, "Unexpected descriptor heap allocation size" );
-        VERIFY(NumSamplerDescriptors == SamplerHeapSpace.GetNumHandles() || NumSamplerDescriptors == 0 && SamplerHeapSpace.GetCpuHandle(0).ptr == 0, "Unexpected descriptor heap allocation size" );
-#endif
-
-        m_CbvSrvUavHeapSpace = std::move(CbcSrvUavHeapSpace);
-        m_SamplerHeapSpace = std::move(SamplerHeapSpace);
-    }
-
-    IVkDescriptorHeap* GetSrvCbvUavDescriptorHeap(){return m_CbvSrvUavHeapSpace.GetDescriptorHeap();}
-    IVkDescriptorHeap* GetSamplerDescriptorHeap()  {return m_SamplerHeapSpace.GetDescriptorHeap();}
-
-    // Returns CPU descriptor handle of a shader visible descriptor heap allocation
-    template<Vk_DESCRIPTOR_HEAP_TYPE HeapType>
-    Vk_CPU_DESCRIPTOR_HANDLE GetShaderVisibleTableCPUDescriptorHandle(Uint32 RootParamInd, Uint32 OffsetFromTableStart = 0)
-    {
-        auto &RootParam = GetDescriptorSet(RootParamInd);
-        VERIFY(HeapType == RootParam.DbgGetHeapType(), "Invalid descriptor heap type");
-
-        Vk_CPU_DESCRIPTOR_HANDLE CPUDescriptorHandle = {0};
-        // Descriptor heap allocation is not assigned for dynamic resources or 
-        // in a special case when resource cache is used to store static 
-        // variable assignments for a shader. It is also not assigned to root views
-        if( RootParam.m_TableStartOffset != InvalidDescriptorOffset )
-        {
-            VERIFY(OffsetFromTableStart < RootParam.m_NumResources, "Offset is out of range");
-            if( HeapType == Vk_DESCRIPTOR_HEAP_TYPE_SAMPLER  )
-            {
-                VERIFY_EXPR(!m_SamplerHeapSpace.IsNull());
-                CPUDescriptorHandle = m_SamplerHeapSpace.GetCpuHandle(RootParam.m_TableStartOffset + OffsetFromTableStart);
-            }
-            else if( HeapType == Vk_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV )
-            {
-                VERIFY_EXPR(!m_CbvSrvUavHeapSpace.IsNull());
-                CPUDescriptorHandle = m_CbvSrvUavHeapSpace.GetCpuHandle(RootParam.m_TableStartOffset + OffsetFromTableStart);
-            }
-            else
-            {
-                UNEXPECTED("Unexpected descriptor heap type");
-            }
-        }
-
-        return CPUDescriptorHandle;
-    }
-
-    // Returns GPU descriptor handle of a shader visible descriptor table
-    template<Vk_DESCRIPTOR_HEAP_TYPE HeapType>
-    Vk_GPU_DESCRIPTOR_HANDLE GetShaderVisibleTableGPUDescriptorHandle(Uint32 RootParamInd, Uint32 OffsetFromTableStart = 0)
-    {
-        auto &RootParam = GetDescriptorSet(RootParamInd);
-        VERIFY(RootParam.m_TableStartOffset != InvalidDescriptorOffset, "GPU descriptor handle must never be requested for dynamic resources");
-        VERIFY(OffsetFromTableStart < RootParam.m_NumResources, "Offset is out of range");
-
-        Vk_GPU_DESCRIPTOR_HANDLE GPUDescriptorHandle = {0};
-        VERIFY( HeapType == RootParam.DbgGetHeapType(), "Invalid descriptor heap type");
-        if( HeapType == Vk_DESCRIPTOR_HEAP_TYPE_SAMPLER )
-        {
-            VERIFY_EXPR(!m_SamplerHeapSpace.IsNull());
-            GPUDescriptorHandle = m_SamplerHeapSpace.GetGpuHandle(RootParam.m_TableStartOffset + OffsetFromTableStart);
-        }
-        else if( HeapType == Vk_DESCRIPTOR_HEAP_TYPE_CBV_SRV_UAV )
-        {
-            VERIFY_EXPR(!m_CbvSrvUavHeapSpace.IsNull());
-            GPUDescriptorHandle = m_CbvSrvUavHeapSpace.GetGpuHandle(RootParam.m_TableStartOffset + OffsetFromTableStart);
-        }
-        else
-        {
-            UNEXPECTED("Unexpected descriptor heap type");
-        }
 
-        return GPUDescriptorHandle;
-    }
-#endif
 #ifdef _DEBUG
     // Only for debug purposes: indicates what types of resources are stored in the cache
     DbgCacheContentType DbgGetContentType()const{return m_DbgContentType;}
@@ -239,14 +139,6 @@ private:
     ShaderResourceCacheVk(ShaderResourceCacheVk&&) = delete;
     ShaderResourceCacheVk& operator = (const ShaderResourceCacheVk&) = delete;
     ShaderResourceCacheVk& operator = (ShaderResourceCacheVk&&) = delete;
-
-#if 0
-    // Allocation in a GPU-visible sampler descriptor heap
-    DescriptorHeapAllocation m_SamplerHeapSpace;
-    
-    // Allocation in a GPU-visible CBV/SRV/UAV descriptor heap
-    DescriptorHeapAllocation m_CbvSrvUavHeapSpace;
-#endif
     
     IMemoryAllocator *m_pAllocator=nullptr; 
     void *m_pMemory = nullptr;
diff --git a/Graphics/GraphicsEngineVulkan/src/DescriptorPoolManager.cpp b/Graphics/GraphicsEngineVulkan/src/DescriptorPoolManager.cpp
index f10a0a74..e9b4dd67 100644
--- a/Graphics/GraphicsEngineVulkan/src/DescriptorPoolManager.cpp
+++ b/Graphics/GraphicsEngineVulkan/src/DescriptorPoolManager.cpp
@@ -28,6 +28,19 @@
 namespace Diligent
 {
 
+void DescriptorPoolAllocation::Release()
+{
+    if (Set != VK_NULL_HANDLE)
+    {
+        VERIFY_EXPR(ParentPoolMgr != nullptr && ParentPool != nullptr);
+        ParentPoolMgr->FreeAllocation(Set, *ParentPool);
+
+        Set = VK_NULL_HANDLE;
+        ParentPoolMgr = nullptr;
+        ParentPool = nullptr;
+    }
+}
+
 void DescriptorPoolManager::CreateNewPool()
 {
     VkDescriptorPoolCreateInfo PoolCI = {};
@@ -40,7 +53,7 @@ void DescriptorPoolManager::CreateNewPool()
     PoolCI.maxSets = m_MaxSets;
     PoolCI.poolSizeCount = static_cast<uint32_t>(m_PoolSizes.size());
     PoolCI.pPoolSizes = m_PoolSizes.data();
-    m_DescriptorPools.emplace_front(m_LogicalDevice, PoolCI);
+    m_DescriptorPools.emplace_front( new VulkanUtilities::VulkanDescriptorPool(m_LogicalDevice, PoolCI) );
 }
 
 DescriptorPoolAllocation DescriptorPoolManager::Allocate(VkDescriptorSetLayout SetLayout)
@@ -52,7 +65,8 @@ DescriptorPoolAllocation DescriptorPoolManager::Allocate(VkDescriptorSetLayout S
     // Try all pools starting from the frontmost
     for(auto it = m_DescriptorPools.begin(); it != m_DescriptorPools.end(); ++it)
     {
-        auto Set = it->AllocateDescriptorSet(SetLayout);
+        auto& Pool = *(*it);
+        auto Set = Pool.AllocateDescriptorSet(SetLayout);
         if(Set != VK_NULL_HANDLE)
         {
             // Move the pool to the front
@@ -60,7 +74,7 @@ DescriptorPoolAllocation DescriptorPoolManager::Allocate(VkDescriptorSetLayout S
             {
                 std::swap(*it, m_DescriptorPools.front());
             }
-            return {Set, *it};
+            return {Set, Pool, *this};
         }
     }
 
@@ -68,19 +82,20 @@ DescriptorPoolAllocation DescriptorPoolManager::Allocate(VkDescriptorSetLayout S
     CreateNewPool();
     LOG_INFO_MESSAGE("Allocated new descriptor pool");
 
-    auto Set = m_DescriptorPools.front().AllocateDescriptorSet(SetLayout);
+    auto &NewPool = *m_DescriptorPools.front();
+    auto Set = NewPool.AllocateDescriptorSet(SetLayout);
     VERIFY(Set != VK_NULL_HANDLE, "Failed to allocate descriptor set");
 
-    return {Set, m_DescriptorPools.front() };
+    return {Set, NewPool, *this };
 }
 
-void DescriptorPoolManager::FreeAllocation(DescriptorPoolAllocation&& Allocation)
+void DescriptorPoolManager::FreeAllocation(VkDescriptorSet Set, VulkanUtilities::VulkanDescriptorPool& Pool)
 {
     std::unique_lock<std::mutex> Lock(m_Mutex, std::defer_lock);
     if (m_IsThreadSafe)
         Lock.lock();
 
-    m_ReleasedAllocations.emplace_back(std::move(Allocation));
+    m_ReleasedAllocations.emplace_back(std::make_pair(Set, &Pool));
 }
 
 void DescriptorPoolManager::DisposeAllocations(uint64_t FenceValue)
@@ -91,8 +106,7 @@ void DescriptorPoolManager::DisposeAllocations(uint64_t FenceValue)
 
     for(auto &Allocation : m_ReleasedAllocations)
     {
-        Allocation.ParentPool.DisposeDescriptorSet(Allocation.Set, FenceValue);
-        Allocation.Set = VK_NULL_HANDLE;
+        Allocation.second->DisposeDescriptorSet(Allocation.first, FenceValue);
     }
     m_ReleasedAllocations.clear();
 }
@@ -104,7 +118,7 @@ void DescriptorPoolManager::ReleaseStaleAllocations(uint64_t LastCompletedFence)
         Lock.lock();
 
     for(auto &Pool : m_DescriptorPools)
-        Pool.ReleaseDiscardedSets(LastCompletedFence);
+        Pool->ReleaseDiscardedSets(LastCompletedFence);
 }
 
 }
diff --git a/Graphics/GraphicsEngineVulkan/src/PipelineLayout.cpp b/Graphics/GraphicsEngineVulkan/src/PipelineLayout.cpp
index 6c7f82cc..78115ff0 100644
--- a/Graphics/GraphicsEngineVulkan/src/PipelineLayout.cpp
+++ b/Graphics/GraphicsEngineVulkan/src/PipelineLayout.cpp
@@ -525,6 +525,31 @@ void PipelineLayout::Finalize(const VulkanUtilities::VulkanLogicalDevice& Logica
 //http://diligentgraphics.com/diligent-engine/architecture/Vk/shader-resource-cache#Initializing-the-Cache-for-Shader-Resource-Binding-Object
 void PipelineLayout::InitResourceCache(RenderDeviceVkImpl *pDeviceVkImpl, ShaderResourceCacheVk& ResourceCache, IMemoryAllocator &CacheMemAllocator)const
 {
+    Uint32 NumSets = 0;
+    std::array<Uint32, 2> SetSizes = {};
+
+    DescriptorPoolAllocation SetAllocation;
+    const auto &StaticAndMutSet = m_LayoutMgr.GetDescriptorSet(SHADER_VARIABLE_TYPE_STATIC);
+    if(StaticAndMutSet.SetIndex >= 0)
+    {
+        NumSets = std::max(NumSets, static_cast<Uint32>(StaticAndMutSet.SetIndex + 1));
+        SetSizes[StaticAndMutSet.SetIndex] = StaticAndMutSet.TotalDescriptors;
+        SetAllocation = pDeviceVkImpl->AllocateDescriptorSet(StaticAndMutSet.VkLayout);
+    }
+
+    const auto &DynamicSet = m_LayoutMgr.GetDescriptorSet(SHADER_VARIABLE_TYPE_DYNAMIC);
+    if(DynamicSet.SetIndex >= 0)
+    {
+        NumSets = std::max(NumSets, static_cast<Uint32>(DynamicSet.SetIndex + 1));
+        SetSizes[DynamicSet.SetIndex] = DynamicSet.TotalDescriptors;
+    }
+  
+    ResourceCache.Initialize(CacheMemAllocator, NumSets, SetSizes.data());
+    if (StaticAndMutSet.SetIndex >= 0)
+    {
+        ResourceCache.GetDescriptorSet(StaticAndMutSet.SetIndex).AssignDescriptorSetAllocation(std::move(SetAllocation));
+    }
+
 #if 0
     // Get root table size for every root index
     // m_RootParams keeps root tables sorted by the array index, not the root index