PipelineStateD3D11Impl: improved D3D11 shader storage

2 files changed, 63 insertions, 60 deletions

diff --git a/Graphics/GraphicsEngineD3D11/include/PipelineStateD3D11Impl.hpp b/Graphics/GraphicsEngineD3D11/include/PipelineStateD3D11Impl.hpp
index f00ec89f..80d64d70 100644
--- a/Graphics/GraphicsEngineD3D11/include/PipelineStateD3D11Impl.hpp
+++ b/Graphics/GraphicsEngineD3D11/include/PipelineStateD3D11Impl.hpp
@@ -71,25 +71,24 @@ public:
     virtual ID3D11InputLayout* DILIGENT_CALL_TYPE GetD3D11InputLayout() override final { return m_pd3d11InputLayout; }
 
     /// Implementation of IPipelineStateD3D11::GetD3D11VertexShader() method.
-    virtual ID3D11VertexShader* DILIGENT_CALL_TYPE GetD3D11VertexShader() override final { return m_pVS; }
+    virtual ID3D11VertexShader* DILIGENT_CALL_TYPE GetD3D11VertexShader() override final { return GetD3D11Shader<ID3D11VertexShader>(VSInd); }
 
     /// Implementation of IPipelineStateD3D11::GetD3D11PixelShader() method.
-    virtual ID3D11PixelShader* DILIGENT_CALL_TYPE GetD3D11PixelShader() override final { return m_pPS; }
+    virtual ID3D11PixelShader* DILIGENT_CALL_TYPE GetD3D11PixelShader() override final { return GetD3D11Shader<ID3D11PixelShader>(PSInd); }
 
     /// Implementation of IPipelineStateD3D11::GetD3D11GeometryShader() method.
-    virtual ID3D11GeometryShader* DILIGENT_CALL_TYPE GetD3D11GeometryShader() override final { return m_pGS; }
+    virtual ID3D11GeometryShader* DILIGENT_CALL_TYPE GetD3D11GeometryShader() override final { return GetD3D11Shader<ID3D11GeometryShader>(GSInd); }
 
     /// Implementation of IPipelineStateD3D11::GetD3D11DomainShader() method.
-    virtual ID3D11DomainShader* DILIGENT_CALL_TYPE GetD3D11DomainShader() override final { return m_pDS; }
+    virtual ID3D11DomainShader* DILIGENT_CALL_TYPE GetD3D11DomainShader() override final { return GetD3D11Shader<ID3D11DomainShader>(DSInd); }
 
     /// Implementation of IPipelineStateD3D11::GetD3D11HullShader() method.
-    virtual ID3D11HullShader* DILIGENT_CALL_TYPE GetD3D11HullShader() override final { return m_pHS; }
+    virtual ID3D11HullShader* DILIGENT_CALL_TYPE GetD3D11HullShader() override final { return GetD3D11Shader<ID3D11HullShader>(HSInd); }
 
     /// Implementation of IPipelineStateD3D11::GetD3D11ComputeShader() method.
-    virtual ID3D11ComputeShader* DILIGENT_CALL_TYPE GetD3D11ComputeShader() override final { return m_pCS; }
+    virtual ID3D11ComputeShader* DILIGENT_CALL_TYPE GetD3D11ComputeShader() override final { return GetD3D11Shader<ID3D11ComputeShader>(CSInd); }
 
-    Uint32      GetNumShaders() const { return m_NumShaders; }
-    SHADER_TYPE GetShaderStageType(Uint32 Index) const;
+    Uint32 GetNumShaders() const { return m_NumShaders; }
 
 #ifdef DILIGENT_DEVELOPMENT
     void DvpVerifySRBResources(class ShaderResourceBindingD3D11Impl* pSRBs[], const D3D11ShaderResourceCounters BaseBindings[], Uint32 NumSRBs) const;
@@ -112,20 +111,27 @@ private:
 
     void ValidateShaderResources(const ShaderD3D11Impl* pShader);
 
+    template <typename D3D11ShaderType>
+    D3D11ShaderType* GetD3D11Shader(Uint32 ShaderInd)
+    {
+        auto idx = m_ShaderIndices[ShaderInd];
+        return idx >= 0 ? static_cast<D3D11ShaderType*>(m_ppd3d11Shaders[idx].p) : nullptr;
+    }
+
 private:
-    std::array<Uint8, 5> m_ShaderTypes = {};
-    Uint8                m_NumShaders  = 0;
+    // ShaderTypeIndex -> index in m_ppd3d11Shaders array
+    std::array<Int8, D3D11ResourceBindPoints::NumShaderTypes> m_ShaderIndices = {-1, -1, -1, -1, -1, -1};
+
+    // The number of shader stages in this pipeline
+    Uint8 m_NumShaders = 0;
 
     CComPtr<ID3D11BlendState>        m_pd3d11BlendState;
     CComPtr<ID3D11RasterizerState>   m_pd3d11RasterizerState;
     CComPtr<ID3D11DepthStencilState> m_pd3d11DepthStencilState;
     CComPtr<ID3D11InputLayout>       m_pd3d11InputLayout;
