gpu: store large compute shared memory in device buffers

This commit is contained in:
nmzik
2026-10-04 22:07:07 +04:00
parent 0ee8800905
commit b0bbaef1eb
17 changed files with 269 additions and 69 deletions
@@ -670,7 +670,9 @@ CompileResult CompileProgram(TranslateResult translated, const CompileOptions& o
IR::EliminateDeadCode(ir.blocks);
IR::CollectShaderInfo(ir, options.input_info);
IR::AllocateBindings(ir, push_data_start_dword);
IR::AllocateBindings(ir, push_data_start_dword,
ir.stage == ShaderType::Compute && options.input_info.compute != nullptr &&
options.input_info.compute->lds_storage);
std::string ir_dump;
if (options.dump_ir) {
ir_dump = MakeIrDump(translated.cfg_dump, ir);
@@ -18,7 +18,7 @@ namespace {
std::abort();
}
void ValidateNativeProgram(const IR::Program& program) {
void ValidateNativeProgram(const IR::Program& program, bool lds_storage) {
using Kind = IR::DescriptorBindingKind;
constexpr auto KindCount = static_cast<size_t>(Kind::Count);
std::array<std::vector<uint32_t>, KindCount> expected;
@@ -53,11 +53,14 @@ void ValidateNativeProgram(const IR::Program& program) {
Expect(Kind::Samplers, Dense(program.info.samplers.size()));
}
auto& buffers = expected[static_cast<size_t>(Kind::Buffers)];
const bool uses_gds = IR::CollectMemoryResources(program, buffers);
const auto shared = IR::CollectMemoryResources(program, buffers);
present[static_cast<size_t>(Kind::Buffers)] = !buffers.empty();
if (uses_gds) {
if (shared.gds) {
Expect(Kind::Gds);
}
if (shared.lds && lds_storage) {
Expect(Kind::SharedMemory);
}
if (program.info.uses_dma) {
Expect(Kind::BdaPagetable);
Expect(Kind::FaultBuffer);
@@ -344,7 +347,8 @@ std::vector<uint32_t> EmitProgram(const IR::Program& program,
!program.shader_info_complete || !program.binding_layout_complete) {
Fail(program, "SPIR-V emitter requires a fully planned native shader program");
}
ValidateNativeProgram(program);
ValidateNativeProgram(program, program.stage == ShaderType::Compute &&
input_info.compute != nullptr && input_info.compute->lds_storage);
IR::ValidateProgram(program, true);
EmitterState state(program, input_info);
const auto* workgroup = ShaderWorkgroupInput(program.stage, input_info);
@@ -607,9 +607,11 @@ void EmitBarrier(EmitterState& state) {
return;
}
const auto memory_scope = tessellation ? spv::ScopeInvocation : spv::ScopeWorkgroup;
const auto semantics = tessellation ? spv::MemorySemanticsMaskNone
: spv::MemorySemanticsAcquireReleaseMask |
spv::MemorySemanticsWorkgroupMemoryMask;
const auto memory = state.lds_storage_class == spv::StorageClassStorageBuffer
? spv::MemorySemanticsUniformMemoryMask
: spv::MemorySemanticsWorkgroupMemoryMask;
const auto semantics = tessellation ? spv::MemorySemanticsMaskNone
: spv::MemorySemanticsAcquireReleaseMask | memory;
state.builder.AddFunction(spv::OpControlBarrier, ConstantU32(state, spv::ScopeWorkgroup),
ConstantU32(state, memory_scope), ConstantU32(state, semantics));
}
@@ -83,7 +83,14 @@ struct EmitterState {
EmitterState(const IR::Program& program_, ShaderStageInputInfo input_info_)
: builder(program_.stage == ShaderType::Mesh ? 0x00010400u : 0x00010300u),
program(program_), input_info(input_info_),
requirements(AnalyzeProgramRequirements(program_)) {}
requirements(AnalyzeProgramRequirements(program_)) {
if (ShaderWorkgroupInput(program.stage, input_info) != nullptr) {
lds_storage_class = spv::StorageClassWorkgroup;
}
if (IR::FindBinding(program.bindings, IR::DescriptorBindingKind::SharedMemory) != nullptr) {
lds_storage_class = spv::StorageClassStorageBuffer;
}
}
Builder builder;
const IR::Program& program;
@@ -123,6 +130,8 @@ struct EmitterState {
uint32_t push_constant_variable = 0;
uint32_t shader_data_storage_variable = 0;
uint32_t flattened_srt_variable = 0;
spv::StorageClass lds_storage_class = spv::StorageClassFunction;
uint32_t lds_base_dwords = 0;
uint32_t lds_variable = 0;
uint32_t lds_u64_variable = 0;
std::array<uint32_t, 2> scratch_variable {};
@@ -1195,16 +1195,22 @@ void EmitSharedAtomic64(ValueEmitContext& ctx, const IR::Inst& inst) {
const auto in_bounds = Binary(state, spv::OpULessThan, TypeBool(state), index,
ConstantU32(state, LdsDwordCount(state) / 2u));
EmitIfCondition(state, in_bounds, [&]() {
auto native_index = index;
auto semantics = spv::MemorySemanticsWorkgroupMemoryMask;
if (state.lds_storage_class == spv::StorageClassStorageBuffer) {
native_index = EmitAddU32(state, index, EmitBinaryU32(
state, spv::OpShiftRightLogical, state.lds_base_dwords, ConstantU32(state, 1)));
semantics = spv::MemorySemanticsUniformMemoryMask;
}
const auto pointer = state.builder.AllocateId();
state.builder.AddFunction(spv::OpAccessChain,
TypePointer(state, spv::StorageClassWorkgroup, TypeScalarU64(state)),
pointer, state.lds_u64_variable, ConstantU32(state, 0), index);
TypePointer(state, state.lds_storage_class, TypeScalarU64(state)),
pointer, state.lds_u64_variable, ConstantU32(state, 0), native_index);
const auto value = Unary(state, spv::OpBitcast, TypeScalarU64(state), ctx.Arg(inst, 1));
state.builder.AddFunction(
SpirvAtomicOpcode(inst.GetOpcode()), TypeScalarU64(state), state.builder.AllocateId(),
pointer, ConstantU32(state, spv::ScopeWorkgroup),
ConstantU32(state, spv::MemorySemanticsAcquireReleaseMask |
spv::MemorySemanticsWorkgroupMemoryMask), value);
ConstantU32(state, spv::MemorySemanticsAcquireReleaseMask | semantics), value);
});
});
}
@@ -121,8 +121,8 @@ void EnsureLdsStorage(EmitterState& state) {
if (state.lds_variable != 0) {
return;
}
if (ShaderWorkgroupInput(state.program.stage, state.input_info) == nullptr) {
EXIT("function LDS was not prepared before SPIR-V function emission\n");
if (state.lds_storage_class != spv::StorageClassWorkgroup) {
EXIT("LDS backing was not prepared before SPIR-V function emission\n");
}
const auto define = [&](uint32_t type, uint32_t bytes) {
const auto array = state.builder.DecoratedType(
@@ -227,12 +227,14 @@ uint32_t EmitMemoryElementInBounds(EmitterState& state, const MemoryResourceAcce
uint32_t EmitMemoryElementPointer(EmitterState& state, const MemoryResourceAccess& access,
uint32_t index) {
if (access.kind == IR::ResourceKind::Lds || access.kind == IR::ResourceKind::Scratch) {
const auto storage_class = access.kind == IR::ResourceKind::Scratch
? spv::StorageClassFunction : state.lds_storage_class;
if (storage_class == spv::StorageClassStorageBuffer) {
return EmitStorageBufferElementPointer(
state, access, EmitAddU32(state, state.lds_base_dwords, index),
TypeStorageBufferElementPointer(state, 32));
}
const auto pointer = state.builder.AllocateId();
const auto storage_class =
access.kind == IR::ResourceKind::Scratch ? spv::StorageClassFunction
: ShaderWorkgroupInput(state.program.stage, state.input_info) != nullptr
? spv::StorageClassWorkgroup
: spv::StorageClassFunction;
if (access.kind == IR::ResourceKind::Lds && state.requirements.shared_int64_atomics) {
state.builder.AddFunction(spv::OpAccessChain, TypeU32ElementPointer(state, storage_class),
pointer, access.object_pointer, ConstantU32(state, 0), index);
@@ -424,7 +426,10 @@ void EmitAtomicMemoryBarrier(EmitterState& state, IR::ResourceKind kind) {
const auto scope = kind == IR::ResourceKind::Lds ? spv::ScopeWorkgroup : spv::ScopeDevice;
const auto memory = [&] {
switch (kind) {
case IR::ResourceKind::Lds: return spv::MemorySemanticsWorkgroupMemoryMask;
case IR::ResourceKind::Lds:
return state.lds_storage_class == spv::StorageClassStorageBuffer
? spv::MemorySemanticsUniformMemoryMask
: spv::MemorySemanticsWorkgroupMemoryMask;
case IR::ResourceKind::Image: return spv::MemorySemanticsImageMemoryMask;
default: return spv::MemorySemanticsUniformMemoryMask;
}
@@ -275,6 +275,17 @@ void DefineDescriptors(EmitterState& state) {
case IR::DescriptorBindingKind::Gds:
state.gds_variable = Define(StorageBufferType(state), "gds");
break;
case IR::DescriptorBindingKind::SharedMemory:
state.lds_variable = Define(StorageBufferType(state), "lds_dwords");
state.builder.AddAnnotation(spv::OpDecorate, state.lds_variable, spv::DecorationCoherent);
if (state.requirements.shared_int64_atomics) {
state.lds_u64_variable = Define(StorageBufferType(state, 64), "lds_qwords");
state.builder.AddAnnotation(spv::OpDecorate, state.lds_u64_variable, spv::DecorationCoherent);
for (const auto variable: {state.lds_variable, state.lds_u64_variable}) {
state.builder.AddAnnotation(spv::OpDecorate, variable, spv::DecorationAliased);
}
}
break;
default: {
EXIT_IF(IR::ImageBindingResourceClass(binding.kind) ==
IR::ImageResourceClass::None);
@@ -408,6 +419,7 @@ uint32_t BuiltInForInput(IR::StageInputKind kind) {
case IR::StageInputKind::BaryCoordSmooth: return spv::BuiltInBaryCoordKHR;
case IR::StageInputKind::BaryCoordNoPerspective: return spv::BuiltInBaryCoordNoPerspKHR;
case IR::StageInputKind::WorkgroupId: return spv::BuiltInWorkgroupId;
case IR::StageInputKind::NumWorkgroups: return spv::BuiltInNumWorkgroups;
case IR::StageInputKind::LocalInvocationId: return spv::BuiltInLocalInvocationId;
case IR::StageInputKind::LocalInvocationIndex: return spv::BuiltInLocalInvocationIndex;
case IR::StageInputKind::GlobalInvocationId: return spv::BuiltInGlobalInvocationId;
@@ -420,15 +432,19 @@ void DefineInputs(EmitterState& state) {
for (const auto& input: state.program.info.inputs) {
state.inputs.push_back({input});
}
const auto add_builtin = [&](IR::StageInputKind kind, uint32_t components,
const char* name) {
if (std::ranges::none_of(state.inputs, [kind](const InputBinding& input) {
return input.kind == kind;
})) {
state.inputs.push_back({{kind, 0, components, name}});
}
};
if (IR::FindBinding(state.program.bindings, IR::DescriptorBindingKind::SharedMemory) != nullptr) {
add_builtin(IR::StageInputKind::WorkgroupId, 3, "gl_WorkGroupID");
add_builtin(IR::StageInputKind::NumWorkgroups, 3, "gl_NumWorkGroups");
}
if (state.lane_count == 2) {
const auto add_builtin = [&](IR::StageInputKind kind, uint32_t components,
const char* name) {
if (std::ranges::none_of(state.inputs, [kind](const InputBinding& input) {
return input.kind == kind;
})) {
state.inputs.push_back({{kind, 0, components, name}});
}
};
add_builtin(IR::StageInputKind::LocalInvocationIndex, 1, "gl_LocalInvocationIndex");
if (std::ranges::any_of(state.inputs, [](const InputBinding& input) {
return input.kind == IR::StageInputKind::GlobalInvocationId;
@@ -459,6 +475,7 @@ void DefineInputs(EmitterState& state) {
case IR::StageInputKind::Layer:
case IR::StageInputKind::SampleId: type = TypeI32(state); break;
case IR::StageInputKind::WorkgroupId:
case IR::StageInputKind::NumWorkgroups:
case IR::StageInputKind::LocalInvocationId:
case IR::StageInputKind::GlobalInvocationId: type = TypeU32Vector(state, 3); break;
case IR::StageInputKind::FragCoord: type = TypeF32Vector(state, 4); break;
@@ -677,7 +694,8 @@ void DefineModule(EmitterState& state) {
state.builder.RequireCapability(spv::CapabilityInt64);
state.builder.RequireCapability(spv::CapabilityInt64Atomics);
}
if (state.requirements.shared_int64_atomics) {
if (state.requirements.shared_int64_atomics &&
state.lds_storage_class == spv::StorageClassWorkgroup) {
state.builder.RequireVersion(0x00010400u);
state.builder.RequireExtension("SPV_KHR_workgroup_memory_explicit_layout");
state.builder.RequireCapability(spv::CapabilityWorkgroupMemoryExplicitLayoutKHR);
@@ -752,6 +752,17 @@ void EmitProgram(EmitterState& state) {
state.builder.AddFunction(spv::OpStore, state.pixel_valid_mask_variable,
ConstantU32(state, 1));
}
if (state.lds_storage_class == spv::StorageClassStorageBuffer && state.lds_variable != 0) {
const auto group_x = EmitInputComponentU32(state, IR::StageInputKind::WorkgroupId, 0);
const auto group_y = EmitInputComponentU32(state, IR::StageInputKind::WorkgroupId, 1);
const auto group_z = EmitInputComponentU32(state, IR::StageInputKind::WorkgroupId, 2);
const auto count_x = EmitInputComponentU32(state, IR::StageInputKind::NumWorkgroups, 0);
const auto count_y = EmitInputComponentU32(state, IR::StageInputKind::NumWorkgroups, 1);
const auto row = EmitAddU32(state, group_y, EmitBinaryU32(state, spv::OpIMul, group_z, count_y));
const auto index = EmitAddU32(state, group_x, EmitBinaryU32(state, spv::OpIMul, row, count_x));
state.lds_base_dwords = EmitBinaryU32(state, spv::OpIMul, index,
ConstantU32(state, LdsDwordCount(state)));
}
EmitMemoryOffsets(state);
if (program.blocks.empty()) {
EmitReturn(ctx);
+3 -1
View File
@@ -199,6 +199,7 @@ enum class StageInputKind {
BaryCoordSmoothCentroid,
BaryCoordNoPerspective,
WorkgroupId,
NumWorkgroups,
LocalInvocationId,
LocalInvocationIndex,
GlobalInvocationId,
@@ -301,11 +302,12 @@ enum class DescriptorBindingKind : uint32_t {
FaultBuffer,
FlattenedSrt,
ShaderData,
SharedMemory,
Count,
};
static_assert(static_cast<uint32_t>(DescriptorBindingKind::Samplers) == 49u);
static_assert(static_cast<uint32_t>(DescriptorBindingKind::Count) == 55u);
static_assert(static_cast<uint32_t>(DescriptorBindingKind::Count) == 56u);
struct PushData {
static constexpr uint32_t DwordCount = 32;
@@ -56,14 +56,15 @@ void AddBinding(BindingLayout& layout, DescriptorBindingKind kind,
} // namespace
bool CollectMemoryResources(const Program& program, std::vector<uint32_t>& buffers) {
SharedMemoryResources CollectMemoryResources(const Program& program, std::vector<uint32_t>& buffers) {
std::array<bool, ShaderInfo::MaxBuffers> live_buffers {};
bool uses_gds = false;
SharedMemoryResources shared;
for (const auto* block: program.blocks) {
for (const auto& inst: *block) {
const auto op = inst.GetOpcode();
if (BufferAccessOf(op) == BufferAccess::None &&
SharedAccessOf(op) == SharedAccess::None) {
SharedAccessOf(op) == SharedAccess::None &&
AddressOpcodeInfoOf(op).access == AddressAccess::None) {
continue;
}
const auto index = inst.Flags<MemoryFlags>().index;
@@ -74,11 +75,13 @@ bool CollectMemoryResources(const Program& program, std::vector<uint32_t>& buffe
if (memory.planning_only) {
continue;
}
shared.lds |= memory.kind == ResourceKind::FlatLocal;
if (SharedAccessOf(op) != SharedAccess::None) {
if (memory.kind != ResourceKind::Lds && memory.kind != ResourceKind::Gds) {
BindingFail("typed shader contains invalid shared-memory metadata");
}
uses_gds |= memory.kind == ResourceKind::Gds;
shared.gds |= memory.kind == ResourceKind::Gds;
shared.lds |= memory.kind == ResourceKind::Lds;
} else if (memory.kind == ResourceKind::Buffer || memory.kind == ResourceKind::ScalarBuffer) {
EXIT_IF(memory.resource >= program.info.buffers.size());
live_buffers.at(memory.resource) = true;
@@ -93,7 +96,7 @@ bool CollectMemoryResources(const Program& program, std::vector<uint32_t>& buffe
buffers.push_back(i);
}
}
return uses_gds;
return shared;
}
bool UsesFlattenedSrt(const Program& program) {
@@ -108,7 +111,7 @@ bool UsesFlattenedSrt(const Program& program) {
std::ranges::any_of(program.info.images, uses_mapping);
}
void AllocateBindings(Program& program, uint32_t push_data_start_dword) {
void AllocateBindings(Program& program, uint32_t push_data_start_dword, bool lds_storage) {
if (!program.shader_info_complete || program.binding_layout_complete) {
EXIT("shader binding layout failed: %s", !program.shader_info_complete
? "shader info is not ready"
@@ -116,7 +119,7 @@ void AllocateBindings(Program& program, uint32_t push_data_start_dword) {
}
BindingLayout next;
std::vector<uint32_t> buffers;
const bool uses_gds = CollectMemoryResources(program, buffers);
const auto shared = CollectMemoryResources(program, buffers);
CollectShaderData(program, next);
next.memory_offset_count = static_cast<uint32_t>(buffers.size());
next.push_data_start_dword =
@@ -160,9 +163,12 @@ void AllocateBindings(Program& program, uint32_t push_data_start_dword) {
}
AddBinding(next, DescriptorBindingKind::Samplers, std::move(resources));
}
if (uses_gds) {
if (shared.gds) {
AddBinding(next, DescriptorBindingKind::Gds);
}
if (shared.lds && lds_storage) {
AddBinding(next, DescriptorBindingKind::SharedMemory);
}
if (program.info.uses_dma) {
AddBinding(next, DescriptorBindingKind::BdaPagetable);
AddBinding(next, DescriptorBindingKind::FaultBuffer);
@@ -5,10 +5,15 @@
namespace Libs::Graphics::ShaderRecompiler::IR {
void AllocateBindings(Program& program, uint32_t push_data_start_dword = 0);
void AllocateBindings(Program& program, uint32_t push_data_start_dword = 0,
bool lds_storage = false);
// Collect live direct buffer indices and report whether the shader accesses GDS.
bool CollectMemoryResources(const Program& program, std::vector<uint32_t>& buffers);
struct SharedMemoryResources {
bool lds = false;
bool gds = false;
};
SharedMemoryResources CollectMemoryResources(const Program& program, std::vector<uint32_t>& buffers);
bool UsesFlattenedSrt(const Program& program);
const DescriptorBinding* FindBinding(const BindingLayout& layout, DescriptorBindingKind kind);