shader: lower wave32 whole-wave UMin and UMax reductions to subgroup reductions

Co-authored-by: Senaxx <michelbeld@gmail.com>
Co-authored-by: Mac <me@itsmemac.com>
This commit is contained in:
nmzik
2026-10-26 01:12:50 +01:00
parent be1c84f4d7
commit 6357740fa1
7 changed files with 102 additions and 60 deletions
@@ -99,28 +99,33 @@ Value GetRealValue(PhiMap& phi_map, Value source, uint32_t lane, uint32_t wave_s
return Value(&phi);
}
// A full-EXEC prefix scan with both half-wave endpoint reads is a reduction.
// A full-EXEC prefix scan with both half-wave endpoint reads is a reduction. Wave64 first joins
// the two rows of each half; a wave32 half is one row.
// Fragment subgroups may have no physical invocation at either endpoint.
class UMinReductionMatcher {
class WaveReductionMatcher {
public:
explicit UMinReductionMatcher(Block* block): block(block) {}
WaveReductionMatcher(Block* block, ValueOpcode operation)
: block(block), operation(operation),
identity(operation == ValueOpcode::UMin32 ? UINT32_MAX : 0u) {}
Value Match(const Inst& minimum) {
const auto* low = Get(minimum.Arg(0), ValueOpcode::ReadLane);
const auto* high = Get(minimum.Arg(1), ValueOpcode::ReadLane);
Value Match(const Inst& reduction, uint32_t wave_size) {
const auto* low = Get(reduction.Arg(0), ValueOpcode::ReadLane);
const auto* high = Get(reduction.Arg(1), ValueOpcode::ReadLane);
if (low == nullptr || high == nullptr || low->Arg(0).Resolve() != high->Arg(0).Resolve()) {
return {};
}
const auto lane = [](Value value, uint32_t expected) {
const auto lane = [&](Value value, uint32_t expected) {
value = value.Resolve();
return value.IsImmediate() && value.GetType() == Type::U32 &&
value.U32() % 64u == expected;
value.U32() % wave_size == expected;
};
if (!((lane(low->Arg(1), 31) && lane(high->Arg(1), 63)) ||
(lane(low->Arg(1), 63) && lane(high->Arg(1), 31)))) {
const auto half = wave_size / 2u;
if (!((lane(low->Arg(1), half - 1u) && lane(high->Arg(1), wave_size - 1u)) ||
(lane(low->Arg(1), wave_size - 1u) && lane(high->Arg(1), half - 1u)))) {
return {};
}
const auto* joined = Get(Unwrap(low->Arg(0)), ValueOpcode::UMin32);
if (wave_size == 32u) return MatchScan(Unwrap(low->Arg(0)));
const auto* joined = Get(Unwrap(low->Arg(0)), operation);
if (joined == nullptr) return {};
for (size_t index = 0; index < 2; index++) {
auto attempt = *this;
@@ -134,26 +139,32 @@ public:
!attempt.FullExec(cross->Arg(3)) || attempt.Unwrap(cross->Arg(0)) != prefix) {
continue;
}
for (const uint32_t shift: {8u, 4u, 2u, 1u}) {
prefix = attempt.MatchRow(prefix, shift);
}
prefix = attempt.Unwrap(prefix);
if (prefix.IsEmpty() || attempt.saved_exec.IsEmpty()) return prefix;
const auto* seed = attempt.Get(prefix, ValueOpcode::SelectU32);
if (seed != nullptr && seed->Arg(2).Resolve() == Value(UINT32_MAX)) {
const auto predicate = seed->Arg(0).Resolve();
const auto* subset = attempt.Get(predicate, ValueOpcode::LogicalAnd);
if (predicate == attempt.saved_exec ||
(subset != nullptr && (subset->Arg(0).Resolve() == attempt.saved_exec ||
subset->Arg(1).Resolve() == attempt.saved_exec))) {
return prefix;
}
}
if (const auto seed = attempt.MatchScan(prefix); !seed.IsEmpty()) return seed;
}
return {};
}
private:
// The row scan's input, whose lanes outside the saved EXEC hold the operation's identity.
Value MatchScan(Value prefix) {
for (const uint32_t shift: {8u, 4u, 2u, 1u}) {
prefix = MatchRow(prefix, shift);
}
prefix = Unwrap(prefix);
if (prefix.IsEmpty() || saved_exec.IsEmpty()) return prefix;
const auto* seed = Get(prefix, ValueOpcode::SelectU32);
if (seed != nullptr && seed->Arg(2).Resolve() == Value(identity)) {
const auto predicate = seed->Arg(0).Resolve();
const auto* subset = Get(predicate, ValueOpcode::LogicalAnd);
if (predicate == saved_exec ||
(subset != nullptr && (subset->Arg(0).Resolve() == saved_exec ||
subset->Arg(1).Resolve() == saved_exec))) {
return prefix;
}
}
return {};
}
Inst* Get(Value value, ValueOpcode opcode) const {
auto* inst = value.Resolve().TryInstruction();
return inst != nullptr && inst->Parent() == block && inst->GetOpcode() == opcode
@@ -192,19 +203,21 @@ private:
Value MatchRow(Value value, uint32_t shift) {
const auto* update = Get(Unwrap(value), ValueOpcode::DppUpdateU32);
if (update == nullptr || !IsRowShift(*update, shift) || !FullExec(update->Arg(2))) return {};
const auto* minimum = Get(update->Arg(0), ValueOpcode::UMin32);
if (minimum == nullptr) return {};
const auto* combined = Get(update->Arg(0), operation);
if (combined == nullptr) return {};
for (size_t index = 0; index < 2; index++) {
const auto* move = Get(minimum->Arg(index), ValueOpcode::DppMoveU32);
const auto* move = Get(combined->Arg(index), ValueOpcode::DppMoveU32);
if (move == nullptr || !IsRowShift(*move, shift) || !FullExec(move->Arg(1))) continue;
const auto previous = Unwrap(minimum->Arg(1 - index));
const auto previous = Unwrap(combined->Arg(1 - index));
if (Unwrap(move->Arg(0)) == previous && Unwrap(update->Arg(1)) == previous) return previous;
}
return {};
}
Block* block;
Value saved_exec;
Block* block;
ValueOpcode operation;
uint32_t identity;
Value saved_exec;
};
} // namespace
@@ -217,18 +230,22 @@ ReadLaneStats EliminateReadLane(Program& program, uint32_t wave_size) {
for (auto* block: program.blocks) {
for (auto& inst: *block) {
if (wave_size == 64u && inst.GetOpcode() == ValueOpcode::UMin32) {
const auto seed = UMinReductionMatcher(block).Match(inst);
const auto opcode = inst.GetOpcode();
if (opcode == ValueOpcode::UMin32 || opcode == ValueOpcode::UMax32) {
const auto seed = WaveReductionMatcher(block, opcode).Match(inst, wave_size);
if (!seed.IsEmpty()) {
const auto position = std::find_if(block->begin(), block->end(),
[&](const Inst& candidate) { return &candidate == &inst; });
auto& reduction = *block->PrependNewInst(position, ValueOpcode::SubgroupUMinU32, {seed});
const auto reduce = opcode == ValueOpcode::UMin32
? ValueOpcode::SubgroupUMinU32
: ValueOpcode::SubgroupUMaxU32;
auto& reduction = *block->PrependNewInst(position, reduce, {seed});
inst.ReplaceUsesWith(Value(&reduction));
stats.rewritten_reductions++;
}
continue;
}
if (inst.GetOpcode() != ValueOpcode::ReadLane) {
if (opcode != ValueOpcode::ReadLane) {
continue;
}
const auto selector = inst.Arg(1).Resolve();