Reworked shader memory alloction for resource layouts in D3D12 pipeline implementation

5 files changed, 68 insertions, 94 deletions

diff --git a/Graphics/GraphicsEngineD3D12/include/PipelineStateD3D12Impl.h b/Graphics/GraphicsEngineD3D12/include/PipelineStateD3D12Impl.h
index 698396b8..e4550356 100644
--- a/Graphics/GraphicsEngineD3D12/include/PipelineStateD3D12Impl.h
+++ b/Graphics/GraphicsEngineD3D12/include/PipelineStateD3D12Impl.h
@@ -66,7 +66,11 @@ public:
     
     const RootSignature& GetRootSignature()const{return m_RootSig;}
     
-    const ShaderResourceLayoutD3D12& GetShaderResLayout(SHADER_TYPE ShaderType)const;
+    const ShaderResourceLayoutD3D12& GetShaderResLayout(Uint32 ShaderInd)const
+    {
+        VERIFY_EXPR(ShaderInd < m_NumShaders);
+        return m_pShaderResourceLayouts[ShaderInd];
+    }
     
     bool dbgContainsShaderResources()const;
 
@@ -79,8 +83,6 @@ public:
 
 private:
 
-    void ParseShaderResourceLayout(IShader *pShader);
-
     /// D3D12 device
     CComPtr<ID3D12PipelineState> m_pd3d12PSO;
     RootSignature m_RootSig;
@@ -89,7 +91,7 @@ private:
     // Must be defined before default SRB
     SRBMemoryAllocator m_SRBMemAllocator;
 
-    ShaderResourceLayoutD3D12* m_pShaderResourceLayouts[6] = {};
+    ShaderResourceLayoutD3D12* m_pShaderResourceLayouts;
 
     // Do not use strong reference to avoid cyclic references
     // Default SRB must be defined after m_SRBMemAllocator
diff --git a/Graphics/GraphicsEngineD3D12/include/ShaderResourceBindingD3D12Impl.h b/Graphics/GraphicsEngineD3D12/include/ShaderResourceBindingD3D12Impl.h
index cf88a372..a87b4b91 100644
--- a/Graphics/GraphicsEngineD3D12/include/ShaderResourceBindingD3D12Impl.h
+++ b/Graphics/GraphicsEngineD3D12/include/ShaderResourceBindingD3D12Impl.h
@@ -47,28 +47,25 @@ public:
                                    bool                           IsPSOInternal);
     ~ShaderResourceBindingD3D12Impl();
 
-    virtual void QueryInterface( const Diligent::INTERFACE_ID &IID, IObject **ppInterface )override;
+    virtual void QueryInterface( const Diligent::INTERFACE_ID &IID, IObject** ppInterface )override;
 
-    virtual void BindResources(Uint32 ShaderFlags, IResourceMapping *pResMapping, Uint32 Flags)override;
+    virtual void BindResources(Uint32 ShaderFlags, IResourceMapping* pResMapping, Uint32 Flags)override;
 
-    virtual IShaderVariable *GetVariable(SHADER_TYPE ShaderType, const char *Name)override;
+    virtual IShaderVariable* GetVariable(SHADER_TYPE ShaderType, const char* Name)override;
 
-    ShaderResourceLayoutD3D12& GetResourceLayout(SHADER_TYPE ResType)
+    ShaderResourceLayoutD3D12& GetResourceLayout(Uint32 ResLayoutInd)
     {
-        auto ShaderInd = GetShaderTypeIndex(ResType);
-        auto ResLayoutInd = m_ResourceLayoutIndex[ShaderInd];
-        VERIFY(ResLayoutInd >= 0, "Shader resource layout is not initialized");
-        VERIFY_EXPR(ResLayoutInd < (Int32)m_NumShaders);
+        VERIFY_EXPR(ResLayoutInd < m_NumShaders);
         return m_pResourceLayouts[ResLayoutInd];
     }
     ShaderResourceCacheD3D12& GetResourceCache(){return m_ShaderResourceCache;}
 
 #ifdef VERIFY_SHADER_BINDINGS
-    void dbgVerifyResourceBindings(const PipelineStateD3D12Impl *pPSO);
+    void dbgVerifyResourceBindings(const PipelineStateD3D12Impl* pPSO);
 #endif
 
     bool StaticResourcesInitialized()const{return m_bStaticResourcesInitialized;}
