shader: fix readlane/writelane for inactive host lanes

src/graphics/shader/recompiler/ShaderRecompiler.cpp

@@ -2,15 +2,16 @@
 
 #include "common/assert.h"
 #include "common/logging/log.h"
+#include "graphics/shader/recompiler/cfg/ShaderCFG.h"
+#include "graphics/shader/recompiler/decompiler/ShaderDecoder.h"
+#include "graphics/shader/recompiler/emitter/SpirvEmitter.h"
 #include "graphics/shader/recompiler/ir/BindingLayout.h"
+#include "graphics/shader/recompiler/ir/ReadLaneElimination.h"
 #include "graphics/shader/recompiler/ir/ResourceMaterialization.h"
 #include "graphics/shader/recompiler/ir/ResourceTracking.h"
 #include "graphics/shader/recompiler/ir/ScalarProvenance.h"
-#include "graphics/shader/recompiler/cfg/ShaderCFG.h"
-#include "graphics/shader/recompiler/decompiler/ShaderDecoder.h"
 #include "graphics/shader/recompiler/ir/ShaderIR.h"
 #include "graphics/shader/recompiler/ir/ShaderInfoCollection.h"
-#include "graphics/shader/recompiler/emitter/SpirvEmitter.h"
 #include "graphics/shader/recompiler/ir/SrtPatcher.h"
 #include "graphics/shader/recompiler/ir/SrtWalker.h"
 
@@ -838,6 +839,11 @@ bool TryRecompile(std::span<const uint32_t> code, const CompileOptions& options,
 	if (!IR::AllocateBindings(ir, layout_options, error)) {
 		return false;
 	}
+	const auto read_lane_stats = IR::EliminateReadLane(ir);
+	if (read_lane_stats.rewritten_reads != 0) {
+		LOGF("%s read-lane elimination: reads=%" PRIu32 " shadow_writes=%" PRIu32 "\n",
+		     GetDumpLabel(options), read_lane_stats.rewritten_reads, read_lane_stats.shadow_writes);
+	}
 	std::string ir_dump;
 	if (options.dump_ir) {
 		ir_dump = MakeIrDump(cfg, ir);

src/graphics/shader/recompiler/emitter/spirvEmitterFlowOps.cpp

@@ -102,7 +102,7 @@ uint32_t EmitWqmLaneU32(EmitterState& state, uint32_t src) {
 		state.builder.AddFunction(
 		    {OpINotEqual, state.bool_type, non_zero, masked, ConstantU32(state, 0)});
 		state.builder.AddFunction({OpSelect, state.uint_type, expanded, non_zero,
-		                            ConstantU32(state, mask), ConstantU32(state, 0)});
+		                           ConstantU32(state, mask), ConstantU32(state, 0)});
 		state.builder.AddFunction({OpBitwiseOr, state.uint_type, combined, ret, expanded});
 		ret = combined;
 	}
@@ -122,8 +122,8 @@ void EmitWqmB64(EmitterState& state, const IR::Instruction& inst) {
 	}
 	const auto ballot = state.builder.AllocateId();
 	state.builder.AddFunction({OpGroupNonUniformBallot, state.vec4_uint_type, ballot,
-	                            ConstantU32(state, ScopeSubgroup),
-	                            EmitLaneMaskOperandActiveBool(state, inst.src[0])});
+	                           ConstantU32(state, ScopeSubgroup),
+	                           EmitLaneMaskOperandActiveBool(state, inst.src[0])});
 	const auto low  = state.builder.AllocateId();
 	const auto high = state.builder.AllocateId();
 	state.builder.AddFunction({OpCompositeExtract, state.uint_type, low, ballot, 0});
@@ -150,8 +150,8 @@ void EmitWqmB64(EmitterState& state, const IR::Instruction& inst) {
 		EmitPerInvocationMask(state, inst.dst, active);
 	} else {
 		const auto result = state.builder.AllocateId();
-		state.builder.AddFunction({OpSelect, state.uint_type, result, active,
-		                            ConstantU32(state, 1), ConstantU32(state, 0)});
+		state.builder.AddFunction({OpSelect, state.uint_type, result, active, ConstantU32(state, 1),
+		                           ConstantU32(state, 0)});
 		EmitStoreU32(state, inst.dst, result);
 		EmitStoreU32(state, OffsetRegisterOperand(inst.dst, 1), ConstantU32(state, 0));
 	}
