Reworked shader memory alloction for resource layouts in D3D12 pipeline implementation

11 files changed, 110 insertions, 128 deletions

diff --git a/Graphics/GraphicsEngine/include/PipelineStateBase.h b/Graphics/GraphicsEngine/include/PipelineStateBase.h
index 77c82cf5..80813fba 100644
--- a/Graphics/GraphicsEngine/include/PipelineStateBase.h
+++ b/Graphics/GraphicsEngine/include/PipelineStateBase.h
@@ -60,8 +60,6 @@ public:
         m_LayoutElements( PSODesc.GraphicsPipeline.InputLayout.NumElements, LayoutElement(), STD_ALLOCATOR_RAW_MEM(LayoutElement, GetRawAllocator(), "Allocator for vector<LayoutElement>" ) ),
         m_NumShaders(0)
     {
-        memset(m_ppShaders, 0, sizeof(m_ppShaders));
-
         if (this->m_Desc.IsComputePipeline)
         {
             const auto &ComputePipeline = PSODesc.ComputePipeline;
@@ -73,7 +71,7 @@ public:
 #define VALIDATE_SHADER_TYPE(Shader, ExpectedType, ShaderName)\
             if (Shader && Shader->GetDesc().ShaderType != ExpectedType)   \
             {                                                   \
-                LOG_ERROR_AND_THROW( GetShaderTypeLiteralName(Shader->GetDesc().ShaderType), " is not valid type for ", ShaderName, " shader" );\
+                LOG_ERROR_AND_THROW( GetShaderTypeLiteralName(Shader->GetDesc().ShaderType), " is not a valid type for ", ShaderName, " shader" );\
             }
             VALIDATE_SHADER_TYPE(ComputePipeline.pCS, SHADER_TYPE_COMPUTE, "compute")
 
@@ -208,6 +206,19 @@ public:
     IShader* const* GetShaders()const{return m_ppShaders;}
     Uint32 GetNumShaders()const{return m_NumShaders;}
 
+    template<typename ShaderType>
+    ShaderType* GetShader(Uint32 ShaderInd)
+    {
+        VERIFY_EXPR(ShaderInd < m_NumShaders);
+        return ValidatedCast<ShaderType>(m_ppShaders[ShaderInd]);
+    }
+    template<typename ShaderType>
+    ShaderType* GetShader(Uint32 ShaderInd)const
+    {
+        VERIFY_EXPR(ShaderInd < m_NumShaders);
+        return ValidatedCast<ShaderType>(m_ppShaders[ShaderInd]);
+    }
+
     // This function only compares shader resource layout hashes, so
     // it can potentially give false negatives
     bool IsIncompatibleWith(IPipelineState *pPSO)const
@@ -230,8 +241,8 @@ protected:
     RefCntAutoPtr<IShader> m_pDS; ///< Strong reference to the domain shader
     RefCntAutoPtr<IShader> m_pHS; ///< Strong reference to the hull shader
     RefCntAutoPtr<IShader> m_pCS; ///< Strong reference to the compute shader
-    IShader *m_ppShaders[5];  ///< Array of pointers to shaders that this PSO uses
-    Uint32 m_NumShaders; ///< Number of shaders that this PSO uses
+    IShader* m_ppShaders[5] = {}; ///< Array of pointers to the shaders used by this PSO
+    Uint32 m_NumShaders = 0;      ///< Number of shaders that this PSO uses
     size_t m_ShaderResourceLayoutHash = 0;///< Hash computed from the shader resource layout
 };
 
diff --git a/Graphics/GraphicsEngine/interface/PipelineState.h b/Graphics/GraphicsEngine/interface/PipelineState.h
index 34962f2e..0c642f3f 100644
--- a/Graphics/GraphicsEngine/interface/PipelineState.h
+++ b/Graphics/GraphicsEngine/interface/PipelineState.h
@@ -62,19 +62,19 @@ struct SampleDesc
 struct GraphicsPipelineDesc
 {
     /// Vertex shader to be used with the pipeline
-    IShader *pVS = nullptr;
+    IShader* pVS = nullptr;
 
     /// Pixel shader to be used with the pipeline
-    IShader *pPS = nullptr;
+    IShader* pPS = nullptr;
 
     /// Domain shader to be used with the pipeline
-    IShader *pDS = nullptr;
+    IShader* pDS = nullptr;
 
     /// Hull shader to be used with the pipeline
-    IShader *pHS = nullptr;
+    IShader* pHS = nullptr;
 
     /// Geometry shader to be used with the pipeline
-    IShader *pGS = nullptr;
+    IShader* pGS = nullptr;
     
     //D3D12_STREAM_OUTPUT_DESC StreamOutput;
     
@@ -136,10 +136,7 @@ struct GraphicsPipelineDesc
 struct ComputePipelineDesc
 {
     /// Compute shader to be used with the pipeline
-    IShader *pCS;
-    ComputePipelineDesc() : 
-        pCS(nullptr)
-    {}
+    IShader* pCS = nullptr;
 };
 
 /// Pipeline state description
diff --git a/Graphics/GraphicsEngineD3D11/include/ShaderD3D11Impl.h b/Graphics/GraphicsEngineD3D11/include/ShaderD3D11Impl.h
index 97ab9ba7..6823dbb0 100644
--- a/Graphics/GraphicsEngineD3D11/include/ShaderD3D11Impl.h
+++ b/Graphics/GraphicsEngineD3D11/include/ShaderD3D11Impl.h
@@ -63,8 +63,8 @@ public:
     ID3DBlob* GetBytecode(){return m_pShaderByteCode;}
 
     ShaderResourceLayoutD3D11& GetStaticResourceLayout(){return m_StaticResLayout;}
-    const std::shared_ptr<const ShaderResourcesD3D11>& GetResources(){return m_pShaderResources;}
-    Uint32 GetShaderTypeIndex(){return m_ShaderTypeIndex;}
+    const std::shared_ptr<const ShaderResourcesD3D11>& GetResources()const{return m_pShaderResources;}
+    Uint32 GetShaderTypeIndex()const{return m_ShaderTypeIndex;}
 
 private:
 
diff --git a/Graphics/GraphicsEngineD3D11/src/PipelineStateD3D11Impl.cpp b/Graphics/GraphicsEngineD3D11/src/PipelineStateD3D11Impl.cpp
index 4241eb31..cd91edef 100644
--- a/Graphics/GraphicsEngineD3D11/src/PipelineStateD3D11Impl.cpp
+++ b/Graphics/GraphicsEngineD3D11/src/PipelineStateD3D11Impl.cpp
@@ -119,7 +119,7 @@ PipelineStateD3D11Impl::PipelineStateD3D11Impl(IReferenceCounters*      pRefCoun
         std::array<size_t, MaxShadersInPipeline> ShaderResCacheDataSizes  = {};
         for (Uint32 s = 0; s < m_NumShaders; ++s)
         {
-            auto* pShader = ValidatedCast<ShaderD3D11Impl>(m_ppShaders[s]);
+            auto* pShader = GetShader<const ShaderD3D11Impl>(s);
             const auto& ShaderResources = *pShader->GetResources();
             std::array<SHADER_VARIABLE_TYPE, 2> AllowedVarTypes = { SHADER_VARIABLE_TYPE_MUTABLE, SHADER_VARIABLE_TYPE_DYNAMIC };
             ShaderResLayoutDataSizes[s] = ShaderResourceLayoutD3D11::GetRequiredMemorySize(ShaderResources, AllowedVarTypes.data(), static_cast<Uint32>(AllowedVarTypes.size()));
@@ -201,8 +201,8 @@ bool PipelineStateD3D11Impl::IsCompatibleWith(const IPipelineState* pPSO)const
 
     for (Uint32 s = 0; s < m_NumShaders; ++s)
     {
-        auto* pShader0 = ValidatedCast<ShaderD3D11Impl>(m_ppShaders[s]);
-        auto* pShader1 = ValidatedCast<ShaderD3D11Impl>(pPSOD3D11->m_ppShaders[s]);
+        auto* pShader0 = GetShader<const ShaderD3D11Impl>(s);
+        auto* pShader1 = pPSOD3D11->GetShader<const ShaderD3D11Impl>(s);
         if (pShader0->GetShaderTypeIndex() != pShader1->GetShaderTypeIndex())
             return false;
         const ShaderResourcesD3D11* pRes0 = pShader0->GetResources().get();
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;
diff --git a/Graphics/GraphicsEngineOpenGL/src/PipelineStateGLImpl.cpp b/Graphics/GraphicsEngineOpenGL/src/PipelineStateGLImpl.cpp
index 8794e91c..4eeebc88 100644
--- a/Graphics/GraphicsEngineOpenGL/src/PipelineStateGLImpl.cpp
+++ b/Graphics/GraphicsEngineOpenGL/src/PipelineStateGLImpl.cpp
@@ -52,7 +52,7 @@ void PipelineStateGLImpl::LinkGLProgram(bool bIsProgramPipelineSupported)
         m_ShaderResourceLayoutHash = 0;
         for (Uint32 Shader = 0; Shader < m_NumShaders; ++Shader)
         {
-            auto *pShaderGL = static_cast<ShaderGLImpl*>(m_ppShaders[Shader]);
+            auto *pShaderGL = GetShader<ShaderGLImpl>(Shader);
             HashCombine(m_ShaderResourceLayoutHash, pShaderGL->m_GlProgObj.GetAllResources().GetHash());
         }
     }
@@ -62,7 +62,7 @@ void PipelineStateGLImpl::LinkGLProgram(bool bIsProgramPipelineSupported)
         m_GLProgram.Create();
         for( Uint32 Shader = 0; Shader < m_NumShaders; ++Shader )
         {
-            auto *pCurrShader = static_cast<ShaderGLImpl*>(m_ppShaders[Shader]);
+            auto *pCurrShader = GetShader<ShaderGLImpl>(Shader);
             glAttachShader( m_GLProgram, pCurrShader->m_GLShaderObj );
             CHECK_GL_ERROR( "glAttachShader() failed" );
         }
@@ -91,7 +91,7 @@ void PipelineStateGLImpl::LinkGLProgram(bool bIsProgramPipelineSupported)
         // Detach shaders from the program object
         for( Uint32 Shader = 0; Shader < m_NumShaders; ++Shader )
         {
-            auto *pCurrShader = static_cast<ShaderGLImpl*>(m_ppShaders[Shader]);
+            auto *pCurrShader = GetShader<ShaderGLImpl>(Shader);
             glDetachShader( m_GLProgram, pCurrShader->m_GLShaderObj );
             CHECK_GL_ERROR( "glDetachShader() failed" );
         }
@@ -101,7 +101,7 @@ void PipelineStateGLImpl::LinkGLProgram(bool bIsProgramPipelineSupported)
         SHADER_VARIABLE_TYPE DefaultVarType = SHADER_VARIABLE_TYPE_STATIC;
         for( Uint32 Shader = 0; Shader < m_NumShaders; ++Shader )
         {
-            auto *pCurrShader = static_cast<ShaderGLImpl*>(m_ppShaders[Shader]);
+            auto *pCurrShader = GetShader<ShaderGLImpl>(Shader);
             const auto& Desc = pCurrShader->GetDesc();
             if (Shader == 0)
             {
@@ -190,7 +190,7 @@ GLObjectWrappers::GLPipelineObj &PipelineStateGLImpl::GetGLProgramPipeline(GLCon
         GLuint Pipeline = it->second;
         for (Uint32 Shader = 0; Shader < m_NumShaders; ++Shader)
         {
-            auto *pCurrShader = static_cast<ShaderGLImpl*>(m_ppShaders[Shader]);
+            auto *pCurrShader = GetShader<ShaderGLImpl>(Shader);
             auto GLShaderBit = ShaderTypeToGLShaderBit(pCurrShader->GetDesc().ShaderType);
             // If the program has an active code for each stage mentioned in set flags,
             // then that code will be used by the pipeline. If program is 0, then the given
diff --git a/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp b/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp
index 7014d51b..07bed5f2 100644
--- a/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp
+++ b/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp
@@ -143,7 +143,7 @@ PipelineStateVkImpl :: PipelineStateVkImpl(IReferenceCounters*      pRefCounters
     for (Uint32 s=0; s < m_NumShaders; ++s)
     {
         new (m_ShaderResourceLayouts + s) ShaderResourceLayoutVk(*this, LogicalDevice, GetRawAllocator());
-        auto *pShaderVk = ValidatedCast<ShaderVkImpl>(m_ppShaders[s]);
+        auto* pShaderVk = GetShader<const ShaderVkImpl>(s);
         ShaderResources[s] = pShaderVk->GetShaderResources();
         ShaderSPIRVs[s] = pShaderVk->GetSPIRV();
     }
@@ -171,7 +171,7 @@ PipelineStateVkImpl :: PipelineStateVkImpl(IReferenceCounters*      pRefCounters
     std::array<VkPipelineShaderStageCreateInfo, MaxShadersInPipeline> ShaderStages = {};
     for (Uint32 s = 0; s < m_NumShaders; ++s)
     {
-        auto *pShader = m_ppShaders[s];
+        auto* pShader = m_ppShaders[s];
         auto ShaderType = pShader->GetDesc().ShaderType;
 
         auto& StageCI = ShaderStages[s];
@@ -441,7 +441,7 @@ void PipelineStateVkImpl::BindShaderResources(IResourceMapping *pResourceMapping
 
 void PipelineStateVkImpl::CreateShaderResourceBinding(IShaderResourceBinding **ppShaderResourceBinding)
 {
-    auto *pRenderDeviceVk = GetDevice<RenderDeviceVkImpl>();
+    auto* pRenderDeviceVk = GetDevice<RenderDeviceVkImpl>();
     auto& SRBAllocator = pRenderDeviceVk->GetSRBAllocator();
     auto pResBindingVk = NEW_RC_OBJ(SRBAllocator, "ShaderResourceBindingVkImpl instance", ShaderResourceBindingVkImpl)(this, false);
     pResBindingVk->QueryInterface(IID_ShaderResourceBinding, reinterpret_cast<IObject**>(ppShaderResourceBinding));
@@ -470,15 +470,15 @@ bool PipelineStateVkImpl::IsCompatibleWith(const IPipelineState *pPSO)const
         {
             for (Uint32 s = 0; s < m_NumShaders; ++s)
             {
-                auto *pShader0 = ValidatedCast<ShaderVkImpl>(m_ppShaders[s]);
-                auto *pShader1 = ValidatedCast<ShaderVkImpl>(pPSOVk->m_ppShaders[s]);
+                auto* pShader0 = GetShader<const ShaderVkImpl>(s);
+                auto* pShader1 = pPSOVk->GetShader<const ShaderVkImpl>(s);
                 if (pShader0->GetDesc().ShaderType != pShader1->GetDesc().ShaderType)
                 {
                     IsCompatibleShaders = false;
                     break;
                 }
-                const auto *pRes0 = pShader0->GetShaderResources().get();
-                const auto *pRes1 = pShader1->GetShaderResources().get();
+                const auto* pRes0 = pShader0->GetShaderResources().get();
+                const auto* pRes1 = pShader1->GetShaderResources().get();
                 if (!pRes0->IsCompatibleWith(*pRes1))
                 {
                     IsCompatibleShaders = false;
@@ -514,11 +514,11 @@ void PipelineStateVkImpl::CommitAndTransitionShaderResources(IShaderResourceBind
 
     // If the shaders contain no resources or static resources only, shader resource binding may be null. 
     // In this case use special internal SRB object
-    auto *pResBindingVkImpl = pShaderResourceBinding ? ValidatedCast<ShaderResourceBindingVkImpl>(pShaderResourceBinding) : m_pDefaultShaderResBinding.get();
+    auto* pResBindingVkImpl = pShaderResourceBinding ? ValidatedCast<ShaderResourceBindingVkImpl>(pShaderResourceBinding) : m_pDefaultShaderResBinding.get();
     
 #ifdef VERIFY_SHADER_BINDINGS
     {
-        auto *pRefPSO = pResBindingVkImpl->GetPipelineState();
+        auto* pRefPSO = pResBindingVkImpl->GetPipelineState();
         if ( IsIncompatibleWith(pRefPSO) )
         {
             LOG_ERROR_MESSAGE("Shader resource binding is incompatible with the pipeline state \"", m_Desc.Name, "\". Operation will be ignored.");
@@ -534,7 +534,7 @@ void PipelineStateVkImpl::CommitAndTransitionShaderResources(IShaderResourceBind
     {
         for (Uint32 s = 0; s < m_NumShaders; ++s)
         {
-            auto *pShaderVk = ValidatedCast<ShaderVkImpl>( m_ppShaders[s] );
+            auto* pShaderVk = GetShader<ShaderVkImpl>(s);
 #ifdef VERIFY_SHADER_BINDINGS
             pShaderVk->DbgVerifyStaticResourceBindings();
 #endif