shader: expose wave32 feature guards through ballot mask simplification

This commit is contained in:
nmzik 2026-10-01 18:50:04 +04:00
parent 6409be2874
commit 35759e48cf
5 changed files with 224 additions and 6 deletions

View File

@ -594,7 +594,7 @@ TranslateResult TranslateProgram(std::span<const uint32_t> code, const CompileOp
GetDumpLabel(options), StageName(options.stage), options.shader_hash,
static_cast<uint64_t>(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<const uint32_t> 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);

View File

@ -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",

View File

@ -4,6 +4,7 @@
#include <algorithm>
#include <bit>
#include <cstdint>
#include <unordered_map>
#include <unordered_set>
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<Inst*>& 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<Inst*>& m_lowered_ancillary;
std::unordered_set<const Inst*> m_visited;
std::unordered_map<const Inst*, Value> m_values;
bool m_grounded = false;
};
} // namespace
void ConstantPropagationPass(const BlockList& blocks) {
void ConstantPropagationPass(const BlockList& blocks, uint32_t wave_size) {
std::unordered_set<Inst*> 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);
}
}

View File

@ -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

View File

@ -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<Block>());
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<ScalarReg>(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();