fix(shader): isolate wave32 operations in wide host subgroups

This commit is contained in:
Nicolas Abosch 2026-10-04 20:59:05 -03:00
parent 9651ec9a92
commit f1d8ec8913
10 changed files with 272 additions and 15 deletions

View File

@ -681,6 +681,14 @@ if(BUILD_TESTING)
set_tests_properties(ray_tracing_miss PROPERTIES ENVIRONMENT_MODIFICATION
"PATH=path_list_prepend:$<TARGET_FILE_DIR:libc>")
endif()
add_test_executable(wave32_wide_subgroup_tests ${LIBS_INCLUDE_DIR}/tests/Wave32WideSubgroup.cpp)
target_link_libraries(wave32_wide_subgroup_tests PRIVATE shader_recompiler_agc_driver libc)
configure_windows_unwind(wave32_wide_subgroup_tests)
add_test(NAME wave32_wide_subgroup COMMAND wave32_wide_subgroup_tests)
if(WIN32)
set_tests_properties(wave32_wide_subgroup PROPERTIES ENVIRONMENT_MODIFICATION
"PATH=path_list_prepend:$<TARGET_FILE_DIR:libc>")
endif()
add_test_executable(vector_compare_i16_u16_tests ${LIBS_INCLUDE_DIR}/tests/VectorCompareI16U16.cpp)
target_link_libraries(vector_compare_i16_u16_tests PRIVATE shader_recompiler_agc_driver libc)
configure_windows_unwind(vector_compare_i16_u16_tests)

View File

