shader: expose wave32 feature guards through ballot mask simplification

src/graphics/shader/recompiler/ShaderRecompiler.cpp

@@ -594,7 +594,7 @@ TranslateResult TranslateProgram(std::span<const uint32_t> code, const CompileOp
 	     GetDumpLabel(options), StageName(options.stage), options.shader_hash,
 	     static_cast<uint64_t>(ir.blocks.size()), phase_ms());
 	IR::RewriteToSsa(ir.blocks);
-	IR::ConstantPropagationPass(ir.blocks);
+	IR::ConstantPropagationPass(ir.blocks, ir.wave_size);
 	IR::ResolveControlFlowIdentities(ir);
 	IR::RemoveIdentities(ir.blocks);
 	IR::EliminateDeadCode(ir.blocks);
@@ -602,7 +602,7 @@ TranslateResult TranslateProgram(std::span<const uint32_t> code, const CompileOp
 	if (read_lane_stats.rewritten_reads != 0) {
 		LOGF("%s read-lane elimination: reads=%" PRIu32 "\n", GetDumpLabel(options),
 		     read_lane_stats.rewritten_reads);
-		IR::ConstantPropagationPass(ir.blocks);
+		IR::ConstantPropagationPass(ir.blocks, ir.wave_size);
 		IR::ResolveControlFlowIdentities(ir);
 		IR::RemoveIdentities(ir.blocks);
 		IR::EliminateDeadCode(ir.blocks);

src/graphics/shader/recompiler/Tessellation.cpp

@@ -217,7 +217,7 @@ void LowerTessellationMemory(IR::Program& program, const CompileOptions& options
 			    2, ActiveAddress(*block, it, it->Arg(2), it->Arg(it->NumArgs() - 1u), resolved));
 		}
 	}
