shader: share repeated raw indirect buffer load bodies

This commit is contained in:
nmzik
2026-10-14 22:17:58 +02:00
parent 6455546f26
commit c1f6838e40
4 changed files with 200 additions and 29 deletions
@@ -115,6 +115,7 @@ struct EmitterState {
uint32_t fault_buffer_variable = 0;
uint32_t bda_pointer_function = 0;
uint32_t indirect_format_function = 0;
std::map<std::array<uint32_t, 3>, uint32_t> indirect_buffer_load_functions;
uint32_t bvh_intersect_function = 0;
uint32_t gds_variable = 0;
uint32_t gds_length = 0;
@@ -524,6 +525,7 @@ void EmitProgram(EmitterState& state);
void DefineGetBdaPointer(EmitterState& state);
void DefineConvertIndirectBufferFormat(ValueEmitContext& ctx);
void DefineIndirectBufferLoads(ValueEmitContext& ctx);
void DefineBvhIntersect(EmitterState& state);
uint32_t GetBdaPointer(EmitterState& state, uint32_t address);
@@ -910,12 +910,22 @@ struct IndirectBufferAccess {
uint32_t raw_index_in_bounds;
};
IndirectBufferAccess PrepareIndirectBuffer(ValueEmitContext& ctx, const IR::Inst& inst) {
auto& state = ctx.state;
std::array<uint32_t, 8> IndirectBufferArgs(ValueEmitContext& ctx, const IR::Inst& inst) {
auto& state = ctx.state;
const auto& handle = *inst.Arg(0).ResolveInstruction();
const auto word1 = ctx.Arg(handle, 1);
const auto records = ctx.Arg(handle, 2);
const auto word3 = ctx.Arg(handle, 3);
auto offset = ctx.Arg(inst, 2);
if (ctx.Memory(inst).offset != 0u) {
const auto immediate = ConstantU32(state, ctx.Memory(inst).offset);
offset = offset == ConstantU32(state, 0) ? immediate
: Binary(state, spv::OpIAdd, TypeU32(state), offset, immediate);
}
return {ctx.Arg(handle, 0), ctx.Arg(handle, 1), ctx.Arg(handle, 2), ctx.Arg(handle, 3),
ctx.Arg(inst, 1), offset, ctx.Arg(inst, 3), BufferLane(state)};
}
IndirectBufferAccess PrepareIndirectBuffer(EmitterState& state,
const std::array<uint32_t, 8>& args, bool formatted) {
const auto& [word0, word1, records, word3, input_index, offset, soffset, lane] = args;
const auto field = [&](uint32_t word, uint32_t first, uint32_t count) {
return EmitBitFieldUExtract(state, word, ConstantU32(state, first), ConstantU32(state, count));
};
@@ -925,14 +935,12 @@ IndirectBufferAccess PrepareIndirectBuffer(ValueEmitContext& ctx, const IR::Inst
const auto stride = field(word1, 16, 14);
const auto swizzle = AndCondition(state, nonzero(stride), nonzero(field(word1, 31, 1)));
const auto index = Binary(
state, spv::OpIAdd, TypeU32(state), ctx.Arg(inst, 1),
Select(state, TypeU32(state), nonzero(field(word3, 23, 1)), BufferLane(state),
state, spv::OpIAdd, TypeU32(state), input_index,
Select(state, TypeU32(state), nonzero(field(word3, 23, 1)), lane,
ConstantU32(state, 0)));
const auto soffset = ctx.Arg(inst, 3);
const auto address = CalculateBufferAddress(state, index, ctx.Arg(inst, 2), soffset,
ctx.Memory(inst).offset, stride, swizzle,
const auto address = CalculateBufferAddress(state, index, offset, soffset, 0u, stride, swizzle,
field(word3, 21, 2));
const auto base = PackU64(state, ctx.Arg(handle, 0), field(word1, 0, 16));
const auto base = PackU64(state, word0, field(word1, 0, 16));
const auto raw_records = Binary(state, spv::OpISub, TypeU32(state), records, soffset);
return {
.address = Binary(state, spv::OpIAdd, TypeU64(state), base,
@@ -941,7 +949,7 @@ IndirectBufferAccess PrepareIndirectBuffer(ValueEmitContext& ctx, const IR::Inst
.stride = stride,
.swizzle = swizzle,
.format = field(word3, 12, 7),
.selector = ctx.Memory(inst).formatted ? field(word3, 0, 3) : 0u,
.selector = formatted ? field(word3, 0, 3) : 0u,
.mode = field(word3, 28, 2),
.records = records,
.index_in_bounds = Binary(state, spv::OpULessThan, TypeBool(state), index, records),
@@ -986,10 +994,9 @@ uint32_t IndirectBufferInBounds(EmitterState& state, const IndirectBufferAccess&
return in_bounds;
}
uint32_t LoadIndirectBuffer(ValueEmitContext& ctx, const IR::Inst& inst, uint32_t components) {
uint32_t LoadIndirectBuffer(ValueEmitContext& ctx, const IndirectBufferAccess& buffer,
uint32_t bits, uint32_t components, bool coherent) {
auto& state = ctx.state;
const auto bits = ctx.Memory(inst).data_bits;
const auto buffer = PrepareIndirectBuffer(ctx, inst);
const auto valid_format = Binary(state, spv::OpINotEqual, TypeBool(state), buffer.format,
ConstantU32(state, 0));
const auto base = Binary(state, spv::OpBitwiseAnd, TypeU64(state), buffer.address,
@@ -1004,7 +1011,7 @@ uint32_t LoadIndirectBuffer(ValueEmitContext& ctx, const IR::Inst& inst, uint32_
IndirectBufferInBounds(state, buffer, component * 4u,
ConstantU32(state, bits / 8u), false)),
[&]() {
const auto value = LoadBdaDword(ctx, address, ctx.Memory(inst).coherent);
const auto value = LoadBdaDword(ctx, address, coherent);
if (bits == 32u) return value;
const auto byte = Binary(
state, spv::OpBitwiseAnd, TypeU32(state),
@@ -1022,7 +1029,7 @@ uint32_t LoadIndirectBuffer(ValueEmitContext& ctx, const IR::Inst& inst, uint32_
uint32_t LoadIndirectFormattedX(ValueEmitContext& ctx, const IR::Inst& inst) {
auto& state = ctx.state;
return EmitValueOrZeroIfCondition(state, ctx.Arg(inst, inst.NumArgs() - 1), [&]() {
const auto buffer = PrepareIndirectBuffer(ctx, inst);
const auto buffer = PrepareIndirectBuffer(state, IndirectBufferArgs(ctx, inst), true);
constexpr auto format_count =
static_cast<uint32_t>(Prospero::BufferFormat::k32_32_32_32Float) + 1u;
const auto geometry = EmitIndexSwitch(
@@ -1087,7 +1094,13 @@ uint32_t LoadBuffer(ValueEmitContext& ctx, const IR::Inst& inst, uint32_t compon
state, ctx.Arg(inst, inst.NumArgs() - 1), TypeU32Composite(state, components),
ConstantU32CompositeZero(state, components), [&]() {
if (mem.kind == IR::ResourceKind::IndirectBuffer) {
return LoadIndirectBuffer(ctx, inst, components);
const auto function = state.indirect_buffer_load_functions.at(
{mem.data_bits, components, uint32_t(mem.coherent)});
const auto args = IndirectBufferArgs(ctx, inst);
const auto result = state.builder.AllocateId();
state.builder.AddFunction(spv::OpFunctionCall, TypeU32Composite(state, components),
result, function, args);
return result;
}
return EmitBufferAccess(ctx, inst, mem, components, [&](const auto& selected, const auto& resource) {
const auto info = Format::GetFormatInfo(
@@ -1184,6 +1197,42 @@ uint32_t GetBdaPointer(EmitterState& state, uint32_t address) {
return result;
}
void DefineIndirectBufferLoads(ValueEmitContext& ctx) {
auto& state = ctx.state;
for (const auto* block: state.program.blocks) {
for (const auto& inst: *block) {
const auto components = IR::BufferComponentCount(inst.GetOpcode());
if (components == 0u || inst.GetOpcode() == IR::ValueOpcode::ReadConstBuffer) continue;
const auto& memory = ctx.Memory(inst);
if (memory.kind != IR::ResourceKind::IndirectBuffer || memory.formatted ||
!memory.SupportsIndirectBufferLoad(inst.GetOpcode())) continue;
auto [entry, inserted] = state.indirect_buffer_load_functions.try_emplace(
std::array {memory.data_bits, components, uint32_t(memory.coherent)}, 0u);
if (!inserted) continue;
const auto function = entry->second = state.builder.AllocateId();
const auto type = TypeU32(state);
const auto result_type = TypeU32Composite(state, components);
const auto signature = state.builder.Type(
spv::OpTypeFunction, result_type, type, type, type, type, type, type, type, type);
const auto name = "load_indirect_buffer_" + std::to_string(memory.data_bits) + "x" +
std::to_string(components) + (memory.coherent ? "_coherent" : "");
state.builder.AddName(function, name.c_str());
state.builder.AddFunction(spv::OpFunction, result_type, function,
spv::FunctionControlMaskNone, signature);
std::array<uint32_t, 8> args;
for (auto& arg: args) {
arg = state.builder.AllocateId();
state.builder.AddFunction(spv::OpFunctionParameter, type, arg);
}
EmitLabel(state, state.builder.AllocateId());
const auto result = LoadIndirectBuffer(ctx, PrepareIndirectBuffer(state, args, false),
memory.data_bits, components, memory.coherent);
state.builder.AddFunction(spv::OpReturnValue, result);
state.builder.AddFunction(spv::OpFunctionEnd);
}
}
}
void DefineConvertIndirectBufferFormat(ValueEmitContext& ctx) {
auto& state = ctx.state;
const auto needed = std::ranges::any_of(state.program.blocks, [&](const IR::Block* block) {
@@ -666,6 +666,7 @@ void EmitProgram(EmitterState& state) {
}
}
DefineGetBdaPointer(state);
DefineIndirectBufferLoads(ctx);
DefineConvertIndirectBufferFormat(ctx);
DefineBvhIntersect(state);
if (program.info.uses_swizzle) {
+130 -11
View File
@@ -30642,7 +30642,7 @@ TestCase BufferLoadDwordx4ZeroesOnlyOutOfBoundsTail() {
return test;
}
TestCase BufferLoadsGpuSelectedDescriptors(u32 width) {
TestCase BufferLoadsGpuSelectedDescriptors(u32 width, u32 wave = 32u) {
const bool scalar = width == 1u;
const bool xyz = width == 3u;
using O = ShaderOpcode;
@@ -30675,6 +30675,14 @@ TestCase BufferLoadsGpuSelectedDescriptors(u32 width) {
: xyz ? "BufferLoadDwordx3GpuSelectedDescriptors"
: "BufferLoadsGpuSelectedDescriptors";
test.initial.resize(2048);
test.compute_info.threads_num[0] = wave;
test.compute_info.threads_num[1] = test.compute_info.threads_num[2] = 1u;
test.compute_info.thread_ids_num = 1;
test.compute_info.wave_size = wave;
test.has_compute_info = true;
test.code.push_back(EncodeVop1(0x01, 25, Vgpr(0)));
AppendSMovLiteral(&test.code, 126, 1u);
AppendSMovLiteral(&test.code, 127, 0u);
for (u32 i = 0; i < std::size(cases); ++i) {
const auto &input = cases[i];
const u32 data_offset = 4096 + i * 128;
@@ -30712,19 +30720,72 @@ TestCase BufferLoadsGpuSelectedDescriptors(u32 width) {
test.expected.push_back(expected == 0 ? 0 : selected * 100 + expected);
}
}
constexpr u32 payload = 6144u, output = 8192u, canary = 0xa5a5a5a5u;
const u32 row = std::size(cases);
const std::array<u32, 4> descriptor{u32(GuestBase + payload), (32u << 16u) | 1u,
wave, 0x5204u | (1u << 23u)};
std::copy(descriptor.begin(), descriptor.end(), test.initial.begin() + 130u + row * 30u);
test.initial[64u + row] = row;
for (u32 lane = 0; lane < wave; ++lane)
for (u32 word = 0; word < 8u; ++word)
test.initial[payload / 4u + lane * 8u + word] = 0x90000000u | (lane << 8u) | word;
test.initial.resize((output + wave * 32u) / 4u, canary);
const auto matrix_results = std::move(test.expected);
test.expected = test.initial;
std::copy(matrix_results.begin(), matrix_results.end(), test.expected.begin());
AppendVMovU32(&test.code, 30, (64u + row) * 4u);
AppendBufferLoadDword(&test.code, 0, 30);
test.code.push_back(EncodeVop1(0x02, 20, Vgpr(0)));
test.code.insert(test.code.end(), {EncodeSop2(0x26, 20, 20, 255), 120u,
EncodeSmem0(0x0a, 8, 0), EncodeSmem1(520, 20)});
AppendSMovLiteral(&test.code, 126, UINT32_MAX);
AppendSMovLiteral(&test.code, 127, wave == 64u ? UINT32_MAX : 0u);
test.code.insert(test.code.end(), {EncodeVop2(0x1a, 24, InlineU32(5), 25),
EncodeVop2(0x25, 24, 255, 24), output});
const u32 opcode = scalar ? 0x0cu : xyz ? 0x0fu : 0x0eu;
for (u32 component = width; component < 4u; ++component)
AppendVMovLiteral(&test.code, 4u + component, canary);
AppendVMovU32(&test.code, 21, 0u);
AppendVMovLiteral(&test.code, 22, UINT32_MAX - 2u);
// U32 byte offset + immediate wraps to four before ADD_TID addressing and bounds.
test.code.insert(test.code.end(), {EncodeMubuf0(opcode, 7u, true, true), EncodeMubuf1(4, 2, 21),
EncodeMubuf0(0x1e, 0u, false, true), EncodeMubuf1(4, 12, 24)});
for (u32 component = 0; component < 4u; ++component)
AppendVMovLiteral(&test.code, 4u + component, canary);
AppendVMovLiteral(&test.code, 21, UINT32_MAX);
AppendVMovU32(&test.code, 22, 0u);
AppendSMovLiteral(&test.code, 126, 0x55555555u);
AppendSMovLiteral(&test.code, 127, wave == 64u ? 0xaaaaaaaau : 0u);
test.code.insert(test.code.end(), {EncodeMubuf0(opcode, 4u, true, true, true), EncodeMubuf1(4, 2, 21)});
AppendSMovLiteral(&test.code, 126, 0u);
AppendSMovLiteral(&test.code, 127, 0u);
test.code.insert(test.code.end(), {EncodeMubuf0(opcode, 0u, true, true, true), EncodeMubuf1(4, 2, 21)});
AppendSMovLiteral(&test.code, 126, UINT32_MAX);
AppendSMovLiteral(&test.code, 127, wave == 64u ? UINT32_MAX : 0u);
test.code.insert(test.code.end(), {EncodeMubuf0(0x1e, 16u, false, true), EncodeMubuf1(4, 12, 24)});
for (u32 lane = 0; lane < wave; ++lane) {
const bool active = (lane & 1u) == (lane < 32u ? 0u : 1u);
for (u32 component = 0; component < width; ++component) {
test.expected[output / 4u + lane * 8u + component] =
0x90000000u | (lane << 8u) | (component + 1u);
test.expected[output / 4u + lane * 8u + 4u + component] = !active ? canary :
lane == 0u ? 0u : 0x90000000u | ((lane - 1u) << 8u) | (component + 1u);
}
}
AppendEnd(&test.code);
test.bda_mappings = {{GuestBase, 0}};
test.opcodes = {O::V_MOV_B32, O::S_MOV_B32, O::BUFFER_LOAD_DWORD,
O::V_READFIRSTLANE_B32, O::S_MUL_I32, O::S_BUFFER_LOAD_DWORDX4,
O::BUFFER_STORE_DWORD, O::S_ENDPGM};
if (scalar) {
test.decoded_counts = {{"BUFFER_LOAD_DWORD ", 2u * std::size(cases)}};
test.decoded_counts = {{"BUFFER_LOAD_DWORD ", 2u * std::size(cases) + 4u}};
} else if (xyz) {
test.opcodes.push_back(O::BUFFER_LOAD_DWORDX3);
} else {
test.opcodes.insert(test.opcodes.end(), {O::BUFFER_LOAD_DWORDX4, O::BUFFER_LOAD_DWORDX2});
}
test.required_spirv = {"OpConvertUToPtr", "PhysicalStorageBuffer"};
test.opcodes.insert(test.opcodes.end(), {O::V_LSHLREV_B32, O::V_ADD_NC_U32, O::BUFFER_STORE_DWORDX4});
test.required_spirv = {"OpConvertUToPtr", "PhysicalStorageBuffer", "Volatile"};
return test;
}
@@ -30936,6 +30997,69 @@ TestCase BufferLoadUbyteGpuSelectedDescriptors64() {
return BufferLoadGpuSelectedDescriptors(64u, ShaderOpcode::BUFFER_LOAD_UBYTE);
}
void CheckSharedIndirectBufferLoads(VulkanHarness &vulkan) {
using namespace ShaderRecompiler::IR;
for (const u32 wave : {32u, 64u}) {
for (const auto &test : {BufferLoadsGpuSelectedDescriptors(1u, wave),
BufferLoadsGpuSelectedDescriptors(3u, wave),
BufferLoadsGpuSelectedDescriptors(4u, wave),
BufferLoadGpuSelectedDescriptors(wave, ShaderOpcode::BUFFER_LOAD_UBYTE),
BufferLoadGpuSelectedDescriptors(wave, ShaderOpcode::BUFFER_LOAD_USHORT)}) {
auto compiled = CompileCase(test, vulkan.SubgroupSize());
struct Profile { u32 components = 0, uses = 0, definitions = 0, calls = 0, reads = 0; };
std::map<std::string, Profile> profiles;
for (const auto *block : compiled.program.blocks)
for (const auto &inst : *block) {
const auto components = BufferComponentCount(inst.GetOpcode());
if (components == 0u || inst.GetOpcode() == ValueOpcode::ReadConstBuffer) continue;
const auto &memory = compiled.program.memory_info[inst.Flags<MemoryFlags>().index];
if (memory.kind != ResourceKind::IndirectBuffer || memory.formatted) continue;
const auto name = "load_indirect_buffer_" + std::to_string(memory.data_bits) + "x" +
std::to_string(components) + (memory.coherent ? "_coherent" : "");
auto &profile = profiles[name];
profile.components = components;
++profile.uses;
}
std::map<u32, std::string> functions;
u32 function = 0, bda = 0, inline_reads = 0;
for (size_t offset = 5; offset < compiled.spirv.size();) {
const auto words = std::span(compiled.spirv).subspan(offset, compiled.spirv[offset] >> 16u);
const auto opcode = static_cast<spv::Op>(words[0] & 0xffffu);
if (opcode == spv::OpName) {
const std::string_view name(reinterpret_cast<const char *>(words.data() + 2));
if (name.starts_with("load_indirect_buffer_")) functions.emplace(words[1], name);
if (name == "get_bda_pointer") bda = words[1];
}
if (opcode == spv::OpFunction) {
function = words[2];
if (functions.contains(function)) ++profiles[functions.at(function)].definitions;
}
if (opcode == spv::OpFunctionCall) {
if (functions.contains(words[3])) ++profiles[functions.at(words[3])].calls;
if (words[3] == bda) {
if (functions.contains(function)) ++profiles[functions.at(function)].reads;
else ++inline_reads;
}
}
offset += words.size();
}
Require(test.name, "raw load function scope", functions.size() == profiles.size() && inline_reads == 0u,
"a raw load profile was unused, missing or still emitted inline");
for (const auto &[name, profile] : profiles)
Require(test.name, name.c_str(), profile.uses != 0u && profile.definitions == 1u &&
profile.calls == profile.uses * (wave == 64u && vulkan.SubgroupSize() == 32u ? 2u : 1u) &&
profile.reads == profile.components,
"raw load helpers duplicated definitions, calls or component reads");
auto buffer = vulkan.CreateStorageBuffer(test.name, test.initial, test.initial.size(), true);
vulkan.Dispatch(test, compiled, buffer);
CompareWords(test, "shared raw loads, lane addressing, EXEC and bounds", test.expected,
vulkan.ReadBuffer(test.name, buffer, test.expected.size()));
vulkan.DestroyBuffer(&buffer);
}
}
std::printf("[compute] %-32s ok\n", "SharedIndirectBufferLoads");
}
TestCase BufferStoreDwordx4DropsOnlyOutOfBoundsTail() {
using O = ShaderOpcode;
@@ -37460,7 +37584,7 @@ void CheckGpuSelectedRecordBufferAccess(VulkanHarness &vulkan) {
// Compile against poison selectors. A separate GPU dispatch writes the actual keys.
auto compiled = CompileCase(test, vulkan.SubgroupSize());
if (table_bytes == 0u || components != 0u)
test.forbidden_spirv = {"convert_indirect_buffer_format"};
test.forbidden_spirv = {"convert_indirect_buffer_format", "load_indirect_buffer_"};
else
test.required_spirv = {"get_bda_pointer", "convert_indirect_buffer_format"};
CheckSpirvEmission(test, compiled.spirv);
@@ -46998,15 +47122,9 @@ int main(int argc, char **argv) {
CheckGpuSelectedRecordBufferAccess(vulkan);
CheckGpuSelectedFormattedRecordStore(vulkan);
CheckRuntimeBufferRecords(vulkan);
RunCase(&vulkan, BufferLoadsGpuSelectedDescriptors());
RunCase(&vulkan, BufferLoadDwordGpuSelectedDescriptors());
RunCase(&vulkan, BufferLoadDwordx3GpuSelectedDescriptors());
CheckSharedIndirectBufferLoads(vulkan);
RunCase(&vulkan, BufferLoadFormatXGpuSelectedDescriptors());
RunCase(&vulkan, BufferLoadFormatXGpuSelectedDescriptors64());
RunCase(&vulkan, BufferLoadUshortGpuSelectedDescriptors());
RunCase(&vulkan, BufferLoadUshortGpuSelectedDescriptors64());
RunCase(&vulkan, BufferLoadUbyteGpuSelectedDescriptors());
RunCase(&vulkan, BufferLoadUbyteGpuSelectedDescriptors64());
RunCase(&vulkan, BufferLoadFormatXRejectsPartialRecord());
RunCase(&vulkan, BufferLoadFormatXyRejectsPartialRecord());
RunCase(&vulkan, BufferLoadDwordx4SnapshotsOverlappingAddress());
@@ -47710,6 +47828,7 @@ int main(int argc, char **argv) {
CheckGpuSelectedFormattedRecordStore(vulkan);
CheckFiniteBufferSelection(vulkan);
CheckRuntimeBufferRecords(vulkan);
CheckSharedIndirectBufferLoads(vulkan);
CheckComputeThreadDimensions(vulkan);
CheckCountedRowImageSelector(vulkan);
CheckPointerImageTable(vulkan);