fixed bug in root signature cache

4 files changed, 62 insertions, 65 deletions

diff --git a/Graphics/GraphicsEngineD3D12/include/RenderDeviceD3D12Impl.hpp b/Graphics/GraphicsEngineD3D12/include/RenderDeviceD3D12Impl.hpp
index b369299f..9a53c556 100644
--- a/Graphics/GraphicsEngineD3D12/include/RenderDeviceD3D12Impl.hpp
+++ b/Graphics/GraphicsEngineD3D12/include/RenderDeviceD3D12Impl.hpp
@@ -172,7 +172,7 @@ public:
                                                               RESOURCE_STATE        InitialState,
                                                               ITopLevelAS**         ppTLAS) override final;
 
-    void CreateRootSignature(const RefCntAutoPtr<class PipelineResourceSignatureD3D12Impl>* ppSignatures, Uint32 SignatureCount, RootSignatureD3D12** ppRootSig);
+    void CreateRootSignature(const RefCntAutoPtr<class PipelineResourceSignatureD3D12Impl>* ppSignatures, Uint32 SignatureCount, size_t Hash, RootSignatureD3D12** ppRootSig);
 
     RootSignatureCacheD3D12& GetRootSignatureCache() { return m_RootSignatureCache; }
 
diff --git a/Graphics/GraphicsEngineD3D12/include/RootSignature.hpp b/Graphics/GraphicsEngineD3D12/include/RootSignature.hpp
index 37e2d486..c7a3e2c4 100644
--- a/Graphics/GraphicsEngineD3D12/include/RootSignature.hpp
+++ b/Graphics/GraphicsEngineD3D12/include/RootSignature.hpp
@@ -31,7 +31,7 @@
 /// Declaration of Diligent::RootSignatureD3D12 class
 #include <array>
 #include <mutex>
-#include <unordered_set>
+#include <unordered_map>
 
 #include "D3D12TypeConversions.hpp"
 #include "ShaderResourceCacheD3D12.hpp"
@@ -44,6 +44,7 @@ namespace Diligent
 
 class RenderDeviceD3D12Impl;
 class PipelineResourceSignatureD3D12Impl;
+class RootSignatureCacheD3D12;
 
 /// Implementation of the Diligent::RootSignature class
 class RootSignatureD3D12 final : public ObjectBase<IObject>
@@ -52,11 +53,10 @@ public:
     RootSignatureD3D12(IReferenceCounters*                                      pRefCounters,
                        RenderDeviceD3D12Impl*                                   pDeviceD3D12Impl,
                        const RefCntAutoPtr<PipelineResourceSignatureD3D12Impl>* ppSignatures,
-                       Uint32                                                   SignatureCount);
+                       Uint32                                                   SignatureCount,
+                       size_t                                                   Hash);
     ~RootSignatureD3D12();
 
-    void Finalize();
-
     size_t GetHash() const { return m_Hash; }
 
     Uint32 GetSignatureCount() const { return m_SignatureCount; }
@@ -93,12 +93,14 @@ public:
     using SignatureArrayType = std::array<RefCntAutoPtr<PipelineResourceSignatureD3D12Impl>, MAX_RESOURCE_SIGNATURES>;
 
 private:
+    void Finalize(RenderDeviceD3D12Impl* pDeviceD3D12Impl);
+
     std::array<Uint16, MAX_RESOURCE_SIGNATURES> m_FirstRootIndex     = {};
     std::array<Uint16, MAX_RESOURCE_SIGNATURES> m_FirstRegisterSpace = {};
 
     Uint32 m_TotalSpacesUsed = 0;
 
-    size_t                       m_Hash = 0;
+    const size_t                 m_Hash;
     CComPtr<ID3D12RootSignature> m_pd3d12RootSignature;
 
     // The number of resource signatures used by this root signature
@@ -106,7 +108,7 @@ private:
     Uint8              m_SignatureCount = 0;
     SignatureArrayType m_Signatures     = {};
 
-    RenderDeviceD3D12Impl* m_pDeviceD3D12Impl;
+    RootSignatureCacheD3D12& m_Cache;
 };
 
 
@@ -159,25 +161,11 @@ public:
     void OnDestroyRootSig(RootSignatureD3D12* pRootSig);
 
 private:
-    struct RootSignatureHash
-    {
-        std::size_t operator()(const RootSignatureD3D12* Key) const noexcept
-        {
-            return Key->GetHash();
-        }
-    };
-
-    struct RootSignatureCompare
-    {
-        bool operator()(const RootSignatureD3D12* lhs, const RootSignatureD3D12* rhs) const noexcept;
-    };
-
-private:
     RenderDeviceD3D12Impl& m_DeviceD3D12Impl;
 
     std::mutex m_RootSigCacheGuard;
 
-    std::unordered_set<RootSignatureD3D12*, RootSignatureHash, RootSignatureCompare> m_RootSigCache;
+    std::unordered_multimap<size_t, RefCntWeakPtr<RootSignatureD3D12>> m_RootSigCache;
 };
 
 } // namespace Diligent
