From 35759e48cf8605df04735244e82b4f708ddddcdc Mon Sep 17 00:00:00 2001 From: nmzik Date: Thu, 1 Oct 2026 18:50:04 +0400 Subject: [PATCH] shader: expose wave32 feature guards through ballot mask simplification --- .../shader/recompiler/ShaderRecompiler.cpp | 4 +- .../shader/recompiler/Tessellation.cpp | 4 +- .../ir/passes/ConstantPropagation.cpp | 101 ++++++++++++++- .../ir/passes/ConstantPropagation.h | 2 +- tests/shaderCfgTests.cpp | 119 ++++++++++++++++++ 5 files changed, 224 insertions(+), 6 deletions(-) diff --git a/src/graphics/shader/recompiler/ShaderRecompiler.cpp b/src/graphics/shader/recompiler/ShaderRecompiler.cpp index 8a0d6cc2..cfafea61 100644 --- a/src/graphics/shader/recompiler/ShaderRecompiler.cpp +++ b/src/graphics/shader/recompiler/ShaderRecompiler.cpp @@ -594,7 +594,7 @@ TranslateResult TranslateProgram(std::span code, const CompileOp GetDumpLabel(options), StageName(options.stage), options.shader_hash, static_cast(ir.blocks.size()), phase_ms()); IR::RewriteToSsa(ir.blocks); - IR::ConstantPropagationPass(ir.blocks); + IR::ConstantPropagationPass(ir.blocks, ir.wave_size); IR::ResolveControlFlowIdentities(ir); IR::RemoveIdentities(ir.blocks); IR::EliminateDeadCode(ir.blocks); @@ -602,7 +602,7 @@ TranslateResult TranslateProgram(std::span code, const CompileOp if (read_lane_stats.rewritten_reads != 0) { LOGF("%s read-lane elimination: reads=%" PRIu32 "\n", GetDumpLabel(options), read_lane_stats.rewritten_reads); - IR::ConstantPropagationPass(ir.blocks); + IR::ConstantPropagationPass(ir.blocks, ir.wave_size); IR::ResolveControlFlowIdentities(ir); IR::RemoveIdentities(ir.blocks); IR::EliminateDeadCode(ir.blocks); diff --git a/src/graphics/shader/recompiler/Tessellation.cpp b/src/graphics/shader/recompiler/Tessellation.cpp index cd436259..8f42333f 100644 --- a/src/graphics/shader/recompiler/Tessellation.cpp +++ b/src/graphics/shader/recompiler/Tessellation.cpp @@ -217,7 +217,7 @@ void LowerTessellationMemory(IR::Program& program, const CompileOptions& options 2, ActiveAddress(*block, it, it->Arg(2), it->Arg(it->NumArgs() - 1u), resolved)); } } - ConstantPropagationPass(program.blocks); + ConstantPropagationPass(program.blocks, program.wave_size); uint32_t reads = 0, writes = 0, factors = 0; for (auto* block: program.blocks) { for (auto it = block->begin(); it != block->end(); ++it) { @@ -317,7 +317,7 @@ void LowerTessellationMemory(IR::Program& program, const CompileOptions& options inst.Invalidate(); } } - ConstantPropagationPass(program.blocks); + ConstantPropagationPass(program.blocks, program.wave_size); RemoveIdentities(program.blocks); EliminateDeadCode(program.blocks); LOGF("%s tessellation lowering: reads=%u writes=%u factors=%u\n", diff --git a/src/graphics/shader/recompiler/ir/passes/ConstantPropagation.cpp b/src/graphics/shader/recompiler/ir/passes/ConstantPropagation.cpp index 1e0c3398..0c38286f 100644 --- a/src/graphics/shader/recompiler/ir/passes/ConstantPropagation.cpp +++ b/src/graphics/shader/recompiler/ir/passes/ConstantPropagation.cpp @@ -4,6 +4,7 @@ #include #include #include +#include #include namespace Libs::Graphics::ShaderRecompiler::IR { @@ -660,12 +661,110 @@ void FoldInstruction(Block& block, Block::iterator instruction, } } +// In wave32, extracting this invocation's bit from a ballot recovers its predicate. +// Keep that identity through scalar EXEC operations and their loop-carried mask Phis. +class LaneMaskProjection { +public: + explicit LaneMaskProjection(std::unordered_set& lowered_ancillary) + : m_lowered_ancillary(lowered_ancillary) {} + + void Fold(Inst& inst) { + if (inst.GetOpcode() != ValueOpcode::INotEqual32 || !Immediate(Arg(inst, 1), 0u)) return; + const auto* bit = Arg(inst, 0).TryInstruction(); + if (bit == nullptr || bit->GetOpcode() != ValueOpcode::BitwiseAnd32 || + !Immediate(Arg(*bit, 1), 1u)) return; + const auto* shift = Arg(*bit, 0).TryInstruction(); + if (shift == nullptr || shift->GetOpcode() != ValueOpcode::ShiftRightLogical32) return; + const auto* index = Arg(*shift, 1).TryInstruction(); + if (index == nullptr || index->GetOpcode() != ValueOpcode::BitwiseAnd32 || + !Immediate(Arg(*index, 1), 31u)) return; + const auto* lane = Arg(*index, 0).TryInstruction(); + if (lane == nullptr || lane->GetOpcode() != ValueOpcode::LaneId) return; + m_visited.clear(); + m_grounded = false; + if (CanProject(Arg(*shift, 0)) && m_grounded) Replace(inst, Project(Arg(*shift, 0))); + } + +private: + static bool Immediate(Value value, uint32_t expected) { + return IsImmediate(value, Type::U32) && value.U32() == expected; + } + + static Value BallotPredicate(const Inst& inst) { + if (inst.GetOpcode() != ValueOpcode::CompositeExtractU32x4 || + !Immediate(Arg(inst, 1), 0u)) return {}; + const auto* source = Arg(inst, 0).TryInstruction(); + return source != nullptr && source->GetOpcode() == ValueOpcode::Ballot + ? Arg(*source, 0) : Value {}; + } + + bool CanProject(Value value) { + if (value.IsImmediate()) { + const bool supported = Immediate(value, 0u) || Immediate(value, UINT32_MAX); + m_grounded |= supported; + return supported; + } + const auto* inst = value.TryInstruction(); + if (inst == nullptr) return false; + if (m_values.contains(inst) || !BallotPredicate(*inst).IsEmpty()) { + m_grounded = true; + return true; + } + if (!m_visited.insert(inst).second) return true; + const auto op = inst->GetOpcode(); + if (op != ValueOpcode::BitwiseAnd32 && op != ValueOpcode::BitwiseOr32 && + op != ValueOpcode::BitwiseNot32 && op != ValueOpcode::SelectU32 && op != ValueOpcode::Phi) + return false; + for (size_t arg = op == ValueOpcode::SelectU32 ? 1u : 0u; arg < inst->NumArgs(); ++arg) { + if (!CanProject(Arg(*inst, arg))) return false; + } + return inst->NumArgs() != 0; + } + + Value Project(Value value) { + if (value.IsImmediate()) return Value(value.U32() != 0u); + auto* source = value.TryInstruction(); + if (const auto found = m_values.find(source); found != m_values.end()) return found->second.Resolve(); + if (const auto predicate = BallotPredicate(*source); !predicate.IsEmpty()) return predicate; + auto* block = source->Parent(); + const auto where = std::find_if(block->begin(), block->end(), + [source](const Inst& candidate) { return &candidate == source; }); + const auto op = source->GetOpcode() == ValueOpcode::BitwiseAnd32 ? ValueOpcode::LogicalAnd + : source->GetOpcode() == ValueOpcode::BitwiseOr32 ? ValueOpcode::LogicalOr + : source->GetOpcode() == ValueOpcode::BitwiseNot32 ? ValueOpcode::LogicalNot + : source->GetOpcode() == ValueOpcode::SelectU32 ? ValueOpcode::SelectU1 + : ValueOpcode::Phi; + auto result = op == ValueOpcode::Phi ? block->PrependNewInst(where, op) + : op == ValueOpcode::LogicalNot + ? block->PrependNewInst(where, op, {Value(false)}) + : op == ValueOpcode::SelectU1 + ? block->PrependNewInst(where, op, {Arg(*source, 0), Value(false), Value(false)}) + : block->PrependNewInst(where, op, {Value(false), Value(false)}); + m_values.emplace(source, Value(&*result)); + if (op == ValueOpcode::Phi) result->SetFlags(Type::U1); + for (size_t arg = op == ValueOpcode::SelectU1 ? 1u : 0u; arg < source->NumArgs(); ++arg) { + const auto projected = Project(Arg(*source, arg)); + if (op == ValueOpcode::Phi) result->AddPhiOperand(source->PhiBlock(arg), projected); + else result->SetArg(arg, projected); + } + FoldInstruction(*block, result, m_lowered_ancillary); + return Value(&*result).Resolve(); + } + + std::unordered_set& m_lowered_ancillary; + std::unordered_set m_visited; + std::unordered_map m_values; + bool m_grounded = false; +}; + } // namespace -void ConstantPropagationPass(const BlockList& blocks) { +void ConstantPropagationPass(const BlockList& blocks, uint32_t wave_size) { std::unordered_set lowered_ancillary; + LaneMaskProjection mask_projection(lowered_ancillary); for (auto* block: blocks) { for (auto inst = block->begin(); inst != block->end(); ++inst) { + if (wave_size == 32u) mask_projection.Fold(*inst); FoldInstruction(*block, inst, lowered_ancillary); } } diff --git a/src/graphics/shader/recompiler/ir/passes/ConstantPropagation.h b/src/graphics/shader/recompiler/ir/passes/ConstantPropagation.h index e158c1ea..6da994cd 100644 --- a/src/graphics/shader/recompiler/ir/passes/ConstantPropagation.h +++ b/src/graphics/shader/recompiler/ir/passes/ConstantPropagation.h @@ -4,6 +4,6 @@ namespace Libs::Graphics::ShaderRecompiler::IR { -void ConstantPropagationPass(const BlockList& blocks); +void ConstantPropagationPass(const BlockList& blocks, uint32_t wave_size = 64); } // namespace Libs::Graphics::ShaderRecompiler::IR diff --git a/tests/shaderCfgTests.cpp b/tests/shaderCfgTests.cpp index 1b504971..b48a27f8 100644 --- a/tests/shaderCfgTests.cpp +++ b/tests/shaderCfgTests.cpp @@ -11551,6 +11551,124 @@ void TestValuePhiValidation() { #endif } +void TestWave32MaskProjection() { + using namespace ShaderRecompiler::IR; + for (const uint32_t wave_size : {32u, 64u}) { + Program program; + program.stage = ShaderType::Compute; + program.wave_size = wave_size; + for (uint32_t id = 0; id < 3; ++id) { + program.block_storage.push_back(std::make_unique()); + program.blocks.push_back(program.block_storage.back().get()); + program.block_info.push_back({.id = id}); + } + auto *entry = program.blocks[0]; + auto *loop = program.blocks[1]; + entry->AddBranch(loop); + loop->AddBranch(loop); + loop->AddBranch(program.blocks[2]); + program.block_info[0].terminator.kind = ShaderRecompiler::CFG::TerminatorKind::Branch; + program.block_info[0].terminator.true_block = 1; + program.block_info[1].terminator.kind = ShaderRecompiler::CFG::TerminatorKind::ConditionalBranch; + program.block_info[1].terminator.true_block = 1; + program.block_info[1].terminator.false_block = 2; + program.block_info[2].terminator.kind = ShaderRecompiler::CFG::TerminatorKind::Return; + IREmitter e(entry); + const auto zero = U32(Value(0u)); + const auto raw = e.GetUserData(static_cast(0)); + const auto active = e.INotEqual(raw, zero); + const auto selected = e.IEqual(raw, U32(Value(7u))); + const auto ballot = [&](IREmitter &ir, U1 predicate, uint32_t component = 0u) { + return ir.CompositeExtract(ir.Emit(ValueOpcode::Ballot, {predicate}), component); + }; + const auto thread_bit = [&](IREmitter &ir, U32 word, bool own_lane = true) { + const auto bit = own_lane + ? ir.BitwiseAnd(U32(ir.Emit(ValueOpcode::LaneId)), U32(Value(31u))) + : U32(Value(3u)); + return ir.INotEqual(ir.BitwiseAnd(ir.ShiftRightLogical(word, bit), U32(Value(1u))), zero); + }; + const auto active_mask = ballot(e, active); + const auto direct = thread_bit(e, active_mask); + const auto arbitrary = thread_bit(e, raw); + const auto mixed = thread_bit(e, e.BitwiseAnd(active_mask, raw)); + const auto mixed_or = thread_bit(e, e.BitwiseOr(active_mask, raw)); + const auto mixed_select = thread_bit(e, e.Select(selected, active_mask, raw)); + const auto different_lane = thread_bit(e, active_mask, false); + const auto high_half = thread_bit(e, ballot(e, active, 1u)); + const auto initial_mask = e.BitwiseAnd(active_mask, ballot(e, selected)); + const auto raw_consumer = e.Emit(ValueOpcode::IAdd32, {initial_mask, Value(3u)}); + const auto initial_predicate = e.LogicalAnd(active, selected); + const auto entry_size = entry->Instructions().size(); + ConstantPropagationPass({entry}, wave_size); + Check(entry->Instructions().size() == entry_size, + "rejected mixed mask graph left unused projected instructions"); + const auto combined_mask = e.BitwiseOr(active_mask, ballot(e, selected)); + const auto combined = thread_bit(e, combined_mask); + const auto gated_mask = e.Select(selected, combined_mask, zero); + const auto gated = thread_bit(e, gated_mask); + const auto rejected = thread_bit(e, e.Select(selected, zero, combined_mask)); + const auto combined_predicate = e.LogicalOr(active, selected); + const auto gated_predicate = e.LogicalAnd(selected, combined_predicate); + const auto rejected_predicate = e.Emit(ValueOpcode::SelectU1, + {selected, Value(false), combined_predicate}); + const auto gated_raw_consumer = e.Emit(ValueOpcode::IAdd32, {gated_mask, Value(3u)}); + auto &mask = loop->AppendNewInst(ValueOpcode::Phi); + mask.SetFlags(Type::U32); + auto &predicate = loop->AppendNewInst(ValueOpcode::Phi); + predicate.SetFlags(Type::U1); + IREmitter body(loop); + const auto matches = body.LogicalAnd(U1(Value(&predicate)), selected); + const auto remaining = body.BitwiseAnd(U32(Value(&mask)), body.BitwiseNot(ballot(body, matches))); + const auto remaining_predicate = body.LogicalAnd(U1(Value(&predicate)), body.LogicalNot(matches)); + const auto remaining_live = thread_bit(body, remaining); + const auto live = thread_bit(body, U32(Value(&mask))); + mask.AddPhiOperand(entry, initial_mask); + mask.AddPhiOperand(loop, remaining); + predicate.AddPhiOperand(entry, initial_predicate); + predicate.AddPhiOperand(loop, remaining_predicate); + program.block_info[1].condition = remaining_predicate; + + ConstantPropagationPass(program.blocks, wave_size); + if (wave_size == 32u) { + Check(direct.Resolve() == active.Resolve(), "wave32 ballot round trip lost its predicate"); + Check(EquivalentValue(program, combined, combined_predicate), + "wave32 OR mask projection changed a lane predicate"); + Check(EquivalentValue(program, gated, gated_predicate) && + EquivalentValue(program, rejected, rejected_predicate), + "wave32 select mask projection changed a scalar gate or selected arm"); + Check(EquivalentValue(program, live.Resolve(), Value(&predicate)), + "projected mask Phi differs from the existing per-lane EXEC Phi"); + Check(EquivalentValue(program, remaining_live.Resolve(), remaining_predicate) && + std::count_if(loop->begin(), loop->end(), [](const Inst &inst) { + return inst.GetOpcode() == ValueOpcode::LogicalAnd; + }) == 3, + "backedge-first projection duplicated the loop's logical mask operation"); + const auto *projected = live.Resolve().TryInstruction(); + Check(projected != nullptr && projected->GetOpcode() == ValueOpcode::Phi && + projected->GetType() == Type::U1 && projected->PhiBlock(0) == entry && + projected->PhiBlock(1) == loop, + "mask projection lost loop Phi type or predecessor order"); + } else { + Check(direct.Resolve().TryInstruction()->GetOpcode() == ValueOpcode::INotEqual32 && + live.Resolve().TryInstruction()->GetOpcode() == ValueOpcode::INotEqual32 && + combined.Resolve().TryInstruction()->GetOpcode() == ValueOpcode::INotEqual32 && + gated.Resolve().TryInstruction()->GetOpcode() == ValueOpcode::INotEqual32 && + rejected.Resolve().TryInstruction()->GetOpcode() == ValueOpcode::INotEqual32, + "wave64 low-half mask was projected onto the wrong invocation"); + } + for (const auto value : {arbitrary, mixed, mixed_or, mixed_select, different_lane, high_half}) { + Check(value.Resolve().TryInstruction()->GetOpcode() == ValueOpcode::INotEqual32, + "mask projection accepted an arbitrary mask, lane or ballot half"); + } + Check(raw_consumer.Resolve().TryInstruction()->Arg(0).Resolve() == initial_mask.Resolve() && + gated_raw_consumer.Resolve().TryInstruction()->Arg(0).Resolve() == gated_mask.Resolve(), + "mask projection changed a raw scalar mask consumer"); + ResolveControlFlowIdentities(program); + RemoveIdentities(program.blocks); + ValidateProgram(program, true); + } +} + void TestU64ShiftConstantPropagation() { using namespace ShaderRecompiler::IR; @@ -14256,6 +14374,7 @@ int main() { TestFinalSsaRejectsRegisterStatePseudos(); #endif TestValuePhiValidation(); + TestWave32MaskProjection(); TestU64ShiftConstantPropagation(); #if KYTY_PLATFORM != KYTY_PLATFORM_WINDOWS TestNativeWideValueValidation();