shader_recompiler: Lower ballot and MBCNT as a pair (#5116)

Trace MaskedBitCount32 inputs back to their ballot source and apply wave64 lowering only when that ballot is in uniform control flow. Keep divergent ballot/MBCNT pairs on native wave32 semantics.
This commit is contained in:
Valdis Bogdāns authored and GitHub committed 2026-09-26 01:33:04 +03:00
1 parent 41f2a42816
commit 2f40f87b5b
2 files changed
+30 -8

No files matched your search

@@ -84,6 +84,23 @@ static std::vector<IR::Block*> FindUniformBlocks(const IR::Program& program) {
return blocks;
}
static IR::Inst* FindBallotForMaskedBitCount(const IR::Inst& mbcnt) {
IR::Value value = mbcnt.Arg(0);
if (value.IsImmediate()) {
return nullptr;
}
IR::Inst* inst = value.Inst();
if (inst->GetOpcode() != IR::Opcode::CompositeExtractU32x2 || inst->Arg(0).IsImmediate()) {
return nullptr;
}
inst = inst->Arg(0).Inst();
if (inst->GetOpcode() != IR::Opcode::UnpackUint2x32 || inst->Arg(0).IsImmediate()) {
return nullptr;
}
inst = inst->Arg(0).Inst();
return inst->GetOpcode() == IR::Opcode::Ballot ? inst : nullptr;
}
void LowerWave64BallotPass(IR::Program& program, const RuntimeInfo& runtime_info,
const Profile& profile) {
if (program.info.hw_stage != HwStage::Compute || profile.subgroup_size == 64) {
@@ -118,9 +135,12 @@ void LowerWave64BallotPass(IR::Program& program, const RuntimeInfo& runtime_info
if (std::ranges::any_of(inst.Uses(), is_unpack)) {
push_worklist(inst);
}
} else if (inst.GetOpcode() == IR::Opcode::GetAttributeU32 &&
inst.Arg(0).Attribute() == IR::Attribute::SubgroupLtMask) {
worklist.push_back(&inst);
} else if (inst.GetOpcode() == IR::Opcode::MaskedBitCount32) {
IR::Inst* const ballot = FindBallotForMaskedBitCount(inst);
if (ballot == nullptr ||
std::ranges::contains(uniform_blocks, ballot->GetParent())) {
worklist.push_back(&inst);
}
}
}
}
@@ -162,13 +182,15 @@ void LowerWave64BallotPass(IR::Program& program, const RuntimeInfo& runtime_info
ir.ShiftLeftLogical(half, ir.Imm32(2u))))};
ir.Barrier();
inst->ReplaceUsesWithAndRemove(value);
} else if (inst->GetOpcode() == IR::Opcode::GetAttributeU32 &&
inst->Arg(0).Attribute() == IR::Attribute::SubgroupLtMask) {
} else if (inst->GetOpcode() == IR::Opcode::MaskedBitCount32) {
const IR::U32 subgroup_invocation_id = ir.BitwiseAnd(invocation_index, ir.Imm32(63));
const IR::U64 mask = ir.ISub(
ir.ShiftLeftLogical(ir.Imm64(u64{1}), subgroup_invocation_id), ir.Imm64(u64{1}));
inst->ReplaceUsesWithAndRemove(
ir.CompositeExtract(ir.UnpackUint2x32(mask), inst->Arg(1).U32()));
const IR::U32 thread_mask{
ir.CompositeExtract(ir.UnpackUint2x32(mask), inst->Arg(2).U1() ? 1u : 0u)};
const IR::U32 masked_value{
ir.BitCount(ir.BitwiseAnd(IR::U32{inst->Arg(0)}, thread_mask))};
inst->ReplaceUsesWithAndRemove(ir.IAdd(masked_value, IR::U32{inst->Arg(1)}));
}
}
}
+1 -1
View File
@@ -116,7 +116,6 @@ IR::Program TranslateProgram(const std::span<const u32>& code, Pools& pools, Inf
Shader::Optimization::LowerUserClipPlanes(program, runtime_info);
Shader::Optimization::PhiSimplificationPass(program);
Shader::Optimization::InverseBallotEliminationPass(program);
Shader::Optimization::LowerHardwareIntrinsics(program);
Shader::IR::DumpProgram(program, info, "pre-lower-phi.");
// Prepare for structurization by clearing flow graph and lowering phis
@@ -138,6 +137,7 @@ IR::Program TranslateProgram(const std::span<const u32>& code, Pools& pools, Inf
Shader::Optimization::SharedMemoryBarrierPass(program, runtime_info, profile);
Shader::Optimization::DeadCodeEliminationPass(program);
Shader::Optimization::LowerWave64BallotPass(program, runtime_info, profile);
Shader::Optimization::LowerHardwareIntrinsics(program);
Shader::Optimization::ConstantPropagationPass(program.post_order_blocks);
Shader::Optimization::DeadCodeEliminationPass(program);
Shader::Optimization::CollectShaderInfoPass(program, profile);