diff --git a/Graphics/GraphicsEngineD3D12/src/RenderDeviceD3D12Impl.cpp b/Graphics/GraphicsEngineD3D12/src/RenderDeviceD3D12Impl.cpp
index fb0f5c43..a7ea6769 100644
--- a/Graphics/GraphicsEngineD3D12/src/RenderDeviceD3D12Impl.cpp
+++ b/Graphics/GraphicsEngineD3D12/src/RenderDeviceD3D12Impl.cpp
@@ -852,9 +852,9 @@ DescriptorHeapAllocation RenderDeviceD3D12Impl::AllocateGPUDescriptors(D3D12_DES
     return m_GPUDescriptorHeaps[Type].Allocate(Count);
 }
 
-void RenderDeviceD3D12Impl::CreateRootSignature(const RefCntAutoPtr<PipelineResourceSignatureD3D12Impl>* ppSignatures, Uint32 SignatureCount, RootSignatureD3D12** ppRootSig)
+void RenderDeviceD3D12Impl::CreateRootSignature(const RefCntAutoPtr<PipelineResourceSignatureD3D12Impl>* ppSignatures, Uint32 SignatureCount, size_t Hash, RootSignatureD3D12** ppRootSig)
 {
-    RootSignatureD3D12* pRootSigD3D12(NEW_RC_OBJ(m_RootSignatureAllocator, "RootSignatureD3D12 instance", RootSignatureD3D12)(this, ppSignatures, SignatureCount));
+    RootSignatureD3D12* pRootSigD3D12(NEW_RC_OBJ(m_RootSignatureAllocator, "RootSignatureD3D12 instance", RootSignatureD3D12)(this, ppSignatures, SignatureCount, Hash));
     pRootSigD3D12->AddRef();
     *ppRootSig = pRootSigD3D12;
 }
diff --git a/Graphics/GraphicsEngineD3D12/src/RootSignature.cpp b/Graphics/GraphicsEngineD3D12/src/RootSignature.cpp
index b4f0e486..4eb16933 100644
--- a/Graphics/GraphicsEngineD3D12/src/RootSignature.cpp
+++ b/Graphics/GraphicsEngineD3D12/src/RootSignature.cpp
@@ -41,10 +41,12 @@ namespace Diligent
 RootSignatureD3D12::RootSignatureD3D12(IReferenceCounters*                                      pRefCounters,
                                        RenderDeviceD3D12Impl*                                   pDeviceD3D12Impl,
                                        const RefCntAutoPtr<PipelineResourceSignatureD3D12Impl>* ppSignatures,
-                                       Uint32                                                   SignatureCount) :
+                                       Uint32                                                   SignatureCount,
+                                       size_t                                                   Hash) :
     ObjectBase<IObject>{pRefCounters},
     m_SignatureCount{static_cast<Uint8>(SignatureCount)},
-    m_pDeviceD3D12Impl{pDeviceD3D12Impl}
+    m_Hash{Hash},
+    m_Cache{pDeviceD3D12Impl->GetRootSignatureCache()}
 {
     VERIFY(m_SignatureCount == SignatureCount, "Signature count (", SignatureCount, ") exceeds maximum representable value");
 
@@ -58,25 +60,15 @@ RootSignatureD3D12::RootSignatureD3D12(IReferenceCounters*
         }
     }
 
-    if (m_SignatureCount > 0)
-    {
-        HashCombine(m_Hash, m_SignatureCount);
-        for (Uint32 i = 0; i < m_SignatureCount; ++i)
-        {
-            if (m_Signatures[i] != nullptr)
-                HashCombine(m_Hash, m_Signatures[i]->GetHash());
-            else
-                HashCombine(m_Hash, 0);
-        }
-    }
+    Finalize(pDeviceD3D12Impl);
 }
 
 RootSignatureD3D12::~RootSignatureD3D12()
 {
-    m_pDeviceD3D12Impl->GetRootSignatureCache().OnDestroyRootSig(this);
+    m_Cache.OnDestroyRootSig(this);
 }
 
-void RootSignatureD3D12::Finalize()
+void RootSignatureD3D12::Finalize(RenderDeviceD3D12Impl* pDeviceD3D12Impl)
 {
     VERIFY(m_pd3d12RootSignature == nullptr, "This root signature is already initialized");
 
@@ -222,14 +214,13 @@ void RootSignatureD3D12::Finalize()
     }
     CHECK_D3D_RESULT_THROW(hr, "Failed to serialize root signature");
 
-    auto* pd3d12Device = m_pDeviceD3D12Impl->GetD3D12Device();
+    auto* pd3d12Device = pDeviceD3D12Impl->GetD3D12Device();
 
     hr = pd3d12Device->CreateRootSignature(0, signature->GetBufferPointer(), signature->GetBufferSize(), __uuidof(m_pd3d12RootSignature), reinterpret_cast<void**>(static_cast<ID3D12RootSignature**>(&m_pd3d12RootSignature)));
     CHECK_D3D_RESULT_THROW(hr, "Failed to create root signature");
 }
 
 