@@ -205,8 +205,7 @@ void EmitSaveexecB32(EmitterState& state, const IR::Instruction& inst) {
 
 	const auto cond = state.builder.AllocateId();
 	const auto scc  = state.builder.AllocateId();
-	state.builder.AddFunction(
-	    {OpINotEqual, state.bool_type, cond, new_low, ConstantU32(state, 0)});
+	state.builder.AddFunction({OpINotEqual, state.bool_type, cond, new_low, ConstantU32(state, 0)});
 	state.builder.AddFunction(
 	    {OpSelect, state.uint_type, scc, cond, ConstantU32(state, 1), ConstantU32(state, 0)});
 	EmitStoreU32(state, SccOperand(), scc);
@@ -269,18 +268,19 @@ void EmitReadFirstLaneU32(EmitterState& state, const IR::Instruction& inst) {
 	const auto first_lane  = state.builder.AllocateId();
 	const auto first_value = state.builder.AllocateId();
 	state.builder.AddFunction({OpGroupNonUniformBallot, state.vec4_uint_type, ballot,
-	                            ConstantU32(state, ScopeSubgroup), active});
+	                           ConstantU32(state, ScopeSubgroup), active});
 	state.builder.AddFunction({OpGroupNonUniformBallotFindLSB, state.uint_type, first_lane,
-	                            ConstantU32(state, ScopeSubgroup), ballot});
+	                           ConstantU32(state, ScopeSubgroup), ballot});
 	state.builder.AddFunction({OpGroupNonUniformShuffle, state.uint_type, first_value,
-	                            ConstantU32(state, ScopeSubgroup), src, first_lane});
+	                           ConstantU32(state, ScopeSubgroup), src, first_lane});
 	EmitStoreU32(state, inst.dst, first_value);
 }
 
 uint32_t EmitLaneIndex(EmitterState& state, const IR::Operand& operand) {
 	const auto lane = state.builder.AllocateId();
+	const auto mask = state.wave_size == 32u ? 31u : 63u;
 	state.builder.AddFunction({OpBitwiseAnd, state.uint_type, lane, EmitValueLoad(state, operand),
-	                            ConstantU32(state, 63)});
+	                           ConstantU32(state, mask)});
 	return lane;
 }
 
@@ -289,7 +289,7 @@ void EmitReadLaneU32(EmitterState& state, const IR::Instruction& inst) {
 	const auto lane  = EmitLaneIndex(state, inst.src[1]);
 	const auto value = state.builder.AllocateId();
 	state.builder.AddFunction({OpGroupNonUniformShuffle, state.uint_type, value,
-	                            ConstantU32(state, ScopeSubgroup), src, lane});
+	                           ConstantU32(state, ScopeSubgroup), src, lane});
 	EmitStoreU32(state, inst.dst, value);
 }
 
@@ -336,10 +336,8 @@ void EmitPermlaneB32(EmitterState& state, const IR::Instruction& inst, bool x16)
 		state.builder.AddFunction(
 		    {OpBitwiseXor, state.uint_type, row_value, row, ConstantU32(state, 16)});
 	}
-	state.builder.AddFunction(
-	    {OpBitwiseAnd, state.uint_type, lane, subid, ConstantU32(state, 15)});
-	state.builder.AddFunction(
-	    {OpBitwiseAnd, state.uint_type, lane8, lane, ConstantU32(state, 7)});
+	state.builder.AddFunction({OpBitwiseAnd, state.uint_type, lane, subid, ConstantU32(state, 15)});
+	state.builder.AddFunction({OpBitwiseAnd, state.uint_type, lane8, lane, ConstantU32(state, 7)});
 	state.builder.AddFunction(
 	    {OpShiftLeftLogical, state.uint_type, shift, lane8, ConstantU32(state, 2)});
 	state.builder.AddFunction(
@@ -350,7 +348,7 @@ void EmitPermlaneB32(EmitterState& state, const IR::Instruction& inst, bool x16)
 	    {OpBitwiseAnd, state.uint_type, index1, index0, ConstantU32(state, 15)});
 	state.builder.AddFunction({OpBitwiseOr, state.uint_type, target, row_value, index1});
 	state.builder.AddFunction({OpGroupNonUniformShuffle, state.uint_type, shuffled,
-	                            ConstantU32(state, ScopeSubgroup), value, target});
+	                           ConstantU32(state, ScopeSubgroup), value, target});
 	uint32_t ret = shuffled;
 	if (!inst.dst.op_sel) {
 		const auto source_active = EmitLaneIndexActiveBool(state, target);
@@ -375,7 +373,7 @@ void EmitBarrier(EmitterState& state, const IR::Instruction& inst) {
 	(void)inst;
 	const auto semantics = MemorySemanticsAcquireRelease | MemorySemanticsWorkgroupMemory;
 	state.builder.AddFunction({OpControlBarrier, ConstantU32(state, ScopeWorkgroup),
-	                            ConstantU32(state, ScopeWorkgroup), ConstantU32(state, semantics)});
+	                           ConstantU32(state, ScopeWorkgroup), ConstantU32(state, semantics)});
 }
 
 } // namespace Libs::Graphics::ShaderRecompiler::Spirv::Emitter

