PSO initialization: updated shader stage initialization

2 files changed, 59 insertions, 39 deletions

diff --git a/Graphics/GraphicsEngine/include/PipelineStateBase.hpp b/Graphics/GraphicsEngine/include/PipelineStateBase.hpp
index df57c671..127d46dd 100644
--- a/Graphics/GraphicsEngine/include/PipelineStateBase.hpp
+++ b/Graphics/GraphicsEngine/include/PipelineStateBase.hpp
@@ -50,9 +50,10 @@ void ValidateGraphicsPipelineCreateInfo(const GraphicsPipelineStateCreateInfo& C
 void ValidateComputePipelineCreateInfo(const ComputePipelineStateCreateInfo& CreateInfo) noexcept(false);
 void ValidateRayTracingPipelineCreateInfo(IRenderDevice* pDevice, Uint32 MaxRecursion, const RayTracingPipelineStateCreateInfo& CreateInfo) noexcept(false);
 
+/// Copies ray tracing shader group names and also initializes the mapping from the group name to its index.
 void CopyRayTracingShaderGroups(std::unordered_map<HashMapStringKey, Uint32, HashMapStringKey::Hasher>& NameToGroupIndex,
                                 const RayTracingPipelineStateCreateInfo&                                CreateInfo,
-                                FixedLinearAllocator&                                                   MemPool) noexcept(false);
+                                FixedLinearAllocator&                                                   MemPool) noexcept;
 
 void CorrectGraphicsPipelineDesc(GraphicsPipelineDesc& GraphicsPipeline) noexcept;
 
@@ -234,7 +235,7 @@ public:
         if (iter != m_pRayTracingPipelineData->NameToGroupIndex.end())
         {
             VERIFY_EXPR(ShaderHandleSize * (iter->second + 1) <= m_pRayTracingPipelineData->ShaderDataSize);
-            std::memcpy(pData, &m_pRayTracingPipelineData->Shaders[ShaderHandleSize * iter->second], ShaderHandleSize);
+            std::memcpy(pData, &m_pRayTracingPipelineData->ShaderHandles[ShaderHandleSize * iter->second], ShaderHandleSize);
             return;
         }
         UNEXPECTED("Can't find shader group with the specified name");
@@ -353,7 +354,7 @@ protected:
         const auto ShaderHandleSize = this->m_pDevice->GetProperties().ShaderGroupHandleSize;
         RTDataSize += ShaderHandleSize * (CreateInfo.GeneralShaderCount + CreateInfo.TriangleHitShaderCount + CreateInfo.ProceduralHitShaderCount);
         // Extra bytes were reserved to avoid compiler errors on zero-sized arrays
-        RTDataSize -= sizeof(RayTracingPipelineData::Shaders);
+        RTDataSize -= sizeof(RayTracingPipelineData::ShaderHandles);
         MemPool.AddSpace(RTDataSize, alignof(RayTracingPipelineData));
     }
 
@@ -369,9 +370,16 @@ protected:
         auto AddShaderStage = [&](IShader* pShader) {
             if (pShader != nullptr)
             {
-                auto ShaderType = pShader->GetDesc().ShaderType;
-                ShaderStages.emplace_back(ShaderType, ValidatedCast<ShaderImplType>(pShader));
-                VERIFY(m_ShaderStageTypes[m_NumShaderStages] == SHADER_TYPE_UNKNOWN, "This shader stage is already initialized.");
+                const auto ShaderType = pShader->GetDesc().ShaderType;
+                ShaderStages.emplace_back(ValidatedCast<ShaderImplType>(pShader));
+                VERIFY(m_ShaderStageTypes[m_NumShaderStages] == SHADER_TYPE_UNKNOWN, "This shader stage has already been initialized.");
+#ifdef DILIGENT_DEBUG
+                for (Uint32 i = 0; i < m_NumShaderStages; ++i)
+                {
+                    VERIFY(m_ShaderStageTypes[i] != ShaderType,
+                           "Shader stage ", GetShaderTypeLiteralName(ShaderType), " has already been initialized in PSO '", this->m_Desc.Name, "'.");
+                }
+#endif
                 m_ShaderStageTypes[m_NumShaderStages++] = ShaderType;
             }
         };
@@ -418,7 +426,7 @@ protected:
         VERIFY_EXPR(CreateInfo.pCS != nullptr);
         VERIFY_EXPR(CreateInfo.pCS->GetDesc().ShaderType == SHADER_TYPE_COMPUTE);
 
