fix(shader): ballot a predicate both wave64 halves share only once

core/libs/tests/Wave32WideSubgroup.cpp

@@ -183,6 +183,47 @@ void Check(std::uint32_t size, const Value& masks) {
     }
 }
 
+std::size_t CountOpcode(std::span<const std::uint32_t> code, spv::Op opcode) {
+    std::size_t found = 0;
+    for (std::size_t offset = 5; offset < code.size(); offset += code[offset] >> 16u) {
+        Require((code[offset] >> 16u) != 0, "malformed SPIR-V instruction");
+        found += static_cast<spv::Op>(code[offset] & 0xffffu) == opcode;
+    }
+    return found;
+}
+
+void CheckTwoLaneBallot(std::uint32_t lowMask, std::uint32_t highMask, bool uniform) {
+    IrProgram program;
+    program.SetWaveSize(64);
+    SpirvEmitterState state(program, {});
+    SpirvValueEmitContext low(state);
+    SpirvValueEmitContext high(state);
+    low.otherHalf = &high;
+    high.otherHalf = &low;
+    high.half = 1u;
+    IrValue predicate(IrOpcode::GetExec, IrType::Bool, 1);
+    if (uniform) predicate.SetImmediateBool(true);
+    const auto lowPredicate = low.Def(&predicate);
+    const auto highPredicate = high.Def(&predicate);
+    const auto ballot = high.Ballot(&predicate);
+    Subgroup subgroup(32);
+    Lanes lowLanes(32), highLanes(32);
+    for (std::uint32_t lane = 0; lane < 32; ++lane) {
+        lowLanes[lane][0] = (lowMask >> lane) & 1u;
+        highLanes[lane][0] = (highMask >> lane) & 1u;
+    }
+    if (!uniform) {
+        subgroup.Set(lowPredicate, lowLanes);
+        subgroup.Set(highPredicate, highLanes);
+    }
+    const auto code = state.module.Finalize();
+    subgroup.Execute(code);
+    Require(CountOpcode(code, spv::OpGroupNonUniformBallot) == (uniform ? 1u : 2u), "a wave64 ballot of one predicate for both halves was not issued once");
+    for (std::uint32_t lane = 0; lane < 32; ++lane) {
+        Require(subgroup.Get(ballot)[lane] == Value{lowMask, highMask, 0, 0}, "wave64 ballot halves mixed up");
+    }
+}
+
 }
 
 int main() {
@@ -192,6 +233,9 @@ int main() {
                 Check(size, masks);
             }
         }
+        CheckTwoLaneBallot(0xffffffffu, 0xffffffffu, true);
+        CheckTwoLaneBallot(0x12345678u, 0xfedcba98u, false);
+        CheckTwoLaneBallot(0u, 0xffffffffu, false);
         std::puts("wave32 wide subgroup tests passed");
         return 0;
     } catch (const std::exception& error) {

core/shader/recompiler/SpirvBackend/src/SpirvEmitter.cpp

@@ -118,15 +118,17 @@ std::uint32_t SpirvValueEmitContext::Ballot(const IrValue* predicate) {
     const auto ballotType = TypeU32Vector(state, 4u);
     const auto scope = ConstantU32(state, spv::ScopeSubgroup);
     const auto low = state.module.AllocateId();
-    EmitSubgroupBallot(state, low, otherHalf == nullptr || half == 0u ? Def(predicate) : otherHalf->Def(predicate));
+    const auto lowPredicate = otherHalf == nullptr || half == 0u ? Def(predicate) : otherHalf->Def(predicate);
+    EmitSubgroupBallot(state, low, lowPredicate);
     if (otherHalf == nullptr) {
         return EmitWaveBallot(state, low);
     }
-    const auto high = state.module.AllocateId();
+    const auto highPredicate = half == 1u ? Def(predicate) : otherHalf->Def(predicate);
+    const auto high = highPredicate == lowPredicate ? low : state.module.AllocateId();
     const auto lowWord = state.module.AllocateId();
     const auto highWord = state.module.AllocateId();
     const auto ballot = state.module.AllocateId();
-    state.module.AddFunction(spv::OpGroupNonUniformBallot, ballotType, high, scope, half == 1u ? Def(predicate) : otherHalf->Def(predicate));
+    if (high != low) state.module.AddFunction(spv::OpGroupNonUniformBallot, ballotType, high, scope, highPredicate);
     state.module.AddFunction(spv::OpCompositeExtract, TypeU32(state), lowWord, low, 0u);
     state.module.AddFunction(spv::OpCompositeExtract, TypeU32(state), highWord, high, 0u);
     state.module.AddFunction(spv::OpCompositeConstruct, ballotType, ballot, lowWord, highWord, ConstantU32(state, 0u), ConstantU32(state, 0u));