Implement RDNA2 per-invocation scratch memory

This commit is contained in:
Senaxx
2026-08-19 17:32:55 +02:00
committed by nmzik
parent 2db0b3cfa7
commit 56a15c2036
20 changed files with 184 additions and 53 deletions
@@ -158,6 +158,11 @@ bool ValidateNativeProgram(const IR::Program& program, std::string* error) {
return Fail(error, "typed address handle has an invalid dense resource");
}
break;
case IR::ValueOpcode::GetScratchResource:
if (inst.NumArgs() != 0 || program.scratch_dwords == 0) {
return Fail(error, "typed scratch handle has invalid shader metadata");
}
break;
case IR::ValueOpcode::GetImageResource:
if (dense >= program.info.images.size()) {
return Fail(error, "typed image handle has an invalid dense resource");
@@ -198,6 +203,18 @@ bool AnalyzeProgramRequirements(IR::Program& program, std::string* error) {
};
for (const auto* block: program.values->blocks) {
for (const auto& inst: *block) {
if (IR::AddressOpcodeInfoOf(inst.GetOpcode()).access != IR::AddressAccess::None) {
const auto memory_index = inst.Flags<IR::MemoryFlags>().index;
if (memory_index >= program.values->memory_info.size()) {
return Fail(error, "address operation has invalid memory metadata");
}
if (program.values->memory_info[memory_index].kind == IR::ResourceKind::Scratch) {
if (program.scratch_dwords == 0) {
return Fail(error, "scratch operation has no per-thread storage");
}
requirements.function_scratch = true;
}
}
if (IR::BufferAccessOf(inst.GetOpcode()) != IR::BufferAccess::None) {
const auto memory_index = inst.Flags<IR::MemoryFlags>().index;
if (memory_index >= program.values->memory_info.size()) {
@@ -535,6 +535,7 @@ bool EmitValueFlow(ValueEmitContext& ctx, const IR::Inst& inst) {
case IR::ValueOpcode::GetSrtResource:
case IR::ValueOpcode::GetBufferResource:
case IR::ValueOpcode::GetAddressResource:
case IR::ValueOpcode::GetScratchResource:
case IR::ValueOpcode::GetImageResource:
case IR::ValueOpcode::GetSamplerResource:
case IR::ValueOpcode::MakeImageAddress: return true;
@@ -321,6 +321,7 @@ struct EmitterState {
uint32_t vsharp_storage_variable = 0;
uint32_t flattened_srt_variable = 0;
uint32_t lds_variable = 0;
uint32_t scratch_variable = 0;
std::array<uint32_t, 14> sampled_image_variables {};
std::array<uint32_t, 10> storage_image_variables {};
uint32_t sampler_variable = 0;
@@ -356,8 +357,8 @@ uint32_t TypeFunction(EmitterState& state);
uint32_t TypeStorageBufferPointer(EmitterState& state);
uint32_t TypeStorageBufferElementPointer(EmitterState& state);
uint32_t TypePushConstantElementPointer(EmitterState& state);
uint32_t TypeLdsArrayPointer(EmitterState& state, uint32_t storage_class);
uint32_t TypeLdsElementPointer(EmitterState& state, uint32_t storage_class);
uint32_t TypeU32ArrayPointer(EmitterState& state, uint32_t storage_class, uint32_t dwords);
uint32_t TypeU32ElementPointer(EmitterState& state, uint32_t storage_class);
inline void EmitLabel(EmitterState& state, uint32_t label) {
state.current_label = label;
@@ -105,12 +105,12 @@ uint32_t AddressByteAddress(ValueEmitContext& ctx, const IR::Inst& inst, const I
}
const auto immediate_low = ConstantU32(state, static_cast<uint32_t>(immediate));
const auto immediate_high = ConstantU32(state, immediate < 0 ? UINT32_MAX : 0u);
const auto base = state.program.info.addresses[mem.resource].specialized_base;
if (mem.kind == IR::ResourceKind::Scratch) {
low = AddU64Low(state, low, high, immediate_low, immediate_high, high);
const auto valid = Binary(state, OpIEqual, TypeBool(state), high, ConstantU32(state, 0));
return Select(state, TypeU32(state), valid, low, ConstantU32(state, UINT32_MAX));
}
const auto base = state.program.info.addresses[mem.resource].specialized_base;
if (mem.kind == IR::ResourceKind::Flat ||
(mem.kind == IR::ResourceKind::Global &&
state.program.info.addresses[mem.resource].unbased)) {
@@ -75,7 +75,8 @@ static void EnsureLdsStorage(EmitterState& state) {
EXIT("function LDS was not prepared before SPIR-V function emission\n");
}
state.lds_variable = state.builder.DefineGlobalVariable(
TypeLdsArrayPointer(state, StorageClassWorkgroup), StorageClassWorkgroup);
TypeU32ArrayPointer(state, StorageClassWorkgroup, LdsDwordCount(state)),
StorageClassWorkgroup);
state.builder.AddName(state.lds_variable, "lds_dwords");
}
@@ -98,10 +99,16 @@ MemoryResourceAccess PrepareMemoryResourceAccess(EmitterState& state, const IR::
}
access.length = state.gds_length;
return access;
case IR::ResourceKind::Scratch:
if (state.scratch_variable == 0) {
EXIT("scratch storage was not prepared before SPIR-V function emission\n");
}
access.object_pointer = state.scratch_variable;
access.length = ConstantU32(state, state.program.scratch_dwords);
return access;
case IR::ResourceKind::ScalarAddress:
case IR::ResourceKind::Flat:
case IR::ResourceKind::Global:
case IR::ResourceKind::Scratch: {
case IR::ResourceKind::Global: {
if (state.address_memory_variable == 0) {
ExitDescriptorBindingFailure(state, IR::DescriptorBindingKind::AddressMemory,
mem.resource,
@@ -155,11 +162,11 @@ uint32_t EmitMemoryElementInBounds(EmitterState& state, const MemoryResourceAcce
uint32_t EmitMemoryElementPointer(EmitterState& state, const MemoryResourceAccess& access,
uint32_t index) {
const auto pointer = state.builder.AllocateId();
if (access.kind == IR::ResourceKind::Lds) {
state.builder.AddFunction({OpAccessChain,
TypeLdsElementPointer(state, state.stage == ShaderType::Compute
? StorageClassWorkgroup
: StorageClassFunction),
if (access.kind == IR::ResourceKind::Lds || access.kind == IR::ResourceKind::Scratch) {
const auto storage_class = access.kind == IR::ResourceKind::Scratch ? StorageClassFunction
: state.stage == ShaderType::Compute ? StorageClassWorkgroup
: StorageClassFunction;
state.builder.AddFunction({OpAccessChain, TypeU32ElementPointer(state, storage_class),
pointer, access.object_pointer, index});
} else {
state.builder.AddFunction({OpAccessChain, TypeStorageBufferElementPointer(state), pointer,
@@ -94,13 +94,13 @@ uint32_t TypePushConstantElementPointer(EmitterState& state) {
return TypePointer(state, StorageClassPushConstant, TypeU32(state));
}
uint32_t TypeLdsArrayPointer(EmitterState& state, uint32_t storage_class) {
const auto count = ConstantU32(state, std::max(LdsDwordCount(state), 1u));
uint32_t TypeU32ArrayPointer(EmitterState& state, uint32_t storage_class, uint32_t dwords) {
const auto count = ConstantU32(state, std::max(dwords, 1u));
const auto array = state.builder.Type(OpTypeArray, {TypeU32(state), count});
return TypePointer(state, storage_class, array);
}
uint32_t TypeLdsElementPointer(EmitterState& state, uint32_t storage_class) {
uint32_t TypeU32ElementPointer(EmitterState& state, uint32_t storage_class) {
return TypePointer(state, storage_class, TypeU32(state));
}
@@ -601,6 +601,9 @@ void DefineModule(EmitterState& state) {
if (state.requirements.function_lds) {
state.lds_variable = state.builder.AllocateId();
}
if (state.requirements.function_scratch) {
state.scratch_variable = state.builder.AllocateId();
}
state.main_func = state.builder.AllocateId();
state.entry_label = state.builder.AllocateId();
@@ -659,6 +662,9 @@ void DefineModule(EmitterState& state) {
if (state.requirements.function_lds) {
state.builder.AddName(state.lds_variable, "lds_dwords");
}
if (state.requirements.function_scratch) {
state.builder.AddName(state.scratch_variable, "scratch_dwords");
}
AddInputAnnotationsAndNames(state);
AddOutputAnnotationsAndNames(state);
AddDescriptorAnnotationsAndNames(state);
@@ -580,8 +580,15 @@ bool EmitValueProgram(EmitterState& state, const IR::ValueProgram& program, std:
{OpFunction, TypeVoid(state), state.main_func, FunctionControlNone, TypeFunction(state)});
EmitLabel(state, state.entry_label);
if (state.requirements.function_lds) {
state.builder.AddFunction({OpVariable, TypeLdsArrayPointer(state, StorageClassFunction),
state.lds_variable, StorageClassFunction});
state.builder.AddFunction(
{OpVariable, TypeU32ArrayPointer(state, StorageClassFunction, LdsDwordCount(state)),
state.lds_variable, StorageClassFunction});
}
if (state.requirements.function_scratch) {
state.builder.AddFunction(
{OpVariable,
TypeU32ArrayPointer(state, StorageClassFunction, state.program.scratch_dwords),
state.scratch_variable, StorageClassFunction});
}
if (state.pixel_valid_mask_variable != 0) {
state.builder.AddFunction({OpVariable,