layers: Fix crash with push constants using spec const

Author: spencer-lunarg

layers/state_tracker/shader_module.cpp

@@ -1823,7 +1823,17 @@ PushConstantVariable::PushConstantVariable(const SHADER_MODULE_STATE& module_sta
         }
     }
     const auto& highest_member = type_struct_info->members[highest_element_index];
-    const uint32_t highest_element_size = module_state.GetTypeBytesSize(highest_member.insn);
+    uint32_t highest_element_size = 0;
+    if (highest_member.insn->Opcode() == spv::OpTypeArray &&
+        module_state.FindDef(highest_member.insn->Word(3))->Opcode() == spv::OpSpecConstant) {
+        // TODO - This is a work-around because currently we only apply SpecConstant for workgroup size
+        // The shader validation needs to be fixed so we handle all cases when spec constant are applied, while still being catious
+        // of the fact that information is not known until pipeline creation (not at shader module creation time)
+        // https://github.com/KhronosGroup/Vulkan-ValidationLayers/issues/5911
+        highest_element_size = module_state.FindDef(highest_member.insn->Word(3))->Word(3);
+    } else {
+        highest_element_size = module_state.GetTypeBytesSize(highest_member.insn);
+    }
     size = (highest_element_size + highest_element_offset) - offset;
 }
 

tests/negative/shader_push_constants.cpp

@@ -456,3 +456,49 @@ TEST_F(NegativeShaderPushConstants, MultipleEntryPoint) {
     pipe.InitState();
     pipe.CreateGraphicsPipeline();
 }
+
+// This is not working because of a bug in the Spec Constant logic
+// https://github.com/KhronosGroup/Vulkan-ValidationLayers/issues/5911
+TEST_F(NegativeShaderPushConstants, DISABLED_SpecConstantSize) {
+    TEST_DESCRIPTION("Use SpecConstant to adjust size of Push Constant Block");
+    ASSERT_NO_FATAL_FAILURE(Init());
+
+    const char *cs_source = R"glsl(
+        #version 460
+        layout (constant_id = 0) const int my_array_size = 1;
+        layout (push_constant) uniform my_buf {
+            float my_array[my_array_size];
+        } pc;
+
+        void main() {
+            float a = pc.my_array[0];
+        }
+    )glsl";
+
+    uint32_t data = 32;
+
+    VkSpecializationMapEntry entry;
+    entry.constantID = 0;
+    entry.offset = 0;
+    entry.size = sizeof(uint32_t);
+
+    VkSpecializationInfo specialization_info = {};
+    specialization_info.mapEntryCount = 1;
+    specialization_info.pMapEntries = &entry;
+    specialization_info.dataSize = sizeof(uint32_t);
+    specialization_info.pData = &data;
+
+    // With spec constant set, this should be 32, not 16
+    VkPushConstantRange push_constant_range = {VK_SHADER_STAGE_COMPUTE_BIT, 0, 16};
+    const VkPipelineLayoutObj pipeline_layout(m_device, {}, {push_constant_range});
+
+    CreateComputePipelineHelper pipe(*this);
+    pipe.InitInfo();
+    pipe.cs_.reset(
+        new VkShaderObj(this, cs_source, VK_SHADER_STAGE_COMPUTE_BIT, SPV_ENV_VULKAN_1_0, SPV_SOURCE_GLSL, &specialization_info));
+    pipe.InitState();
+    pipe.pipeline_layout_ = VkPipelineLayoutObj(m_device, {}, {push_constant_range});
+    m_errorMonitor->SetDesiredFailureMsg(VK_DEBUG_REPORT_ERROR_BIT_EXT, "VUID-VkComputePipelineCreateInfo-layout-07987");
+    pipe.CreateComputePipeline();
+    m_errorMonitor->VerifyFound();
+}

tests/positive/shader_push_constants.cpp

@@ -745,3 +745,30 @@ TEST_F(PositiveShaderPushConstants, MultipleStructs) {
     m_commandBuffer->EndRenderPass();
     m_commandBuffer->end();
 }
+
+TEST_F(PositiveShaderPushConstants, SpecConstantSize) {
+    TEST_DESCRIPTION("Use SpecConstant to adjust size of Push Constant Block");
+    ASSERT_NO_FATAL_FAILURE(Init());
+
+    const char *cs_source = R"glsl(
+        #version 460
+        layout (constant_id = 2) const int my_array_size = 1;
+        layout (push_constant) uniform my_buf {
+            float my_array[my_array_size];
+        } pc;
+
+        void main() {
+            float a = pc.my_array[0];
+        }
+    )glsl";
+
+    VkPushConstantRange push_constant_range = {VK_SHADER_STAGE_COMPUTE_BIT, 0, 32};
+    const VkPipelineLayoutObj pipeline_layout(m_device, {}, {push_constant_range});
+
+    CreateComputePipelineHelper pipe(*this);
+    pipe.InitInfo();
+    pipe.cs_.reset(new VkShaderObj(this, cs_source, VK_SHADER_STAGE_COMPUTE_BIT));
+    pipe.InitState();
+    pipe.pipeline_layout_ = VkPipelineLayoutObj(m_device, {}, {push_constant_range});
+    pipe.CreateComputePipeline();
+}