Merge pull request #2450 from oneandonlydean/up639/half-wave-vertex-exec

fix(shader): match half-wave reductions in vertex programs
This commit is contained in:
Dean 2026-10-10 17:47:33 +13:00 committed by GitHub
commit 29fe66dfd1
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 25 additions and 4 deletions

View File

@ -807,6 +807,23 @@ void verifyHalfWaveReduction() {
const auto wide = "half-wave reduction: a wave64 scan reduced by SPIR-V opcode " + std::to_string(reduction.reduce) + " did not read lanes 31 and 63 as reductions of host invocations 0-31 and 32-63 on 64 lanes";
require(halves(pixel(scan(reduction.identity, reduction.vector, reduction.scalar), 64u, true), reduction.reduce, reduction.identity), wide.c_str());
}
const auto vertex = [](std::vector<std::uint32_t> code, std::uint32_t subgroupSize) {
*std::find(code.begin(), code.end(), 0xf800180fu) = 0xf80008cfu;
RecompileRequest request{};
request.shader = {ShaderStage::Vertex, 0x10000u, code, 0, {}};
request.context.waveSize = 64;
request.context.vertex = ShaderVertexStageInfo{};
request.target.vulkanVersion = 0x00401000u;
request.target.spirvVersion = 0x00010300u;
request.target.subgroupSize = subgroupSize;
request.target.supportedCapabilities = capabilities;
request.target.fragmentShaderBarycentricEnabled = false;
request.layout.pushConstantSizeBytes = 128;
request.useCache = false;
return Recompile(request).spirv;
};
require(reductions(vertex(scan(0xffffffffu, 0x26000000u, 0x83800000u), 32u), spv::OpGroupNonUniformUMin) == 1u, "half-wave reduction: the wave64 vertex UMin scan, whose entry EXEC is every invocation, did not read lane 31 as a subgroup UMin on 32 lanes");
require(halves(vertex(scan(0x80000000u, 0x24000000u, 0x84000000u), 64u), spv::OpGroupNonUniformSMax, 0x80000000u), "half-wave reduction: the wave64 vertex SMax scan did not read lanes 31 and 63 as reductions of host invocations 0-31 and 32-63 on 64 lanes");
const auto umin = scan(0xffffffffu, 0x26000000u, 0x83800000u);
const auto patched = [&umin](std::initializer_list<std::pair<std::size_t, std::uint32_t>> words) {
auto code = umin;

View File

@ -297,6 +297,8 @@ bool AllLanesBit(const IrValue* bit) {
class EntryLaneWalk {
public:
explicit EntryLaneWalk(IrShaderStage stage) : everyInvocationEnters(stage == IrShaderStage::Vertex || stage == IrShaderStage::Local || stage == IrShaderStage::TessellationControl || stage == IrShaderStage::TessellationEvaluation) {}
bool ZeroBit(const IrValue* bit) {
bit = Resolved(bit);
if (bit == nullptr || !Spend()) return false;
@ -317,7 +319,8 @@ private:
bool NotHelper(const IrValue* predicate) {
predicate = Resolved(predicate);
if (predicate == nullptr || predicate->HasImmediate() || !Spend()) return false;
if (predicate == nullptr || !Spend()) return false;
if (predicate->HasImmediate()) return everyInvocationEnters && predicate->Type() == IrType::Bool && predicate->ImmediateBool();
switch (predicate->Opcode()) {
case IrOpcode::LogicalAnd: return NotHelper(predicate->Argument(0)) || NotHelper(predicate->Argument(1));
case IrOpcode::IEqual32: return EitherOrder(predicate, [](const IrValue* builtin, const IrValue* zero) { return IsU32(zero, 0u) && Is(builtin, IrOpcode::GetBuiltin) && IsU32(builtin->Argument(0), static_cast<std::uint32_t>(StageInputKind::HelperInvocation)) && IsU32(builtin->Argument(1), 0u); });
@ -362,6 +365,7 @@ private:
return result;
}
bool everyInvocationEnters;
std::unordered_set<const IrValue*> phis;
std::uint32_t budget = 512u;
};
@ -393,7 +397,7 @@ bool CrossesRowsOf(const IrValue* value, const IrValue* scan) {
return flags.x16 && !flags.fetchInactive && Resolved(permlane->Argument(0)) == scan && IsU32(permlane->Argument(1), 0xffffffffu) && IsU32(permlane->Argument(2), 0xffffffffu) && AllLanesBit(permlane->Argument(3));
}
std::optional<HalfWaveScan> MatchHalfWaveScan(const IrValue* value) {
std::optional<HalfWaveScan> MatchHalfWaveScan(const IrValue* value, IrShaderStage stage) {
const auto* combined = Unmasked(Resolved(value));
const auto reduction = std::find_if(WaveReductions.begin(), WaveReductions.end(), [&](const WaveReduction& candidate) { return Is(combined, candidate.opcode); });
const IrValue* scan = nullptr;
@ -403,14 +407,14 @@ std::optional<HalfWaveScan> MatchHalfWaveScan(const IrValue* value) {
if (scan == nullptr) return std::nullopt;
}
const auto* source = Unmasked(scan);
if (!Is(source, IrOpcode::SelectU32) || !IsU32(source->Argument(2), reduction->identity) || !EntryLaneWalk{}.ZeroBit(source->Argument(0))) return std::nullopt;
if (!Is(source, IrOpcode::SelectU32) || !IsU32(source->Argument(2), reduction->identity) || !EntryLaneWalk{stage}.ZeroBit(source->Argument(0))) return std::nullopt;
return HalfWaveScan{&*reduction, source};
}
std::optional<std::uint32_t> EmitHalfWaveReduction(SpirvValueEmitContext& ctx, const IrValue& inst, std::uint32_t lane) {
auto& state = ctx.state;
if (state.program.WaveSize() != 64u || (lane & 31u) != 31u) return std::nullopt;
const auto scan = MatchHalfWaveScan(inst.Argument(0));
const auto scan = MatchHalfWaveScan(inst.Argument(0), state.program.Resources().stage);
if (!scan) return std::nullopt;
const auto read = "v_readlane_b32 of lane " + std::to_string(lane) + " of a wave64 half-wave reduction scan";
const auto& capabilities = state.supportedCapabilities;