src/graphics/shader/recompiler/ir/ReadLaneElimination.cpp

@@ -0,0 +1,323 @@
+#include "graphics/shader/recompiler/ir/ReadLaneElimination.h"
+
+#include "graphics/shader/recompiler/ir/SrtWalker.h"
+
+#include <algorithm>
+#include <iterator>
+#include <map>
+#include <set>
+#include <utility>
+
+namespace Libs::Graphics::ShaderRecompiler::IR {
+
+namespace {
+
+constexpr uint32_t FirstTemporaryScalarRegister = 128;
+
+struct LaneKey {
+	uint32_t reg  = 0;
+	uint32_t lane = 0;
+
+	auto operator<=>(const LaneKey&) const = default;
+};
+
+using LaneSet = std::set<LaneKey>;
+
+bool PairDwordOpcode(Opcode op) {
+	switch (op) {
+		case Opcode::MoveU64:
+		case Opcode::WqmB64:
+		case Opcode::SaveexecB64:
+		case Opcode::BitwiseAndU64:
+		case Opcode::BitwiseAndNotU64:
+		case Opcode::BitwiseOrU64:
+		case Opcode::BitwiseOrNotU64:
+		case Opcode::BitwiseXorU64:
+		case Opcode::BitwiseNandU64:
+		case Opcode::BitwiseNorU64:
+		case Opcode::BitwiseXnorU64:
+		case Opcode::BitwiseNotU64:
+		case Opcode::BitFieldMaskU64:
+		case Opcode::BitFieldExtractU64:
+		case Opcode::BitReplicateB64B32:
+		case Opcode::ShiftLeftLogicalU64:
+		case Opcode::ShiftRightLogicalU64:
+		case Opcode::SelectU64: return true;
+		default: return false;
+	}
+}
+
+bool ResolveLane(const Program& program, const Instruction& inst, uint32_t source_index,
+                 uint32_t& lane) {
+	if (source_index >= inst.src_count || (program.wave_size != 32 && program.wave_size != 64)) {
+		return false;
+	}
+	const auto& selector = inst.src[source_index];
+	if (selector.kind == OperandKind::ImmediateU32) {
+		lane = selector.imm % program.wave_size;
+		return true;
+	}
+	uint32_t folded = 0;
+	if (!FoldScalarConstant(program.provenance, inst.scalar_sources[source_index], folded)) {
+		return false;
+	}
+	lane = folded % program.wave_size;
+	return true;
+}
+
+bool UniformWriteSource(const Instruction& inst) {
+	if (inst.src_count == 0) {
+		return false;
+	}
+	const auto& source = inst.src[0];
+	if (source.kind == OperandKind::ImmediateU32 || source.kind == OperandKind::PcRelativeU32) {
+		return true;
+	}
+	return source.kind == OperandKind::Register &&
+	       (source.reg.file == RegisterFile::Scalar || source.reg.file == RegisterFile::Scc ||
+	        source.reg.file == RegisterFile::M0);
+}
+
+bool WriteLaneKey(const Program& program, const Instruction& inst, LaneKey& key) {
+	if (inst.op != Opcode::WriteLaneU32 || inst.dst.kind != OperandKind::Register ||
+	    inst.dst.reg.file != RegisterFile::Vector || !UniformWriteSource(inst)) {
+		return false;
+	}
+	uint32_t lane = 0;
+	if (!ResolveLane(program, inst, 1, lane)) {
+		return false;
+	}
+	key = {inst.dst.reg.index, lane};
+	return true;
+}
+
+bool ReadLaneKey(const Program& program, const Instruction& inst, LaneKey& key) {
+	if (inst.op != Opcode::ReadLaneU32 || inst.src_count < 2 ||
+	    inst.src[0].kind != OperandKind::Register || inst.src[0].reg.file != RegisterFile::Vector) {
+		return false;
+	}
+	uint32_t lane = 0;
+	if (!ResolveLane(program, inst, 1, lane)) {
+		return false;
+	}
+	key = {inst.src[0].reg.index, lane};
+	return true;
+}
+
+void InvalidateRegister(LaneSet& valid, uint32_t reg) {
+	const auto first = valid.lower_bound({reg, 0});
+	const auto last  = valid.lower_bound({reg + 1u, 0});
+	valid.erase(first, last);
+}
+
+void ApplyInstruction(const Program& program, const Instruction& inst, LaneSet& valid) {
+	if (inst.op == Opcode::WriteLaneU32 && inst.dst.kind == OperandKind::Register &&
+	    inst.dst.reg.file == RegisterFile::Vector) {
+		LaneKey key;
+		if (WriteLaneKey(program, inst, key)) {
+			valid.insert(key);
+			return;
+		}
+		uint32_t lane = 0;
+		if (ResolveLane(program, inst, 1, lane)) {
+			valid.erase({inst.dst.reg.index, lane});
+		} else {
+			InvalidateRegister(valid, inst.dst.reg.index);
+		}
+		return;
+	}
+
+	if (inst.op == Opcode::MoveRelDestU32 && inst.dst.kind == OperandKind::Register &&
+	    inst.dst.reg.file == RegisterFile::Vector) {
+		valid.clear();
+		return;
+	}
+
+	if (inst.dst.kind == OperandKind::Register && inst.dst.reg.file == RegisterFile::Vector) {
+		uint32_t dwords = std::max(inst.memory.data_dwords, 1u);
+		if (PairDwordOpcode(inst.op) || inst.op == Opcode::UMadU64U32) {
+			dwords = std::max(dwords, 2u);
+		}
+		for (uint32_t i = 0; i < dwords && inst.dst.reg.index <= UINT32_MAX - i; i++) {
+			InvalidateRegister(valid, inst.dst.reg.index + i);
+		}
+	}
+	if (inst.dst2.kind == OperandKind::Register && inst.dst2.reg.file == RegisterFile::Vector) {
+		InvalidateRegister(valid, inst.dst2.reg.index);
+	}
+}
+
+LaneSet TransferBlock(const Program& program, const BasicBlock& block, LaneSet state) {
+	for (const auto& inst: block.instructions) {
+		ApplyInstruction(program, inst, state);
+	}
+	return state;
+}
+
+LaneSet Intersect(const LaneSet& left, const LaneSet& right) {
+	LaneSet result;
+	std::set_intersection(left.begin(), left.end(), right.begin(), right.end(),
+	                      std::inserter(result, result.end()));
+	return result;
+}
+
+uint32_t NextTemporaryScalarRegister(const Program& program) {
+	uint32_t   next     = FirstTemporaryScalarRegister;
+	const auto consider = [&next](const Operand& operand) {
+		if (operand.kind == OperandKind::Register && operand.reg.file == RegisterFile::Scalar &&
+		    operand.reg.index >= next && operand.reg.index != UINT32_MAX) {
+			next = operand.reg.index + 1u;
+		}
+	};
+	for (const auto& block: program.blocks) {
+		for (const auto& inst: block.instructions) {
+			consider(inst.dst);
+			consider(inst.dst2);
+			for (uint32_t i = 0; i < inst.src_count; i++) {
+				consider(inst.src[i]);
+			}
+		}
+	}
+	return next;
+}
+
+Operand ScalarRegisterOperand(uint32_t reg) {
+	Operand operand;
+	operand.kind      = OperandKind::Register;
+	operand.reg.file  = RegisterFile::Scalar;
+	operand.reg.index = reg;
+	return operand;
+}
+
+Instruction ShadowWrite(const Instruction& write, uint32_t temporary) {
+	Instruction shadow;
+	shadow.pc        = write.pc;
+	shadow.op        = Opcode::MoveU32;
+	shadow.dst       = ScalarRegisterOperand(temporary);
+	shadow.src[0]    = write.src[0];
+	shadow.src_count = 1;
+	return shadow;
+}
+
+Instruction ShadowRead(const Instruction& read, uint32_t temporary) {
+	Instruction rewritten;
+	rewritten.pc        = read.pc;
+	rewritten.op        = Opcode::MoveU32;
+	rewritten.dst       = read.dst;
+	rewritten.src[0]    = ScalarRegisterOperand(temporary);
+	rewritten.src_count = 1;
+	return rewritten;
+}
+
+} // namespace
+
+ReadLaneEliminationStats EliminateReadLane(Program& program) {
+	ReadLaneEliminationStats stats;
+	if (program.blocks.empty() || (program.wave_size != 32 && program.wave_size != 64)) {
+		return stats;
+	}
+
+	LaneSet universe;
+	for (const auto& block: program.blocks) {
+		for (const auto& inst: block.instructions) {
+			LaneKey key;
+			if (WriteLaneKey(program, inst, key)) {
+				universe.insert(key);
+			}
+		}
+	}
+	if (universe.empty()) {
+		return stats;
+	}
+
+	const size_t         block_count = program.blocks.size();
+	std::vector<LaneSet> entry(block_count, universe);
+	std::vector<LaneSet> exit(block_count, universe);
+	entry[0].clear();
+	for (size_t block = 0; block < block_count; block++) {
+		exit[block] = TransferBlock(program, program.blocks[block], entry[block]);
+	}
+
+	bool changed = true;
+	while (changed) {
+		changed = false;
+		for (size_t block_index = 0; block_index < block_count; block_index++) {
+			LaneSet     next_entry;
+			const auto& block = program.blocks[block_index];
+			if (block_index != 0 && !block.predecessors.empty()) {
+				next_entry = universe;
+				for (const auto predecessor: block.predecessors) {
+					if (predecessor >= block_count) {
+						next_entry.clear();
+						break;
+					}
+					next_entry = Intersect(next_entry, exit[predecessor]);
+				}
+			}
+			auto next_exit = TransferBlock(program, block, next_entry);
+			if (next_entry != entry[block_index] || next_exit != exit[block_index]) {
+				entry[block_index] = std::move(next_entry);
+				exit[block_index]  = std::move(next_exit);
+				changed            = true;
+			}
+		}
+	}
+
+	LaneSet forwarded;
+	for (size_t block_index = 0; block_index < block_count; block_index++) {
+		auto state = entry[block_index];
+		for (const auto& inst: program.blocks[block_index].instructions) {
+			LaneKey key;
+			if (ReadLaneKey(program, inst, key) && state.contains(key)) {
+				forwarded.insert(key);
+			}
+			ApplyInstruction(program, inst, state);
+		}
+	}
+	if (forwarded.empty()) {
+		return stats;
+	}
+
+	std::map<LaneKey, uint32_t> temporaries;
+	auto                        next_temporary = NextTemporaryScalarRegister(program);
+	for (const auto& key: forwarded) {
+		if (next_temporary == UINT32_MAX) {
+			return {};
+		}
+		temporaries.emplace(key, next_temporary++);
+	}
+
+	for (size_t block_index = 0; block_index < block_count; block_index++) {
+		const auto original  = std::move(program.blocks[block_index].instructions);
+		auto&      rewritten = program.blocks[block_index].instructions;
+		rewritten.clear();
+		rewritten.reserve(original.size() + temporaries.size());
+		auto state = entry[block_index];
+		for (const auto& inst: original) {
+			LaneKey read_key;
+			if (ReadLaneKey(program, inst, read_key) && state.contains(read_key)) {
+				const auto temporary = temporaries.find(read_key);
+				if (temporary != temporaries.end()) {
+					rewritten.push_back(ShadowRead(inst, temporary->second));
+					stats.rewritten_reads++;
+					ApplyInstruction(program, inst, state);
+					continue;
+				}
+			}
+
+			rewritten.push_back(inst);
+			LaneKey write_key;
+			if (WriteLaneKey(program, inst, write_key)) {
+				const auto temporary = temporaries.find(write_key);
+				if (temporary != temporaries.end()) {
+					rewritten.push_back(ShadowWrite(inst, temporary->second));
+					stats.shadow_writes++;
+				}
+			}
+			ApplyInstruction(program, inst, state);
+		}
+	}
+	return stats;
+}
+
+} // namespace Libs::Graphics::ShaderRecompiler::IR