-        ShaderStages.emplace_back(SHADER_TYPE_COMPUTE, ValidatedCast<ShaderImplType>(CreateInfo.pCS));
+        ShaderStages.emplace_back(ValidatedCast<ShaderImplType>(CreateInfo.pCS));
         m_ShaderStageTypes[m_NumShaderStages++] = SHADER_TYPE_COMPUTE;
 
         VERIFY_EXPR(!ShaderStages.empty() && ShaderStages.size() == m_NumShaderStages);
@@ -433,15 +441,13 @@ protected:
 
         std::unordered_set<IShader*> UniqueShaders;
 
-        auto AddShaderStage = [&ShaderStages, &UniqueShaders](IShader* pShader) {
+        auto AddShader = [&ShaderStages, &UniqueShaders](IShader* pShader) {
             if (pShader != nullptr && UniqueShaders.insert(pShader).second)
             {
-                auto  ShaderType = pShader->GetDesc().ShaderType;
-                auto  StageInd   = GetShaderTypePipelineIndex(ShaderType, PIPELINE_TYPE_RAY_TRACING);
-                auto& Stage      = ShaderStages[StageInd];
+                const auto ShaderType = pShader->GetDesc().ShaderType;
+                const auto StageInd   = GetShaderTypePipelineIndex(ShaderType, PIPELINE_TYPE_RAY_TRACING);
+                auto&      Stage      = ShaderStages[StageInd];
                 Stage.Append(ValidatedCast<ShaderImplType>(pShader));
-                VERIFY_EXPR(Stage.Type == SHADER_TYPE_UNKNOWN || Stage.Type == ShaderType);
-                Stage.Type = ShaderType;
             }
         };
 
@@ -450,18 +456,18 @@ protected:
 
         for (Uint32 i = 0; i < CreateInfo.GeneralShaderCount; ++i)
         {
-            AddShaderStage(CreateInfo.pGeneralShaders[i].pShader);
+            AddShader(CreateInfo.pGeneralShaders[i].pShader);
         }
         for (Uint32 i = 0; i < CreateInfo.TriangleHitShaderCount; ++i)
         {
-            AddShaderStage(CreateInfo.pTriangleHitShaders[i].pClosestHitShader);
-            AddShaderStage(CreateInfo.pTriangleHitShaders[i].pAnyHitShader);
+            AddShader(CreateInfo.pTriangleHitShaders[i].pClosestHitShader);
+            AddShader(CreateInfo.pTriangleHitShaders[i].pAnyHitShader);
         }
         for (Uint32 i = 0; i < CreateInfo.ProceduralHitShaderCount; ++i)
         {
-            AddShaderStage(CreateInfo.pProceduralHitShaders[i].pIntersectionShader);
-            AddShaderStage(CreateInfo.pProceduralHitShaders[i].pClosestHitShader);
-            AddShaderStage(CreateInfo.pProceduralHitShaders[i].pAnyHitShader);
+            AddShader(CreateInfo.pProceduralHitShaders[i].pIntersectionShader);
+            AddShader(CreateInfo.pProceduralHitShaders[i].pClosestHitShader);
+            AddShader(CreateInfo.pProceduralHitShaders[i].pAnyHitShader);
         }
 
         if (ShaderStages[GetShaderTypePipelineIndex(SHADER_TYPE_RAY_GEN, PIPELINE_TYPE_RAY_TRACING)].Count() == 0)
@@ -479,7 +485,7 @@ protected:
                 continue;
             }
 
-            VERIFY(m_ShaderStageTypes[m_NumShaderStages] == SHADER_TYPE_UNKNOWN, "This shader stage is already initialized.");
+            VERIFY(m_ShaderStageTypes[m_NumShaderStages] == SHADER_TYPE_UNKNOWN, "This shader stage has already been initialized.");
             m_ShaderStageTypes[m_NumShaderStages++] = iter->Type;
             ++iter;
         }
@@ -637,7 +643,7 @@ protected:
         const auto ShaderDataSize   = ShaderHandleSize * (CreateInfo.GeneralShaderCount + CreateInfo.TriangleHitShaderCount + CreateInfo.ProceduralHitShaderCount);
         RTDataSize += ShaderDataSize;
         // Extra bytes were reserved to avoid compiler errors on zero-sized arrays
