Fix wave64 LDS synchronization on subgroup32 hosts

src/graphics/host_gpu/renderer/pipeline/shaders.cpp

@@ -1019,8 +1019,7 @@ void CreatePipelineInternal(GraphicContext& graphics, PipelineCache::ComputePipe
 	comp_shader_stage_info.pSpecializationInfo = nullptr;
 	EXIT_IF(!input_info.stage);
 	const auto wave_size = input_info.stage.program->wave_size;
-	if (graphics.subgroup_size_control_enabled &&
-	    (graphics.required_subgroup_size_stages & vk::ShaderStageFlagBits::eCompute) &&
+	if (graphics.compute_subgroup_size_control_enabled &&
 	    wave_size >= graphics.min_subgroup_size && wave_size <= graphics.max_subgroup_size) {
 		comp_subgroup_size.sType =
 		    vk::StructureType::ePipelineShaderStageRequiredSubgroupSizeCreateInfo;

src/graphics/host_gpu/renderer/renderCompute.cpp

@@ -210,7 +210,9 @@ void RenderExecutor::DispatchDirect(uint64_t submit_id, RenderCommandBuffer& buf
 
 	ShaderComputeInputInfo    input_info {};
 	std::span<const uint32_t> cs_shader;
-	if (!ShaderCompileInfoCS(cs_regs, sh_regs, input_info, cs_shader)) {
+	if (!ShaderCompileInfoCS(cs_regs, sh_regs,
+	                         !m_context.GetGraphics().compute_wave64_supported, input_info,
+	                         cs_shader)) {
 		EXIT("ShaderCompileInfoCS failed for dispatch with CS shader 0x%016" PRIx64 "\n",
 		     cs_regs.cs_regs.data_addr);
 	}

src/graphics/host_gpu/vulkanInstance.h

@@ -25,7 +25,8 @@ struct VulkanInstance {
 	VmaAllocator                       allocator                         = nullptr;
 	bool                               memory_budget_ext_enabled         = false;
 	bool                               rt_extensions_enabled             = false;
-	bool                               subgroup_size_control_enabled     = false;
+	bool                               compute_subgroup_size_control_enabled = false;
+	bool                               compute_wave64_supported          = false;
 	bool                               sample_rate_shading_enabled       = false;
 	uint32_t                           subgroup_size                     = 0;
 	uint32_t                           min_subgroup_size                 = 0;

src/graphics/presentation/window/vulkanWindow.cpp

@@ -493,14 +493,19 @@ static void VulkanInitSubgroupSizeControl(vk::PhysicalDevice physical_device,
 	graphics.min_subgroup_size             = subgroup_size_control.minSubgroupSize;
 	graphics.max_subgroup_size             = subgroup_size_control.maxSubgroupSize;
 	graphics.required_subgroup_size_stages = subgroup_size_control.requiredSubgroupSizeStages;
-	graphics.subgroup_size_control_enabled = features13.subgroupSizeControl == VK_TRUE &&
-	                                         subgroup_size_control.minSubgroupSize <= 64 &&
-	                                         subgroup_size_control.maxSubgroupSize >= 64;
+	graphics.compute_subgroup_size_control_enabled =
+	    features13.subgroupSizeControl == VK_TRUE &&
+	    (graphics.required_subgroup_size_stages & vk::ShaderStageFlagBits::eCompute) &&
+	    subgroup_size_control.minSubgroupSize <= 64 &&
+	    subgroup_size_control.maxSubgroupSize >= 64;
+	graphics.compute_wave64_supported =
+	    graphics.subgroup_size == 64u || graphics.compute_subgroup_size_control_enabled;
 
-	LOGF("Vulkan subgroup: default=%u min=%u max=%u stages=0x%08x size_control=%s\n",
+	LOGF("Vulkan subgroup: default=%u min=%u max=%u stages=0x%08x size_control=%s wave64=%s\n",
 	     graphics.subgroup_size, graphics.min_subgroup_size, graphics.max_subgroup_size,
 	     static_cast<vk::ShaderStageFlags::MaskType>(graphics.required_subgroup_size_stages),
-	     graphics.subgroup_size_control_enabled ? "true" : "false");
+	     graphics.compute_subgroup_size_control_enabled ? "true" : "false",
+	     graphics.compute_wave64_supported ? "true" : "false");
 }
 
 static vk::Device VulkanCreateDevice(vk::PhysicalDevice physical_device, const VulkanExtensions& r,
@@ -645,8 +650,9 @@ static vk::Device VulkanCreateDevice(vk::PhysicalDevice physical_device, const V
 		robustness2.nullDescriptor      = supported_robustness2.nullDescriptor;
 	}
 
-	const bool subgroup_size_control_enabled = graphics.subgroup_size_control_enabled &&
-	                                           supported_features13.subgroupSizeControl == VK_TRUE;
+	const bool subgroup_size_control_enabled =
+	    graphics.compute_subgroup_size_control_enabled &&
+	    supported_features13.subgroupSizeControl == VK_TRUE;
 
 	auto features13 = required_features13;
 #if defined(__APPLE__)

src/graphics/shader/recompiler/ShaderRecompiler.cpp

@@ -15,6 +15,7 @@
 #include "graphics/shader/recompiler/ir/passes/ResourceMaterialization.h"
 #include "graphics/shader/recompiler/ir/passes/ResourceTracking.h"
 #include "graphics/shader/recompiler/ir/passes/ShaderInfoCollection.h"
+#include "graphics/shader/recompiler/ir/passes/SharedMemoryBarrier.h"
 #include "graphics/shader/recompiler/ir/passes/SrtWalker.h"
 #include "graphics/shader/recompiler/ir/passes/SsaRewrite.h"
 
@@ -741,6 +742,14 @@ bool TryRecompile(std::span<const uint32_t> code, const CompileOptions& options,
 		IR::RemoveIdentities(ir.values->blocks);
 		IR::EliminateDeadCode(ir.values->blocks);
 	}
+	if (options.stage == ShaderType::Compute) {
+		const auto lds_barriers =
+		    IR::InsertSharedMemoryBarriers(*ir.values, ir.wave_size, *compute);
+		if (lds_barriers.inserted_barriers != 0) {
+			LOGF("%s wave64 LDS synchronization: barriers=%" PRIu32 "\n", GetDumpLabel(options),
+			     lds_barriers.inserted_barriers);
+		}
+	}
 	std::string srt_error;
 	if (!IR::BuildSrtPlan(ir, &srt_error)) {
 		LOGF("%s SRT planning failed: %s\n", GetDumpLabel(options), srt_error.c_str());

src/graphics/shader/recompiler/ir/passes/SharedMemoryBarrier.cpp

@@ -0,0 +1,153 @@
+#include "graphics/shader/recompiler/ir/passes/SharedMemoryBarrier.h"
+
+#include "graphics/shader/recompiler/ir/ValueProgram.h"
+
+#include <algorithm>
+#include <queue>
+#include <unordered_map>
+#include <unordered_set>
+
+namespace Libs::Graphics::ShaderRecompiler::IR {
+namespace {
+
+enum class LdsPhase { None, Read, Write };
+
+LdsPhase GetLdsPhase(const ValueProgram& program, const Inst& inst) {
+	const auto access = SharedAccessOf(inst.GetOpcode());
+	if (access != SharedAccess::Read && access != SharedAccess::Write) {
+		return LdsPhase::None;
+	}
+	const auto index = inst.Flags<MemoryFlags>().index;
+	if (index >= program.memory_info.size() ||
+	    program.memory_info[index].kind != ResourceKind::Lds) {
+		return LdsPhase::None;
+	}
+	return access == SharedAccess::Read ? LdsPhase::Read : LdsPhase::Write;
+}
+
+bool DependsOnInvocation(Value value) {
+	std::queue<Value>         queue;
+	std::unordered_set<Inst*> visited;
+	queue.push(value);
+	while (!queue.empty()) {
+		value = queue.front().Resolve();
+		queue.pop();
+		auto* inst = value.TryInstruction();
+		if (inst == nullptr || !visited.insert(inst).second) {
+			continue;
+		}
+		if (inst->GetOpcode() == ValueOpcode::LaneId) {
+			return true;
+		}
+		if (inst->GetOpcode() == ValueOpcode::GetBuiltin) {
+			const auto kind = static_cast<StageInputKind>(inst->Arg(0).U32());
+			if (kind == StageInputKind::LocalInvocationId ||
+			    kind == StageInputKind::LocalInvocationIndex ||
+			    kind == StageInputKind::GlobalInvocationId) {
+				return true;
+			}
+		}
+		for (size_t index = 0; index < inst->NumArgs(); index++) {
+			queue.push(inst->Arg(index));
+		}
+	}
+	return false;
+}
+
+uint32_t InsertBarrier(Block& block, Block::iterator position) {
+	block.PrependNewInst(position, ValueOpcode::Barrier);
+	return 1;
+}
+
+uint32_t InsertBlockBarriers(ValueProgram& program, Block& block) {
+	uint32_t inserted = 0;
+	LdsPhase phase    = LdsPhase::None;
+	for (auto inst = block.begin(); inst != block.end(); ++inst) {
+		if (inst->GetOpcode() == ValueOpcode::Barrier) {
+			phase = LdsPhase::None;
+			continue;
+		}
+		const auto next = GetLdsPhase(program, *inst);
+		if (next == LdsPhase::None) {
+			continue;
+		}
+		if (phase != LdsPhase::None && phase != next) {
+			inserted += InsertBarrier(block, inst);
+		}
+		phase = next;
+	}
+	if (phase != LdsPhase::None) {
+		block.AppendNewInst(ValueOpcode::Barrier);
+		inserted++;
+	}
+	return inserted;
+}
+
+uint32_t InsertMergeBarrier(Block& block) {
+	auto position = std::ranges::find_if(
+	    block, [](const Inst& inst) { return inst.GetOpcode() != ValueOpcode::Phi; });
+	if (position != block.end() && position->GetOpcode() == ValueOpcode::Barrier) {
+		return 0;
+	}
+	return InsertBarrier(block, position);
+}
+
+} // namespace
+
+SharedMemoryBarrierStats InsertSharedMemoryBarriers(ValueProgram& program, uint32_t wave_size,
+	                                                 const ShaderComputeInputInfo& compute_info) {
+	SharedMemoryBarrierStats stats;
+	const auto threadgroup_size = compute_info.threads_num[0] * compute_info.threads_num[1] *
+	                              compute_info.threads_num[2];
+	if (wave_size != 64u || !compute_info.needs_lds_barriers ||
+	    compute_info.lds_size_dwords == 0u || threadgroup_size != 64u) {
+		return stats;
+	}
+	if (program.blocks.size() != program.block_info.size()) {
+		return stats;
+	}
+	for (const auto& info: program.block_info) {
+		if (info.terminator.kind == CFG::TerminatorKind::ConditionalBranch &&
+		    info.terminator.merge_block == UINT32_MAX && DependsOnInvocation(info.condition)) {
+			return stats;
+		}
+	}
+
+	std::unordered_map<uint32_t, Block*> blocks;
+	for (size_t index = 0; index < program.blocks.size(); index++) {
+		blocks.emplace(program.block_info[index].id, program.blocks[index]);
+	}
+
+	uint32_t                               divergence_depth = 0;
+	std::unordered_map<uint32_t, uint32_t> divergence_ends;
+	std::unordered_set<uint32_t>           synchronized_merges;
+	for (size_t index = 0; index < program.blocks.size(); index++) {
+		const auto id = program.block_info[index].id;
+		if (const auto end = divergence_ends.find(id); end != divergence_ends.end()) {
+			EXIT_IF(end->second > divergence_depth);
+			divergence_depth -= end->second;
+		}
+
+		if (divergence_depth == 0) {
+			stats.inserted_barriers += InsertBlockBarriers(program, *program.blocks[index]);
+		}
+
+		const auto& info = program.block_info[index];
+		const auto& term = info.terminator;
+		if (term.kind != CFG::TerminatorKind::ConditionalBranch || term.merge_block == UINT32_MAX ||
+		    !DependsOnInvocation(info.condition)) {
+			continue;
+		}
+		if (divergence_depth == 0 && synchronized_merges.insert(term.merge_block).second) {
+			const auto merge = blocks.find(term.merge_block);
+			if (merge != blocks.end()) {
+				stats.inserted_barriers += InsertMergeBarrier(*merge->second);
+			}
+		}
+		divergence_depth++;
+		divergence_ends[term.merge_block]++;
+	}
+	return stats;
+}
+
+} // namespace Libs::Graphics::ShaderRecompiler::IR

src/graphics/shader/recompiler/ir/passes/SharedMemoryBarrier.h

@@ -0,0 +1,15 @@
+#pragma once
+
+#include "graphics/shader/recompiler/ir/ShaderIR.h"
+
+namespace Libs::Graphics::ShaderRecompiler::IR {
+
+struct SharedMemoryBarrierStats {
+	uint32_t inserted_barriers = 0;
+};
+
+[[nodiscard]] SharedMemoryBarrierStats
+InsertSharedMemoryBarriers(ValueProgram& program, uint32_t wave_size,
+	                       const ShaderComputeInputInfo& compute_info);
+
+} // namespace Libs::Graphics::ShaderRecompiler::IR

src/graphics/shader/shader.cpp

@@ -1212,10 +1212,12 @@ bool ShaderCompileInfoPS(const HW::PixelShaderInfo& regs, const HW::ShaderRegist
 }
 
 bool ShaderCompileInfoCS(const HW::ComputeShaderInfo& regs, const HW::ShaderRegisters& sh,
-                         ShaderComputeInputInfo& info, std::span<const uint32_t>& spirv) {
+                         bool needs_lds_barriers, ShaderComputeInputInfo& info,
+                         std::span<const uint32_t>& spirv) {
 	spirv = {};
 
 	ShaderGetStaticInputInfoCS(regs, sh, info);
+	info.needs_lds_barriers = needs_lds_barriers;
 	const auto shader_hash = regs.cs_regs.data_addr;
 	const auto program_id  = ShaderGetIdCS(regs, info, false);
 	const auto key = MakeShaderStageProgramKey(ShaderType::Compute, shader_hash, program_id);
@@ -1344,9 +1346,11 @@ void ShaderDbgDumpInputInfo(const ShaderComputeInputInfo& info) {
 	     "\t thread_ids_num     = %d\n"
 	     "\t wave_size          = %u\n"
 	     "\t lds_size_dwords    = %u\n"
+	     "\t needs_lds_barriers = %s\n"
 	     "\t threads_num        = {%u, %u, %u}\n"
 	     "\t tg_size_en         = %s\n",
 	     info.workgroup_register, info.thread_ids_num, info.wave_size, info.lds_size_dwords,
+	     info.needs_lds_barriers ? "true" : "false",
 	     info.threads_num[0], info.threads_num[1], info.threads_num[2],
 	     info.tg_size_en ? "true" : "false");
 	LOGF("\t threadgroup_id     = {%s, %s, %s}\n", info.group_id[0] ? "true" : "false",
@@ -1742,6 +1746,7 @@ ShaderId ShaderGetIdCS(const HW::ComputeShaderInfo& regs, const ShaderComputeInp
 	ret.ids.push_back(input_info.thread_ids_num);
 	ret.ids.push_back(input_info.lds_size_dwords);
 	ret.ids.push_back(input_info.scratch_size_dwords);
+	ret.ids.push_back(static_cast<uint32_t>(input_info.needs_lds_barriers));
 
 	for (int i = 0; i < 3; i++) {
 		ret.ids.push_back(input_info.threads_num[i]);

src/graphics/shader/shader.h

@@ -107,6 +107,7 @@ struct ShaderComputeInputInfo {
 	uint32_t           scratch_size_dwords        = 0;
 	bool               group_id[3]                = {false, false, false};
 	bool               dispatch_thread_dimensions = false;
+	bool               needs_lds_barriers          = false;
 	uint32_t           wave_size                  = 64;
 	int                thread_ids_num             = 0;
 	int                workgroup_register         = 0;
@@ -268,7 +269,8 @@ bool ShaderCompileInfoPS(const HW::PixelShaderInfo& regs, const HW::ShaderRegist
                          std::span<const Prospero::ColorComponentMapping, 8> target_export_mapping,
                          ShaderPixelInputInfo& input_info, std::span<const uint32_t>& spirv);
 bool ShaderCompileInfoCS(const HW::ComputeShaderInfo& regs, const HW::ShaderRegisters& sh,
-                         ShaderComputeInputInfo& input_info, std::span<const uint32_t>& spirv);
+                         bool needs_lds_barriers, ShaderComputeInputInfo& input_info,
+                         std::span<const uint32_t>& spirv);
 bool ShaderCompileSpirvVS(const HW::VertexShaderInfo& regs, const HW::ShaderRegisters& sh,
                           ShaderVertexInputInfo& input_info, std::vector<uint32_t>& spirv);
 bool ShaderCompileSpirvPS(const HW::PixelShaderInfo& regs, const HW::ShaderRegisters& sh,