Few minor improvements to RT pipeline initialization

3 files changed, 32 insertions, 21 deletions

diff --git a/Graphics/GraphicsEngine/include/PipelineStateBase.hpp b/Graphics/GraphicsEngine/include/PipelineStateBase.hpp
index 2c46f314..79033337 100644
--- a/Graphics/GraphicsEngine/include/PipelineStateBase.hpp
+++ b/Graphics/GraphicsEngine/include/PipelineStateBase.hpp
@@ -51,9 +51,9 @@ void ValidateComputePipelineCreateInfo(const ComputePipelineStateCreateInfo& Cre
 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;
+void CopyRTShaderGroupNames(std::unordered_map<HashMapStringKey, Uint32, HashMapStringKey::Hasher>& NameToGroupIndex,
+                            const RayTracingPipelineStateCreateInfo&                                CreateInfo,
+                            FixedLinearAllocator&                                                   MemPool) noexcept;
 
 void CorrectGraphicsPipelineDesc(GraphicsPipelineDesc& GraphicsPipeline) noexcept;
 
@@ -246,7 +246,7 @@ public:
         VERIFY_EXPR(m_pRayTracingPipelineData != nullptr);
 
         const auto ShaderHandleSize = m_pRayTracingPipelineData->ShaderHandleSize;
-        VERIFY_EXPR(ShaderHandleSize <= DataSize);
+        VERIFY(ShaderHandleSize <= DataSize, "DataSize (", DataSize, ") must be at least as large as the shader handle size (", ShaderHandleSize, ").");
 
         if (Name == nullptr || Name[0] == '\0')
         {
@@ -262,7 +262,7 @@ public:
             std::memcpy(pData, &m_pRayTracingPipelineData->ShaderHandles[ShaderHandleSize * iter->second], ShaderHandleSize);
             return;
         }
-        UNEXPECTED("Can't find shader group with the specified name");
+        UNEXPECTED("Can't find shader group '", Name, "'.");
     }
 
 protected:
@@ -657,7 +657,7 @@ protected:
                                 FixedLinearAllocator&                    MemPool) noexcept
     {
         TNameToGroupIndexMap NameToGroupIndex;
-        CopyRayTracingShaderGroups(NameToGroupIndex, CreateInfo, MemPool);
+        CopyRTShaderGroupNames(NameToGroupIndex, CreateInfo, MemPool);
 
         CopyResourceLayout(CreateInfo.PSODesc.ResourceLayout, this->m_Desc.ResourceLayout, MemPool);
 
@@ -761,12 +761,20 @@ protected:
     struct RayTracingPipelineData
     {
         RayTracingPipelineDesc Desc;
-        TNameToGroupIndexMap   NameToGroupIndex;
+
+        // Mapping from the shader group name to its index in the pipeline.
+        // It is used to find the shader handle in ShaderHandles array.
+        TNameToGroupIndexMap NameToGroupIndex;
 
         Uint32 ShaderHandleSize = 0;
         Uint32 ShaderDataSize   = 0;
 
-        Uint8 ShaderHandles[sizeof(void*)] = {}; // The actual array size will be ShaderDataSize
+        // Array of shader handles for every group in the pipeline.
+        // The handles will be copied to the SBT using NameToGroupIndex to find
+        // handles by group name.
+        // The actual array size will be determined at run time and will be stored
+        // in ShaderDataSize.
+        Uint8 ShaderHandles[sizeof(void*)] = {};
     };
     static_assert(offsetof(RayTracingPipelineData, ShaderHandles) % sizeof(void*) == 0, "ShaderHandles member is expected to be sizeof(void*)-aligned");
 
diff --git a/Graphics/GraphicsEngine/src/PipelineStateBase.cpp b/Graphics/GraphicsEngine/src/PipelineStateBase.cpp
index f99b3a0f..2f0297bf 100644
--- a/Graphics/GraphicsEngine/src/PipelineStateBase.cpp
+++ b/Graphics/GraphicsEngine/src/PipelineStateBase.cpp
@@ -278,7 +278,7 @@ void ValidateRayTracingPipelineCreateInfo(IRenderDevice* pDevice, Uint32 MaxRecu
         if (GroupName == nullptr)
             LOG_PSO_ERROR_AND_THROW(MemberName, "[", GroupInd, "].Name must not be null.");
 
-        if (*GroupName == 0)
+        if (GroupName[0] == '\0')
             LOG_PSO_ERROR_AND_THROW(MemberName, "[", GroupInd, "].Name must not be empty.");
 
         const bool IsNewName = GroupNames.emplace(HashMapStringKey{GroupName}).second;
@@ -335,9 +335,9 @@ void ValidateRayTracingPipelineCreateInfo(IRenderDevice* pDevice, Uint32 MaxRecu
     }
 }
 
-void CopyRayTracingShaderGroups(std::unordered_map<HashMapStringKey, Uint32, HashMapStringKey::Hasher>& NameToGroupIndex,
-                                const RayTracingPipelineStateCreateInfo&                                CreateInfo,
-                                FixedLinearAllocator&                                                   MemPool) noexcept
+void CopyRTShaderGroupNames(std::unordered_map<HashMapStringKey, Uint32, HashMapStringKey::Hasher>& NameToGroupIndex,
+                            const RayTracingPipelineStateCreateInfo&                                CreateInfo,
+                            FixedLinearAllocator&                                                   MemPool) noexcept
 {
     Uint32 GroupIndex = 0;
 
diff --git a/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp b/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp
index 50f3bb9e..5427e257 100644
--- a/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp
+++ b/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp
@@ -395,7 +395,7 @@ std::vector<VkRayTracingShaderGroupCreateInfoKHR> BuildRTShaderGroupDescription(
                         return idx;
                 }
                 UNEXPECTED("Unable to find shader '", pShader->GetDesc().Name, "' in the shader stage. This should never happen and is a bug.");
-                return ~0U;
+                return VK_SHADER_UNUSED_KHR;
             }
             else
             {
@@ -403,7 +403,7 @@ std::vector<VkRayTracingShaderGroupCreateInfoKHR> BuildRTShaderGroupDescription(
             }
         }
         UNEXPECTED("Unable to find corresponding shader stage for shader '", pShader->GetDesc().Name, "'. This should never happen and is a bug.");
-        return ~0U;
+        return VK_SHADER_UNUSED_KHR;
     };
 
     std::vector<VkRayTracingShaderGroupCreateInfoKHR> ShaderGroups;