-    CComPtr<ID3D11VertexShader>      m_pVS;
-    CComPtr<ID3D11PixelShader>       m_pPS;
-    CComPtr<ID3D11GeometryShader>    m_pGS;
-    CComPtr<ID3D11DomainShader>      m_pDS;
-    CComPtr<ID3D11HullShader>        m_pHS;
-    CComPtr<ID3D11ComputeShader>     m_pCS;
+
+    using D3D11ShaderAutoPtrType             = CComPtr<ID3D11DeviceChild>;
+    D3D11ShaderAutoPtrType* m_ppd3d11Shaders = nullptr; // Shader array indexed by m_ShaderIndices[]
 
 #ifdef DILIGENT_DEVELOPMENT
     // Shader resources for all shaders in all shader stages in the pipeline.
@@ -136,9 +142,4 @@ private:
 #endif
 };
 
-__forceinline SHADER_TYPE GetShaderStageType(const ShaderD3D11Impl* pShader)
-{
-    return pShader->GetDesc().ShaderType;
-}
-
 } // namespace Diligent
diff --git a/Graphics/GraphicsEngineD3D11/src/PipelineStateD3D11Impl.cpp b/Graphics/GraphicsEngineD3D11/src/PipelineStateD3D11Impl.cpp
index 7d5fcc52..492ba5ae 100644
--- a/Graphics/GraphicsEngineD3D11/src/PipelineStateD3D11Impl.cpp
+++ b/Graphics/GraphicsEngineD3D11/src/PipelineStateD3D11Impl.cpp
@@ -37,8 +37,10 @@
 
 namespace Diligent
 {
+
 namespace
 {
+
 void VerifyResourceMerge(const PipelineStateDesc&        PSODesc,
                          const D3DShaderResourceAttribs& ExistingRes,
                          const D3DShaderResourceAttribs& NewResAttribs)
@@ -64,6 +66,11 @@ void VerifyResourceMerge(const PipelineStateDesc&        PSODesc,
 #undef LOG_RESOURCE_MERGE_ERROR_AND_THROW
 }
 
+__forceinline SHADER_TYPE GetShaderStageType(const ShaderD3D11Impl* pShader)
+{
+    return pShader->GetDesc().ShaderType;
+}
+
 } // namespace
 
 
@@ -263,45 +270,50 @@ void PipelineStateD3D11Impl::InitInternalObjects(const PSOCreateInfoType&
     ExtractShaders<ShaderD3D11Impl>(CreateInfo, Shaders);
 
     m_NumShaders = static_cast<Uint8>(Shaders.size());
-
-    for (Uint32 s = 0; s < Shaders.size(); ++s)
+    for (Uint32 s = 0; s < m_NumShaders; ++s)
     {
-        auto ShaderType  = Shaders[s]->GetDesc().ShaderType;
-        m_ShaderTypes[s] = static_cast<Uint8>(GetShaderTypeIndex(ShaderType));
-        VERIFY_EXPR(ShaderType == GetShaderTypeFromIndex(m_ShaderTypes[s]));
+        const auto ShaderType    = Shaders[s]->GetDesc().ShaderType;
+        const auto ShaderTypeIdx = GetShaderTypeIndex(ShaderType);
+        VERIFY_EXPR(m_ShaderIndices[ShaderTypeIdx] < 0);
+        m_ShaderIndices[ShaderTypeIdx] = static_cast<Int8>(s);
     }
 
     FixedLinearAllocator MemPool{GetRawAllocator()};
 
     ReserveSpaceForPipelineDesc(CreateInfo, MemPool);
+    MemPool.AddSpace<D3D11ShaderAutoPtrType>(m_NumShaders);
 
     MemPool.Reserve();
 
     InitializePipelineDesc(CreateInfo, MemPool);
+    m_ppd3d11Shaders = MemPool.ConstructArray<D3D11ShaderAutoPtrType>(m_NumShaders);
 
     InitResourceLayouts(CreateInfo, Shaders, ByteCodes);
 
     auto* pDeviceD3D11 = GetDevice()->GetD3D11Device();
-    for (Uint32 s = 0; s < Shaders.size(); ++s)
+    for (Uint32 s = 0; s < m_NumShaders; ++s)
     {
-        auto        ShaderType = Shaders[s]->GetDesc().ShaderType;
+        const auto  ShaderType = Shaders[s]->GetDesc().ShaderType;
         const auto& pByteCode  = ByteCodes[s];
         switch (ShaderType)
         {
-#define CREATE_SHADER(SHADER_NAME, ShaderName, pShader)                                                                                   \
-    case SHADER_TYPE_##SHADER_NAME:                                                                                                       \
-    {                                                                                                                                     \
-        HRESULT hr = pDeviceD3D11->Create##ShaderName##Shader(pByteCode->GetBufferPointer(), pByteCode->GetBufferSize(), NULL, &pShader); \
-        CHECK_D3D_RESULT_THROW(hr, "Failed to create D3D11 shader");                                                                      \
-        break;                                                                                                                            \
+#define CREATE_SHADER(SHADER_NAME, ShaderName)                                                                                         \
+    case SHADER_TYPE_##SHADER_NAME:                                                                                                    \
+    {                                                                                                                                  \
+        CComPtr<ID3D11##ShaderName##Shader> pShader;                                                                                   \
+                                                                                                                                       \
+        auto hr = pDeviceD3D11->Create##ShaderName##Shader(pByteCode->GetBufferPointer(), pByteCode->GetBufferSize(), NULL, &pShader); \
+        CHECK_D3D_RESULT_THROW(hr, "Failed to create D3D11 shader");                                                                   \
+        m_ppd3d11Shaders[s] = pShader;                                                                                                 \
+        break;                                                                                                                         \
     }
             // clang-format off
-            CREATE_SHADER(VERTEX,   Vertex,   m_pVS)
-            CREATE_SHADER(PIXEL,    Pixel,    m_pPS)
-            CREATE_SHADER(GEOMETRY, Geometry, m_pGS)
-            CREATE_SHADER(DOMAIN,   Domain,   m_pDS)
-            CREATE_SHADER(HULL,     Hull,     m_pHS)
-            CREATE_SHADER(COMPUTE,  Compute,  m_pCS)
+            CREATE_SHADER(VERTEX,   Vertex)
+            CREATE_SHADER(PIXEL,    Pixel)
+            CREATE_SHADER(GEOMETRY, Geometry)
+            CREATE_SHADER(DOMAIN,   Domain)
+            CREATE_SHADER(HULL,     Hull)
+            CREATE_SHADER(COMPUTE,  Compute)
             // clang-format on
             default: LOG_ERROR_AND_THROW("Unknown shader type");
         }
@@ -320,7 +332,7 @@ PipelineStateD3D11Impl::PipelineStateD3D11Impl(IReferenceCounters*
         std::vector<CComPtr<ID3DBlob>> ByteCodes;
         InitInternalObjects(CreateInfo, ByteCodes);
 
-        if (m_pVS == nullptr)
+        if (GetD3D11VertexShader() == nullptr)
             LOG_ERROR_AND_THROW("Vertex shader is null");
 
         auto& GraphicsPipeline = GetGraphicsPipelineDesc();
@@ -384,19 +396,14 @@ PipelineStateD3D11Impl::~PipelineStateD3D11Impl()
 
 void PipelineStateD3D11Impl::Destruct()
 {
-    m_ShaderTypes = {};
-    m_NumShaders  = 0;
-
-    m_pd3d11BlendState        = nullptr;
-    m_pd3d11RasterizerState   = nullptr;
-    m_pd3d11DepthStencilState = nullptr;
-    m_pd3d11InputLayout       = nullptr;
-    m_pVS                     = nullptr;
-    m_pPS                     = nullptr;
-    m_pGS                     = nullptr;
-    m_pDS                     = nullptr;
-    m_pHS                     = nullptr;
-    m_pCS                     = nullptr;
+    m_pd3d11BlendState.Release();
+    m_pd3d11RasterizerState.Release();
+    m_pd3d11DepthStencilState.Release();
+    m_pd3d11InputLayout.Release();
+    for (Uint32 s = 0; s < m_NumShaders; ++s)
+        m_ppd3d11Shaders[s].~D3D11ShaderAutoPtrType();
+    m_ppd3d11Shaders = nullptr;
+    m_ShaderIndices.fill(-1);
 
     TPipelineStateBase::Destruct();
 }
@@ -410,17 +417,12 @@ bool PipelineStateD3D11Impl::IsCompatibleWith(const IPipelineState* pPSO) const
         return false;
 
     const auto& rhs = *ValidatedCast<const PipelineStateD3D11Impl>(pPSO);
-    if (m_NumShaders != rhs.m_NumShaders || m_ShaderTypes != rhs.m_ShaderTypes)
+    if (m_ActiveShaderStages != rhs.m_ActiveShaderStages)
         return false;
 
     return true;
 }
 
-SHADER_TYPE PipelineStateD3D11Impl::GetShaderStageType(Uint32 Index) const
-{
-    return GetShaderTypeFromIndex(m_ShaderTypes[Index]);
-}
-
 void PipelineStateD3D11Impl::ValidateShaderResources(const ShaderD3D11Impl* pShader)
 {
     const auto& pShaderResources = pShader->GetShaderResources();