@ -53,11 +53,10 @@ std::vector<std::array<float, 4>> Triangles() {
return vertices;
}
void Draw(AgcDriver::VulkanDevice& device, std::uint32_t waveSize) {
void Draw(AgcDriver::VulkanDevice& device, std::uint32_t waveSize, const ShaderRecompiler::SpirvTarget& target) {
static const auto triangles = Triangles();
Pixels.fill(std::byte{0});
Lanes.fill(Sentinel);
const auto target = device.Target();
std::vector<std::uint32_t> vertexUserData(4, 0u);
const auto vertexBuffer = BufferDescriptor(triangles.data(), 16u, static_cast<std::uint32_t>(triangles.size()));
@ -156,10 +155,18 @@ int main() {
}
for (const auto waveSize : {64u, 32u}) {
for (std::uint32_t pass = 0; pass < 4; ++pass) {
Draw(*device, waveSize);
Draw(*device, waveSize, device->Target());
Check(waveSize);
}
}
if (device->Target().subgroupSize == 32u) {
auto wide = device->Target();
wide.subgroupSize = 64u;
for (std::uint32_t pass = 0; pass < 4; ++pass) {
Draw(*device, 32u, wide);
Check(32u);
}
}
std::puts("pixel helper lane tests passed");
return 0;
} catch (const std::exception& error) {

View File

@ -70,12 +70,22 @@ int main() {
const auto device = OpenVulkanTestDevice();
if (!device) return VulkanTestSkipped;
const auto target = device->ComputeTarget(WaveSize);
if (target.subgroupSize != WaveSize) {
if (target.subgroupSize < WaveSize) {
std::printf("skipped, the device runs wave32 programs on %u-wide subgroups\n", target.subgroupSize);
return VulkanTestSkipped;
}
Run(*device, target);
Check();
if (device->Target().subgroupSize != target.subgroupSize && device->Target().subgroupSize >= WaveSize) {
Run(*device, device->Target());
Check();
}
if (device->Target().subgroupSize == WaveSize) {
auto wide = device->Target();
wide.subgroupSize = 64;
Run(*device, wide);
Check();
}
std::puts("wave32 subgroup tests passed");
return 0;
} catch (const std::exception& error) {

View File

@ -0,0 +1,201 @@
#include "SpirvBackend/SpirvEmitter.hpp"
#include "SpirvBackend/SpirvEmitterInstructions.hpp"
#include "SpirvBackend/SpirvMemory/SpirvSubgroup.hpp"
#include "SpirvBackend/SpirvMemory/SpirvConstants.hpp"
#include <array>
#include <bit>
#include <cstdio>
#include <stdexcept>
#include <span>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
using namespace ShaderRecompiler;
namespace {
using Value = std::array<std::uint32_t, 4>;
using Lanes = std::vector<Value>;
void Require(bool condition, const char* reason) {
if (!condition) throw std::runtime_error(reason);
}
class Subgroup {
public:
explicit Subgroup(std::uint32_t size) : size(size) {}
void Set(std::uint32_t id, Lanes lanes) {
values[id] = std::move(lanes);
}
const Lanes& Get(std::uint32_t id) const {
return values.at(id);
}
void Execute(std::span<const std::uint32_t> code) {
for (std::size_t offset = 5; offset < code.size();) {
const auto count = code[offset] >> 16u;
Require(count != 0 && count <= code.size() - offset, "malformed SPIR-V instruction");
const auto words = code.subspan(offset, count);
const auto op = static_cast<spv::Op>(words[0] & 0xffffu);
offset += count;
switch (op) {
case spv::OpTypeInt:
case spv::OpTypeBool:
case spv::OpTypeVector:
continue;
case spv::OpStore:
values[words[1]] = values.at(words[2]);
continue;
default:
break;
}
auto& result = values[words[2]];
result.resize(size);
for (std::uint32_t lane = 0; lane < size; ++lane) {
const auto arg = [&](std::size_t index) -> const Value& { return values.at(words[index]).at(lane); };
auto& value = result[lane];
switch (op) {
case spv::OpConstant: value[0] = words[3]; break;
case spv::OpConstantTrue: value[0] = 1; break;
case spv::OpConstantFalse: value[0] = 0; break;
case spv::OpLoad: value = arg(3); break;
case spv::OpBitwiseAnd: value[0] = arg(3)[0] & arg(4)[0]; break;
case spv::OpBitwiseOr: value[0] = arg(3)[0] | arg(4)[0]; break;
case spv::OpBitwiseXor: value[0] = arg(3)[0] ^ arg(4)[0]; break;
case spv::OpIAdd: value[0] = arg(3)[0] + arg(4)[0]; break;
case spv::OpIMul: value[0] = arg(3)[0] * arg(4)[0]; break;
case spv::OpShiftRightLogical: value[0] = arg(3)[0] >> arg(4)[0]; break;
case spv::OpShiftLeftLogical: value[0] = arg(3)[0] << arg(4)[0]; break;
case spv::OpINotEqual: value[0] = arg(3)[0] != arg(4)[0]; break;
case spv::OpIEqual: value[0] = arg(3)[0] == arg(4)[0]; break;
case spv::OpULessThan: value[0] = arg(3)[0] < arg(4)[0]; break;
case spv::OpLogicalAnd: value[0] = arg(3)[0] && arg(4)[0]; break;
case spv::OpSelect: value = arg(3)[0] ? arg(4) : arg(5); break;
case spv::OpCompositeExtract: value[0] = arg(3).at(words[4]); break;
case spv::OpVectorExtractDynamic: value[0] = arg(3).at(arg(4)[0]); break;
case spv::OpCompositeConstruct:
for (std::size_t component = 0; component < count - 3u; ++component) value.at(component) = arg(3u + component)[0];
break;
case spv::OpGroupNonUniformBallot:
value.fill(0);
for (std::uint32_t source = 0; source < size; ++source) {
if (values.at(words[4])[source][0]) value[source / 32u] |= 1u << (source % 32u);
}
break;
case spv::OpGroupNonUniformBallotFindLSB:
value[0] = 0xffffffffu;
for (std::uint32_t component = 0; component < 4; ++component) {
if (arg(4)[component] != 0) {
value[0] = component * 32u + std::countr_zero(arg(4)[component]);
break;
}
}
break;
case spv::OpGroupNonUniformShuffle:
Require(arg(5)[0] < size, "shuffle escaped the host subgroup");
value = values.at(words[4]).at(arg(5)[0]);
break;
default:
throw std::runtime_error("unexpected opcode in subgroup test: " + std::to_string(op));
}
}
}
}
private:
std::uint32_t size;
std::unordered_map<std::uint32_t, Lanes> values;
};
void Check(std::uint32_t size, const Value& masks) {
IrProgram program;
program.SetWaveSize(32);
SpirvEmitterState state(program, {});
state.splitSubgroup = size > 32;
state.subgroupLocalInvocationIdVariable = state.module.AllocateId();
SpirvValueEmitContext context(state);
context.scratchU32Variable = state.module.AllocateId();
IrValue source(IrOpcode::GetAttribute, IrType::U32, 1);
IrValue predicate(IrOpcode::GetExec, IrType::Bool, 2);
IrValue requestedLane(IrOpcode::GetAttribute, IrType::U32, 3);
IrValue first(IrOpcode::ReadFirstLane, IrType::U32, 4);
first.AddArgument(&source);
first.AddArgument(&predicate);
IrValue read(IrOpcode::ReadLane, IrType::U32, 5);
read.AddArgument(&source);
read.AddArgument(&requestedLane);
IrValue control(IrOpcode::Void, IrType::U32, 6);
control.SetImmediateU32(0x7c1fu);
IrValue swizzle(IrOpcode::SwizzleU32, IrType::U32, 7);
swizzle.AddArgument(&source);
swizzle.AddArgument(&control);
swizzle.AddArgument(&predicate);
IrValue replacement(IrOpcode::Void, IrType::U32, 8);
replacement.SetImmediateU32(0x12345678u);
IrValue writeTarget(IrOpcode::Void, IrType::U32, 9);
writeTarget.SetImmediateU32(7u);
IrValue write(IrOpcode::WriteLane, IrType::U32, 10);
write.AddArgument(&replacement);
write.AddArgument(&writeTarget);
write.AddArgument(&source);
const auto sourceId = context.Def(&source);
const auto predicateId = context.Def(&predicate);
const auto requestedId = context.Def(&requestedLane);
const auto logicalId = EmitSubgroupLocalInvocationId(state);
const auto ballot = context.Ballot(&predicate);
const auto readFirst = EmitReadFirstLane(context, first);
const auto readLane = EmitReadLane(context, read);
const auto active = EmitBallotLaneActiveBool(state, ballot, requestedId);
const auto quad = EmitDppQuadPermTargetLane(state, logicalId, 0x1bu);
const auto quadValue = context.Shuffle(read, 0, quad.lane);
const auto hardwareActive = EmitSubgroupLaneActiveBool(state, requestedId);
const auto swizzled = EmitSwizzleU32(context, swizzle);
const auto written = EmitWriteLane(context, write);
Subgroup subgroup(size);
Lanes host(size), sources(size), predicates(size), requested(size);
for (std::uint32_t lane = 0; lane < size; ++lane) {
host[lane][0] = lane;
sources[lane][0] = 0x1000u + lane;
predicates[lane][0] = (masks[lane / 32u] >> (lane % 32u)) & 1u;
requested[lane][0] = 31u - lane % 32u;
}
subgroup.Set(state.subgroupLocalInvocationIdVariable, host);
subgroup.Set(sourceId, sources);
subgroup.Set(predicateId, predicates);
subgroup.Set(requestedId, requested);
subgroup.Execute(state.module.Finalize());
for (std::uint32_t lane = 0; lane < size; ++lane) {
const auto mask = masks[lane / 32u];
const auto base = lane / 32u * 32u;
Require(subgroup.Get(logicalId)[lane][0] == lane % 32u, "lane id escaped its wave32");
Require(subgroup.Get(ballot)[lane] == Value{mask, 0, 0, 0}, "ballot mixed independent wave32s");
Require(subgroup.Get(readFirst)[lane][0] == 0x1000u + base + (mask == 0 ? 0u : std::countr_zero(mask)), "readfirstlane read a different wave32");
Require(subgroup.Get(readLane)[lane][0] == 0x1000u + base + requested[lane][0], "readlane read a different wave32");
Require(subgroup.Get(active)[lane][0] == ((mask >> requested[lane][0]) & 1u), "source lane activity came from another wave32");
Require(subgroup.Get(quadValue)[lane][0] == 0x1000u + (lane & ~3u) + 3u - lane % 4u, "DPP quad permutation crossed waves");
Require(subgroup.Get(hardwareActive)[lane][0] == 1u, "upper wave32 lanes were treated as inactive");
Require(subgroup.Get(swizzled)[lane][0] == (((mask >> requested[lane][0]) & 1u) != 0 ? 0x1000u + base + requested[lane][0] : 0u), "DS swizzle read another wave32 or ignored source EXEC");
Require(subgroup.Get(written)[lane][0] == (lane % 32u == 7u ? 0x12345678u : 0x1000u + lane), "writelane missed the upper wave32");
}
}
}
int main() {
try {
for (std::uint32_t size : {32u, 64u, 128u}) {
for (const Value& masks : {Value{0xffffffffu, 0xffffffffu, 0xffffffffu, 0xffffffffu}, Value{1, 0x80000000u, 0x100, 0x4000}, Value{0x55555555u, 0xaaaaaaaau, 0x12345678u, 0xfedcba98u}, Value{0, 1, 0, 0x80000000u}, Value{1, 0, 0x80000000u, 0}, Value{0, 0, 0, 0}}) {
Check(size, masks);
}
}
std::puts("wave32 wide subgroup tests passed");
return 0;
} catch (const std::exception& error) {
std::fprintf(stderr, "%s\n", error.what());
return 1;
}
}

View File

@ -82,6 +82,7 @@ struct SpirvEmitterState {
SpirvRequirements requirements;
std::uint32_t laneCount = 1;
bool splitSubgroup = false;
std::unordered_set<const IrValue*> sharedLaneValues;
std::uint32_t laneHalf = 0;
// The target's SPIR-V version and what the device accepts, for capabilities an emitter adds

View File

@ -7,6 +7,8 @@
namespace ShaderRecompiler {
std::uint32_t EmitSubgroupLocalInvocationId(SpirvEmitterState& state);
std::uint32_t EmitHostSubgroupLane(SpirvEmitterState& state, std::uint32_t lane);
std::uint32_t EmitWaveBallot(SpirvEmitterState& state, std::uint32_t ballot);
DppTargetLane EmitDppQuadPermTargetLane(SpirvEmitterState& state, std::uint32_t subid, std::uint32_t control);
DppTargetLane EmitDppRowShiftTargetLane(SpirvEmitterState& state, std::uint32_t subid, std::uint32_t amount, bool left);
DppTargetLane EmitDppRowRotateRightTargetLane(SpirvEmitterState& state, std::uint32_t subid, std::uint32_t amount);

View File

@ -1,4 +1,5 @@
#include "SpirvBackend/SpirvEmitterInstructions.hpp"
#include "SpirvBackend/SpirvMemory/SpirvSubgroup.hpp"
#include "IntermediateRepresentation/IrValue.hpp"
#include <spirv/unified1/GLSL.std.450.h>
#include <spirv/unified1/spirv.hpp>
@ -201,10 +202,11 @@ std::uint32_t EmitDsMaskedLaneRead(SpirvEmitterState& state, std::uint32_t sourc
if (state.laneCount == 2) {
lane = Binary(state, spv::OpBitwiseAnd, TypeU32(state), lane, ConstantU32(state, 31u));
}
const auto physicalLane = EmitHostSubgroupLane(state, lane);
const auto shuffled = state.module.AllocateId();
state.module.AddFunction(spv::OpGroupNonUniformShuffle, TypeU32(state), shuffled, ConstantU32(state, spv::ScopeSubgroup), source, lane);
state.module.AddFunction(spv::OpGroupNonUniformShuffle, TypeU32(state), shuffled, ConstantU32(state, spv::ScopeSubgroup), source, physicalLane);
const auto sourceExec = state.module.AllocateId();
state.module.AddFunction(spv::OpGroupNonUniformShuffle, TypeBool(state), sourceExec, ConstantU32(state, spv::ScopeSubgroup), exec, lane);
state.module.AddFunction(spv::OpGroupNonUniformShuffle, TypeBool(state), sourceExec, ConstantU32(state, spv::ScopeSubgroup), exec, physicalLane);
const auto sourceActive = Binary(state, spv::OpLogicalAnd, TypeBool(state), sourceExec, EmitSubgroupLaneActiveBool(state, lane));
return Select(state, TypeU32(state), sourceActive, shuffled, ConstantU32(state, 0u));
}

View File

@ -4,6 +4,7 @@
#include "SpirvBackend/SpirvEmitterState.hpp"
#include "SpirvBackend/SpirvFlowEmitter.hpp"
#include "SpirvBackend/SpirvMemory/SpirvModuleSetup.hpp"
#include "SpirvBackend/SpirvMemory/SpirvSubgroup.hpp"
#include <spirv/unified1/GLSL.std.450.h>
#include <spirv/unified1/spirv.hpp>
#include <algorithm>
@ -119,7 +120,7 @@ std::uint32_t SpirvValueEmitContext::Ballot(const IrValue* predicate) {
const auto low = state.module.AllocateId();
state.module.AddFunction(spv::OpGroupNonUniformBallot, ballotType, low, scope, otherHalf == nullptr || half == 0u ? Def(predicate) : otherHalf->Def(predicate));
if (otherHalf == nullptr) {
return low;
return EmitWaveBallot(state, low);
}
const auto high = state.module.AllocateId();
const auto lowWord = state.module.AllocateId();
@ -158,7 +159,7 @@ std::uint32_t SpirvValueEmitContext::Shuffle(const IrValue& inst, std::size_t in
const auto scope = ConstantU32(state, spv::ScopeSubgroup);
const auto low = state.module.AllocateId();
if (otherHalf == nullptr) {
state.module.AddFunction(spv::OpGroupNonUniformShuffle, type, low, scope, Arg(inst, index), lane);
state.module.AddFunction(spv::OpGroupNonUniformShuffle, type, low, scope, Arg(inst, index), EmitHostSubgroupLane(state, lane));
return low;
}
const auto physicalLane = EmitBinaryU32(state, spv::OpBitwiseAnd, lane, ConstantU32(state, 31u));
@ -247,6 +248,8 @@ std::vector<std::uint32_t> SpirvEmitter::Emit(const IrProgram& program, const Sh
state.supportedCapabilities = target.supportedCapabilities;
state.supportedExtensions = target.supportedExtensions;
state.nonConstantImageOffsets = target.nonConstantImageOffsets;
state.splitSubgroup = program.WaveSize() == 32u && target.subgroupSize > 32u;
if (state.splitSubgroup && (state.requirements.subgroupBallot || state.requirements.subgroupShuffle)) state.requirements.subgroupLocalInvocationId = true;
const auto* workgroup = ShaderWorkgroupInputFor(state);
state.laneCount = workgroup != nullptr && program.WaveSize() == 64u && workgroup->hostSubgroupSize == 32u ? 2u : 1u;
if (state.laneCount == 2u) state.sharedLaneValues = WaveUniformValues(program);

View File

@ -1,5 +1,6 @@
#include "SpirvBackend/SpirvBda.hpp"
#include "SpirvBackend/SpirvEmitterInstructions.hpp"
#include "SpirvBackend/SpirvMemory/SpirvSubgroup.hpp"
#include "SpirvBackend/SpirvBufferFormat.hpp"
#include <spirv/unified1/spirv.hpp>
#include <algorithm>
@ -1082,7 +1083,7 @@ std::uint32_t AppendConsume(SpirvValueEmitContext& ctx, const IrValue& inst, boo
return value;
});
const auto result = state.module.AllocateId();
state.module.AddFunction(spv::OpGroupNonUniformShuffle, TypeU32(state), result, ConstantU32(state, spv::ScopeSubgroup), atomic, sourceLane);
state.module.AddFunction(spv::OpGroupNonUniformShuffle, TypeU32(state), result, ConstantU32(state, spv::ScopeSubgroup), atomic, EmitHostSubgroupLane(state, sourceLane));
return result;
}

View File

@ -17,17 +17,39 @@ namespace ShaderRecompiler
throw std::runtime_error("SPIR-V module emission failed: " + reason);
}
std::uint32_t HostInvocationId(SpirvEmitterState& state) {
if (state.subgroupLocalInvocationIdVariable == 0) {
FailEmit("SubgroupLocalInvocationId was not declared before function emission");
}
const auto value = state.module.AllocateId();
state.module.AddFunction(spv::OpLoad, TypeU32(state), value, state.subgroupLocalInvocationIdVariable);
return value;
}
}
std::uint32_t EmitSubgroupLocalInvocationId(SpirvEmitterState& state) {
if (state.subgroupLocalInvocationIdVariable == 0) {
FailEmit("SubgroupLocalInvocationId was not declared before function emission");
}
const auto value = state.module.AllocateId();
state.module.AddFunction(spv::OpLoad, TypeU32(state), value, state.subgroupLocalInvocationIdVariable);
const auto value = HostInvocationId(state);
if (state.splitSubgroup) return EmitBinaryU32(state, spv::OpBitwiseAnd, value, ConstantU32(state, 31u));
return state.laneHalf == 0 ? value : EmitAddU32(state, value, ConstantU32(state, 32));
}
std::uint32_t EmitHostSubgroupLane(SpirvEmitterState& state, std::uint32_t lane) {
if (!state.splitSubgroup) return lane;
const auto base = EmitBinaryU32(state, spv::OpBitwiseAnd, HostInvocationId(state), ConstantU32(state, ~31u));
return EmitAddU32(state, base, lane);
}
std::uint32_t EmitWaveBallot(SpirvEmitterState& state, std::uint32_t ballot) {
if (!state.splitSubgroup) return ballot;
const auto word = EmitBinaryU32(state, spv::OpShiftRightLogical, HostInvocationId(state), ConstantU32(state, 5u));
const auto mask = state.module.AllocateId();
const auto result = state.module.AllocateId();
state.module.AddFunction(spv::OpVectorExtractDynamic, TypeU32(state), mask, ballot, word);
state.module.AddFunction(spv::OpCompositeConstruct, TypeU32Vector(state, 4u), result, mask, ConstantU32(state, 0u), ConstantU32(state, 0u), ConstantU32(state, 0u));
return result;
}
DppTargetLane EmitDppQuadPermTargetLane(SpirvEmitterState& state, std::uint32_t subid, std::uint32_t control) {
constexpr std::uint32_t laneBits = 2u;
constexpr std::uint32_t laneMask = (1u << laneBits) - 1u;
@ -157,6 +179,6 @@ namespace ShaderRecompiler
std::uint32_t EmitSubgroupLaneActiveBool(SpirvEmitterState& state, std::uint32_t lane) {
const auto activeBallot = state.module.AllocateId();
state.module.AddFunction(spv::OpGroupNonUniformBallot, TypeU32Vector(state, 4), activeBallot, ConstantU32(state, spv::ScopeSubgroup), ConstantBool(state, true));
return EmitBallotLaneActiveBool(state, activeBallot, lane);
return EmitBallotLaneActiveBool(state, EmitWaveBallot(state, activeBallot), lane);
}
}