-        RTDataSize -= sizeof(RayTracingPipelineData::Shaders);
+        RTDataSize -= sizeof(RayTracingPipelineData::ShaderHandles);
 
         this->m_pRayTracingPipelineData = static_cast<RayTracingPipelineData*>(MemPool.Allocate(RTDataSize, alignof(RayTracingPipelineData)));
         new (this->m_pRayTracingPipelineData) RayTracingPipelineData{};
@@ -736,9 +742,9 @@ protected:
         Uint32 ShaderHandleSize = 0;
         Uint32 ShaderDataSize   = 0;
 
-        Uint8 Shaders[sizeof(void*)] = {}; // The actual array size will be ShaderDataSize
+        Uint8 ShaderHandles[sizeof(void*)] = {}; // The actual array size will be ShaderDataSize
     };
-    static_assert(offsetof(RayTracingPipelineData, Shaders) % sizeof(void*) == 0, "Shaders member is expected to be sizeof(void*)-aligned");
+    static_assert(offsetof(RayTracingPipelineData, ShaderHandles) % sizeof(void*) == 0, "ShaderHandles member is expected to be sizeof(void*)-aligned");
 
     union
     {
diff --git a/Graphics/GraphicsEngine/src/PipelineStateBase.cpp b/Graphics/GraphicsEngine/src/PipelineStateBase.cpp
index 8a24ea32..cb25f257 100644
--- a/Graphics/GraphicsEngine/src/PipelineStateBase.cpp
+++ b/Graphics/GraphicsEngine/src/PipelineStateBase.cpp
@@ -27,6 +27,10 @@
 
 #include "PipelineStateBase.hpp"
 
+#include <unordered_set>
+
+#include "HashUtils.hpp"
+
 namespace Diligent
 {
 
@@ -265,6 +269,8 @@ void ValidateRayTracingPipelineCreateInfo(IRenderDevice* pDevice, Uint32 MaxRecu
         LOG_PSO_ERROR_AND_THROW("MaxRecursionDepth must not exceed the ", MaxRecursion);
     }
 
+    std::unordered_set<HashMapStringKey, HashMapStringKey::Hasher> GroupNames;
+
     for (Uint32 i = 0; i < CreateInfo.GeneralShaderCount; ++i)
     {
         const auto& Group = CreateInfo.pGeneralShaders[i];
@@ -273,6 +279,10 @@ void ValidateRayTracingPipelineCreateInfo(IRenderDevice* pDevice, Uint32 MaxRecu
         if (Group.Name == nullptr)
             LOG_PSO_ERROR_AND_THROW("pGeneralShaders[", i, "].Name must not be null.");
 
+        const bool IsNewName = GroupNames.emplace(HashMapStringKey{Group.Name}).second;
+        if (!IsNewName)
+            LOG_PSO_ERROR_AND_THROW("pGeneralShaders[", i, "].Name ('", Group.Name, "') has already been assigned to another group. All group names must be unique.");
+
         switch (Group.pShader->GetDesc().ShaderType)
         {
             case SHADER_TYPE_RAY_GEN:
@@ -291,6 +301,10 @@ void ValidateRayTracingPipelineCreateInfo(IRenderDevice* pDevice, Uint32 MaxRecu
         if (Group.Name == nullptr)
             LOG_PSO_ERROR_AND_THROW("pTriangleHitShaders[", i, "].Name must not be null.");
 
+        const bool IsNewName = GroupNames.emplace(HashMapStringKey{Group.Name}).second;
+        if (!IsNewName)
+            LOG_PSO_ERROR_AND_THROW("pTriangleHitShaders[", i, "].Name ('", Group.Name, "') has already been assigned to another group. All group names must be unique.");
+
         VALIDATE_SHADER_TYPE(Group.pClosestHitShader, SHADER_TYPE_RAY_CLOSEST_HIT, "ray tracing triangle closest hit.");
 
         if (Group.pAnyHitShader != nullptr)
@@ -305,6 +319,10 @@ void ValidateRayTracingPipelineCreateInfo(IRenderDevice* pDevice, Uint32 MaxRecu
         if (Group.Name == nullptr)
             LOG_PSO_ERROR_AND_THROW("pProceduralHitShaders[", i, "].Name must not be null.");
 
+        const bool IsNewName = GroupNames.emplace(HashMapStringKey{Group.Name}).second;
+        if (!IsNewName)
+            LOG_PSO_ERROR_AND_THROW("pProceduralHitShaders[", i, "].Name ('", Group.Name, "') has already been assigned to another group. All group names must be unique.");
+
         VALIDATE_SHADER_TYPE(Group.pIntersectionShader, SHADER_TYPE_RAY_INTERSECTION, "ray tracing procedural intersection.");
 
         if (Group.pClosestHitShader != nullptr)
@@ -316,34 +334,30 @@ void ValidateRayTracingPipelineCreateInfo(IRenderDevice* pDevice, Uint32 MaxRecu
 
 void CopyRayTracingShaderGroups(std::unordered_map<HashMapStringKey, Uint32, HashMapStringKey::Hasher>& NameToGroupIndex,
                                 const RayTracingPipelineStateCreateInfo&                                CreateInfo,
-                                FixedLinearAllocator&                                                   MemPool) noexcept(false)
+                                FixedLinearAllocator&                                                   MemPool) noexcept
 {
-    const auto& PSODesc    = CreateInfo.PSODesc;
-    Uint32      GroupIndex = 0;
+    Uint32 GroupIndex = 0;
 
     for (Uint32 i = 0; i < CreateInfo.GeneralShaderCount; ++i)
     {
-        const auto* Name         = CreateInfo.pGeneralShaders[i].Name;
-        const bool  IsUniqueName = NameToGroupIndex.emplace(HashMapStringKey{MemPool.CopyString(Name)}, GroupIndex++).second;
-        if (!IsUniqueName)
-            LOG_PSO_ERROR_AND_THROW("pGeneralShaders[", i, "].Name ('", Name, "') must be unique.");
+        const auto* Name      = CreateInfo.pGeneralShaders[i].Name;
+        const bool  IsNewName = NameToGroupIndex.emplace(HashMapStringKey{MemPool.CopyString(Name)}, GroupIndex++).second;
+        VERIFY(IsNewName, "All group names must be unique. ValidateRayTracingPipelineCreateInfo() should've caught this error.");
     }
     for (Uint32 i = 0; i < CreateInfo.TriangleHitShaderCount; ++i)
     {
-        const auto* Name         = CreateInfo.pTriangleHitShaders[i].Name;
-        const bool  IsUniqueName = NameToGroupIndex.emplace(HashMapStringKey{MemPool.CopyString(Name)}, GroupIndex++).second;
-        if (!IsUniqueName)
-            LOG_PSO_ERROR_AND_THROW("pTriangleHitShaders[", i, "].Name ('", Name, "') must be unique.");
+        const auto* Name      = CreateInfo.pTriangleHitShaders[i].Name;
+        const bool  IsNewName = NameToGroupIndex.emplace(HashMapStringKey{MemPool.CopyString(Name)}, GroupIndex++).second;
+        VERIFY(IsNewName, "All group names must be unique. ValidateRayTracingPipelineCreateInfo() should've caught this error.");
     }
     for (Uint32 i = 0; i < CreateInfo.ProceduralHitShaderCount; ++i)
     {
-        const auto* Name         = CreateInfo.pProceduralHitShaders[i].Name;
-        const bool  IsUniqueName = NameToGroupIndex.emplace(HashMapStringKey{MemPool.CopyString(Name)}, GroupIndex++).second;
-        if (!IsUniqueName)
-            LOG_PSO_ERROR_AND_THROW("pProceduralHitShaders[", i, "].Name ('", Name, "') must be unique.");
+        const auto* Name      = CreateInfo.pProceduralHitShaders[i].Name;
+        const bool  IsNewName = NameToGroupIndex.emplace(HashMapStringKey{MemPool.CopyString(Name)}, GroupIndex++).second;
+        VERIFY(IsNewName, "All group names must be unique. ValidateRayTracingPipelineCreateInfo() should've caught this error.");
     }
 
-    VERIFY_EXPR(Uint32{CreateInfo.GeneralShaderCount} + Uint32{CreateInfo.TriangleHitShaderCount} + Uint32{CreateInfo.ProceduralHitShaderCount} == GroupIndex);
+    VERIFY_EXPR(CreateInfo.GeneralShaderCount + CreateInfo.TriangleHitShaderCount + CreateInfo.ProceduralHitShaderCount == GroupIndex);
 }
 
 #undef VALIDATE_SHADER_TYPE