@@ -428,11 +428,11 @@ std::vector<VkRayTracingShaderGroupCreateInfoKHR> BuildRTShaderGroupDescription(
             VERIFY(Iter != NameToGroupIndex.end(),
                    "Can't find general shader '", GeneralShader.Name,
                    "'. This looks to be a bug as NameToGroupIndex is initialized by "
-                   "CopyRayTracingShaderGroups that processes the same general shaders.");
+                   "CopyRTShaderGroupNames() that processes the same general shaders.");
             VERIFY(Iter->second == ShaderGroups.size(),
                    "General shader group '", GeneralShader.Name, "' index mismatch: (", Iter->second, " != ", ShaderGroups.size(),
                    "). This looks to be a bug as NameToGroupIndex is initialized by "
-                   "CopyRayTracingShaderGroups that processes the same shaders in the same order.");
+                   "CopyRTShaderGroupNames() that processes the same shaders in the same order.");
         }
 #endif
 
@@ -458,11 +458,11 @@ std::vector<VkRayTracingShaderGroupCreateInfoKHR> BuildRTShaderGroupDescription(
             VERIFY(Iter != NameToGroupIndex.end(),
                    "Can't find triangle hit group '", TriHitShader.Name,
                    "'. This looks to be a bug as NameToGroupIndex is initialized by "
-                   "CopyRayTracingShaderGroups that processes the same hit groups.");
+                   "CopyRTShaderGroupNames() that processes the same hit groups.");
             VERIFY(Iter->second == ShaderGroups.size(),
                    "Triangle hit group '", TriHitShader.Name, "' index mismatch: (", Iter->second, " != ", ShaderGroups.size(),
                    "). This looks to be a bug as NameToGroupIndex is initialized by "
-                   "CopyRayTracingShaderGroups that processes the same hit groups in the same order.");
+                   "CopyRTShaderGroupNames() that processes the same hit groups in the same order.");
         }
 #endif
 
@@ -488,11 +488,11 @@ std::vector<VkRayTracingShaderGroupCreateInfoKHR> BuildRTShaderGroupDescription(
             VERIFY(Iter != NameToGroupIndex.end(),
                    "Can't find procedural hit group '", ProcHitShader.Name,
                    "'. This looks to be a bug as NameToGroupIndex is initialized by "
-                   "CopyRayTracingShaderGroups that processes the same hit groups.");
+                   "CopyRTShaderGroupNames() that processes the same hit groups.");
             VERIFY(Iter->second == ShaderGroups.size(),
                    "Procedural hit group '", ProcHitShader.Name, "' index mismatch: (", Iter->second, " != ", ShaderGroups.size(),
                    "). This looks to be a bug as NameToGroupIndex is initialized by "
-                   "CopyRayTracingShaderGroups that processes the same hit groups in the same order.");
+                   "CopyRTShaderGroupNames() that processes the same hit groups in the same order.");
         }
 #endif
 
@@ -697,7 +697,7 @@ PipelineStateVkImpl::TShaderStages PipelineStateVkImpl::InitInternalObjects(
     InitResourceLayouts(CreateInfo, ShaderStages);
 
     // Create shader modules and initialize shader stages
-    InitPipelineShaderStages(GetDevice()->GetLogicalDevice(), ShaderStages, ShaderModules, vkShaderStages);
+    InitPipelineShaderStages(LogicalDevice, ShaderStages, ShaderModules, vkShaderStages);
 
     return ShaderStages;
 }
@@ -768,6 +768,9 @@ PipelineStateVkImpl::PipelineStateVkImpl(IReferenceCounters*
 
         CreateRayTracingPipeline(pDeviceVk, vkShaderStages, ShaderGroups, m_PipelineLayout, m_Desc, GetRayTracingPipelineDesc(), m_Pipeline);
 
+        VERIFY(m_pRayTracingPipelineData->NameToGroupIndex.size() == ShaderGroups.size(),
+               "The size of NameToGroupIndex map does not match the actual number of groups in the pipeline. This is a bug.");
+        // Get shader group handles from the PSO.
         auto err = LogicalDevice.GetRayTracingShaderGroupHandles(m_Pipeline, 0, static_cast<uint32_t>(ShaderGroups.size()), m_pRayTracingPipelineData->ShaderDataSize, m_pRayTracingPipelineData->ShaderHandles);
         DEV_CHECK_ERR(err == VK_SUCCESS, "Failed to get shader group handles");
         (void)err;