shader: keep fragment layer values under lane masks

Co-authored-by: shadowbeat070 <70843316+shadowbeat070@users.noreply.github.com>

src/graphics/shader/recompiler/ShaderRecompiler.cpp

@@ -651,6 +651,9 @@ TranslateResult TranslateProgram(std::span<const uint32_t> code, const CompileOp
 	IR::ResolveControlFlowIdentities(ir);
 	IR::RemoveIdentities(ir.blocks);
 	IR::EliminateDeadCode(ir.blocks);
+	IR::EliminateMaskedValues(ir);
+	IR::RemoveIdentities(ir.blocks);
+	IR::EliminateDeadCode(ir.blocks);
 	const auto read_lane_stats = IR::EliminateReadLane(ir, ir.wave_size);
 	if (read_lane_stats.rewritten_reads != 0) {
 		LOGF("%s read-lane elimination: reads=%" PRIu32 "\n", GetDumpLabel(options),

src/graphics/shader/recompiler/backend/spirv/SpirvBuilder.cpp

@@ -32,6 +32,10 @@ uint32_t Builder::AllocateId() {
 	return m_next_id++;
 }
 
+void Builder::RequireVersion(uint32_t version) {
+	m_version = std::max(m_version, version);
+}
+
 void Builder::RequireCapability(uint32_t capability) {
 	if (m_required_capabilities.insert(capability).second) {
 		AddCapability({capability});

src/graphics/shader/recompiler/backend/spirv/SpirvBuilder.h

@@ -28,6 +28,7 @@ public:
 	KYTY_CLASS_DEFAULT_COPY(Builder);
 
 	uint32_t AllocateId();
+	void     RequireVersion(uint32_t version);
 	void     RequireCapability(uint32_t capability);
 	void     RequireExtension(const char* name);
 	uint32_t Import(const char* name);

src/graphics/shader/recompiler/backend/spirv/spirvEmitterFlow.cpp

@@ -71,7 +71,8 @@ uint32_t EmitBuiltinU32(ValueEmitContext& ctx, IR::StageInputKind kind, uint32_t
 		    {OpSelect, TypeU32(state), bits, value, ConstantU32(state, 1), ConstantU32(state, 0)});
 		return bits;
 	}
-	if (kind == IR::StageInputKind::VertexIndex || kind == IR::StageInputKind::InstanceIndex) {
+	if (kind == IR::StageInputKind::VertexIndex || kind == IR::StageInputKind::InstanceIndex ||
+	    kind == IR::StageInputKind::Layer) {
 		const auto value = state.builder.AllocateId();
 		const auto bits  = state.builder.AllocateId();
 		state.builder.AddFunction({OpLoad, TypeI32(state), value, variable});

src/graphics/shader/recompiler/backend/spirv/spirvEmitterInternal.h

@@ -50,8 +50,8 @@ enum : uint32_t {
 	CapabilityGroupNonUniform                = 61,
 	CapabilityGroupNonUniformBallot          = 64,
 	CapabilityGroupNonUniformShuffle         = 65,
+	CapabilityShaderLayer                    = 69,
 	CapabilitySignedZeroInfNanPreserve       = 4466,
-	CapabilityShaderViewportIndexLayerEXT    = 5254,
 	CapabilityFragmentBarycentricKHR         = 5284,
 	CapabilityComputeDerivativeGroupQuadsKHR = 5288,
 	CapabilityPhysicalStorageBufferAddresses = 5347,

src/graphics/shader/recompiler/backend/spirv/spirvEmitterModule.cpp

@@ -518,6 +518,7 @@ uint32_t BuiltInForInput(IR::StageInputKind kind) {
 		case IR::StageInputKind::InstanceIndex: return BuiltInInstanceIndex;
 		case IR::StageInputKind::FragCoord: return BuiltInFragCoord;
 		case IR::StageInputKind::FrontFacing: return BuiltInFrontFacing;
+		case IR::StageInputKind::Layer: return BuiltInLayer;
 		case IR::StageInputKind::BaryCoordSmooth: return BuiltInBaryCoordKHR;
 		case IR::StageInputKind::BaryCoordNoPerspective: return BuiltInBaryCoordNoPerspKHR;
 		case IR::StageInputKind::WorkgroupId: return BuiltInWorkgroupId;
@@ -541,6 +542,9 @@ void AddInputAnnotationsAndNames(EmitterState& state) {
 	}
 	for (const auto& input: state.inputs) {
 		state.builder.AddName(input.variable_id, input.debug_name.c_str());
+		if (input.kind == IR::StageInputKind::Layer) {
+			state.builder.AddAnnotation({OpDecorate, input.variable_id, DecorationFlat});
+		}
 		if (input.kind == IR::StageInputKind::Parameter) {
 			const auto flat = PixelParameterIsFlat(state, input.location);
 			if (input.per_vertex) {
@@ -710,9 +714,9 @@ void DefineModule(EmitterState& state) {
 	if (state.cull_distance_variable != 0) {
 		state.builder.RequireCapability(CapabilityCullDistance);
 	}
-	if (state.layer_variable != 0) {
-		state.builder.RequireCapability(CapabilityShaderViewportIndexLayerEXT);
-		state.builder.RequireExtension("SPV_EXT_shader_viewport_index_layer");
+	if (state.layer_variable != 0 || InputVariableForKind(state, IR::StageInputKind::Layer) != 0) {
+		state.builder.RequireVersion(0x00010500u);
+		state.builder.RequireCapability(CapabilityShaderLayer);
 	}
 	if (state.requirements.image_gather_extended) {
 		state.builder.RequireCapability(CapabilityImageGatherExtended);
@@ -805,6 +809,7 @@ void DefineModule(EmitterState& state) {
 		switch (input.kind) {
 			case IR::StageInputKind::VertexIndex:
 			case IR::StageInputKind::InstanceIndex:
+			case IR::StageInputKind::Layer:
 				ptr_type = TypePointer(state, StorageClassInput, TypeI32(state));
 				break;
 			case IR::StageInputKind::WorkgroupId:

src/graphics/shader/recompiler/frontend/translate/Control.cpp

@@ -259,13 +259,20 @@ void Translator::S_MOV_B64(const Decoder::Instruction& inst) {
 }
 
 void Translator::S_WQM_B64(const Decoder::Instruction& inst) {
-	const auto mask_valid = ReadMaskValid(inst.src0);
+	const auto source_mask = ReadMask(inst.src0);
+	const auto mask_valid  = ReadMaskValid(inst.src0);
 	const auto result =
 	    IR::U64(ir.Emit(IR::ValueOpcode::WqmU64, {ReadOperand(inst.src0, IR::Type::U64)}));
+	// WQM activates whole quads and cannot clear an already active lane.
+	const auto widened_mask = ir.LogicalOr(source_mask, ThreadBit(ExtractU64(result)));
 	WriteOperand(DestinationOperand(inst), result);
-	if (inst.dst.kind == Decoder::OperandKind::Sgpr) {
+	if (inst.dst.kind == Decoder::OperandKind::ExecLo) {
+		ir.SetExec(widened_mask);
+	} else if (inst.dst.kind == Decoder::OperandKind::VccLo) {
+		ir.SetVcc(widened_mask);
+	} else if (inst.dst.kind == Decoder::OperandKind::Sgpr) {
 		const auto dst = static_cast<IR::ScalarReg>(inst.dst.reg);
-		ir.SetThreadBitScalarReg(dst, ThreadBit(ExtractU64(result)));
+		ir.SetThreadBitScalarReg(dst, widened_mask);
 		ir.SetScalarMaskTag(dst, mask_valid);
 	}
 	ir.SetScc(IR::U1(ir.Emit(IR::ValueOpcode::INotEqual64, {result, IR::Value(uint64_t {0})})));

src/graphics/shader/recompiler/frontend/translate/Translate.cpp

@@ -1137,9 +1137,13 @@ IR::Program TranslateProgram(const Decoder::Program& decoded, const CFG::Graph&
 				                      builtin(IR::StageInputKind::FragCoord, 3));
 			}
 			if (ps->ps_front_face) {
-				entry_ir.SetVectorReg(static_cast<IR::VectorReg>(reg),
+				entry_ir.SetVectorReg(static_cast<IR::VectorReg>(reg++),
 				                      builtin(IR::StageInputKind::FrontFacing));
 			}
+			if (ps->ps_ancillary) {
+				entry_ir.SetVectorReg(static_cast<IR::VectorReg>(reg),
+				                      builtin(IR::StageInputKind::PackedAncillary));
+			}
 		} else if (options.stage == ShaderType::Vertex) {
 			entry_ir.SetVectorReg(static_cast<IR::VectorReg>(5),
 			                      builtin(IR::StageInputKind::VertexIndex));

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

@@ -160,6 +160,8 @@ enum class StageInputKind {
 	InstanceIndex,
 	FragCoord,
 	FrontFacing,
+	PackedAncillary,
+	Layer,
 	BaryCoordSmooth,
 	BaryCoordNoPerspective,
 	WorkgroupId,

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

@@ -1,8 +1,10 @@
 #include "graphics/shader/recompiler/ir/passes/ConstantPropagation.h"
+#include "graphics/shader/recompiler/ir/ShaderIR.h"
 
 #include <algorithm>
 #include <bit>
 #include <cstdint>
+#include <unordered_map>
 
 namespace Libs::Graphics::ShaderRecompiler::IR {
 namespace {
@@ -116,6 +118,87 @@ bool FoldSelect(Inst& inst) {
 	return false;
 }
 
+bool IsSelect(ValueOpcode opcode) {
+	return opcode == ValueOpcode::SelectU1 || opcode == ValueOpcode::SelectU32 ||
+	       opcode == ValueOpcode::SelectF32;
+}
+
+bool Implies(Value condition, Value required, uint32_t depth = 0) {
+	condition = condition.Resolve();
+	if (condition == required) return true;
+	const auto* inst = condition.TryInstruction();
+	return depth < 64u && inst != nullptr && inst->GetOpcode() == ValueOpcode::LogicalAnd &&
+	       (Implies(inst->Arg(0), required, depth + 1u) ||
+	        Implies(inst->Arg(1), required, depth + 1u));
+}
+
+bool IsAcyclic(const BlockList& blocks) {
+	std::unordered_map<const Block*, bool> completed;
+	const auto visit = [&](auto&& self, const Block* block) -> bool {
+		const auto [entry, added] = completed.emplace(block, false);
+		if (!added) return entry->second;
+		for (const auto* successor: block->ImmSuccessors()) {
+			if (!self(self, successor)) return false;
+		}
+		completed[block] = true;
+		return true;
+	};
+	for (const auto* block: blocks) {
+		if (!visit(visit, block)) return false;
+	}
+	return true;
+}
+
+class MaskedUses {
+public:
+	MaskedUses(const Program& program, Value condition)
+	    : program(program), condition(condition.Resolve()) {}
+
+	bool OnlyWhenActive(const Inst& inst) {
+		const auto [entry, added] = completed.emplace(&inst, false);
+		if (!added) return entry->second;
+		for (const auto& use: inst.Uses()) {
+			const auto& user = *use.user;
+			const auto  op   = user.GetOpcode();
+			if (IsSelect(op)) {
+				if (use.operand == 0) return false;
+				if (use.operand == 1 && Implies(user.Arg(0), condition)) continue;
+			} else if (op == ValueOpcode::SetAttribute) {
+				if (use.operand == 0 && Implies(user.Arg(1), condition)) continue;
+				return false;
+			} else if (!IsLaneLocal(user, use.operand)) {
+				return false;
+			}
+			if (!OnlyWhenActive(user)) return false;
+		}
+		completed[&inst] = true;
+		return true;
+	}
+
+private:
+	bool IsLaneLocal(const Inst& inst, size_t operand) const {
+		const auto op = inst.GetOpcode();
+		// These opcode-table groups contain only construction, conversion and arithmetic.
+		if ((op >= ValueOpcode::BitCastU16F16 && op <= ValueOpcode::BitFieldSExtract) ||
+		    (op >= ValueOpcode::WqmU64 && op <= ValueOpcode::FPFract32) ||
+		    op == ValueOpcode::Identity || op == ValueOpcode::Phi ||
+		    op == ValueOpcode::MakeImageAddress) {
+			return true;
+		}
+		if (op != ValueOpcode::ImageSampleRaw || operand != 2) return false;
+		// Explicit LOD or gradients keep coordinates local; implicit derivatives read neighbors.
+		const auto index = inst.Flags<MemoryFlags>().index;
+		return index < program.memory_info.size() &&
+		       (program.memory_info[index].image_sample_flags &
+		        (Decoder::ImageSampleFlagLod | Decoder::ImageSampleFlagLevelZero |
+		         Decoder::ImageSampleFlagDerivative)) != 0;
+	}
+
+	const Program&                       program;
+	Value                                condition;
+	std::unordered_map<const Inst*, bool> completed;
+};
+
 bool FoldPhi(Inst& inst) {
 	Value same;
 	for (size_t index = 0; index < inst.NumArgs(); index++) {
@@ -192,7 +275,8 @@ bool FoldCompositeExtract(Inst& inst, ValueOpcode construct, size_t components)
 	return false;
 }
 
-void FoldInstruction(Inst& inst) {
+void FoldInstruction(Block& block, Block::iterator instruction) {
+	auto& inst = *instruction;
 	switch (inst.GetOpcode()) {
 		case ValueOpcode::Phi: FoldPhi(inst); return;
 		case ValueOpcode::SelectU1:
@@ -223,6 +307,21 @@ void FoldInstruction(Inst& inst) {
 			const auto value  = Arg(inst, 0);
 			const auto offset = Arg(inst, 1);
 			const auto count  = Arg(inst, 2);
+			const auto* source = value.TryInstruction();
+			if (source != nullptr && source->GetOpcode() == ValueOpcode::GetBuiltin &&
+			    source->Arg(0) == Value(static_cast<uint32_t>(StageInputKind::PackedAncillary)) &&
+			    IsImmediate(offset, Type::U32) && IsImmediate(count, Type::U32) &&
+			    offset.U32() >= 16u && offset.U32() < 27u && count.U32() != 0u &&
+			    count.U32() <= 27u - offset.U32()) {
+				// Preserve the extraction (including sign extension) while exposing only its
+				// used hardware field. Other live ancillary fields remain unsupported.
+				const auto layer = block.PrependNewInst(
+				    instruction, ValueOpcode::GetBuiltin,
+				    {Value(static_cast<uint32_t>(StageInputKind::Layer)), Value(0u)});
+				inst.SetArg(0, Value(&*layer));
+				inst.SetArg(1, Value(offset.U32() - 16u));
+				return;
+			}
 			if (!IsImmediate(value, Type::U32) || !IsImmediate(offset, Type::U32) ||
 			    !IsImmediate(count, Type::U32) || offset.U32() > 32u ||
 			    count.U32() > 32u - offset.U32()) {
@@ -584,8 +683,22 @@ void FoldInstruction(Inst& inst) {
 
 void ConstantPropagationPass(const BlockList& blocks) {
 	for (auto* block: blocks) {
-		for (auto& inst: block->Instructions()) {
-			FoldInstruction(inst);
+		for (auto inst = block->begin(); inst != block->end(); ++inst) {
+			FoldInstruction(*block, inst);
+		}
+	}
+}
+
+void EliminateMaskedValues(Program& program) {
+	// A loop-carried value may have been produced under a previous iteration's predicate.
+	if (!IsAcyclic(program.blocks)) return;
+	for (auto block = program.blocks.rbegin(); block != program.blocks.rend(); ++block) {
+		auto& instructions = (*block)->Instructions();
+		for (auto inst = instructions.rbegin(); inst != instructions.rend(); ++inst) {
+			if (IsSelect(inst->GetOpcode()) &&
+			    MaskedUses(program, inst->Arg(0)).OnlyWhenActive(*inst)) {
+				Replace(*inst, inst->Arg(1));
+			}
 		}
 	}
 }

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

@@ -4,6 +4,9 @@
 
 namespace Libs::Graphics::ShaderRecompiler::IR {
 
+struct Program;
+
 void ConstantPropagationPass(const BlockList& blocks);
+void EliminateMaskedValues(Program& program);
 
 } // namespace Libs::Graphics::ShaderRecompiler::IR

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

@@ -115,6 +115,13 @@ void ValidateValueReferences(const Program& program, const ShaderInfoOptions& op
 					const auto kind      = static_cast<StageInputKind>(inst.Arg(0).U32());
 					const auto component = inst.Arg(1).U32();
 					switch (kind) {
+						case StageInputKind::PackedAncillary:
+							return Fail("packed pixel ancillary input has an unsupported live use");
+						case StageInputKind::Layer:
+							if (program.stage != ShaderType::Pixel || component != 0u) {
+								return Fail("typed layer input is invalid");
+							}
+							break;
 						case StageInputKind::VertexIndex:
 						case StageInputKind::InstanceIndex:
 						case StageInputKind::FrontFacing:
@@ -251,6 +258,7 @@ void CollectBuiltinInputs(const Program& program, ShaderInfo& info) {
 				case StageInputKind::FrontFacing:
 					AddInput(info, kind, 0, 1, "gl_FrontFacing");
 					break;
+				case StageInputKind::Layer: AddInput(info, kind, 0, 1, "gl_Layer"); break;
 				case StageInputKind::BaryCoordSmooth:
 					AddInput(info, kind, 0, 3, "gl_BaryCoordKHR");
 					break;
@@ -269,6 +277,7 @@ void CollectBuiltinInputs(const Program& program, ShaderInfo& info) {
 				case StageInputKind::GlobalInvocationId:
 					AddInput(info, kind, 0, 3, "gl_GlobalInvocationID");
 					break;
+				case StageInputKind::PackedAncillary:
 				case StageInputKind::Parameter: break;
 			}
 		}

src/graphics/shader/shader.cpp

@@ -699,6 +699,7 @@ static void ShaderGetStaticInputInfoPS(
 	ps_info.ps_pos_z                     = (active_inputs & 0x00000400u) != 0;
 	ps_info.ps_pos_w                     = (active_inputs & 0x00000800u) != 0;
 	ps_info.ps_front_face                = (active_inputs & 0x00001000u) != 0;
+	ps_info.ps_ancillary                 = (active_inputs & 0x00002000u) != 0;
 	ps_info.ps_sample_shading            = (active_inputs & 0x00000011u) != 0;
 	ps_info.ps_no_perspective            = (sh.ps_input_ena & sh.ps_input_addr & 0x00000020u) != 0;
 	ps_info.ps_pixel_kill_enable         = sh.db_shader_control.shader_kill_enable;
@@ -808,6 +809,7 @@ void BuildStageStaticKey(const ShaderPixelInputInfo& info, std::vector<uint32_t>
 	key.push_back(static_cast<uint32_t>(info.ps_pos_z));
 	key.push_back(static_cast<uint32_t>(info.ps_pos_w));
 	key.push_back(static_cast<uint32_t>(info.ps_front_face));
+	key.push_back(static_cast<uint32_t>(info.ps_ancillary));
 	key.push_back(static_cast<uint32_t>(info.ps_no_perspective));
 	key.push_back(static_cast<uint32_t>(info.ps_pixel_kill_enable));
 	key.push_back(static_cast<uint32_t>(info.ps_depth_export_enable));
@@ -985,6 +987,7 @@ void ShaderDbgDumpInputInfo(const ShaderPixelInputInfo& info) {
 	     "\t ps_pos_z             = %s\n"
 	     "\t ps_pos_w             = %s\n"
 	     "\t ps_front_face        = %s\n"
+	     "\t ps_ancillary         = %s\n"
 	     "\t ps_sample_shading    = %s\n"
 	     "\t ps_no_perspective    = %s\n"
 	     "\t ps_pixel_kill_enable = %s\n"
@@ -994,6 +997,7 @@ void ShaderDbgDumpInputInfo(const ShaderPixelInputInfo& info) {
 	     info.ps_perspective_center_vgpr, info.ps_pos_x ? "true" : "false",
 	     info.ps_pos_y ? "true" : "false", info.ps_pos_z ? "true" : "false",
 	     info.ps_pos_w ? "true" : "false", info.ps_front_face ? "true" : "false",
+	     info.ps_ancillary ? "true" : "false",
 	     info.ps_sample_shading ? "true" : "false", info.ps_no_perspective ? "true" : "false",
 	     info.ps_pixel_kill_enable ? "true" : "false", info.ps_early_z ? "true" : "false",
 	     info.ps_execute_on_noop ? "true" : "false");

src/graphics/shader/shader.h

@@ -125,6 +125,7 @@ struct ShaderPixelInputInfo {
 	bool                                           ps_pos_z                     = false;
 	bool                                           ps_pos_w                     = false;
 	bool                                           ps_front_face                = false;
+	bool                                           ps_ancillary                 = false;
 	bool                                           ps_no_perspective            = false;
 	bool                                           ps_pixel_kill_enable         = false;
 	bool                                           ps_depth_export_enable       = false;