-
 LocalRootSignatureD3D12::LocalRootSignatureD3D12(const char* pCBName, Uint32 ShaderRecordSize) :
     m_pName{pCBName},
     m_ShaderRecordSize{ShaderRecordSize}
@@ -281,11 +272,12 @@ bool LocalRootSignatureD3D12::Create(ID3D12Device* pDevice, Uint32 RegisterSpace
 }
 
 
-
-bool RootSignatureCacheD3D12::RootSignatureCompare::operator()(const RootSignatureD3D12* lhs, const RootSignatureD3D12* rhs) const noexcept
+namespace
+{
+bool RootSignatureCompare(const RootSignatureD3D12* lhs, const RefCntAutoPtr<PipelineResourceSignatureD3D12Impl>* ppSignatures, Uint32 SignatureCount) noexcept
 {
     const Uint32 LSigCount = lhs->GetSignatureCount();
-    const Uint32 RSigCount = rhs->GetSignatureCount();
+    const Uint32 RSigCount = SignatureCount;
 
     if (LSigCount != RSigCount)
         return false;
@@ -293,7 +285,7 @@ bool RootSignatureCacheD3D12::RootSignatureCompare::operator()(const RootSignatu
     for (Uint32 i = 0; i < LSigCount; ++i)
     {
         auto* pLSig = lhs->GetSignature(i);
-        auto* pRSig = rhs->GetSignature(i);
+        auto* pRSig = ppSignatures[i].RawPtr();
 
         if (pLSig == pRSig)
             continue;
@@ -306,6 +298,7 @@ bool RootSignatureCacheD3D12::RootSignatureCompare::operator()(const RootSignatu
     }
     return true;
 }
+} // namespace
 
 RootSignatureCacheD3D12::RootSignatureCacheD3D12(RenderDeviceD3D12Impl& DeviceD3D12Impl) :
     m_DeviceD3D12Impl{DeviceD3D12Impl}
@@ -319,38 +312,54 @@ RootSignatureCacheD3D12::~RootSignatureCacheD3D12()
 
 RefCntAutoPtr<RootSignatureD3D12> RootSignatureCacheD3D12::GetRootSig(const RefCntAutoPtr<PipelineResourceSignatureD3D12Impl>* ppSignatures, Uint32 SignatureCount)
 {
-    RefCntAutoPtr<RootSignatureD3D12> pNewRootSig;
-    m_DeviceD3D12Impl.CreateRootSignature(ppSignatures, SignatureCount, &pNewRootSig);
-
-    if (pNewRootSig == nullptr)
-        return {};
-
-    RefCntAutoPtr<RootSignatureD3D12> Result;
-    bool                              Inserted = false;
+    size_t Hash = 0;
+    if (SignatureCount)
     {
-        std::lock_guard<std::mutex> Lock{m_RootSigCacheGuard};
+        HashCombine(Hash, SignatureCount);
+        for (Uint32 i = 0; i < SignatureCount; ++i)
+        {
+            if (ppSignatures[i] != nullptr)
+            {
+                VERIFY(ppSignatures[i]->GetDesc().BindingIndex == i, "Signature placed to another binding index");
+                HashCombine(Hash, ppSignatures[i]->GetHash());
+            }
+            else
+                HashCombine(Hash, 0);
+        }
+    }
 
-        auto IterAndFlag = m_RootSigCache.insert(pNewRootSig.RawPtr());
-        Inserted         = IterAndFlag.second;
+    std::lock_guard<std::mutex> Lock{m_RootSigCacheGuard};
 
-        if (Inserted)
+    auto Range = m_RootSigCache.equal_range(Hash);
+    for (auto Iter = Range.first; Iter != Range.second; ++Iter)
+    {
+        if (auto Ptr = Iter->second.Lock())
         {
-            pNewRootSig->Finalize();
-            Result = std::move(pNewRootSig);
+            if (RootSignatureCompare(Ptr, ppSignatures, SignatureCount))
+                return Ptr;
         }
-        else
-            Result = *IterAndFlag.first;
     }
-    return Result;
+
+    RefCntAutoPtr<RootSignatureD3D12> pNewRootSig;
+    m_DeviceD3D12Impl.CreateRootSignature(ppSignatures, SignatureCount, Hash, &pNewRootSig);
+
+    m_RootSigCache.emplace(Hash, pNewRootSig);
+    return pNewRootSig;
 }
 
 void RootSignatureCacheD3D12::OnDestroyRootSig(RootSignatureD3D12* pRootSig)
 {
     std::lock_guard<std::mutex> Lock{m_RootSigCacheGuard};
 
-    auto Iter = m_RootSigCache.find(pRootSig);
-    if (Iter != m_RootSigCache.end() && *Iter == pRootSig)
-        m_RootSigCache.erase(Iter);
+    auto Range = m_RootSigCache.equal_range(pRootSig->GetHash());
+
+    for (auto Iter = Range.first; Iter != Range.second;)
+    {
+        if (!Iter->second.IsValid())
+            Iter = m_RootSigCache.erase(Iter);
+        else
+            ++Iter;
+    }
 }
 
 } // namespace Diligent