-	ConstantPropagationPass(program.blocks);
+	ConstantPropagationPass(program.blocks, program.wave_size);
 	uint32_t reads = 0, writes = 0, factors = 0;
 	for (auto* block: program.blocks) {
 		for (auto it = block->begin(); it != block->end(); ++it) {
@@ -317,7 +317,7 @@ void LowerTessellationMemory(IR::Program& program, const CompileOptions& options
 			inst.Invalidate();
 		}
 	}
-	ConstantPropagationPass(program.blocks);
+	ConstantPropagationPass(program.blocks, program.wave_size);
 	RemoveIdentities(program.blocks);
 	EliminateDeadCode(program.blocks);
 	LOGF("%s tessellation lowering: reads=%u writes=%u factors=%u\n",

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

@@ -4,6 +4,7 @@
 #include <algorithm>
 #include <bit>
 #include <cstdint>
+#include <unordered_map>
 #include <unordered_set>
 
 namespace Libs::Graphics::ShaderRecompiler::IR {
@@ -660,12 +661,110 @@ void FoldInstruction(Block& block, Block::iterator instruction,
 	}
 }
 
+// In wave32, extracting this invocation's bit from a ballot recovers its predicate.
+// Keep that identity through scalar EXEC operations and their loop-carried mask Phis.
+class LaneMaskProjection {
+public:
+	explicit LaneMaskProjection(std::unordered_set<Inst*>& lowered_ancillary)
+	    : m_lowered_ancillary(lowered_ancillary) {}
+
+	void Fold(Inst& inst) {
+		if (inst.GetOpcode() != ValueOpcode::INotEqual32 || !Immediate(Arg(inst, 1), 0u)) return;
+		const auto* bit = Arg(inst, 0).TryInstruction();
+		if (bit == nullptr || bit->GetOpcode() != ValueOpcode::BitwiseAnd32 ||
+		    !Immediate(Arg(*bit, 1), 1u)) return;
+		const auto* shift = Arg(*bit, 0).TryInstruction();
+		if (shift == nullptr || shift->GetOpcode() != ValueOpcode::ShiftRightLogical32) return;
+		const auto* index = Arg(*shift, 1).TryInstruction();
+		if (index == nullptr || index->GetOpcode() != ValueOpcode::BitwiseAnd32 ||
+		    !Immediate(Arg(*index, 1), 31u)) return;
+		const auto* lane = Arg(*index, 0).TryInstruction();
+		if (lane == nullptr || lane->GetOpcode() != ValueOpcode::LaneId) return;
+		m_visited.clear();
+		m_grounded = false;
+		if (CanProject(Arg(*shift, 0)) && m_grounded) Replace(inst, Project(Arg(*shift, 0)));
+	}
+
+private:
+	static bool Immediate(Value value, uint32_t expected) {
+		return IsImmediate(value, Type::U32) && value.U32() == expected;
+	}
+
+	static Value BallotPredicate(const Inst& inst) {
+		if (inst.GetOpcode() != ValueOpcode::CompositeExtractU32x4 ||
+		    !Immediate(Arg(inst, 1), 0u)) return {};
+		const auto* source = Arg(inst, 0).TryInstruction();
+		return source != nullptr && source->GetOpcode() == ValueOpcode::Ballot
+		           ? Arg(*source, 0) : Value {};
+	}
+
+	bool CanProject(Value value) {
+		if (value.IsImmediate()) {
+			const bool supported = Immediate(value, 0u) || Immediate(value, UINT32_MAX);
+			m_grounded |= supported;
+			return supported;
+		}
+		const auto* inst = value.TryInstruction();
+		if (inst == nullptr) return false;
+		if (m_values.contains(inst) || !BallotPredicate(*inst).IsEmpty()) {
+			m_grounded = true;
+			return true;
+		}
+		if (!m_visited.insert(inst).second) return true;
+		const auto op = inst->GetOpcode();
+		if (op != ValueOpcode::BitwiseAnd32 && op != ValueOpcode::BitwiseOr32 &&
+		    op != ValueOpcode::BitwiseNot32 && op != ValueOpcode::SelectU32 && op != ValueOpcode::Phi)
+			return false;
+		for (size_t arg = op == ValueOpcode::SelectU32 ? 1u : 0u; arg < inst->NumArgs(); ++arg) {
+			if (!CanProject(Arg(*inst, arg))) return false;
+		}
+		return inst->NumArgs() != 0;
+	}
+
+	Value Project(Value value) {
+		if (value.IsImmediate()) return Value(value.U32() != 0u);
+		auto* source = value.TryInstruction();
+		if (const auto found = m_values.find(source); found != m_values.end()) return found->second.Resolve();
+		if (const auto predicate = BallotPredicate(*source); !predicate.IsEmpty()) return predicate;
+		auto* block = source->Parent();
+		const auto where = std::find_if(block->begin(), block->end(),
+		                               [source](const Inst& candidate) { return &candidate == source; });
+		const auto op = source->GetOpcode() == ValueOpcode::BitwiseAnd32 ? ValueOpcode::LogicalAnd
+		              : source->GetOpcode() == ValueOpcode::BitwiseOr32 ? ValueOpcode::LogicalOr
+		              : source->GetOpcode() == ValueOpcode::BitwiseNot32 ? ValueOpcode::LogicalNot
+		              : source->GetOpcode() == ValueOpcode::SelectU32 ? ValueOpcode::SelectU1
+		                                                               : ValueOpcode::Phi;
+		auto result = op == ValueOpcode::Phi ? block->PrependNewInst(where, op)
+		            : op == ValueOpcode::LogicalNot
+		                ? block->PrependNewInst(where, op, {Value(false)})
+		            : op == ValueOpcode::SelectU1
+		                ? block->PrependNewInst(where, op, {Arg(*source, 0), Value(false), Value(false)})
+		                : block->PrependNewInst(where, op, {Value(false), Value(false)});
+		m_values.emplace(source, Value(&*result));
+		if (op == ValueOpcode::Phi) result->SetFlags(Type::U1);
+		for (size_t arg = op == ValueOpcode::SelectU1 ? 1u : 0u; arg < source->NumArgs(); ++arg) {
+			const auto projected = Project(Arg(*source, arg));
+			if (op == ValueOpcode::Phi) result->AddPhiOperand(source->PhiBlock(arg), projected);
+			else result->SetArg(arg, projected);
+		}
+		FoldInstruction(*block, result, m_lowered_ancillary);
+		return Value(&*result).Resolve();
+	}
+
+	std::unordered_set<Inst*>& m_lowered_ancillary;
+	std::unordered_set<const Inst*> m_visited;
+	std::unordered_map<const Inst*, Value> m_values;
+	bool m_grounded = false;
+};
+
 } // namespace
 
-void ConstantPropagationPass(const BlockList& blocks) {
+void ConstantPropagationPass(const BlockList& blocks, uint32_t wave_size) {
 	std::unordered_set<Inst*> lowered_ancillary;
+	LaneMaskProjection mask_projection(lowered_ancillary);
 	for (auto* block: blocks) {
 		for (auto inst = block->begin(); inst != block->end(); ++inst) {
+			if (wave_size == 32u) mask_projection.Fold(*inst);
 			FoldInstruction(*block, inst, lowered_ancillary);
 		}
 	}

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

@@ -4,6 +4,6 @@
 
 namespace Libs::Graphics::ShaderRecompiler::IR {
 
-void ConstantPropagationPass(const BlockList& blocks);
+void ConstantPropagationPass(const BlockList& blocks, uint32_t wave_size = 64);
 
 } // namespace Libs::Graphics::ShaderRecompiler::IR

tests/shaderCfgTests.cpp

@@ -11551,6 +11551,124 @@ void TestValuePhiValidation() {
 #endif
 }
 
+void TestWave32MaskProjection() {
+  using namespace ShaderRecompiler::IR;
+  for (const uint32_t wave_size : {32u, 64u}) {
+    Program program;
+    program.stage = ShaderType::Compute;
+    program.wave_size = wave_size;
+    for (uint32_t id = 0; id < 3; ++id) {
+      program.block_storage.push_back(std::make_unique<Block>());
+      program.blocks.push_back(program.block_storage.back().get());
+      program.block_info.push_back({.id = id});
+    }
+    auto *entry = program.blocks[0];
+    auto *loop = program.blocks[1];
+    entry->AddBranch(loop);
+    loop->AddBranch(loop);
+    loop->AddBranch(program.blocks[2]);
+    program.block_info[0].terminator.kind = ShaderRecompiler::CFG::TerminatorKind::Branch;
+    program.block_info[0].terminator.true_block = 1;
+    program.block_info[1].terminator.kind = ShaderRecompiler::CFG::TerminatorKind::ConditionalBranch;
+    program.block_info[1].terminator.true_block = 1;
+    program.block_info[1].terminator.false_block = 2;
+    program.block_info[2].terminator.kind = ShaderRecompiler::CFG::TerminatorKind::Return;
+    IREmitter e(entry);
+    const auto zero = U32(Value(0u));
+    const auto raw = e.GetUserData(static_cast<ScalarReg>(0));
+    const auto active = e.INotEqual(raw, zero);
+    const auto selected = e.IEqual(raw, U32(Value(7u)));
+    const auto ballot = [&](IREmitter &ir, U1 predicate, uint32_t component = 0u) {
+      return ir.CompositeExtract(ir.Emit(ValueOpcode::Ballot, {predicate}), component);
+    };
+    const auto thread_bit = [&](IREmitter &ir, U32 word, bool own_lane = true) {
+      const auto bit = own_lane
+          ? ir.BitwiseAnd(U32(ir.Emit(ValueOpcode::LaneId)), U32(Value(31u)))
+          : U32(Value(3u));
+      return ir.INotEqual(ir.BitwiseAnd(ir.ShiftRightLogical(word, bit), U32(Value(1u))), zero);
+    };
+    const auto active_mask = ballot(e, active);
+    const auto direct = thread_bit(e, active_mask);
+    const auto arbitrary = thread_bit(e, raw);
+    const auto mixed = thread_bit(e, e.BitwiseAnd(active_mask, raw));
+    const auto mixed_or = thread_bit(e, e.BitwiseOr(active_mask, raw));
+    const auto mixed_select = thread_bit(e, e.Select(selected, active_mask, raw));
+    const auto different_lane = thread_bit(e, active_mask, false);
+    const auto high_half = thread_bit(e, ballot(e, active, 1u));
+    const auto initial_mask = e.BitwiseAnd(active_mask, ballot(e, selected));
+    const auto raw_consumer = e.Emit(ValueOpcode::IAdd32, {initial_mask, Value(3u)});
+    const auto initial_predicate = e.LogicalAnd(active, selected);
+    const auto entry_size = entry->Instructions().size();
+    ConstantPropagationPass({entry}, wave_size);
+    Check(entry->Instructions().size() == entry_size,
+          "rejected mixed mask graph left unused projected instructions");
+    const auto combined_mask = e.BitwiseOr(active_mask, ballot(e, selected));
+    const auto combined = thread_bit(e, combined_mask);
+    const auto gated_mask = e.Select(selected, combined_mask, zero);
+    const auto gated = thread_bit(e, gated_mask);
+    const auto rejected = thread_bit(e, e.Select(selected, zero, combined_mask));
+    const auto combined_predicate = e.LogicalOr(active, selected);
+    const auto gated_predicate = e.LogicalAnd(selected, combined_predicate);
+    const auto rejected_predicate = e.Emit(ValueOpcode::SelectU1,
+                                          {selected, Value(false), combined_predicate});
+    const auto gated_raw_consumer = e.Emit(ValueOpcode::IAdd32, {gated_mask, Value(3u)});
+    auto &mask = loop->AppendNewInst(ValueOpcode::Phi);
+    mask.SetFlags(Type::U32);
+    auto &predicate = loop->AppendNewInst(ValueOpcode::Phi);
+    predicate.SetFlags(Type::U1);
+    IREmitter body(loop);
+    const auto matches = body.LogicalAnd(U1(Value(&predicate)), selected);
+    const auto remaining = body.BitwiseAnd(U32(Value(&mask)), body.BitwiseNot(ballot(body, matches)));
+    const auto remaining_predicate = body.LogicalAnd(U1(Value(&predicate)), body.LogicalNot(matches));
+    const auto remaining_live = thread_bit(body, remaining);
+    const auto live = thread_bit(body, U32(Value(&mask)));
+    mask.AddPhiOperand(entry, initial_mask);
+    mask.AddPhiOperand(loop, remaining);
+    predicate.AddPhiOperand(entry, initial_predicate);
+    predicate.AddPhiOperand(loop, remaining_predicate);
+    program.block_info[1].condition = remaining_predicate;
+
+    ConstantPropagationPass(program.blocks, wave_size);
+    if (wave_size == 32u) {
+      Check(direct.Resolve() == active.Resolve(), "wave32 ballot round trip lost its predicate");
+      Check(EquivalentValue(program, combined, combined_predicate),
+            "wave32 OR mask projection changed a lane predicate");
+      Check(EquivalentValue(program, gated, gated_predicate) &&
+                EquivalentValue(program, rejected, rejected_predicate),
+            "wave32 select mask projection changed a scalar gate or selected arm");
+      Check(EquivalentValue(program, live.Resolve(), Value(&predicate)),
+            "projected mask Phi differs from the existing per-lane EXEC Phi");
+      Check(EquivalentValue(program, remaining_live.Resolve(), remaining_predicate) &&
+                std::count_if(loop->begin(), loop->end(), [](const Inst &inst) {
+                  return inst.GetOpcode() == ValueOpcode::LogicalAnd;
+                }) == 3,
+            "backedge-first projection duplicated the loop's logical mask operation");
+      const auto *projected = live.Resolve().TryInstruction();
+      Check(projected != nullptr && projected->GetOpcode() == ValueOpcode::Phi &&
+                projected->GetType() == Type::U1 && projected->PhiBlock(0) == entry &&
+                projected->PhiBlock(1) == loop,
+            "mask projection lost loop Phi type or predecessor order");
+    } else {
+      Check(direct.Resolve().TryInstruction()->GetOpcode() == ValueOpcode::INotEqual32 &&
+                live.Resolve().TryInstruction()->GetOpcode() == ValueOpcode::INotEqual32 &&
+                combined.Resolve().TryInstruction()->GetOpcode() == ValueOpcode::INotEqual32 &&
+                gated.Resolve().TryInstruction()->GetOpcode() == ValueOpcode::INotEqual32 &&
+                rejected.Resolve().TryInstruction()->GetOpcode() == ValueOpcode::INotEqual32,
+            "wave64 low-half mask was projected onto the wrong invocation");
+    }
+    for (const auto value : {arbitrary, mixed, mixed_or, mixed_select, different_lane, high_half}) {
+      Check(value.Resolve().TryInstruction()->GetOpcode() == ValueOpcode::INotEqual32,
+            "mask projection accepted an arbitrary mask, lane or ballot half");
+    }
+    Check(raw_consumer.Resolve().TryInstruction()->Arg(0).Resolve() == initial_mask.Resolve() &&
+              gated_raw_consumer.Resolve().TryInstruction()->Arg(0).Resolve() == gated_mask.Resolve(),
+          "mask projection changed a raw scalar mask consumer");
+    ResolveControlFlowIdentities(program);
+    RemoveIdentities(program.blocks);
+    ValidateProgram(program, true);
+  }
+}
+
 void TestU64ShiftConstantPropagation() {
   using namespace ShaderRecompiler::IR;
 
@@ -14256,6 +14374,7 @@ int main() {
   TestFinalSsaRejectsRegisterStatePseudos();
 #endif
   TestValuePhiValidation();
+  TestWave32MaskProjection();
   TestU64ShiftConstantPropagation();
 #if KYTY_PLATFORM != KYTY_PLATFORM_WINDOWS
   TestNativeWideValueValidation();