shader: add native 64-bit LDS OR with shared typed storage views

This commit is contained in:
nmzik
2026-09-29 18:41:22 +04:00
parent 35c4dca2ce
commit 6799ecbc5b
14 changed files with 221 additions and 16 deletions
@@ -520,16 +520,20 @@ static vk::Device VulkanCreateDevice(GraphicContext& graphics,
#endif
depth_clip_control.depthClipControl = VK_TRUE;
auto features12 = WindowContext::RequiredVulkan12Features();
features12.pNext = &depth_clip_control;
const bool workgroup_layout_extension =
HasExtension(device_extensions, VK_KHR_WORKGROUP_MEMORY_EXPLICIT_LAYOUT_EXTENSION_NAME);
vk::PhysicalDeviceWorkgroupMemoryExplicitLayoutFeaturesKHR supported_workgroup_layout {};
vk::PhysicalDeviceVulkan12Features supported_features12 {};
supported_features12.pNext = workgroup_layout_extension ? &supported_workgroup_layout : nullptr;
vk::PhysicalDeviceVulkan13Features supported_features13 {};
supported_features13.pNext = &supported_features12;
const auto robustness2_ext_enabled =
HasExtension(device_extensions, VK_EXT_ROBUSTNESS_2_EXTENSION_NAME);
vk::PhysicalDeviceRobustness2FeaturesEXT supported_robustness2 {};
if (robustness2_ext_enabled) {
supported_robustness2.pNext = supported_features13.pNext;
supported_features13.pNext = &supported_robustness2;
}
@@ -557,6 +561,22 @@ static vk::Device VulkanCreateDevice(GraphicContext& graphics,
supported_features2.pNext = &provoking_vertex;
}
physical_device.getFeatures2(&supported_features2);
auto features12 = WindowContext::RequiredVulkan12Features();
features12.shaderSharedInt64Atomics = supported_features12.shaderSharedInt64Atomics;
vk::PhysicalDeviceWorkgroupMemoryExplicitLayoutFeaturesKHR workgroup_layout {};
workgroup_layout.workgroupMemoryExplicitLayout =
supported_workgroup_layout.workgroupMemoryExplicitLayout;
workgroup_layout.pNext = &depth_clip_control;
features12.pNext = workgroup_layout_extension ? static_cast<void*>(&workgroup_layout)
: static_cast<void*>(&depth_clip_control);
if (!features12.shaderSharedInt64Atomics || !workgroup_layout.workgroupMemoryExplicitLayout) {
Log::WriteToConsoleAndLog(fmt::format(
"WARNING: Native 64-bit LDS atomics are unavailable: shaderSharedInt64Atomics={}, "
"workgroupMemoryExplicitLayout={}. Continuing with the selected Vulkan device.\n",
features12.shaderSharedInt64Atomics != VK_FALSE,
workgroup_layout.workgroupMemoryExplicitLayout != VK_FALSE));
}
graphics.mesh_shader_enabled = mesh_extension && supported_mesh.meshShader;
vk::PhysicalDeviceSubgroupSizeControlProperties subgroup_size_control {};
@@ -1049,6 +1069,7 @@ void WindowContext::CreateVulkan() {
for (const auto* extension: {VK_EXT_ROBUSTNESS_2_EXTENSION_NAME,
VK_EXT_PROVOKING_VERTEX_EXTENSION_NAME,
VK_EXT_MESH_SHADER_EXTENSION_NAME,
VK_KHR_WORKGROUP_MEMORY_EXPLICIT_LAYOUT_EXTENSION_NAME,
VK_EXT_DEPTH_RANGE_UNRESTRICTED_EXTENSION_NAME}) {
if (HasExtension(available_extensions, extension)) {
device_extensions.push_back(extension);
@@ -229,6 +229,12 @@ Emitter::SpirvRequirements Emitter::AnalyzeProgramRequirements(const IR::Program
if (kind != IR::ResourceKind::Lds && kind != IR::ResourceKind::Gds) {
Fail(program, "shared operation has invalid resource kind");
}
if (inst.GetOpcode() == IR::ValueOpcode::SharedAtomicOr64) {
if (kind != IR::ResourceKind::Lds || program.stage != ShaderType::Compute) {
Fail(program, "64-bit shared atomics require compute LDS");
}
requirements.shared_int64_atomics = true;
}
if (program.stage != ShaderType::Compute && program.stage != ShaderType::Mesh &&
kind == IR::ResourceKind::Lds) {
requirements.function_lds = true;
@@ -222,6 +222,7 @@ void EmitLoadMemory(ValueEmitContext& ctx, const IR::Inst& inst
void EmitStoreMemory(ValueEmitContext& ctx, const IR::Inst& inst);
uint32_t EmitAtomic32(ValueEmitContext& ctx, const IR::Inst& inst);
uint32_t EmitBufferAtomic64(ValueEmitContext& ctx, const IR::Inst& inst);
void EmitSharedAtomicOr64(ValueEmitContext& ctx, const IR::Inst& inst);
uint32_t EmitBufferFloatAtomic(ValueEmitContext& ctx, const IR::Inst& inst);
void EmitSharedFloatAtomic(ValueEmitContext& ctx, const IR::Inst& inst);
uint32_t EmitSharedIncDec(ValueEmitContext& ctx, const IR::Inst& inst);
@@ -67,6 +67,7 @@ struct SpirvRequirements {
bool function_scratch = false;
bool pixel_valid_mask = false;
bool buffer_int64_atomics = false;
bool shared_int64_atomics = false;
bool coherent_buffers = false;
bool float64 = false;
};
@@ -101,6 +102,7 @@ struct EmitterState {
uint32_t shader_data_storage_variable = 0;
uint32_t flattened_srt_variable = 0;
uint32_t lds_variable = 0;
uint32_t lds_u64_variable = 0;
std::array<uint32_t, 2> scratch_variable {};
std::array<uint32_t, IR::ImageBindingCount> image_variables {};
uint32_t sampler_variable = 0;
@@ -376,6 +378,7 @@ Prospero::BufferFormat StorageBufferFormat(const EmitterState& state, const IR::
void EmitMemoryOffsets(EmitterState& state);
uint32_t LdsDwordCount(const EmitterState& state);
void EnsureLdsStorage(EmitterState& state);
struct MemoryResourceAccess {
IR::ResourceKind kind = IR::ResourceKind::None;
@@ -548,6 +548,7 @@ spv::Op SpirvAtomicOpcode(IR::ValueOpcode opcode) {
case IR::ValueOpcode::SharedAtomicAnd32: return spv::OpAtomicAnd;
case IR::ValueOpcode::BufferAtomicOr32:
case IR::ValueOpcode::BufferAtomicOr64:
case IR::ValueOpcode::SharedAtomicOr64:
case IR::ValueOpcode::SharedAtomicOr32: return spv::OpAtomicOr;
case IR::ValueOpcode::BufferAtomicXor32:
case IR::ValueOpcode::SharedAtomicXor32: return spv::OpAtomicXor;
@@ -1064,6 +1065,32 @@ uint32_t EmitBufferAtomic64(ValueEmitContext& ctx, const IR::Inst& inst) {
});
}
void EmitSharedAtomicOr64(ValueEmitContext& ctx, const IR::Inst& inst) {
auto& state = ctx.state;
const auto& mem = ctx.Memory(inst);
EnsureLdsStorage(state);
EmitIfCondition(state, ctx.Arg(inst, 2), [&]() {
const auto address = Binary(state, spv::OpBitwiseAnd, TypeU32(state),
ByteAddress(ctx, inst, mem), ConstantU32(state, 0xfff8u));
const auto index = Binary(state, spv::OpShiftRightLogical, TypeU32(state), address,
ConstantU32(state, 3u));
const auto in_bounds = Binary(state, spv::OpULessThan, TypeBool(state), index,
ConstantU32(state, LdsDwordCount(state) / 2u));
EmitIfCondition(state, in_bounds, [&]() {
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);
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);
});
});
}
uint32_t EmitBufferFloatAtomic(ValueEmitContext& ctx, const IR::Inst& inst) {
const auto& mem = ctx.Memory(inst);
const bool max_value = inst.GetOpcode() == IR::ValueOpcode::BufferAtomicFMax32;
@@ -63,16 +63,34 @@ uint32_t LdsDwordCount(const EmitterState& state) {
return workgroup != nullptr ? workgroup->lds_size_dwords : 8192u;
}
static void EnsureLdsStorage(EmitterState& state) {
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");
}
state.lds_variable = state.builder.DefineGlobalVariable(
TypeU32ArrayPointer(state, spv::StorageClassWorkgroup, LdsDwordCount(state)),
spv::StorageClassWorkgroup);
const auto define = [&](uint32_t type, uint32_t bytes) {
const auto array = state.builder.DecoratedType(
spv::OpTypeArray, {{spv::OpDecorate, {spv::DecorationArrayStride, bytes}}}, type,
ConstantU32(state, std::max(LdsDwordCount(state) * 4u / bytes, 1u)));
const auto block = state.builder.DecoratedType(
spv::OpTypeStruct, {{spv::OpMemberDecorate, {0, spv::DecorationOffset, 0}},
{spv::OpDecorate, {spv::DecorationBlock}}}, array);
const auto variable = state.builder.DefineGlobalVariable(
TypePointer(state, spv::StorageClassWorkgroup, block), spv::StorageClassWorkgroup);
state.builder.AddAnnotation(spv::OpDecorate, variable, spv::DecorationAliased);
return variable;
};
if (state.requirements.shared_int64_atomics) {
state.lds_variable = define(TypeU32(state), 4u);
state.lds_u64_variable = define(TypeScalarU64(state), 8u);
state.builder.AddName(state.lds_u64_variable, "lds_qwords");
} else {
state.lds_variable = state.builder.DefineGlobalVariable(
TypeU32ArrayPointer(state, spv::StorageClassWorkgroup, LdsDwordCount(state)),
spv::StorageClassWorkgroup);
}
state.builder.AddName(state.lds_variable, "lds_dwords");
}
@@ -167,6 +185,11 @@ uint32_t EmitMemoryElementPointer(EmitterState& state, const MemoryResourceAcces
: 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);
return pointer;
}
state.builder.AddFunction(spv::OpAccessChain, TypeU32ElementPointer(state, storage_class),
pointer, access.object_pointer, index);
return pointer;
@@ -626,10 +626,15 @@ void DefineModule(EmitterState& state) {
state.builder.RequireCapability(spv::CapabilityPhysicalStorageBufferAddresses);
state.builder.RequireExtension("SPV_KHR_physical_storage_buffer");
}
if (state.requirements.buffer_int64_atomics) {
if (state.requirements.buffer_int64_atomics || state.requirements.shared_int64_atomics) {
state.builder.RequireCapability(spv::CapabilityInt64);
state.builder.RequireCapability(spv::CapabilityInt64Atomics);
}
if (state.requirements.shared_int64_atomics) {
state.builder.RequireVersion(0x00010400u);
state.builder.RequireExtension("SPV_KHR_workgroup_memory_explicit_layout");
state.builder.RequireCapability(spv::CapabilityWorkgroupMemoryExplicitLayoutKHR);
}
if (state.clip_distance_variable != 0) {
state.builder.RequireCapability(spv::CapabilityClipDistance);
}
@@ -106,6 +106,7 @@ constexpr MemoryOpcodeInfo DS_OPCODE_LIST[] = {
{0x39u, Opcode::DS_READ_I8, 1, 8, true}, {0x3au, Opcode::DS_READ_U8, 1, 8},
{0x3bu, Opcode::DS_READ_I16, 1, 16, true}, {0x3cu, Opcode::DS_READ_U16, 1, 16},
{0x3du, Opcode::DS_CONSUME, 1, 32}, {0x3eu, Opcode::DS_APPEND, 1, 32},
{0x4au, Opcode::DS_OR_B64, 2, 32},
{0x4du, Opcode::DS_WRITE_B64, 2, 32}, {0x4eu, Opcode::DS_WRITE2_B64, 4, 32},
{0x4fu, Opcode::DS_WRITE2ST64_B64, 4, 32}, {0x76u, Opcode::DS_READ_B64, 2, 32},
{0x77u, Opcode::DS_READ2_B64, 4, 32}, {0x78u, Opcode::DS_READ2ST64_B64, 4, 32},
@@ -185,6 +186,7 @@ bool IsDsAtomicOpcode(Opcode opcode) {
case Opcode::DS_AND_B32:
case Opcode::DS_AND_RTN_B32:
case Opcode::DS_OR_B32:
case Opcode::DS_OR_B64:
case Opcode::DS_OR_RTN_B32:
case Opcode::DS_XOR_B32:
case Opcode::DS_XOR_RTN_B32:
@@ -681,6 +681,7 @@ std::string InstructionToString(const Instruction& inst) {
case Opcode::DS_AND_B32:
case Opcode::DS_AND_RTN_B32:
case Opcode::DS_OR_B32:
case Opcode::DS_OR_B64:
case Opcode::DS_OR_RTN_B32:
case Opcode::DS_XOR_B32:
case Opcode::DS_XOR_RTN_B32:
@@ -528,6 +528,7 @@ enum class Opcode {
DS_AND_B32,
DS_AND_RTN_B32,
DS_OR_B32,
DS_OR_B64,
DS_OR_RTN_B32,
DS_XOR_B32,
DS_XOR_RTN_B32,
@@ -570,7 +570,10 @@ void Translator::DS_ATOMIC(const Decoder::Instruction& inst, IR::ValueOpcode opc
bool returns_value) {
const auto memory = MemoryInfoFromDecoded(inst);
const auto address = ReadU32(MemorySourceAt(inst, 1));
const auto result = ir.Emit(opcode, {address, ReadU32(MemorySourceAt(inst, 0)), ir.GetExec()},
const auto data_src = MemorySourceAt(inst, 0);
const IR::Value value =
memory.data_dwords == 2u ? IR::Value(ReadU64(data_src)) : IR::Value(ReadU32(data_src));
const auto result = ir.Emit(opcode, {address, value, ir.GetExec()},
AddMemoryInfo(memory, inst.pc));
if (returns_value) {
WriteOperand(inst.dst, result);
@@ -1003,6 +1006,8 @@ void Translator::EmitMemory(const Decoder::Instruction& inst) {
return DS_ATOMIC(inst, IR::ValueOpcode::SharedAtomicAnd32, true);
case Decoder::Opcode::DS_OR_B32:
return DS_ATOMIC(inst, IR::ValueOpcode::SharedAtomicOr32, false);
case Decoder::Opcode::DS_OR_B64:
return DS_ATOMIC(inst, IR::ValueOpcode::SharedAtomicOr64, false);
case Decoder::Opcode::DS_OR_RTN_B32:
return DS_ATOMIC(inst, IR::ValueOpcode::SharedAtomicOr32, true);
case Decoder::Opcode::DS_XOR_B32:
@@ -173,6 +173,7 @@ SharedAccess SharedAccessOf(ValueOpcode opcode) {
case ValueOpcode::SharedAtomicUMax32:
case ValueOpcode::SharedAtomicAnd32:
case ValueOpcode::SharedAtomicOr32:
case ValueOpcode::SharedAtomicOr64:
case ValueOpcode::SharedAtomicXor32: return SharedAccess::Atomic;
case ValueOpcode::DataAppend: return SharedAccess::Append;
case ValueOpcode::DataConsume: return SharedAccess::Consume;
@@ -182,6 +183,7 @@ SharedAccess SharedAccessOf(ValueOpcode opcode) {
uint32_t SharedComponentCount(ValueOpcode opcode) {
switch (opcode) {
case ValueOpcode::SharedAtomicOr64:
case ValueOpcode::LoadSharedU32x2:
case ValueOpcode::WriteSharedU32x2: return 2u;
case ValueOpcode::LoadSharedU32x3:
@@ -276,6 +276,7 @@ VALUE_OPCODE(SharedAtomicSMax32, U32, U32, U32, U1)
VALUE_OPCODE(SharedAtomicUMax32, U32, U32, U32, U1)
VALUE_OPCODE(SharedAtomicAnd32, U32, U32, U32, U1)
VALUE_OPCODE(SharedAtomicOr32, U32, U32, U32, U1)
VALUE_OPCODE(SharedAtomicOr64, Void, U32, U64, U1)
VALUE_OPCODE(SharedAtomicXor32, U32, U32, U32, U1)
VALUE_OPCODE(DataAppend, U32, U32, U1)
VALUE_OPCODE(DataConsume, U32, U32, U1)