From b0bbaef1eb6e1dc8ead73b4535fc8021f0fe6837 Mon Sep 17 00:00:00 2001 From: nmzik Date: Sun, 4 Oct 2026 22:07:07 +0400 Subject: [PATCH] gpu: store large compute shared memory in device buffers --- .../renderer/pipeline/descriptors.cpp | 7 +- .../host_gpu/renderer/pipeline/descriptors.h | 1 + .../renderer/pipeline/pipelineCache.cpp | 15 +-- .../host_gpu/renderer/renderCompute.cpp | 38 +++++- .../shader/recompiler/ShaderRecompiler.cpp | 4 +- .../recompiler/backend/spirv/SpirvEmitter.cpp | 12 +- .../backend/spirv/spirvEmitterFlow.cpp | 8 +- .../backend/spirv/spirvEmitterInternal.h | 11 +- .../backend/spirv/spirvEmitterMemory.cpp | 14 +- .../spirv/spirvEmitterMemoryHelpers.cpp | 21 +-- .../backend/spirv/spirvEmitterModule.cpp | 36 ++++-- .../backend/spirv/spirvEmitterProgram.cpp | 11 ++ src/graphics/shader/recompiler/ir/ShaderIR.h | 4 +- .../recompiler/ir/passes/BindingLayout.cpp | 22 ++-- .../recompiler/ir/passes/BindingLayout.h | 11 +- src/graphics/shader/shader.h | 1 + tests/ShaderRecompilerComputeTests.cpp | 122 ++++++++++++++++-- 17 files changed, 269 insertions(+), 69 deletions(-) diff --git a/src/graphics/host_gpu/renderer/pipeline/descriptors.cpp b/src/graphics/host_gpu/renderer/pipeline/descriptors.cpp index 318c7465..142e3638 100644 --- a/src/graphics/host_gpu/renderer/pipeline/descriptors.cpp +++ b/src/graphics/host_gpu/renderer/pipeline/descriptors.cpp @@ -68,7 +68,8 @@ vk::DescriptorType NativeDescriptorType(BindingKind kind) { case BindingKind::BdaPagetable: case BindingKind::FaultBuffer: case BindingKind::FlattenedSrt: - case BindingKind::ShaderData: return vk::DescriptorType::eStorageBuffer; + case BindingKind::ShaderData: + case BindingKind::SharedMemory: return vk::DescriptorType::eStorageBuffer; case BindingKind::Count: EXIT("invalid native descriptor binding kind"); } EXIT("invalid native descriptor binding kind"); @@ -779,6 +780,7 @@ void RenderExecutor::PrepareBindings(const ShaderStageRuntime& runtime, prepared.gds = {nullptr, 0, VK_WHOLE_SIZE}; prepared.flattened_srt = {}; prepared.shader_data_buffer = {}; + prepared.shared_memory = {}; prepared.images.resize(program.info.images.size()); prepared.samplers.clear(); prepared.shader_data.clear(); @@ -1120,12 +1122,15 @@ void RenderExecutor::CommitBindings(CommandBuffer& buffer, } case BindingKind::FlattenedSrt: case BindingKind::ShaderData: + case BindingKind::SharedMemory: case BindingKind::Gds: { const vk::DescriptorBufferInfo* view = &descriptors.gds; if (binding.kind == BindingKind::FlattenedSrt) { view = &descriptors.flattened_srt; } else if (binding.kind == BindingKind::ShaderData) { view = &descriptors.shader_data_buffer; + } else if (binding.kind == BindingKind::SharedMemory) { + view = &descriptors.shared_memory; } EXIT_IF(view->buffer == nullptr); m_descriptor_buffers.push_back(*view); diff --git a/src/graphics/host_gpu/renderer/pipeline/descriptors.h b/src/graphics/host_gpu/renderer/pipeline/descriptors.h index 3bfebcb6..766b24dd 100644 --- a/src/graphics/host_gpu/renderer/pipeline/descriptors.h +++ b/src/graphics/host_gpu/renderer/pipeline/descriptors.h @@ -44,6 +44,7 @@ struct PreparedBindings { vk::DescriptorBufferInfo gds {nullptr, 0, VK_WHOLE_SIZE}; vk::DescriptorBufferInfo flattened_srt; vk::DescriptorBufferInfo shader_data_buffer; + vk::DescriptorBufferInfo shared_memory; std::vector shader_data; }; diff --git a/src/graphics/host_gpu/renderer/pipeline/pipelineCache.cpp b/src/graphics/host_gpu/renderer/pipeline/pipelineCache.cpp index 2eca0762..e0479cce 100644 --- a/src/graphics/host_gpu/renderer/pipeline/pipelineCache.cpp +++ b/src/graphics/host_gpu/renderer/pipeline/pipelineCache.cpp @@ -680,19 +680,8 @@ ShaderProgram PipelineCache::GetComputeProgram(const HW::ComputeShaderInfo& regs ShaderComputeInputInfo& input_info) { input_info.host_subgroup_size = m_graphics.SupportsComputeWave64() ? 64u : 32u; const auto params = PrepareProgram(regs, sh, input_info); - // Use one effective size for the cache key, LDS declaration, and access bounds. - const auto max_lds_dwords = - m_graphics.GetPhysicalDeviceProperties().limits.maxComputeSharedMemorySize / 4u; - if (input_info.lds_size_dwords > max_lds_dwords) { - static std::atomic_bool warned = false; - if (!warned.exchange(true, std::memory_order_relaxed)) { - PipelineCacheLog("GPU warning: game compute shader requests {} bytes of LDS, but " - "the Vulkan device limit is {} bytes. Clamping LDS; rendering may " - "be incorrect.", - input_info.lds_size_dwords * 4u, max_lds_dwords * 4u); - } - } - input_info.lds_size_dwords = std::min(input_info.lds_size_dwords, max_lds_dwords); + input_info.lds_storage = input_info.lds_size_dwords * 4u > + m_graphics.GetPhysicalDeviceProperties().limits.maxComputeSharedMemorySize; uint32_t push_data_cursor = 0; return m_program_cache->Get(params, input_info, push_data_cursor); } diff --git a/src/graphics/host_gpu/renderer/renderCompute.cpp b/src/graphics/host_gpu/renderer/renderCompute.cpp index ef8ab03a..daec7af7 100644 --- a/src/graphics/host_gpu/renderer/renderCompute.cpp +++ b/src/graphics/host_gpu/renderer/renderCompute.cpp @@ -20,6 +20,7 @@ #include "graphics/host_gpu/vulkanCommon.h" #include "graphics/shader/recompiler/ir/ShaderIR.h" #include "graphics/shader/recompiler/ir/passes/ResourceMaterialization.h" +#include "graphics/shader/recompiler/ir/passes/BindingLayout.h" #include "graphics/shader/shader.h" #include "kernel/eventQueue.h" #include "kernel/pthread.h" @@ -195,6 +196,35 @@ bool RenderExecutor::TryConsumeComputeImageClear(const ShaderComputeInputInfo& i return true; } +static void BindSharedMemory(RenderContext& context, ShaderComputeInputInfo& input, + PreparedBindings& bindings, uint64_t indirect_args = 0) { + if (ShaderRecompiler::IR::FindBinding(input.stage.program->bindings, + ShaderRecompiler::IR::DescriptorBindingKind::SharedMemory) == nullptr) { + return; + } + auto& cache = context.GetBufferCache(); + if (indirect_args != 0) { + cache.ReadMemory(indirect_args, sizeof(vk::DispatchIndirectCommand)); + std::memcpy(input.workgroup_counts, reinterpret_cast(indirect_args), + sizeof(input.workgroup_counts)); + } + // LDS has no contents to preserve between dispatches. The existing shader hazard + // barriers also order other users of this GPU-only utility buffer. + auto& storage = cache.GetUtilityBuffer(MemoryUsage::DeviceLocal); + const auto limit = std::min(storage.Size(), + context.GetGraphics().GetPhysicalDeviceProperties().limits.maxStorageBufferRange); + uint64_t size = sizeof(uint32_t); + if (std::ranges::find(input.workgroup_counts, 0u) == std::end(input.workgroup_counts)) { + size = uint64_t {input.lds_size_dwords} * sizeof(uint32_t); + EXIT_IF(size == 0 || size > limit); + for (const auto count: input.workgroup_counts) { + EXIT_IF(size > limit / count); + size *= count; + } + } + bindings.shared_memory = {storage.Handle(), 0, size}; +} + void RenderExecutor::DispatchDirect(uint64_t submit_id, CommandBuffer& buffer, uint32_t thread_group_x, uint32_t thread_group_y, uint32_t thread_group_z, uint32_t mode) { @@ -317,13 +347,15 @@ void RenderExecutor::DispatchDirect(uint64_t submit_id, CommandBuffer& buffer, m_context.PrepareBda(); } RebindImages(bindings); + BindSharedMemory(m_context, input_info, bindings); RebindBuffers(bindings); auto vk_buffer = buffer.Handle(); PreparedBindings* descriptor_stage = &bindings; CommitBindings(buffer, vk::PipelineBindPoint::eCompute, pipeline, std::span {&descriptor_stage, 1u}); - bool has_storage_writes = HasShaderBufferWrites(input_info.stage); + bool has_storage_writes = bindings.shared_memory.buffer != nullptr || + HasShaderBufferWrites(input_info.stage); has_storage_writes = std::any_of(program.info.images.begin(), program.info.images.end(), [](const auto& image) { @@ -430,6 +462,7 @@ void RenderExecutor::DispatchIndirect(uint64_t submit_id, CommandBuffer& buffer, if (program.info.uses_dma) { m_context.PrepareBda(); } + BindSharedMemory(m_context, input_info, bindings, args_addr); RebindImages(bindings); // Acquiring arguments can merge cache buffers; finalize shader bindings afterward. const auto [args_buffer, args_offset] = m_context.GetBufferCache().ObtainBuffer( @@ -440,7 +473,8 @@ void RenderExecutor::DispatchIndirect(uint64_t submit_id, CommandBuffer& buffer, CommitBindings(buffer, vk::PipelineBindPoint::eCompute, pipeline, std::span {&descriptor_stage, 1u}); const auto vk_buffer = buffer.Handle(); - const bool has_storage_writes = HasShaderBufferWrites(input_info.stage) || + const bool has_storage_writes = bindings.shared_memory.buffer != nullptr || + HasShaderBufferWrites(input_info.stage) || std::any_of(program.info.images.begin(), program.info.images.end(), [](const auto& image) { return image.written && image.resource_class == ShaderRecompiler::IR::ImageResourceClass::Storage; diff --git a/src/graphics/shader/recompiler/ShaderRecompiler.cpp b/src/graphics/shader/recompiler/ShaderRecompiler.cpp index 9b74c9a9..cb07edfd 100644 --- a/src/graphics/shader/recompiler/ShaderRecompiler.cpp +++ b/src/graphics/shader/recompiler/ShaderRecompiler.cpp @@ -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); diff --git a/src/graphics/shader/recompiler/backend/spirv/SpirvEmitter.cpp b/src/graphics/shader/recompiler/backend/spirv/SpirvEmitter.cpp index 7bb8f4c1..45d71531 100644 --- a/src/graphics/shader/recompiler/backend/spirv/SpirvEmitter.cpp +++ b/src/graphics/shader/recompiler/backend/spirv/SpirvEmitter.cpp @@ -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(Kind::Count); std::array, 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(Kind::Buffers)]; - const bool uses_gds = IR::CollectMemoryResources(program, buffers); + const auto shared = IR::CollectMemoryResources(program, buffers); present[static_cast(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 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); diff --git a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterFlow.cpp b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterFlow.cpp index 2a1d1ee1..9f8b42d6 100644 --- a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterFlow.cpp +++ b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterFlow.cpp @@ -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)); } diff --git a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterInternal.h b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterInternal.h index 91c155f2..0724a4a8 100644 --- a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterInternal.h +++ b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterInternal.h @@ -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 scratch_variable {}; diff --git a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterMemory.cpp b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterMemory.cpp index e8afe445..02f7dc72 100644 --- a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterMemory.cpp +++ b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterMemory.cpp @@ -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); }); }); } diff --git a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterMemoryHelpers.cpp b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterMemoryHelpers.cpp index de07dd0a..bb449a93 100644 --- a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterMemoryHelpers.cpp +++ b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterMemoryHelpers.cpp @@ -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; } diff --git a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterModule.cpp b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterModule.cpp index 0b9ce260..ea5e82db 100644 --- a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterModule.cpp +++ b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterModule.cpp @@ -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); diff --git a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterProgram.cpp b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterProgram.cpp index ac78acd7..03ba346c 100644 --- a/src/graphics/shader/recompiler/backend/spirv/spirvEmitterProgram.cpp +++ b/src/graphics/shader/recompiler/backend/spirv/spirvEmitterProgram.cpp @@ -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); diff --git a/src/graphics/shader/recompiler/ir/ShaderIR.h b/src/graphics/shader/recompiler/ir/ShaderIR.h index 625b877a..694cd94b 100644 --- a/src/graphics/shader/recompiler/ir/ShaderIR.h +++ b/src/graphics/shader/recompiler/ir/ShaderIR.h @@ -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(DescriptorBindingKind::Samplers) == 49u); -static_assert(static_cast(DescriptorBindingKind::Count) == 55u); +static_assert(static_cast(DescriptorBindingKind::Count) == 56u); struct PushData { static constexpr uint32_t DwordCount = 32; diff --git a/src/graphics/shader/recompiler/ir/passes/BindingLayout.cpp b/src/graphics/shader/recompiler/ir/passes/BindingLayout.cpp index 38dc2391..a9067dbd 100644 --- a/src/graphics/shader/recompiler/ir/passes/BindingLayout.cpp +++ b/src/graphics/shader/recompiler/ir/passes/BindingLayout.cpp @@ -56,14 +56,15 @@ void AddBinding(BindingLayout& layout, DescriptorBindingKind kind, } // namespace -bool CollectMemoryResources(const Program& program, std::vector& buffers) { +SharedMemoryResources CollectMemoryResources(const Program& program, std::vector& buffers) { std::array 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().index; @@ -74,11 +75,13 @@ bool CollectMemoryResources(const Program& program, std::vector& 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& 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 buffers; - const bool uses_gds = CollectMemoryResources(program, buffers); + const auto shared = CollectMemoryResources(program, buffers); CollectShaderData(program, next); next.memory_offset_count = static_cast(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); diff --git a/src/graphics/shader/recompiler/ir/passes/BindingLayout.h b/src/graphics/shader/recompiler/ir/passes/BindingLayout.h index c7ca51b9..566e2f04 100644 --- a/src/graphics/shader/recompiler/ir/passes/BindingLayout.h +++ b/src/graphics/shader/recompiler/ir/passes/BindingLayout.h @@ -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& buffers); +struct SharedMemoryResources { + bool lds = false; + bool gds = false; +}; + +SharedMemoryResources CollectMemoryResources(const Program& program, std::vector& buffers); bool UsesFlattenedSrt(const Program& program); const DescriptorBinding* FindBinding(const BindingLayout& layout, DescriptorBindingKind kind); diff --git a/src/graphics/shader/shader.h b/src/graphics/shader/shader.h index 3c54a28e..e86a5229 100644 --- a/src/graphics/shader/shader.h +++ b/src/graphics/shader/shader.h @@ -151,6 +151,7 @@ struct ShaderComputeInputInfo: ShaderWorkgroupInputInfo { uint32_t workgroup_counts[3] = {0, 0, 0}; bool group_id[3] = {false, false, false}; bool dispatch_thread_dimensions = false; + bool lds_storage = false; int thread_ids_num = 0; int workgroup_register = 0; bool tg_size_en = false; diff --git a/tests/ShaderRecompilerComputeTests.cpp b/tests/ShaderRecompilerComputeTests.cpp index 9a4c152f..b7f01de4 100644 --- a/tests/ShaderRecompilerComputeTests.cpp +++ b/tests/ShaderRecompilerComputeTests.cpp @@ -10016,6 +10016,79 @@ public: "a retained descriptor load used stale dimensions or the wrong record stride"); } } + + // Exchange values across two waves in the last 512 bytes of each + // workgroup's 64-KiB LDS, on a host restricted to 48-KiB native shared memory. + constexpr auto lds_output = base + 0xb0000; + constexpr auto lds_args = base + 0xc0000; + constexpr u32 lds_words = 8 * 128; + std::vector lds_shader{ + EncodeSop2(0x1e, 8, 5, InlineU32(1)), + EncodeSop2(0x1e, 9, 6, InlineU32(2)), + EncodeSop2(0x02, 8, 8, 9), EncodeSop2(0x02, 8, 8, 4)}; + AppendVop3(&lds_shader, 0x346, 2, 8, InlineU32(7), Vgpr(0)); + lds_shader.push_back(EncodeVop2(0x1a, 5, InlineU32(2), 2)); + lds_shader.push_back(EncodeVop2(0x25, 2, InlineU32(1), 2)); + lds_shader.push_back(EncodeVop2(0x1a, 1, InlineU32(2), 0)); + AppendVMovLiteral(&lds_shader, 7, 64 * 1024 - 128 * 4); + lds_shader.insert(lds_shader.end(), { + EncodeVop2(0x25, 1, Vgpr(7), 1), + EncodeDs0(0x0d), EncodeDs1(0, 2, 1), EncodeSopp(0x0a, 0), + EncodeVop2(0x1d, 3, InlineU32(64), 0), + EncodeVop2(0x1a, 3, InlineU32(2), 3), EncodeVop2(0x25, 3, Vgpr(7), 3), + EncodeDs0(0x36), EncodeDs1(4, 0, 3), + EncodeMubuf0(0x1c), EncodeMubuf1(4, 0, 5)}); + AppendEnd(&lds_shader); + ShaderMapUserData(reinterpret_cast(lds_shader.data()), + {.type = Prospero::ShaderBinaryType::kCs, + .code_size_bytes = static_cast(lds_shader.size() * sizeof(u32))}); + auto &host_lds_limit = context.GetGraphics().physical_device_properties.limits.maxComputeSharedMemorySize; + const auto saved_lds_limit = host_lds_limit; + host_lds_limit = std::min(host_lds_limit, 48u * 1024u); + for (u32 dispatch = 0; dispatch < 4; dispatch++) { + cache.FillBuffer(lds_output, lds_words * sizeof(u32), sentinel, false); + const bool indirect = dispatch == 1 || dispatch == 2; + const bool empty = dispatch == 2; + if (indirect) { + const auto groups = empty ? std::array{4096u, 0u, 1u} : std::array{2u, 2u, 2u}; + shaders.SetCsShader({.data_addr = reinterpret_cast(writer.data()), + .num_thread_x = 1, .num_thread_y = 1, .num_thread_z = 1, + .wave_size = 64, .user_sgpr = 15}); + for (u32 axis = 0; axis < 3; axis++) { + set_buffer(axis * 4, lds_args + axis * sizeof(u32), sizeof(u32)); + shaders.SetCsUserSgpr(12 + axis, groups[axis], HW::UserSgprType::Unknown); + } + processor.DispatchDirect(1, 1, 1, 0x41u); + Require(name, "GPU-authored LDS dispatch arguments", + cache.HasGpuDirtyBytes(lds_args, 3 * sizeof(u32)), + "indirect LDS dimensions were already visible on the CPU"); + } + shaders.SetCsShader({.data_addr = reinterpret_cast(lds_shader.data()), + .num_thread_x = 128, .num_thread_y = 1, .num_thread_z = 1, + .wave_size = 64, .user_sgpr = 4, .tgid_x_en = true, + .tgid_y_en = true, .tgid_z_en = true, .lds_size = 128}); + set_buffer(0, lds_output, lds_words * sizeof(u32)); + if (indirect) { + const std::array packet{static_cast(lds_args), + static_cast(lds_args >> 32u), 0x41u}; + Require(name, "indirect storage LDS dispatch", + CpOpDispatchIndirect(processor, 0xc0021600u, packet.data(), 0, 0) == 3, + "the GPU-authored LDS dispatch was not consumed"); + } else { + processor.DispatchDirect(2, 2, 2, 0x41u); + } + cache.ReadMemory(lds_output, lds_words * sizeof(u32)); + std::array actual{}; + Require(name, "storage LDS readback", + LibKernel::Memory::TryReadBacking(lds_output, actual.data(), sizeof(actual)), + "LDS result buffer could not be read back"); + for (u32 word = 0; word < actual.size(); word++) { + const auto expected = empty ? sentinel : (word ^ 64u) + 1u; + Require(name, "storage LDS upper range, isolation and barrier", actual[word] == expected, + "64-KiB LDS lost upper bytes, workgroup isolation, or cross-wave writes"); + } + } + host_lds_limit = saved_lds_limit; context.UnmapMemory(base, allocation_size); scheduler.Finish(); }); @@ -14858,12 +14931,28 @@ public: std::vector sampler_infos; Buffer flattened_buffer; Buffer user_data_buffer; + Buffer shared_buffer; vk::DescriptorBufferInfo flattened_info{}; vk::DescriptorBufferInfo user_data_info{}; + vk::DescriptorBufferInfo shared_info{}; vk::DescriptorBufferInfo gds_info{}; vk::DescriptorBufferInfo bda_pagetable_info{}; vk::DescriptorBufferInfo fault_buffer_info{}; + if (Binding(Kind::SharedMemory) != nullptr) { + const auto bytes = uint64_t{test.compute_info.lds_size_dwords} * sizeof(u32) * + test.dispatch_x * test.dispatch_y * test.dispatch_z; + shared_buffer = CreateHostBuffer(test.name, bytes, vk::BufferUsageFlagBits::eStorageBuffer, {}); + shared_info = {shared_buffer.buffer, 0, shared_buffer.size}; + vk::WriteDescriptorSet write{}; + write.dstSet = descriptor_set; + write.dstBinding = Native(Kind::SharedMemory); + write.descriptorCount = 1; + write.descriptorType = vk::DescriptorType::eStorageBuffer; + write.pBufferInfo = &shared_info; + writes.push_back(write); + } + const bool uses_bda = Binding(Kind::BdaPagetable) != nullptr; Require(test.name, "dispatch", uses_bda == (Binding(Kind::FaultBuffer) != nullptr), @@ -15176,6 +15265,7 @@ public: if (user_data_buffer.buffer != nullptr) { DestroyBuffer(&user_data_buffer); } + DestroyBuffer(&shared_buffer); m_device.destroyDescriptorPool(descriptor_pool, nullptr); m_device.destroyPipeline(pipeline, nullptr); m_device.destroyPipelineLayout(pipeline_layout, nullptr); @@ -34252,24 +34342,27 @@ void CheckComputeLdsLimit(VulkanHarness &vulkan) { regs.cs_regs.num_thread_z = 1; regs.cs_regs.wave_size = 32; PipelineCache cache(graphics); - ShaderProgram at_limit; + ShaderProgram previous; + uint16_t previous_units = 0; for (const auto units : {static_cast(limit_units - 1u), limit_units, - uint16_t{112}, uint16_t{128}}) { + uint16_t{112}, uint16_t{128}, uint16_t{128}}) { regs.cs_regs.lds_size = units; ShaderComputeInputInfo input{}; const auto program = cache.GetComputeProgram(regs, {}, input); - const auto expected = units < limit_units ? units * 128u : limit / 4u; - Require(name, "device allocation", input.lds_size_dwords == expected, - "compute LDS allocation did not respect the device limit"); + Require(name, "guest allocation and device backing", + input.lds_size_dwords == units * 128u && + input.lds_storage == (units > limit_units), + "compute LDS was truncated or chose backing that exceeds the device limit"); Require(name, "pipeline creation", cache.GetComputePipeline(input, program).pipeline != nullptr, - "clamped LDS shader did not create a Vulkan compute pipeline"); - if (units == limit_units) { - at_limit = program; - } else if (units > limit_units) { - Require(name, "effective-size cache reuse", program.id == at_limit.id, - "equivalent clamped LDS allocations compiled separate shaders"); + "full-size LDS shader did not create a Vulkan compute pipeline"); + if (previous.id != 0) { + Require(name, "guest-size cache identity", + (program.id == previous.id) == (units == previous_units), + "different LDS sizes shared a shader, or repeating a size recompiled it"); } + previous = program; + previous_units = units; } std::printf("[gpu] %-32s ok\n", name); } @@ -39129,6 +39222,13 @@ int main(int argc, char **argv) { RunCase(&vulkan, DsWideLdsPartialBounds()); RunCase(&vulkan, DsAtomic64Bounds(false)); RunCase(&vulkan, DsAtomic64Bounds(true)); + for (auto test : {DsWideLdsPartialBounds(), DsAtomic64Bounds(false), DsAtomic64Bounds(true)}) { + test.compute_info.lds_storage = true; + std::erase(test.required_spirv, "WorkgroupMemoryExplicitLayoutKHR"); + test.required_spirv.push_back("StorageBuffer"); + RunCase(&vulkan, test); + } + vulkan.CheckNativeIndirectDispatch(); return 0; } if (argc == 2 && std::strcmp(argv[1], "--new-opcodes-only") == 0) {