Vulkan: fix handling of push constant blocks of different sizes across shader stages

Author: TheMostDiligent

Graphics/GraphicsEngineVulkan/src/PipelineStateVkImpl.cpp

@@ -566,23 +566,10 @@ MergedPushConstantMapType MergePushConstants(const PipelineStateVkImpl::TShaderS
                 const SPIRVShaderResourceAttribs& PCAttribs  = ShaderResources.GetPushConstant(i);
                 MergedPushConstantInfo&           MergedPC   = MergedPushConstants[PCAttribs.Name];
                 const Uint32                      BufferSize = PCAttribs.GetConstantBufferSize();
-                if (MergedPC.Stages == SHADER_TYPE_UNKNOWN)
-                {
-                    MergedPC.Stages = Stage.Type;
-                    MergedPC.Size   = BufferSize;
-                }
-                else if (MergedPC.Size == BufferSize)
-                {
-                    MergedPC.Stages |= Stage.Type;
-                }
-                else
-                {
-                    LOG_ERROR_AND_THROW("Push constant '", PCAttribs.Name,
-                                        "' is shared between multiple shaders in pipeline '", (PSOName ? PSOName : ""),
-                                        "', but its size varies (",
-                                        MergedPC.Size, " vs ", BufferSize,
-                                        "). A push constant shared between multiple shaders must have the same size in all shaders.");
-                }
+
+                // Combine push constants from all stages
+                MergedPC.Stages |= Stage.Type;
+                MergedPC.Size = std::max(MergedPC.Size, BufferSize);
             }
         }
     }
@@ -682,7 +669,7 @@ PipelineResourceSignatureDescWrapper PipelineStateVkImpl::GetDefaultResourceSign
                         auto it = MergedPushConstants.find(VarDesc.Name);
                         if (it != MergedPushConstants.end())
                         {
-                            VERIFY_EXPR(Attribs.GetConstantBufferSize() == it->second.Size);
+                            VERIFY_EXPR(Attribs.GetConstantBufferSize() <= it->second.Size);
                             VarDesc.ShaderStages = it->second.Stages;
                         }
                         else

Tests/DiligentCoreAPITest/src/InlineConstantsTest.cpp

@@ -108,7 +108,7 @@ const std::string VulkanPushConstants_VS{
 struct PushConstants_t
 {
     float4 g_Positions[6];
-    float4 g_Colors[3];
+    // Make VS push constant size smaller than PS
 };
 
 [[vk::push_constant]] ConstantBuffer<PushConstants_t> PushConstants;
@@ -121,8 +121,12 @@ struct PSInput
 
 void main(uint VertexId : SV_VertexId, out PSInput PSIn)
 {
+    float4 Colors[3];
+    Colors[0] = float4(1.0, 0.0, 0.0, 1.0);
+    Colors[1] = float4(0.0, 1.0, 0.0, 1.0);
+    Colors[2] = float4(0.0, 0.0, 1.0, 1.0);
     PSIn.Pos = PushConstants.g_Positions[VertexId];
-    PSIn.Color = PushConstants.g_Colors[VertexId % 3];
+    PSIn.Color = Colors[VertexId % 3];
 }
 )"};
 
@@ -131,7 +135,7 @@ const std::string VulkanPushConstants_PS{
 struct PushConstants_t
 {
     float4 g_Positions[6];
-    float4 g_Colors[3];
+    float4 g_Colors[3]; // Make PS push constant size larger than VS
 };
 
 [[vk::push_constant]] ConstantBuffer<PushConstants_t> PushConstants;
@@ -144,8 +148,6 @@ struct PSInput
 
 float4 main(in PSInput PSIn) : SV_Target
 {
-    // Use push constants from PS to apply color modulation
-    // This ensures both VS and PS access the same PushConstants
     return float4(
         PSIn.Color.r * PushConstants.g_Colors[0].r,
         PSIn.Color.g * PushConstants.g_Colors[1].g,