PSO initialization: updated shader stage initialization

3 files changed, 20 insertions, 4 deletions

diff --git a/Graphics/GraphicsEngineVulkan/include/ShaderResourceLayoutVk.hpp b/Graphics/GraphicsEngineVulkan/include/ShaderResourceLayoutVk.hpp
index 41da59cb..f0b17529 100644
--- a/Graphics/GraphicsEngineVulkan/include/ShaderResourceLayoutVk.hpp
+++ b/Graphics/GraphicsEngineVulkan/include/ShaderResourceLayoutVk.hpp
@@ -110,7 +110,7 @@ public:
     struct ShaderStageInfo
     {
         ShaderStageInfo() {}
-        ShaderStageInfo(SHADER_TYPE Stage, const ShaderVkImpl* pShader);
+        ShaderStageInfo(const ShaderVkImpl* pShader);
 
         void   Append(const ShaderVkImpl* pShader);
         size_t Count() const;
diff --git a/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp b/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp
index 6ef0e47c..42970686 100644
--- a/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp
+++ b/Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp
@@ -798,7 +798,7 @@ PipelineStateVkImpl::PipelineStateVkImpl(IReferenceCounters*
 
         CreateRayTracingPipeline(pDeviceVk, vkShaderStages, ShaderGroups, m_PipelineLayout, m_Desc, GetRayTracingPipelineDesc(), m_Pipeline);
 
-        auto err = LogicalDevice.GetRayTracingShaderGroupHandles(m_Pipeline, 0, static_cast<uint32_t>(ShaderGroups.size()), m_pRayTracingPipelineData->ShaderDataSize, &m_pRayTracingPipelineData->Shaders[0]);
+        auto err = LogicalDevice.GetRayTracingShaderGroupHandles(m_Pipeline, 0, static_cast<uint32_t>(ShaderGroups.size()), m_pRayTracingPipelineData->ShaderDataSize, &m_pRayTracingPipelineData->ShaderHandles[0]);
         VERIFY(err == VK_SUCCESS, "Failed to get shader group handles");
         (void)err;
     }
diff --git a/Graphics/GraphicsEngineVulkan/src/ShaderResourceLayoutVk.cpp b/Graphics/GraphicsEngineVulkan/src/ShaderResourceLayoutVk.cpp
index 8439ef06..d23eb131 100644
--- a/Graphics/GraphicsEngineVulkan/src/ShaderResourceLayoutVk.cpp
+++ b/Graphics/GraphicsEngineVulkan/src/ShaderResourceLayoutVk.cpp
@@ -95,8 +95,8 @@ static SHADER_RESOURCE_VARIABLE_TYPE FindShaderVariableType(SHADER_TYPE
 }
 
 
-ShaderResourceLayoutVk::ShaderStageInfo::ShaderStageInfo(SHADER_TYPE Stage, const ShaderVkImpl* pShader) :
-    Type{Stage},
+ShaderResourceLayoutVk::ShaderStageInfo::ShaderStageInfo(const ShaderVkImpl* pShader) :
+    Type{pShader->GetDesc().ShaderType},
     Shaders{pShader},
     SPIRVs{pShader->GetSPIRV()}
 {
@@ -104,6 +104,22 @@ ShaderResourceLayoutVk::ShaderStageInfo::ShaderStageInfo(SHADER_TYPE Stage, cons
 
 void ShaderResourceLayoutVk::ShaderStageInfo::Append(const ShaderVkImpl* pShader)
 {
+    VERIFY_EXPR(pShader != nullptr);
+    VERIFY(std::find(Shaders.begin(), Shaders.end(), pShader) == Shaders.end(),
+           "Shader '", pShader->GetDesc().Name, "' already exists in the stage. Shaders must be deduplicated.");
+
+    const auto NewShaderType = pShader->GetDesc().ShaderType;
+    if (Type == SHADER_TYPE_UNKNOWN)
+    {
+        VERIFY_EXPR(Shaders.empty() && SPIRVs.empty());
+        Type = NewShaderType;
+    }
+    else
+    {
+        VERIFY(Type == NewShaderType, "The type (", GetShaderTypeLiteralName(NewShaderType),
+               ") of shader '", pShader->GetDesc().Name, "' being added to the stage is incosistent with the stage type (",
+               GetShaderTypeLiteralName(Type), ").");
+    }
     Shaders.push_back(pShader);
     SPIRVs.push_back(pShader->GetSPIRV());
 }