src/graphics/shader/recompiler/ir/ReadLaneElimination.h

@@ -0,0 +1,20 @@
+#ifndef EMULATOR_INCLUDE_EMULATOR_GRAPHICS_SHADER_RECOMPILER_READLANEELIMINATION_H_
+#define EMULATOR_INCLUDE_EMULATOR_GRAPHICS_SHADER_RECOMPILER_READLANEELIMINATION_H_
+
+#include "graphics/shader/recompiler/ir/ShaderIR.h"
+
+namespace Libs::Graphics::ShaderRecompiler::IR {
+
+struct ReadLaneEliminationStats {
+	uint32_t rewritten_reads = 0;
+	uint32_t shadow_writes   = 0;
+};
+
+// Replaces fixed-lane ReadLane operations that are reached by a matching WriteLane on every
+// control-flow path. A synthetic scalar register snapshots the value at WriteLane execution time,
+// so the rewrite remains valid when the source SGPR is subsequently overwritten.
+[[nodiscard]] ReadLaneEliminationStats EliminateReadLane(Program& program);
+
+} // namespace Libs::Graphics::ShaderRecompiler::IR
+
+#endif /* EMULATOR_INCLUDE_EMULATOR_GRAPHICS_SHADER_RECOMPILER_READLANEELIMINATION_H_ */