A number of corrections for PSO refactoring

7 files changed, 100 insertions, 105 deletions

diff --git a/Graphics/GraphicsEngineVulkan/include/PipelineStateVkImpl.hpp b/Graphics/GraphicsEngineVulkan/include/PipelineStateVkImpl.hpp
index 303496bb..2185dabd 100644
--- a/Graphics/GraphicsEngineVulkan/include/PipelineStateVkImpl.hpp
+++ b/Graphics/GraphicsEngineVulkan/include/PipelineStateVkImpl.hpp
@@ -56,7 +56,6 @@ class PipelineStateVkImpl final : public PipelineStateBase<IPipelineStateVk, Ren
 {
 public:
     using TPipelineStateBase = PipelineStateBase<IPipelineStateVk, RenderDeviceVkImpl>;
-    using ShaderSPIRVs_t     = std::vector<std::vector<uint32_t>>;
 
     PipelineStateVkImpl(IReferenceCounters* pRefCounters, RenderDeviceVkImpl* pDeviceVk, const PipelineStateCreateInfo& CreateInfo);
     ~PipelineStateVkImpl();
@@ -105,7 +104,7 @@ public:
 
     const ShaderResourceLayoutVk& GetShaderResLayout(Uint32 ShaderInd) const
     {
-        VERIFY_EXPR(ShaderInd < m_NumShaderTypes);
+        VERIFY_EXPR(ShaderInd < GetNumShaderStages());
         return m_ShaderResourceLayouts[ShaderInd];
     }
 
@@ -128,19 +127,19 @@ public:
 private:
     const ShaderResourceLayoutVk& GetStaticShaderResLayout(Uint32 ShaderInd) const
     {
-        VERIFY_EXPR(ShaderInd < m_NumShaderTypes);
-        return m_ShaderResourceLayouts[m_NumShaderTypes + ShaderInd];
+        VERIFY_EXPR(ShaderInd < GetNumShaderStages());
+        return m_ShaderResourceLayouts[GetNumShaderStages() + ShaderInd];
     }
 
     const ShaderResourceCacheVk& GetStaticResCache(Uint32 ShaderInd) const
     {
-        VERIFY_EXPR(ShaderInd < m_NumShaderTypes);
+        VERIFY_EXPR(ShaderInd < GetNumShaderStages());
         return m_StaticResCaches[ShaderInd];
     }
 
     ShaderVariableManagerVk& GetStaticVarMgr(Uint32 ShaderInd) const
     {
-        VERIFY_EXPR(ShaderInd < m_NumShaderTypes);
+        VERIFY_EXPR(ShaderInd < GetNumShaderStages());
         return m_StaticVarsMgrs[ShaderInd];
     }
 
@@ -156,7 +155,7 @@ private:
 
     // Resource layout index in m_ShaderResourceLayouts array for every shader stage,
     // indexed by the shader type pipeline index (returned by GetShaderTypePipelineIndex)
-    std::array<Int8, MAX_SHADERS_IN_PIPELINE> m_ResourceLayoutIndex;
+    std::array<Int8, MAX_SHADERS_IN_PIPELINE> m_ResourceLayoutIndex = {-1, -1, -1, -1, -1};
 
     bool m_HasStaticResources    = false;
     bool m_HasNonStaticResources = false;
diff --git a/Graphics/GraphicsEngineVulkan/include/ShaderResourceBindingVkImpl.hpp b/Graphics/GraphicsEngineVulkan/include/ShaderResourceBindingVkImpl.hpp
index 85230b02..026bb9da 100644
--- a/Graphics/GraphicsEngineVulkan/include/ShaderResourceBindingVkImpl.hpp
+++ b/Graphics/GraphicsEngineVulkan/include/ShaderResourceBindingVkImpl.hpp
@@ -81,7 +81,7 @@ private:
 
     // Resource layout index in m_ShaderResourceCache array for every shader stage,
     // indexed by the shader type pipeline index (returned by GetShaderTypePipelineIndex)
-    std::array<Int8, MAX_SHADERS_IN_PIPELINE> m_ResourceLayoutIndex;
+    std::array<Int8, MAX_SHADERS_IN_PIPELINE> m_ResourceLayoutIndex = {-1, -1, -1, -1, -1};
 
     bool  m_bStaticResourcesInitialized = false;
     Uint8 m_NumShaders                  = 0;
diff --git a/Graphics/GraphicsEngineVulkan/include/ShaderResourceLayoutVk.hpp b/Graphics/GraphicsEngineVulkan/include/ShaderResourceLayoutVk.hpp
index c6432f4a..62ff38b5 100644
--- a/Graphics/GraphicsEngineVulkan/include/ShaderResourceLayoutVk.hpp
+++ b/Graphics/GraphicsEngineVulkan/include/ShaderResourceLayoutVk.hpp
@@ -110,13 +110,23 @@
 namespace Diligent
 {
 
+class ShaderVkImpl;
+
 /// Diligent::ShaderResourceLayoutVk class
 // sizeof(ShaderResourceLayoutVk)==56 (MS compiler, x64)
 class ShaderResourceLayoutVk
 {
 public:
-    using ShaderStages_t = std::vector<std::pair<SHADER_TYPE, IShader*>>;
-    using ShaderSPIRVs_t = std::vector<std::vector<uint32_t>>;
+    struct ShaderStageInfo
+    {
+        ShaderStageInfo(SHADER_TYPE         _Type,
+                        const ShaderVkImpl* _pShader);
+
+        const SHADER_TYPE         Type;
+        const ShaderVkImpl* const pShader;
+        std::vector<uint32_t>     SPIRV;
+    };
+    using TShaderStages = std::vector<ShaderStageInfo>;
 
     ShaderResourceLayoutVk(const VulkanUtilities::VulkanLogicalDevice& LogicalDevice) :
         m_LogicalDevice{LogicalDevice}
@@ -134,7 +144,7 @@ public:
 
     // This method is called by PipelineStateVkImpl class instance to initialize static
     // shader resource layout and the cache
-    void InitializeStaticResourceLayout(IShader*                          pShader,
+    void InitializeStaticResourceLayout(const ShaderVkImpl*               pShader,
                                         IMemoryAllocator&                 LayoutDataAllocator,
                                         const PipelineResourceLayoutDesc& ResourceLayoutDesc,
                                         ShaderResourceCacheVk&            StaticResourceCache);
@@ -142,11 +152,10 @@ public:
     // This method is called by PipelineStateVkImpl class instance to initialize resource
     // layouts for all shader stages in the pipeline.
     static void Initialize(IRenderDevice*                    pRenderDevice,
-                           const ShaderStages_t&             ShaderStages,
+                           TShaderStages&                    ShaderStages,
                            ShaderResourceLayoutVk            Layouts[],
                            IMemoryAllocator&                 LayoutDataAllocator,
                            const PipelineResourceLayoutDesc& ResourceLayoutDesc,
-                           ShaderSPIRVs_t&                   SPIRVs,
                            class PipelineLayout&             PipelineLayout,
                            bool                              VerifyVariables,
                            bool                              VerifyStaticSamplers);
@@ -281,7 +290,7 @@ public:
 
 #ifdef DILIGENT_DEVELOPMENT
     bool        dvpVerifyBindings(const ShaderResourceCacheVk& ResourceCache) const;
-    static void dvpVerifyResourceLayoutDesc(const ShaderStages_t&             ShaderStages,
+    static void dvpVerifyResourceLayoutDesc(const TShaderStages&              ShaderStages,
                                             const PipelineResourceLayoutDesc& ResourceLayoutDesc,
                                             bool                              VerifyVariables,
                                             bool                              VerifyStaticSamplers);
@@ -349,7 +358,7 @@ private:
         return m_NumResources[SHADER_RESOURCE_VARIABLE_TYPE_NUM_TYPES];
     }
 
-    void AllocateMemory(IShader*                             pShader,
+    void AllocateMemory(const ShaderVkImpl*                  pShader,
                         IMemoryAllocator&                    Allocator,
                         const PipelineResourceLayoutDesc&    ResourceLayoutDesc,
                         const SHADER_RESOURCE_VARIABLE_TYPE* AllowedVarTypes,
diff --git a/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp b/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp
index 881196ec..fbc5fae3 100644
--- a/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp
+++ b/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp
@@ -149,26 +149,21 @@ static bool StripReflection(std::vector<uint32_t>& SPIRV)
 #endif
 }
 
-static void InitializeShaderStages(const VulkanUtilities::VulkanLogicalDevice&        LogicalDevice,
-                                   const PipelineStateVkImpl::ShaderStages_t&         ShaderStages,
-                                   PipelineStateVkImpl::ShaderSPIRVs_t&               ShaderSPIRVs,
-                                   std::vector<VulkanUtilities::ShaderModuleWrapper>& ShaderModules,
-                                   std::vector<VkPipelineShaderStageCreateInfo>&      Stages)
+static void InitPipelineShaderStages(const VulkanUtilities::VulkanLogicalDevice&        LogicalDevice,
+                                     ShaderResourceLayoutVk::TShaderStages&             ShaderStages,
+                                     std::vector<VulkanUtilities::ShaderModuleWrapper>& vkShaderModules,
+                                     std::vector<VkPipelineShaderStageCreateInfo>&      vkPipelineShaderStages)
 {
-    VERIFY_EXPR(ShaderStages.size() == ShaderSPIRVs.size());
-
     for (size_t s = 0; s < ShaderStages.size(); ++s)
     {
-        auto*      pShaderVk  = ValidatedCast<ShaderVkImpl>(ShaderStages[s].second);
-        auto&      SPIRV      = ShaderSPIRVs[s];
-        const auto ShaderType = ShaderStages[s].first;
+        auto& StageInfo = ShaderStages[s];
 
         VkPipelineShaderStageCreateInfo StageCI = {};
 
         StageCI.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
         StageCI.pNext = nullptr;
         StageCI.flags = 0; //  reserved for future use
-        StageCI.stage = ShaderTypeToVkShaderStageFlagBit(ShaderType);
+        StageCI.stage = ShaderTypeToVkShaderStageFlagBit(StageInfo.Type);
 
         VkShaderModuleCreateInfo ShaderModuleCI = {};
 
@@ -179,22 +174,22 @@ static void InitializeShaderStages(const VulkanUtilities::VulkanLogicalDevice&
         // We have to strip reflection instructions to fix the follownig validation error:
         //     SPIR-V module not valid: DecorateStringGOOGLE requires one of the following extensions: SPV_GOOGLE_decorate_string
         // Optimizer also performs validation and may catch problems with the byte code.
-        if (!StripReflection(SPIRV))
-            LOG_ERROR("Failed to strip reflection information from shader '", pShaderVk->GetDesc().Name, "'. This may indicate a problem with the byte code.");
+        if (!StripReflection(StageInfo.SPIRV))
+            LOG_ERROR("Failed to strip reflection information from shader '", StageInfo.pShader->GetDesc().Name, "'. This may indicate a problem with the byte code.");
 
-        ShaderModuleCI.codeSize = SPIRV.size() * sizeof(uint32_t);
-        ShaderModuleCI.pCode    = SPIRV.data();
+        ShaderModuleCI.codeSize = StageInfo.SPIRV.size() * sizeof(uint32_t);
+        ShaderModuleCI.pCode    = StageInfo.SPIRV.data();
 
-        ShaderModules.push_back(LogicalDevice.CreateShaderModule(ShaderModuleCI, pShaderVk->GetDesc().Name));
+        vkShaderModules.push_back(LogicalDevice.CreateShaderModule(ShaderModuleCI, StageInfo.pShader->GetDesc().Name));
 
-        StageCI.module              = ShaderModules.back();
-        StageCI.pName               = pShaderVk->GetEntryPoint();
+        StageCI.module              = vkShaderModules.back();
+        StageCI.pName               = StageInfo.pShader->GetEntryPoint();
         StageCI.pSpecializationInfo = nullptr;
 
-        Stages.push_back(StageCI);
+        vkPipelineShaderStages.push_back(StageCI);
     }
 
-    VERIFY_EXPR(ShaderModules.size() == Stages.size());
+    VERIFY_EXPR(vkShaderModules.size() == vkPipelineShaderStages.size());
 }
 
 
@@ -417,47 +412,41 @@ PipelineStateVkImpl::PipelineStateVkImpl(IReferenceCounters*            pRefCoun
 
     const auto& LogicalDevice = pDeviceVk->GetLogicalDevice();
 
-    ShaderStages_t ShaderStages;
-    ShaderSPIRVs_t ShaderSPIRVs;
-    ExtractShaders(ShaderStages);
-
-    ShaderSPIRVs.resize(ShaderStages.size());
-    for (size_t s = 0; s < ShaderSPIRVs.size(); ++s)
-    {
-        auto* pShaderVk = ValidatedCast<ShaderVkImpl>(ShaderStages[s].second);
-        ShaderSPIRVs[s] = pShaderVk->GetSPIRV();
-    }
+    ShaderResourceLayoutVk::TShaderStages ShaderStages;
+    ExtractShaders<ShaderVkImpl>(ShaderStages);
 
     // clang-format off
     static_assert((sizeof(ShaderResourceLayoutVk)  % sizeof(void*)) == 0, "sizeof(ShaderResourceLayoutVk) is expected to be a multiple of sizeof(void*)");
     static_assert((sizeof(ShaderResourceCacheVk)   % sizeof(void*)) == 0, "sizeof(ShaderResourceCacheVk) is expected to be a multiple of sizeof(void*)");
     static_assert((sizeof(ShaderVariableManagerVk) % sizeof(void*)) == 0, "sizeof(ShaderVariableManagerVk) is expected to be a multiple of sizeof(void*)");
     // clang-format on
-    const auto  MemSize = (sizeof(ShaderResourceLayoutVk) * 2 + sizeof(ShaderResourceCacheVk) + sizeof(ShaderVariableManagerVk)) * GetNumShaderTypes();
+    const auto  MemSize = (sizeof(ShaderResourceLayoutVk) * 2 + sizeof(ShaderResourceCacheVk) + sizeof(ShaderVariableManagerVk)) * GetNumShaderStages();
     auto* const pRawMem =
         ALLOCATE_RAW(GetRawAllocator(), "Raw memory for ShaderResourceLayoutVk, ShaderResourceCacheVk, and ShaderVariableManagerVk arrays", MemSize);
 
     m_ShaderResourceLayouts = reinterpret_cast<ShaderResourceLayoutVk*>(pRawMem);
-    m_StaticResCaches       = reinterpret_cast<ShaderResourceCacheVk*>(m_ShaderResourceLayouts + GetNumShaderTypes() * 2);
-    m_StaticVarsMgrs        = reinterpret_cast<ShaderVariableManagerVk*>(m_StaticResCaches + GetNumShaderTypes());
+    m_StaticResCaches       = reinterpret_cast<ShaderResourceCacheVk*>(m_ShaderResourceLayouts + GetNumShaderStages() * 2);
+    m_StaticVarsMgrs        = reinterpret_cast<ShaderVariableManagerVk*>(m_StaticResCaches + GetNumShaderStages());
 
     for (size_t s = 0; s < ShaderStages.size(); ++s)
     {
+        auto&      StageInfo     = ShaderStages[s];
+        const auto ShaderType    = StageInfo.Type;
+        const auto ShaderTypeInd = GetShaderTypePipelineIndex(ShaderType, m_Desc.PipelineType);
+
         new (m_ShaderResourceLayouts + s) ShaderResourceLayoutVk{LogicalDevice};
 
-        const auto ShaderType                = ShaderStages[s].first;
-        auto&      Shaders                   = ShaderStages[s].second;
-        const auto ShaderTypeInd             = GetShaderTypePipelineIndex(ShaderType, m_Desc.PipelineType);
+
         m_ResourceLayoutIndex[ShaderTypeInd] = static_cast<Int8>(s);
 
         auto* pStaticResLayout = new (m_ShaderResourceLayouts + ShaderStages.size() + s) ShaderResourceLayoutVk{LogicalDevice};
         auto* pStaticResCache  = new (m_StaticResCaches + s) ShaderResourceCacheVk{ShaderResourceCacheVk::DbgCacheContentType::StaticShaderResources};
-        pStaticResLayout->InitializeStaticResourceLayout(Shaders, GetRawAllocator(), m_Desc.ResourceLayout, m_StaticResCaches[s]);
+        pStaticResLayout->InitializeStaticResourceLayout(StageInfo.pShader, GetRawAllocator(), m_Desc.ResourceLayout, m_StaticResCaches[s]);
 
         new (m_StaticVarsMgrs + s) ShaderVariableManagerVk{*this, *pStaticResLayout, GetRawAllocator(), nullptr, 0, *pStaticResCache};
     }
     ShaderResourceLayoutVk::Initialize(pDeviceVk, ShaderStages, m_ShaderResourceLayouts, GetRawAllocator(),
-                                       m_Desc.ResourceLayout, ShaderSPIRVs, m_PipelineLayout,
+                                       m_Desc.ResourceLayout, m_PipelineLayout,
                                        (CreateInfo.Flags & PSO_CREATE_FLAG_IGNORE_MISSING_VARIABLES) == 0,
                                        (CreateInfo.Flags & PSO_CREATE_FLAG_IGNORE_MISSING_STATIC_SAMPLERS) == 0);
     m_PipelineLayout.Finalize(LogicalDevice);
@@ -465,7 +454,7 @@ PipelineStateVkImpl::PipelineStateVkImpl(IReferenceCounters*            pRefCoun
     if (m_Desc.SRBAllocationGranularity > 1)
     {
         std::array<size_t, MAX_SHADERS_IN_PIPELINE> ShaderVariableDataSizes = {};
-        for (Uint32 s = 0; s < GetNumShaderTypes(); ++s)
+        for (Uint32 s = 0; s < GetNumShaderStages(); ++s)
         {
             const SHADER_RESOURCE_VARIABLE_TYPE AllowedVarTypes[] = {SHADER_RESOURCE_VARIABLE_TYPE_MUTABLE, SHADER_RESOURCE_VARIABLE_TYPE_DYNAMIC};
 
@@ -477,13 +466,13 @@ PipelineStateVkImpl::PipelineStateVkImpl(IReferenceCounters*            pRefCoun
         auto   DescriptorSetSizes = m_PipelineLayout.GetDescriptorSetSizes(NumSets);
         auto   CacheMemorySize    = ShaderResourceCacheVk::GetRequiredMemorySize(NumSets, DescriptorSetSizes.data());
 
-        m_SRBMemAllocator.Initialize(m_Desc.SRBAllocationGranularity, GetNumShaderTypes(), ShaderVariableDataSizes.data(), 1, &CacheMemorySize);
+        m_SRBMemAllocator.Initialize(m_Desc.SRBAllocationGranularity, GetNumShaderStages(), ShaderVariableDataSizes.data(), 1, &CacheMemorySize);
     }
 
     // Create shader modules and initialize shader stages
     std::vector<VkPipelineShaderStageCreateInfo>      VkShaderStages;
     std::vector<VulkanUtilities::ShaderModuleWrapper> ShaderModules;
-    InitializeShaderStages(LogicalDevice, ShaderStages, ShaderSPIRVs, ShaderModules, VkShaderStages);
+    InitPipelineShaderStages(LogicalDevice, ShaderStages, ShaderModules, VkShaderStages);
 
     // Create pipeline
     switch (m_Desc.PipelineType)
@@ -498,7 +487,7 @@ PipelineStateVkImpl::PipelineStateVkImpl(IReferenceCounters*            pRefCoun
 
     m_HasStaticResources    = false;
     m_HasNonStaticResources = false;
-    for (Uint32 s = 0; s < GetNumShaderTypes(); ++s)
+    for (Uint32 s = 0; s < GetNumShaderStages(); ++s)
     {
         const auto& Layout = m_ShaderResourceLayouts[s];
         if (Layout.GetResourceCount(SHADER_RESOURCE_VARIABLE_TYPE_STATIC) != 0)
@@ -518,12 +507,12 @@ PipelineStateVkImpl::~PipelineStateVkImpl()
     m_PipelineLayout.Release(m_pDevice, m_Desc.CommandQueueMask);
 
     auto& RawAllocator = GetRawAllocator();
-    for (Uint32 s = 0; s < GetNumShaderTypes() * 2; ++s)
+    for (Uint32 s = 0; s < GetNumShaderStages() * 2; ++s)
     {
         m_ShaderResourceLayouts[s].~ShaderResourceLayoutVk();
     }
 
-    for (Uint32 s = 0; s < GetNumShaderTypes(); ++s)
+    for (Uint32 s = 0; s < GetNumShaderStages(); ++s)
     {
         m_StaticResCaches[s].~ShaderResourceCacheVk();
         m_StaticVarsMgrs[s].DestroyVariables(GetRawAllocator());
@@ -563,14 +552,14 @@ bool PipelineStateVkImpl::IsCompatibleWith(const IPipelineState* pPSO) const
 #ifdef DILIGENT_DEBUG
     {
         bool IsCompatibleShaders = true;
-        if (GetNumShaderTypes() != pPSOVk->GetNumShaderTypes())
+        if (GetNumShaderStages() != pPSOVk->GetNumShaderStages())
             IsCompatibleShaders = false;
 
         if (IsCompatibleShaders)
         {
-            for (Uint32 s = 0; s < GetNumShaderTypes(); ++s)
+            for (Uint32 s = 0; s < GetNumShaderStages(); ++s)
             {
-                if (GetShaderTypes()[s] != pPSOVk->GetShaderTypes()[s])
+                if (GetShaderStageType(s) != pPSOVk->GetShaderStageType(s))
                 {
                     IsCompatibleShaders = false;
                     break;
@@ -636,7 +625,7 @@ void PipelineStateVkImpl::CommitAndTransitionShaderResources(IShaderResourceBind
     auto& ResourceCache = pResBindingVkImpl->GetResourceCache();
 
 #ifdef DILIGENT_DEVELOPMENT
-    for (Uint32 s = 0; s < GetNumShaderTypes(); ++s)
+    for (Uint32 s = 0; s < GetNumShaderStages(); ++s)
     {
         m_ShaderResourceLayouts[s].dvpVerifyBindings(ResourceCache);
     }
@@ -671,7 +660,7 @@ void PipelineStateVkImpl::CommitAndTransitionShaderResources(IShaderResourceBind
             // Allocate vulkan descriptor set for dynamic resources
             DynamicDescrSet = pCtxVkImpl->AllocateDynamicDescriptorSet(DynamicDescriptorSetVkLayout, DynamicDescrSetName);
             // Commit all dynamic resource descriptors
-            for (Uint32 s = 0; s < GetNumShaderTypes(); ++s)
+            for (Uint32 s = 0; s < GetNumShaderStages(); ++s)
             {
                 const auto& Layout = m_ShaderResourceLayouts[s];
                 if (Layout.GetResourceCount(SHADER_RESOURCE_VARIABLE_TYPE_DYNAMIC) != 0)
@@ -688,7 +677,7 @@ void PipelineStateVkImpl::CommitAndTransitionShaderResources(IShaderResourceBind
 
 void PipelineStateVkImpl::BindStaticResources(Uint32 ShaderFlags, IResourceMapping* pResourceMapping, Uint32 Flags)
 {
-    for (Uint32 s = 0; s < GetNumShaderTypes(); ++s)
+    for (Uint32 s = 0; s < GetNumShaderStages(); ++s)
     {
         auto ShaderType = GetStaticShaderResLayout(s).GetShaderType();
         if ((ShaderType & ShaderFlags) != 0)
@@ -732,7 +721,7 @@ IShaderResourceVariable* PipelineStateVkImpl::GetStaticVariableByIndex(SHADER_TY
 
 void PipelineStateVkImpl::InitializeStaticSRBResources(ShaderResourceCacheVk& ResourceCache) const
 {
-    for (Uint32 s = 0; s < GetNumShaderTypes(); ++s)
+    for (Uint32 s = 0; s < GetNumShaderStages(); ++s)
     {
         const auto& StaticResLayout = GetStaticShaderResLayout(s);
         const auto& StaticResCache  = GetStaticResCache(s);
@@ -742,7 +731,7 @@ void PipelineStateVkImpl::InitializeStaticSRBResources(ShaderResourceCacheVk& Re
         {
             LOG_ERROR_MESSAGE("Static resources in SRB of PSO '", GetDesc().Name,
                               "' will not be successfully initialized because not all static resource bindings in shader '",
-                              GetShaderTypeLiteralName(GetShaderTypes()[s]),
+                              GetShaderTypeLiteralName(GetShaderStageType(s)),
                               "' are valid. Please make sure you bind all static resources to PSO before calling InitializeStaticResources() "
                               "directly or indirectly by passing InitStaticResources=true to CreateShaderResourceBinding() method.");
         }
diff --git a/Graphics/GraphicsEngineVulkan/src/ShaderResourceBindingVkImpl.cpp b/Graphics/GraphicsEngineVulkan/src/ShaderResourceBindingVkImpl.cpp
index 75f3bb37..91008811 100644
--- a/Graphics/GraphicsEngineVulkan/src/ShaderResourceBindingVkImpl.cpp
+++ b/Graphics/GraphicsEngineVulkan/src/ShaderResourceBindingVkImpl.cpp
@@ -49,8 +49,7 @@ ShaderResourceBindingVkImpl::ShaderResourceBindingVkImpl(IReferenceCounters*  pR
 {
     m_ResourceLayoutIndex.fill(-1);
 
-    auto* pShaderTypes = pPSO->GetShaderTypes();
-    m_NumShaders       = static_cast<decltype(m_NumShaders)>(pPSO->GetNumShaderTypes());
+    m_NumShaders = static_cast<decltype(m_NumShaders)>(pPSO->GetNumShaderStages());
 
     auto* pRenderDeviceVkImpl = pPSO->GetDevice();
     // This will only allocate memory and initialize descriptor sets in the resource cache
@@ -62,7 +61,7 @@ ShaderResourceBindingVkImpl::ShaderResourceBindingVkImpl(IReferenceCounters*  pR
 
     for (Uint32 s = 0; s < m_NumShaders; ++s)
     {
-        auto ShaderInd = GetShaderTypePipelineIndex(pShaderTypes[s], pPSO->GetDesc().PipelineType);
+        auto ShaderInd = GetShaderTypePipelineIndex(pPSO->GetShaderStageType(s), pPSO->GetDesc().PipelineType);
 
         m_ResourceLayoutIndex[ShaderInd] = static_cast<Int8>(s);
 
diff --git a/Graphics/GraphicsEngineVulkan/src/ShaderResourceLayoutVk.cpp b/Graphics/GraphicsEngineVulkan/src/ShaderResourceLayoutVk.cpp
index a7410210..8b00c394 100644
--- a/Graphics/GraphicsEngineVulkan/src/ShaderResourceLayoutVk.cpp
+++ b/Graphics/GraphicsEngineVulkan/src/ShaderResourceLayoutVk.cpp
@@ -93,6 +93,14 @@ static SHADER_RESOURCE_VARIABLE_TYPE FindShaderVariableType(SHADER_TYPE
     }
 }
 
+ShaderResourceLayoutVk::ShaderStageInfo::ShaderStageInfo(SHADER_TYPE         _Type,
+                                                         const ShaderVkImpl* _pShader) :
+    Type{_Type},
+    pShader{_pShader},
+    SPIRV{pShader->GetSPIRV()}
+{
+}
+
 
 ShaderResourceLayoutVk::~ShaderResourceLayoutVk()
 {
@@ -103,7 +111,7 @@ ShaderResourceLayoutVk::~ShaderResourceLayoutVk()
         GetImmutableSampler(s).~ImmutableSamplerPtrType();
 }
 
-void ShaderResourceLayoutVk::AllocateMemory(IShader*                             pShader,
+void ShaderResourceLayoutVk::AllocateMemory(const ShaderVkImpl*                  pShader,
                                             IMemoryAllocator&                    Allocator,
                                             const PipelineResourceLayoutDesc&    ResourceLayoutDesc,
                                             const SHADER_RESOURCE_VARIABLE_TYPE* AllowedVarTypes,
@@ -112,16 +120,14 @@ void ShaderResourceLayoutVk::AllocateMemory(IShader*
 {
     VERIFY(!m_ResourceBuffer, "Memory has already been initialized");
     VERIFY_EXPR(!m_pResources);
+    m_pResources = pShader->GetShaderResources();
 
-    const Uint32 AllowedTypeBits = GetAllowedTypeBits(AllowedVarTypes, NumAllowedTypes);
-    const auto   ShaderType      = pShader->GetDesc().ShaderType;
+    const auto ShaderType = pShader->GetDesc().ShaderType;
+    VERIFY_EXPR(m_pResources->GetShaderType() == ShaderType);
     // Count the number of resources to allocate all needed memory
     {
-        auto*       pShaderVk             = ValidatedCast<ShaderVkImpl>(pShader);
-        auto        pResources            = pShaderVk->GetShaderResources();
-        const auto* CombinedSamplerSuffix = pResources->GetCombinedSamplerSuffix();
-        VERIFY_EXPR(pResources->GetShaderType() == ShaderType);
-        m_pResources = pResources;
+        const Uint32 AllowedTypeBits       = GetAllowedTypeBits(AllowedVarTypes, NumAllowedTypes);
+        const auto*  CombinedSamplerSuffix = m_pResources->GetCombinedSamplerSuffix();
         m_pResources->ProcessResources(
             [&](const SPIRVShaderResourceAttribs& ResAttribs, Uint32) //
             {
@@ -172,11 +178,11 @@ void ShaderResourceLayoutVk::AllocateMemory(IShader*
 }
 
 
-Uint32 FindAssignedSampler(const ShaderResourceLayoutVk&     Layout,
-                           const SPIRVShaderResources&       Resources,
-                           const SPIRVShaderResourceAttribs& SepImg,
-                           Uint32                            CurrResourceCount,
-                           SHADER_RESOURCE_VARIABLE_TYPE     ImgVarType)
+static Uint32 FindAssignedSampler(const ShaderResourceLayoutVk&     Layout,
+                                  const SPIRVShaderResources&       Resources,
+                                  const SPIRVShaderResourceAttribs& SepImg,
+                                  Uint32                            CurrResourceCount,
+                                  SHADER_RESOURCE_VARIABLE_TYPE     ImgVarType)
 {
     using VkResource = ShaderResourceLayoutVk::VkResource;
     VERIFY_EXPR(SepImg.Type == SPIRVShaderResourceAttribs::ResourceType::SeparateImage);
@@ -211,7 +217,7 @@ Uint32 FindAssignedSampler(const ShaderResourceLayoutVk&     Layout,
 }
 
 
-void ShaderResourceLayoutVk::InitializeStaticResourceLayout(IShader*                          pShader,
+void ShaderResourceLayoutVk::InitializeStaticResourceLayout(const ShaderVkImpl*               pShader,
                                                             IMemoryAllocator&                 LayoutDataAllocator,
                                                             const PipelineResourceLayoutDesc& ResourceLayoutDesc,
                                                             ShaderResourceCacheVk&            StaticResourceCache)
@@ -277,7 +283,7 @@ void ShaderResourceLayoutVk::InitializeStaticResourceLayout(IShader*
 }
 
 #ifdef DILIGENT_DEVELOPMENT
-void ShaderResourceLayoutVk::dvpVerifyResourceLayoutDesc(const ShaderStages_t&             ShaderStages,
+void ShaderResourceLayoutVk::dvpVerifyResourceLayoutDesc(const TShaderStages&              ShaderStages,
                                                          const PipelineResourceLayoutDesc& ResourceLayoutDesc,
                                                          bool                              VerifyVariables,
                                                          bool                              VerifyStaticSamplers)
@@ -290,11 +296,11 @@ void ShaderResourceLayoutVk::dvpVerifyResourceLayoutDesc(const ShaderStages_t&
             const auto  ShaderType = Stages & static_cast<SHADER_TYPE>(~(static_cast<Uint32>(Stages) - 1));
             const char* ShaderName = nullptr;
 
-            for (size_t s = 0; s < ShaderStages.size(); ++s)
+            for (const auto& StageInfo : ShaderStages)
             {
-                if ((Stages & ShaderStages[s].first) != 0)
+                if ((Stages & StageInfo.Type) != 0)
                 {
-                    ShaderName = ShaderStages[s].second->GetDesc().Name;
+                    ShaderName = StageInfo.pShader->GetDesc().Name;
                     break;
                 }
             }
@@ -334,8 +340,7 @@ void ShaderResourceLayoutVk::dvpVerifyResourceLayoutDesc(const ShaderStages_t&
             bool VariableFound = false;
             for (size_t s = 0; s < ShaderStages.size() && !VariableFound; ++s)
             {
-                const auto* pShaderVk = ValidatedCast<ShaderVkImpl>(ShaderStages[s].second);
-                const auto& Resources = *pShaderVk->GetShaderResources();
+                const auto& Resources = *ShaderStages[s].pShader->GetShaderResources();
                 if ((VarDesc.ShaderStages & Resources.GetShaderType()) != 0)
                 {
                     for (Uint32 res = 0; res < Resources.GetTotalResources() && !VariableFound; ++res)
@@ -368,8 +373,7 @@ void ShaderResourceLayoutVk::dvpVerifyResourceLayoutDesc(const ShaderStages_t&
             bool SamplerFound = false;
             for (size_t s = 0; s < ShaderStages.size() && !SamplerFound; ++s)
             {
-                const auto* pShaderVk = ValidatedCast<ShaderVkImpl>(ShaderStages[s].second);
-                const auto& Resources = *pShaderVk->GetShaderResources();
+                const auto& Resources = *ShaderStages[s].pShader->GetShaderResources();
                 if ((StSamDesc.ShaderStages & Resources.GetShaderType()) == 0)
                     continue;
 
@@ -407,11 +411,10 @@ void ShaderResourceLayoutVk::dvpVerifyResourceLayoutDesc(const ShaderStages_t&
 #endif
 
 void ShaderResourceLayoutVk::Initialize(IRenderDevice*                    pRenderDevice,
-                                        const ShaderStages_t&             ShaderStages,
+                                        TShaderStages&                    ShaderStages,
                                         ShaderResourceLayoutVk            Layouts[],
                                         IMemoryAllocator&                 LayoutDataAllocator,
                                         const PipelineResourceLayoutDesc& ResourceLayoutDesc,
-                                        ShaderSPIRVs_t&                   SPIRVs,
                                         class PipelineLayout&             PipelineLayout,
                                         bool                              VerifyVariables,
                                         bool                              VerifyStaticSamplers)
@@ -427,7 +430,8 @@ void ShaderResourceLayoutVk::Initialize(IRenderDevice*                    pRende
 
     for (size_t s = 0; s < ShaderStages.size(); ++s)
     {
-        Layouts[s].AllocateMemory(ShaderStages[s].second, LayoutDataAllocator, ResourceLayoutDesc, AllowedVarTypes, NumAllowedTypes, AllocateImmutableSamplers);
+        Layouts[s].AllocateMemory(ShaderStages[s].pShader, LayoutDataAllocator, ResourceLayoutDesc,
+                                  AllowedVarTypes, NumAllowedTypes, AllocateImmutableSamplers);
     }
 
     //VERIFY_EXPR(NumShaders <= MAX_SHADERS_IN_PIPELINE);
@@ -475,7 +479,7 @@ void ShaderResourceLayoutVk::Initialize(IRenderDevice*                    pRende
             }
         }
 
-        auto& ShaderSPIRV = SPIRVs[ShaderInd];
+        auto& ShaderSPIRV = ShaderStages[ShaderInd].SPIRV;
         PipelineLayout.AllocateResourceSlot(Attribs, VarType, vkImmutableSampler, Resources.GetShaderType(), DescriptorSet, Binding, CacheOffset, ShaderSPIRV);
         VERIFY(DescriptorSet <= std::numeric_limits<decltype(VkResource::DescriptorSet)>::max(), "Descriptor set (", DescriptorSet, ") excceeds maximum representable value");
         VERIFY(Binding <= std::numeric_limits<decltype(VkResource::Binding)>::max(), "Binding (", Binding, ") excceeds maximum representable value");
@@ -498,9 +502,8 @@ void ShaderResourceLayoutVk::Initialize(IRenderDevice*                    pRende
     // First process uniform buffers for all shader stages to make sure all UBs go first in every descriptor set
     for (size_t s = 0; s < ShaderStages.size(); ++s)
     {
-        auto& pShader   = ShaderStages[s].second;
         auto& Layout    = Layouts[s];
-        auto* pShaderVk = ValidatedCast<ShaderVkImpl>(pShader);
+        auto* pShaderVk = ShaderStages[s].pShader;
         auto& Resources = *pShaderVk->GetShaderResources();
         for (Uint32 n = 0; n < Resources.GetNumUBs(); ++n)
         {
@@ -516,10 +519,8 @@ void ShaderResourceLayoutVk::Initialize(IRenderDevice*                    pRende
     // Second, process all storage buffers
     for (size_t s = 0; s < ShaderStages.size(); ++s)
     {
-        auto& pShader   = ShaderStages[s].second;
         auto& Layout    = Layouts[s];
-        auto* pShaderVk = ValidatedCast<ShaderVkImpl>(pShader);
-        auto& Resources = *pShaderVk->GetShaderResources();
+        auto& Resources = *ShaderStages[s].pShader->GetShaderResources();
         for (Uint32 n = 0; n < Resources.GetNumSBs(); ++n)
         {
             const auto& SB      = Resources.GetSB(n);
@@ -534,10 +535,8 @@ void ShaderResourceLayoutVk::Initialize(IRenderDevice*                    pRende
     // Finally, process all other resource types
     for (size_t s = 0; s < ShaderStages.size(); ++s)
     {
-        auto& pShader   = ShaderStages[s].second;
         auto& Layout    = Layouts[s];
-        auto* pShaderVk = ValidatedCast<ShaderVkImpl>(pShader);
-        auto& Resources = *pShaderVk->GetShaderResources();
+        auto& Resources = *ShaderStages[s].pShader->GetShaderResources();
         // clang-format off
         Resources.ProcessResources(
             [&](const SPIRVShaderResourceAttribs& UB, Uint32)
diff --git a/Graphics/GraphicsEngineVulkan/src/VulkanTypeConversions.cpp b/Graphics/GraphicsEngineVulkan/src/VulkanTypeConversions.cpp
index 95261295..55bb7d8f 100644
--- a/Graphics/GraphicsEngineVulkan/src/VulkanTypeConversions.cpp
+++ b/Graphics/GraphicsEngineVulkan/src/VulkanTypeConversions.cpp
@@ -1515,7 +1515,7 @@ VkAccessFlags AccessFlagsToVkAccessFlags(ACCESS_FLAGS AccessFlags)
 VkShaderStageFlagBits ShaderTypeToVkShaderStageFlagBit(SHADER_TYPE ShaderType)
 {
     static_assert(SHADER_TYPE_LAST == SHADER_TYPE_MESH, "Please update the switch below to handle the new shader type");
-    VERIFY((ShaderType & (ShaderType - 1)) == 0, "More than one shader type specified");
+    VERIFY(IsPowerOfTwo(Uint32{ShaderType}), "More than one shader type is specified");
     switch (ShaderType)
     {
         // clang-format off