vk: fix push constant implementation and test (#9588)

All stages share a push constant "block". We need to order
the fragment stage bytes to after the vertex stage if they
both exists.

FIXES=453776664

filament/backend/src/vulkan/VulkanHandles.cpp

@@ -265,24 +265,31 @@ void VulkanDescriptorSet::addNewSet(VkDescriptorSet vkSet, OnRecycle&& onRecycle
 
 PushConstantDescription::PushConstantDescription(backend::Program const& program) {
     mRangeCount = 0;
-    for (auto stage : { ShaderStage::VERTEX, ShaderStage::FRAGMENT, ShaderStage::COMPUTE }) {
+    uint32_t offset = 0;
+
+    // The range is laid out so that the vertex constants are defined as the first set of bytes,
+    // followed by fragment and compute. This means we need to keep track of the offset for each
+    // stage. We do the bookeeping in mDescriptions.
+    for (auto stage: { ShaderStage::VERTEX, ShaderStage::FRAGMENT, ShaderStage::COMPUTE }) {
         auto const& constants = program.getPushConstants(stage);
         if (constants.empty()) {
             continue;
         }
 
+        auto& description = mDescriptions[(uint8_t) stage];
         // We store the type of the constant for type-checking when writing.
-        auto& types = mTypes[(uint8_t) stage];
-        types.reserve(constants.size());
-        std::for_each(constants.cbegin(), constants.cend(), [&types] (Program::PushConstant t) {
-            types.push_back(t.type);
-        });
+        description.types.reserve(constants.size());
+        std::for_each(constants.cbegin(), constants.cend(),
+                [&description](Program::PushConstant t) { description.types.push_back(t.type); });
 
+        uint32_t const constantsSize = (uint32_t) constants.size() * ENTRY_SIZE;
         mRanges[mRangeCount++] = {
             .stageFlags = getVkStage(stage),
-            .offset = 0,
-            .size = (uint32_t) constants.size() * ENTRY_SIZE,
+            .offset = offset,
+            .size = constantsSize,
         };
+        description.offset = offset;
+        offset += constantsSize;
     }
 }
 
@@ -290,7 +297,10 @@ void PushConstantDescription::write(VkCommandBuffer cmdbuf, VkPipelineLayout lay
         backend::ShaderStage stage, uint8_t index, backend::PushConstantVariant const& value) {
 
     uint32_t binaryValue = 0;
-    UTILS_UNUSED_IN_RELEASE auto const& types = mTypes[(uint8_t) stage];
+    auto const& description = mDescriptions[(uint8_t) stage];
+    UTILS_UNUSED_IN_RELEASE auto const& types = description.types;
+    uint32_t const offset = description.offset;
+
     if (std::holds_alternative<bool>(value)) {
         assert_invariant(types[index] == ConstantType::BOOL);
         bool const bval = std::get<bool>(value);
@@ -304,7 +314,8 @@ void PushConstantDescription::write(VkCommandBuffer cmdbuf, VkPipelineLayout lay
         int const ival = std::get<int>(value);
         binaryValue = *reinterpret_cast<uint32_t const*>(&ival);
     }
-    vkCmdPushConstants(cmdbuf, layout, getVkStage(stage), index * ENTRY_SIZE, ENTRY_SIZE,
+
+    vkCmdPushConstants(cmdbuf, layout, getVkStage(stage), offset + index * ENTRY_SIZE, ENTRY_SIZE,
             &binaryValue);
 }
 

filament/backend/src/vulkan/VulkanHandles.h

@@ -245,7 +245,13 @@ struct PushConstantDescription {
 private:
     static constexpr uint32_t ENTRY_SIZE = sizeof(uint32_t);
 
-    utils::FixedCapacityVector<backend::ConstantType> mTypes[Program::SHADER_TYPE_COUNT];
+    struct ConstantDescription {
+        utils::FixedCapacityVector<backend::ConstantType> types;
+        uint32_t offset = 0;
+    };
+
+    // Describes the constants in each shader stage.
+    ConstantDescription mDescriptions[Program::SHADER_TYPE_COUNT];
     VkPushConstantRange mRanges[Program::SHADER_TYPE_COUNT];
     uint32_t mRangeCount;
 };

filament/backend/src/vulkan/VulkanPipelineLayoutCache.cpp

@@ -33,7 +33,7 @@ VkPipelineLayout VulkanPipelineLayoutCache::getLayout(
 
     // build the push constant layout key
     uint32_t const pushConstantRangeCount = program->getPushConstantRangeCount();
-    auto const& pushConstantRanges = program->getPushConstantRanges();
+    auto pushConstantRanges = program->getPushConstantRanges();
     if (pushConstantRangeCount > 0) {
         assert_invariant(pushConstantRangeCount <= Program::SHADER_TYPE_COUNT);
         for (uint8_t i = 0; i < pushConstantRangeCount; ++i) {

filament/backend/test/test_PushConstants.cpp

@@ -72,7 +72,13 @@ void main() {
 static const char* const triangleFs = R"(#version 450 core
 
 layout(push_constant) uniform Constants {
+#if defined(TARGET_VULKAN_ENVIRONMENT)
+    // offset here accounts for the size of the push constants in the vertex stage.  Vulkan has one
+    // block of memory for all stages to share.
+    layout(offset=16) float red;
+#else
     float red;
+#endif
     bool padding;       // test correct bool padding
     float green;
     float blue;
@@ -104,7 +110,6 @@ void initPushConstants() {
 }
 
 TEST_F(BackendTest, PushConstants) {
-    SKIP_IF(SkipEnvironment(OperatingSystem::CI, Backend::VULKAN), "see b/453776664");
     SKIP_IF(Backend::WEBGPU, "Push constants not supported on WebGPU");
 
     initPushConstants();

libs/filamat/src/shaders/CodeGenerator.cpp

@@ -963,6 +963,10 @@ utils::io::sstream& CodeGenerator::generateSpecializationConstant(utils::io::sst
     return out;
 }
 
+// Note that we've only introduced push constants to the vertex stage.  If we want to add push
+// constants to the fragment stage, in vulkan, we would have to offset the definition of the field
+// by the size of the constant struct in the vertex stage. This is due to vulkan having essentially
+// one block of memory for push constants that is shared across all stages).
 utils::io::sstream& CodeGenerator::generatePushConstants(utils::io::sstream& out,
         MaterialBuilder::PushConstantList const& pushConstants, size_t const layoutLocation) const {
     if (UTILS_UNLIKELY(pushConstants.empty())) {