Merge pull request #4162 from Celegans12/fix/shader-wave64-lane-reductions

fix(shader): match half-wave reductions over relatively addressed vector registers
This commit is contained in:
Dean 2026-10-12 23:54:08 +13:00 committed by GitHub
commit 0e61bedf9a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 41 additions and 7 deletions

View File

@ -1341,6 +1341,13 @@ void verifyHalfWaveReduction() {
divergent.insert(divergent.begin() + 3, 0x88fe1a7eu);
divergent.insert(divergent.begin() + 2, {0xbe860480u, 0x7d821688u, 0xbe9a246au, 0xbf880001u});
require(reductions(pixel(divergent, 32u, true), spv::OpGroupNonUniformUMin) == 1u, "half-wave reduction: keys masked inside a divergent branch did not read lane 31 as a subgroup UMin");
auto windowed = umin;
windowed.insert(windowed.begin() + 23, 0x7e008514u);
require(reductions(pixel(windowed, 32u, true), spv::OpGroupNonUniformUMin) == 1u, "half-wave reduction: a UMin scan of registers in a v_movreld_b32 window did not read lane 31 as a subgroup UMin on 32 lanes");
require(halves(pixel(windowed, 64u, true), spv::OpGroupNonUniformUMin, 0xffffffffu), "half-wave reduction: a UMin scan of registers in a v_movreld_b32 window did not read lanes 31 and 63 as reductions of host invocations 0-31 and 32-63 on 64 lanes");
auto overwritten = umin;
overwritten.insert(overwritten.begin() + 17, {0x7ef80500u, 0x7e008514u});
expectFailure([&] { static_cast<void>(pixel(overwritten, 32u, true)); }, lane63, "half-wave reduction: a scan of window registers read lane 63 after a v_movreld_b32 at a runtime M0 that may overwrite them");
const char* arithmetic = "v_readlane_b32 of lane 31 of a wave64 half-wave reduction scan needs subgroup arithmetic";
expectFailure([&] { static_cast<void>(pixel(umin, 32u, false)); }, arithmetic, "half-wave reduction: a device without subgroup arithmetic read lane 31 on 32 lanes");
expectFailure([&] { static_cast<void>(pixel(umin, 64u, false)); }, arithmetic, "half-wave reduction: a device without subgroup arithmetic read lane 31 on 64 lanes");

View File

@ -194,13 +194,16 @@ int main() {
}};
for (const auto& layout : layouts) {
for (const auto& reduction : reductions) {
const auto code = PixelCode(reduction);
std::uint32_t partial = 0;
for (std::uint32_t pass = 0; pass < 4; ++pass) {
Draw(*device, code, layout);
partial += Check(reduction, target.subgroupSize);
for (const bool windowed : {false, true}) {
auto code = PixelCode(reduction);
if (windowed) code.insert(std::find(code.begin(), code.end(), 0x7e0e02f2u), 0x7e008514u);
std::uint32_t partial = 0;
for (std::uint32_t pass = 0; pass < 4; ++pass) {
Draw(*device, code, layout);
partial += Check(reduction, target.subgroupSize);
}
std::printf("%s%s, laid out for %u lanes on %u-lane subgroups: %u partial waves over 4 draws\n", reduction.name, windowed ? " in a v_movreld_b32 window" : "", layout.subgroupSize, target.subgroupSize, partial);
}
std::printf("%s, laid out for %u lanes on %u-lane subgroups: %u partial waves over 4 draws\n", reduction.name, layout.subgroupSize, target.subgroupSize, partial);
}
}
std::puts("pixel wave reduction tests passed");

View File

@ -1,5 +1,6 @@
#include "SpirvBackend/SpirvEmitterInstructions.hpp"
#include "SpirvBackend/SpirvMemory/SpirvSubgroup.hpp"
#include "IntermediateRepresentation/IrBlock.hpp"
#include "IntermediateRepresentation/IrValue.hpp"
#include <spirv/unified1/GLSL.std.450.h>
#include <spirv/unified1/spirv.hpp>
@ -227,8 +228,31 @@ struct HalfWaveScan {
const IrValue* source = nullptr;
};
const IrValue* StoredWindowValue(const IrValue* load) {
const auto* slot = load->Argument(0)->Resolve();
const auto* block = load->Parent();
if (slot == nullptr || !slot->HasImmediate() || block == nullptr) return nullptr;
const auto& instructions = block->Instructions();
auto at = std::find(instructions.begin(), instructions.end(), load);
if (at == instructions.end()) return nullptr;
while (at != instructions.begin()) {
const auto* previous = *--at;
if (previous->HasImmediate() || previous->Opcode() != IrOpcode::StoreVectorWindow) continue;
const auto* stored = previous->Argument(0)->Resolve();
if (stored == nullptr || !stored->HasImmediate()) return nullptr;
if (stored->ImmediateU32() == slot->ImmediateU32()) return previous->Argument(1)->Resolve();
}
return nullptr;
}
const IrValue* Resolved(const IrValue* value) {
return value != nullptr ? value->Resolve() : nullptr;
value = value != nullptr ? value->Resolve() : nullptr;
while (value != nullptr && !value->HasImmediate() && value->Opcode() == IrOpcode::LoadVectorWindow) {
const auto* stored = StoredWindowValue(value);
if (stored == nullptr) break;
value = stored;
}
return value;
}
bool Is(const IrValue* value, IrOpcode opcode) {