[Fix, Test] (ShaderTranspiler): relocate a late length constant so an offset atomic-counter block still flattens

MobileGL/MG_Test/ShaderTranspiler/FlattenAtomicCounterBlockTest.cpp

@@ -21,6 +21,7 @@
 #include <cstring>
 #include <map>
 #include <string>
+#include <utility>
 #include <vector>
 
 using namespace MobileGL;
@@ -156,6 +157,53 @@ void main() {
     g_out.value[0] = 1u;
 }
 )";
+
+// glslang emits constants in FIRST-USE order, so a shader that does not use the flattened
+// array's length until after it has declared the counter block leaves that constant BELOW the
+// block. The pass needs the length to build `uint[length]` immediately before the block (SPIR-V
+// forbids forward type references), and it used to decline the whole block in that case - which
+// left the offsets in place and made SPIRV-Cross refuse the stage outright:
+//
+//     Push constant block cannot be expressed as neither std430 nor std140.
+//
+// That is KHR-GL43.compute_shader.pipeline-compute-chain: its first kernel declares two counters
+// at offset 8 (so the flattened array is 4 elements) and first uses the value 4 after the block,
+// so the kernel never reached the driver and every buffer, image and counter it writes kept its
+// initial value. Here `i < 4u` is what puts `uint 4` below the block; the ordering assertion
+// below is the fixture's own latch, so a future glslang that emits constants differently reports
+// a stale fixture rather than silently testing nothing.
+constexpr const char* kLateLengthConstantCounters = R"(#version 430 core
+layout(local_size_x = 1) in;
+layout(binding = 1, offset = 8) uniform atomic_uint g_counter[2];
+layout(std430, binding = 0) buffer Output { uint value[]; } g_out;
+void main() {
+    uint i = atomicCounterIncrement(g_counter[1]);
+    if (i < 4u) { g_out.value[0] = i; }
+}
+)";
+
+// Index of the first OpConstant of type uint with value |value|, and of struct |structId|, in
+// the module's instruction order. -1 when absent.
+std::pair<Int64, Int64> UintConstantAndStructOrder(const Vector<Uint32>& spirv, Uint32 structId,
+                                                   Uint32 value) {
+    Int64 index = 0, constantIndex = -1, structIndex = -1;
+    Uint32 uintTypeId = 0;
+    ForEachInstruction(spirv, [&](spv::Op opcode, const Uint32* words, Uint32 wordCount) {
+        if (opcode == spv::Op::OpTypeInt && wordCount >= 4u && words[2] == 32u && words[3] == 0u) {
+            uintTypeId = words[1];
+        }
+        if (opcode == spv::Op::OpConstant && wordCount >= 4u && words[1] == uintTypeId &&
+            words[3] == value && constantIndex < 0) {
+            constantIndex = index;
+        }
+        if (opcode == spv::Op::OpTypeStruct && wordCount >= 2u && words[1] == structId) {
+            structIndex = index;
+        }
+        ++index;
+    });
+    return {constantIndex, structIndex};
+}
+
 } // namespace
 
 TEST(FlattenAtomicCounterBlockPass, MovesTheBlockToOffsetZeroAndGrowsTheArray) {
@@ -212,3 +260,44 @@ TEST(FlattenAtomicCounterBlockPass, IsIdempotent) {
     ASSERT_TRUE(ShaderCompiler::FlattenAtomicCounterBlockOffsetsForEssl(once, twice, true));
     EXPECT_EQ(twice, once);
 }
+
+// The block must still flatten when the module already declares the flattened array's length
+// constant BELOW the block. The pass relocates that constant instead of declining; declining
+// left the offsets in place and cost the whole stage its transpile.
+TEST(FlattenAtomicCounterBlockPass, FlattensWhenTheLengthConstantIsDeclaredAfterTheBlock) {
+    const Vector<Uint32> input = CompileCompute(kLateLengthConstantCounters);
+    ASSERT_FALSE(input.empty());
+
+    const Uint32 structId = FindAtomicCounterBlockStructId(input);
+    ASSERT_NE(structId, 0u);
+    ASSERT_EQ(MemberOffsetOf(input, structId, 0u), 8);
+
+    // The fixture's precondition, asserted rather than assumed: two counters at offset 8 need a
+    // 4-element array, and this shader's `uint 4` really does sit below the block.
+    const auto [constantIndex, structIndex] = UintConstantAndStructOrder(input, structId, 4u);
+    ASSERT_GE(constantIndex, 0) << "fixture is stale: the module no longer declares a uint 4";
+    ASSERT_GE(structIndex, 0);
+    ASSERT_GT(constantIndex, structIndex)
+        << "fixture is stale: `uint 4` is no longer declared after the counter block, so this "
+           "test would pass without exercising the relocation at all";
+
+    Vector<Uint32> output;
+    ASSERT_TRUE(ShaderCompiler::FlattenAtomicCounterBlockOffsetsForEssl(input, output, true));
+    ASSERT_FALSE(output.empty());
+    ASSERT_NE(output, input) << "the block was declined; the offsets are still in the module and "
+                                "SPIRV-Cross will refuse the stage";
+
+    const Uint32 outStructId = FindAtomicCounterBlockStructId(output);
+    ASSERT_EQ(outStructId, structId);
+    EXPECT_EQ(MemberCountOf(output, outStructId), 1u);
+    EXPECT_EQ(MemberOffsetOf(output, outStructId, 0u), 0);
+    EXPECT_EQ(ArrayLengthOf(output, MemberTypeOf(output, outStructId, 0u)), 4);
+    // The relocation moved a definition; the module has to still be well-ordered.
+    EXPECT_TRUE(Validates(output));
+
+    // The symptom the CTS case actually failed on: with the block declined this throws.
+    MG_Util::ShaderTranspiler::SpvcSession session(
+        output, MG_Util::ShaderTranspiler::SessionUsageBit::Transpile);
+    auto essl = ShaderCompiler::DecompileShader(session);
+    EXPECT_TRUE(essl) << "ESSL transpile failed: " << (essl ? String{} : essl.error().log);
+}

MobileGL/MG_Util/ShaderTranspiler/SpirvPasses/FlattenAtomicCounterBlockPass.cpp

@@ -325,9 +325,31 @@ namespace MobileGL {
                 // appending it at the end of the section would make the module invalid. A
                 // duplicate OpTypeArray is legal (SPIR-V 2.8 exempts aggregates from the
                 // uniqueness rule, and so does spirv-val), so no search for an existing one is
-                // needed; the LENGTH CONSTANT is not exempt, and if the module already declares
-                // it after the block there is nowhere legal to put the array - the block is then
-                // declined and keeps today's behaviour. Returns 0 for that.
+                // needed; the LENGTH CONSTANT is not exempt, so when the module already declares
+                // it the pass has to work with the one instruction that exists.
+                //
+                // That instruction is not always in a usable place. GetDefiningInstruction only
+                // honours `position` when it MINTS the constant; when the module already has one
+                // it hands back the existing instruction wherever it happens to sit, and glslang
+                // emits constants in first-use order, so a shader whose first use of the value is
+                // below the counter block declares it below the block. The flattened array would
+                // then forward-reference its own length.
+                //
+                // KHR-GL43.compute_shader.pipeline-compute-chain is exactly that shader: two
+                // counters at offset 8 need a 4-element array, and its `%uint_4` is first used by
+                // a later declaration, so it lands AFTER gl_AtomicCounterBlock_1. Declining there
+                // - which is what this used to do - left the offsets in place, and SPIRV-Cross
+                // then refused the whole stage with "Push constant block cannot be expressed as
+                // neither std430 nor std140", so the chain's first kernel never reached the
+                // driver and every resource it writes stayed at its initial value.
+                //
+                // Moving the constant UP to just before the block is always legal, which is why
+                // this is a relocation and not a second declaration: an OpConstant's only operand
+                // is its result TYPE, and that type already precedes the block (it is the element
+                // type of the counter array the block declares). Every existing use sits after
+                // the constant's old position and therefore after its new one too, so no use is
+                // left dangling - moving a definition earlier in the types/constants section
+                // cannot invalidate anything. Ordering is all that changes; def-use is untouched.
                 uint32_t CreateCounterArrayTypeBefore(IRContext* context, Instruction* structType,
                                                       uint32_t uintTypeId, uint32_t length) {
                     auto* constantMgr = context->get_constant_mgr();
@@ -341,7 +363,12 @@ namespace MobileGL {
                     if (position == context->types_values_end()) return 0;
                     Instruction* lengthInst = constantMgr->GetDefiningInstruction(lengthConstant, 0, &position);
                     if (lengthInst == nullptr) return 0;
-                    if (!DeclaredBefore(context, lengthInst->result_id(), structType->result_id())) return 0;
+                    if (!DeclaredBefore(context, lengthInst->result_id(), structType->result_id())) {
+                        // Pre-existing constant, declared below the block. Relocate it; see above
+                        // for why that is sound. InsertBefore unlinks it from its current spot
+                        // first, so this is a move rather than an aliasing second entry.
+                        lengthInst->InsertBefore(structType);
+                    }
 
                     const uint32_t arrayTypeId = context->TakeNextId();
                     if (arrayTypeId == 0) return 0;