gpu: store large compute shared memory in device buffers

2026-10-20 10:56:58 +02:00 · 2026-10-04 22:07:07 +04:00
parent 0ee8800905
commit b0bbaef1eb
17 changed files with 269 additions and 69 deletions
@@ -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);
@@ -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<uint32_t>                 shader_data;
 };

@@ -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);
 }
@@ -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<const void*>(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<uint64_t>(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;
@@ -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);
@@ -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);
@@ -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;