-    void InitializeStaticResources(const PipelineStateD3D12Impl *pPSO);
+    void InitializeStaticResources(const PipelineStateD3D12Impl* pPSO);
 
 private:
 
@@ -77,7 +74,7 @@ private:
     // Resource layout index in m_ResourceLayouts[] array for every shader stage
     Int8 m_ResourceLayoutIndex[6] = {-1, -1, -1, -1, -1, -1};
     bool m_bStaticResourcesInitialized = false;
-    Uint32 m_NumShaders = 0;
+    const Uint32 m_NumShaders = 0;
 };
 
 }
diff --git a/Graphics/GraphicsEngineD3D12/include/ShaderResourceLayoutD3D12.h b/Graphics/GraphicsEngineD3D12/include/ShaderResourceLayoutD3D12.h
index 79f58da7..e1ee36b6 100644
--- a/Graphics/GraphicsEngineD3D12/include/ShaderResourceLayoutD3D12.h
+++ b/Graphics/GraphicsEngineD3D12/include/ShaderResourceLayoutD3D12.h
@@ -333,6 +333,8 @@ public:
 
     IObject& GetOwner(){return m_Owner;}
 
+    const ShaderResourcesD3D12& GetShaderResources(){return *m_pResources;}
+
 private:
     void InitVariablesHashMap();
 
diff --git a/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp b/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
index d4dcf2f6..0ff30050 100644
--- a/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
+++ b/Graphics/GraphicsEngineD3D12/src/PipelineStateD3D12Impl.cpp
@@ -59,24 +59,6 @@ private:
     std::array<D3D12_PRIMITIVE_TOPOLOGY_TYPE, PRIMITIVE_TOPOLOGY_NUM_TOPOLOGIES> m_Map;
 };
 
-void PipelineStateD3D12Impl::ParseShaderResourceLayout(IShader* pShader)
-{
-    VERIFY_EXPR(pShader);
-
-    auto ShaderType = pShader->GetDesc().ShaderType;
-    auto ShaderInd = GetShaderTypeIndex(ShaderType);
-    auto *pShaderD3D12 = ValidatedCast<ShaderD3D12Impl>(pShader);
-    
-    VERIFY(m_pShaderResourceLayouts[ShaderInd] == nullptr, "Shader resource layout has already been initialized");
-
-    auto pDeviceD3D12Impl = ValidatedCast<RenderDeviceD3D12Impl>(pShaderD3D12->GetDevice());
-    auto &ShaderResLayoutAllocator = GetRawAllocator();
-
-    auto *pRawMem = ALLOCATE(ShaderResLayoutAllocator, "Raw memory for ShaderResourceLayoutD3D12", sizeof(ShaderResourceLayoutD3D12));
-    m_pShaderResourceLayouts[ShaderInd] = new (pRawMem) ShaderResourceLayoutD3D12(*this, GetRawAllocator());
-    m_pShaderResourceLayouts[ShaderInd]->Initialize(pDeviceD3D12Impl->GetD3D12Device(), pShaderD3D12->GetShaderResources(), GetRawAllocator(), nullptr, 0, nullptr, &m_RootSig);
-}
-
 PipelineStateD3D12Impl :: PipelineStateD3D12Impl(IReferenceCounters*      pRefCounters,
                                                  RenderDeviceD3D12Impl*   pDeviceD3D12,
                                                  const PipelineStateDesc& PipelineDesc) : 
@@ -86,9 +68,23 @@ PipelineStateD3D12Impl :: PipelineStateD3D12Impl(IReferenceCounters*      pRefCo
     m_pDefaultShaderResBinding(nullptr, STDDeleter<ShaderResourceBindingD3D12Impl, FixedBlockMemoryAllocator>(pDeviceD3D12->GetSRBAllocator()) )
 {
     auto pd3d12Device = pDeviceD3D12->GetD3D12Device();
+    
+    m_RootSig.AllocateStaticSamplers( GetShaders(), GetNumShaders() );
+
+    auto& ShaderResLayoutAllocator = GetRawAllocator();
+    auto* pShaderResLayoutRawMem = ALLOCATE(ShaderResLayoutAllocator, "Raw memory for ShaderResourceLayoutD3D12", sizeof(ShaderResourceLayoutD3D12) * m_NumShaders);
+    m_pShaderResourceLayouts = reinterpret_cast<ShaderResourceLayoutD3D12*>(pShaderResLayoutRawMem);
+    for (Uint32 s=0; s < m_NumShaders; ++s)
+    {
+        auto* pShaderD3D12 = GetShader<ShaderD3D12Impl>(s);
+        new (m_pShaderResourceLayouts+s) ShaderResourceLayoutD3D12(*this, GetRawAllocator());
+        m_pShaderResourceLayouts[s].Initialize(pDeviceD3D12->GetD3D12Device(), pShaderD3D12->GetShaderResources(), GetRawAllocator(), nullptr, 0, nullptr, &m_RootSig);
+    }
+    m_RootSig.Finalize(pd3d12Device);
+
     if (PipelineDesc.IsComputePipeline)
     {
-        auto &ComputePipeline = PipelineDesc.ComputePipeline;
+        auto& ComputePipeline = PipelineDesc.ComputePipeline;
 
         if( ComputePipeline.pCS == nullptr )
             LOG_ERROR_AND_THROW("Compute shader is not set in the pipeline desc");
@@ -111,8 +107,6 @@ PipelineStateD3D12Impl :: PipelineStateD3D12Impl(IReferenceCounters*      pRefCo
         // The only valid bit is D3D12_PIPELINE_STATE_FLAG_TOOL_DEBUG, which can only be set on WARP devices.
         d3d12PSODesc.Flags = D3D12_PIPELINE_STATE_FLAG_NONE;
 
-        ParseShaderResourceLayout(ComputePipeline.pCS);
-        m_RootSig.Finalize(pd3d12Device);
         d3d12PSODesc.pRootSignature = m_RootSig.GetD3D12RootSignature();
 
         HRESULT hr = pd3d12Device->CreateComputePipelineState(&d3d12PSODesc, __uuidof(ID3D12PipelineState), reinterpret_cast<void**>( static_cast<ID3D12PipelineState**>(&m_pd3d12PSO)) );
@@ -123,34 +117,26 @@ PipelineStateD3D12Impl :: PipelineStateD3D12Impl(IReferenceCounters*      pRefCo
     {
         const auto& GraphicsPipeline = PipelineDesc.GraphicsPipeline;
         D3D12_GRAPHICS_PIPELINE_STATE_DESC d3d12PSODesc = {};
-
-        m_RootSig.AllocateStaticSamplers( GetShaders(), GetNumShaders() );
             
-#define INIT_SHADER(VarName, ExpectedType)\
-        if (GraphicsPipeline.p##VarName)                                            \
-        {                                                                           \
-            auto ShaderType = GraphicsPipeline.p##VarName->GetDesc().ShaderType;    \
-            if( ShaderType != ExpectedType )                                        \
-                LOG_ERROR_AND_THROW( GetShaderTypeLiteralName(ShaderType), " shader is provided while ", GetShaderTypeLiteralName(ExpectedType), " is expected");\
-            auto *pByteCode = ValidatedCast<ShaderD3D12Impl>(GraphicsPipeline.p##VarName)->GetShaderByteCode(); \
-            d3d12PSODesc.VarName.pShaderBytecode = pByteCode->GetBufferPointer();   \
-            d3d12PSODesc.VarName.BytecodeLength = pByteCode->GetBufferSize();       \
-            ParseShaderResourceLayout(GraphicsPipeline.p##VarName);                 \
-        }                                                                           \
-        else                                                                        \
-        {                                                                           \
-            d3d12PSODesc.VarName.pShaderBytecode = nullptr;                         \
-            d3d12PSODesc.VarName.BytecodeLength = 0;                                \
+        for (Uint32 s=0; s < m_NumShaders; ++s)
+        {
+            auto* pShaderD3D12 = GetShader<ShaderD3D12Impl>(s);
+            auto ShaderType = pShaderD3D12->GetDesc().ShaderType;
+            D3D12_SHADER_BYTECODE *pd3d12ShaderBytecode = nullptr;
+            switch(ShaderType)
+            {
+                case SHADER_TYPE_VERTEX:   pd3d12ShaderBytecode = &d3d12PSODesc.VS; break;
+                case SHADER_TYPE_PIXEL:    pd3d12ShaderBytecode = &d3d12PSODesc.PS; break;
+                case SHADER_TYPE_GEOMETRY: pd3d12ShaderBytecode = &d3d12PSODesc.GS; break;
+                case SHADER_TYPE_HULL:     pd3d12ShaderBytecode = &d3d12PSODesc.HS; break;
+                case SHADER_TYPE_DOMAIN:   pd3d12ShaderBytecode = &d3d12PSODesc.DS; break;
+                default: UNEXPECTED("Unexpected shader type");
+            }
+            auto *pByteCode = pShaderD3D12->GetShaderByteCode();
+            pd3d12ShaderBytecode->pShaderBytecode = pByteCode->GetBufferPointer();
+            pd3d12ShaderBytecode->BytecodeLength  = pByteCode->GetBufferSize();
         }
 
-        INIT_SHADER(VS, SHADER_TYPE_VERTEX);
-        INIT_SHADER(PS, SHADER_TYPE_PIXEL);
-        INIT_SHADER(GS, SHADER_TYPE_GEOMETRY);
-        INIT_SHADER(DS, SHADER_TYPE_DOMAIN);
-        INIT_SHADER(HS, SHADER_TYPE_HULL);
-#undef INIT_SHADER
-
-        m_RootSig.Finalize(pd3d12Device);
         d3d12PSODesc.pRootSignature = m_RootSig.GetD3D12RootSignature();
         
         memset(&d3d12PSODesc.StreamOutput, 0, sizeof(d3d12PSODesc.StreamOutput));
@@ -220,14 +206,14 @@ PipelineStateD3D12Impl :: PipelineStateD3D12Impl(IReferenceCounters*      pRefCo
         for (Uint32 s = 0; s < m_NumShaders; ++s)
         {
             std::array<SHADER_VARIABLE_TYPE, 3> AllowedVarTypes = { SHADER_VARIABLE_TYPE_STATIC, SHADER_VARIABLE_TYPE_MUTABLE, SHADER_VARIABLE_TYPE_DYNAMIC };
-            ShaderResLayoutDataSizes[s] = ShaderResourceLayoutD3D12::GetRequiredMemorySize(*m_pShaderResourceLayouts[s], AllowedVarTypes.data(), static_cast<Uint32>(AllowedVarTypes.size()));
+            ShaderResLayoutDataSizes[s] = ShaderResourceLayoutD3D12::GetRequiredMemorySize(m_pShaderResourceLayouts[s], AllowedVarTypes.data(), static_cast<Uint32>(AllowedVarTypes.size()));
         }
 
         auto CacheMemorySize = m_RootSig.GetResourceCacheRequiredMemSize();
         m_SRBMemAllocator.Initialize(PipelineDesc.SRBAllocationGranularity, m_NumShaders, ShaderResLayoutDataSizes.data(), 1, &CacheMemorySize);
     }
 
-    auto &SRBAllocator = pDeviceD3D12->GetSRBAllocator();
+    auto& SRBAllocator = pDeviceD3D12->GetSRBAllocator();
     // Default shader resource binding must be initialized after resource layouts are parsed!
     m_pDefaultShaderResBinding.reset( NEW_RC_OBJ(SRBAllocator, "ShaderResourceBindingD3D12Impl instance", ShaderResourceBindingD3D12Impl, this)(this, true) );
 
@@ -236,15 +222,12 @@ PipelineStateD3D12Impl :: PipelineStateD3D12Impl(IReferenceCounters*      pRefCo
 
 PipelineStateD3D12Impl::~PipelineStateD3D12Impl()
 {
-    auto &ShaderResLayoutAllocator = GetRawAllocator();
-    for(Int32 l = 0; l < _countof(m_pShaderResourceLayouts); ++l)
+    auto& ShaderResLayoutAllocator = GetRawAllocator();
+    for(Uint32 s = 0; s < m_NumShaders; ++s)
     {
-        if (m_pShaderResourceLayouts[l] != nullptr)
-        {
-            m_pShaderResourceLayouts[l]->~ShaderResourceLayoutD3D12();
-            ShaderResLayoutAllocator.Free(m_pShaderResourceLayouts[l]);
-        }
+        m_pShaderResourceLayouts[s].~ShaderResourceLayoutD3D12();
     }
+    ShaderResLayoutAllocator.Free(m_pShaderResourceLayouts);
 
     // D3D12 object can only be destroyed when it is no longer used by the GPU
     auto *pDeviceD3D12Impl = ValidatedCast<RenderDeviceD3D12Impl>(GetDevice());
@@ -272,7 +255,7 @@ void PipelineStateD3D12Impl::BindShaderResources(IResourceMapping* pResourceMapp
 void PipelineStateD3D12Impl::CreateShaderResourceBinding(IShaderResourceBinding** ppShaderResourceBinding)
 {
     auto *pRenderDeviceD3D12 = ValidatedCast<RenderDeviceD3D12Impl>( GetDevice() );
-    auto &SRBAllocator = pRenderDeviceD3D12->GetSRBAllocator();
+    auto& SRBAllocator = pRenderDeviceD3D12->GetSRBAllocator();
     auto pResBindingD3D12 = NEW_RC_OBJ(SRBAllocator, "ShaderResourceBindingD3D12Impl instance", ShaderResourceBindingD3D12Impl)(this, false);
     pResBindingD3D12->QueryInterface(IID_ShaderResourceBinding, reinterpret_cast<IObject**>(ppShaderResourceBinding));
 }
@@ -300,8 +283,8 @@ bool PipelineStateD3D12Impl::IsCompatibleWith(const IPipelineState* pPSO)const
         {
             for (Uint32 s = 0; s < m_NumShaders; ++s)
             {
-                auto *pShader0 = ValidatedCast<ShaderD3D12Impl>(m_ppShaders[s]);
-                auto *pShader1 = ValidatedCast<ShaderD3D12Impl>(pPSOD3D12->m_ppShaders[s]);
+                auto* pShader0 = GetShader<const ShaderD3D12Impl>(s);
+                auto* pShader1 = pPSOD3D12->GetShader<const ShaderD3D12Impl>(s);
                 if (pShader0->GetDesc().ShaderType != pShader1->GetDesc().ShaderType)
                 {
                     IsCompatibleShaders = false;
@@ -325,13 +308,6 @@ bool PipelineStateD3D12Impl::IsCompatibleWith(const IPipelineState* pPSO)const
     return IsSameRootSignature;
 }
 
-const ShaderResourceLayoutD3D12& PipelineStateD3D12Impl::GetShaderResLayout(SHADER_TYPE ShaderType)const 
-{
-    auto ShaderInd = GetShaderTypeIndex(ShaderType);
-    VERIFY_EXPR(m_pShaderResourceLayouts[ShaderInd] != nullptr);
-    return *m_pShaderResourceLayouts[ShaderInd];
-}
-
 ShaderResourceCacheD3D12* PipelineStateD3D12Impl::CommitAndTransitionShaderResources(IShaderResourceBinding* pShaderResourceBinding, 
                                                                                      CommandContext&         Ctx,
                                                                                      bool                    CommitResources,
@@ -370,7 +346,7 @@ ShaderResourceCacheD3D12* PipelineStateD3D12Impl::CommitAndTransitionShaderResou
 #endif
 
     auto *pDeviceD3D12Impl = ValidatedCast<RenderDeviceD3D12Impl>( GetDevice() );
-    auto &ResourceCache = pResBindingD3D12Impl->GetResourceCache();
+    auto& ResourceCache = pResBindingD3D12Impl->GetResourceCache();
     if(CommitResources)
     {
         if(m_Desc.IsComputePipeline)
diff --git a/Graphics/GraphicsEngineD3D12/src/ShaderResourceBindingD3D12Impl.cpp b/Graphics/GraphicsEngineD3D12/src/ShaderResourceBindingD3D12Impl.cpp
index 26e37b48..c40f0a78 100644
--- a/Graphics/GraphicsEngineD3D12/src/ShaderResourceBindingD3D12Impl.cpp
+++ b/Graphics/GraphicsEngineD3D12/src/ShaderResourceBindingD3D12Impl.cpp
@@ -34,10 +34,10 @@ ShaderResourceBindingD3D12Impl::ShaderResourceBindingD3D12Impl(IReferenceCounter
                                                                PipelineStateD3D12Impl* pPSO,
                                                                bool                    IsPSOInternal) :
     TBase( pRefCounters, pPSO, IsPSOInternal ),
-    m_ShaderResourceCache(ShaderResourceCacheD3D12::DbgCacheContentType::SRBResources)
+    m_ShaderResourceCache(ShaderResourceCacheD3D12::DbgCacheContentType::SRBResources),
+    m_NumShaders(pPSO->GetNumShaders())
 {
     auto* ppShaders = pPSO->GetShaders();
-    m_NumShaders = pPSO->GetNumShaders();
 
     auto* pRenderDeviceD3D12Impl = ValidatedCast<RenderDeviceD3D12Impl>(pPSO->GetDevice());
     auto& ResCacheDataAllocator = pPSO->GetSRBMemoryAllocator().GetResourceCacheDataAllocator(0);
@@ -51,12 +51,12 @@ ShaderResourceBindingD3D12Impl::ShaderResourceBindingD3D12Impl(IReferenceCounter
         auto* pShader = ppShaders[s];
         auto ShaderType = pShader->GetDesc().ShaderType;
         auto ShaderInd = GetShaderTypeIndex(ShaderType);
-        
+
         auto& ShaderResLayoutDataAllocator = pPSO->GetSRBMemoryAllocator().GetShaderVariableDataAllocator(s);
 
         // http://diligentgraphics.com/diligent-engine/architecture/d3d12/shader-resource-layout#Initializing-Resource-Layouts-in-a-Shader-Resource-Binding-Object
         std::array<SHADER_VARIABLE_TYPE, 3> AllowedVarTypes = { SHADER_VARIABLE_TYPE_STATIC, SHADER_VARIABLE_TYPE_MUTABLE, SHADER_VARIABLE_TYPE_DYNAMIC };
-        const auto &SrcLayout = pPSO->GetShaderResLayout(ShaderType);
+        const auto& SrcLayout = pPSO->GetShaderResLayout(s);
         new (m_pResourceLayouts + s) ShaderResourceLayoutD3D12(*this, SrcLayout, ShaderResLayoutDataAllocator, AllowedVarTypes.data(), static_cast<Uint32>(AllowedVarTypes.size()), m_ShaderResourceCache);
 
         m_ResourceLayoutIndex[ShaderInd] = static_cast<Int8>(s);
@@ -75,15 +75,12 @@ IMPLEMENT_QUERY_INTERFACE( ShaderResourceBindingD3D12Impl, IID_ShaderResourceBin
 
 void ShaderResourceBindingD3D12Impl::BindResources(Uint32 ShaderFlags, IResourceMapping* pResMapping, Uint32 Flags)
 {
-    for (auto ShaderInd = 0; ShaderInd <= CSInd; ++ShaderInd )
+    for (Uint32 s = 0; s < m_NumShaders; ++s )
     {
-        if (ShaderFlags & GetShaderTypeFromIndex(ShaderInd))
+        const auto& ShaderRes = m_pResourceLayouts[s].GetShaderResources();
+        if (ShaderFlags & ShaderRes.GetShaderType())
         {
-            auto ResLayoutInd = m_ResourceLayoutIndex[ShaderInd];
-            if(ResLayoutInd >= 0)
-            {
-                m_pResourceLayouts[ResLayoutInd].BindResources(pResMapping, Flags, &m_ShaderResourceCache);
-            }
+            m_pResourceLayouts[s].BindResources(pResMapping, Flags, &m_ShaderResourceCache);
         }
     }
 }
@@ -133,8 +130,8 @@ void ShaderResourceBindingD3D12Impl::InitializeStaticResources(const PipelineSta
 #ifdef VERIFY_SHADER_BINDINGS
         pShader->DbgVerifyStaticResourceBindings();
 #endif
-        auto &ConstResLayout = pShader->GetConstResLayout();
-        GetResourceLayout(pShader->GetDesc().ShaderType).CopyStaticResourceDesriptorHandles(ConstResLayout);
+        auto& ConstResLayout = pShader->GetConstResLayout();
+        GetResourceLayout(s).CopyStaticResourceDesriptorHandles(ConstResLayout);
     }
 
     m_bStaticResourcesInitialized = true;