shader: propagate late constants through affected IR users

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

@@ -22,6 +22,59 @@ void Replace(Inst& inst, Value value) {
 	inst.ReplaceUsesWith(value.Resolve());
 }
 
+class FoldWorklist {
+public:
+	struct Location {
+		Block::iterator instruction;
+		uint64_t visited = 0;
+		bool queued = false;
+	};
+
+	void Register(Block::iterator instruction) {
+		locations.emplace(&*instruction, Location {instruction});
+		Queue(&*instruction);
+	}
+
+	void Queue(Inst* inst) {
+		const auto found = locations.find(inst);
+		if (found != locations.end() && !found->second.queued) {
+			found->second.queued = true;
+			pending.push_back(inst);
+		}
+	}
+
+	void QueueUsers(const std::vector<Use>& users) {
+		++generation;
+		affected.clear();
+		for (const auto& use: users) affected.push_back(use.user);
+		// Pattern folds inspect nested operands even when intermediate operations
+		// stay dynamic. Visit consumers once per change, including cyclic Phi edges.
+		for (size_t i = 0; i < affected.size(); ++i) {
+			auto* inst = affected[i];
+			const auto found = locations.find(inst);
+			if (found == locations.end() || found->second.visited == generation) continue;
+			found->second.visited = generation;
+			Queue(inst);
+			for (const auto& use: inst->Uses()) affected.push_back(use.user);
+		}
+	}
+
+	Block::iterator Prepend(Block& block, Block::iterator before, ValueOpcode opcode,
+	                         std::initializer_list<Value> args = {}) {
+		const auto result = block.PrependNewInst(before, opcode, args);
+		Register(result);
+		return result;
+	}
+
+	std::unordered_map<Inst*, Location> locations;
+	std::vector<Inst*> pending;
+	std::unordered_set<Inst*> lowered_ancillary;
+
+private:
+	std::vector<Inst*> affected;
+	uint64_t generation = 0;
+};
+
 template <typename Function>
 bool FoldU32(Inst& inst, Function function) {
 	const auto lhs = Arg(inst, 0);
@@ -196,7 +249,7 @@ bool FoldCompositeExtract(Inst& inst, ValueOpcode construct, size_t components)
 }
 
 void FoldInstruction(Block& block, Block::iterator instruction,
-                      std::unordered_set<Inst*>& lowered_ancillary) {
+	                      FoldWorklist& worklist) {
 	auto& inst = *instruction;
 	switch (inst.GetOpcode()) {
 		case ValueOpcode::StoreBufferU32: {
@@ -208,7 +261,7 @@ void FoldInstruction(Block& block, Block::iterator instruction,
 		case ValueOpcode::SelectU1:
 			if (!FoldSelect(inst) && IsImmediate(Arg(inst, 2), Type::U1) &&
 			    !Arg(inst, 2).U1()) {
-				auto result = block.PrependNewInst(instruction, ValueOpcode::LogicalAnd,
+				auto result = worklist.Prepend(block, instruction, ValueOpcode::LogicalAnd,
 				                                   {Arg(inst, 0), Arg(inst, 1)});
 				Replace(inst, Value(&*result));
 			}
@@ -224,8 +277,8 @@ void FoldInstruction(Block& block, Block::iterator instruction,
 				case 0u: Replace(inst, Value(false)); break;
 				case 0x3ffu: Replace(inst, Value(true)); break;
 				case 3u: {
-					const auto result = block.PrependNewInst(
+					const auto result = worklist.Prepend(
-					    instruction, ValueOpcode::FPIsNan32, {Arg(inst, 0)});
+					    block, instruction, ValueOpcode::FPIsNan32, {Arg(inst, 0)});
 					Replace(inst, Value(&*result));
 					break;
 				}
@@ -280,12 +333,12 @@ void FoldInstruction(Block& block, Block::iterator instruction,
 					if (offset.U32() >= field.start && offset.U32() < field.end &&
 					    count.U32() <= field.end - offset.U32()) {
 						// Preserve extraction and sign extension while exposing only the used field.
-						const auto input = block.PrependNewInst(
+						const auto input = worklist.Prepend(
-						    instruction, ValueOpcode::GetBuiltin,
+						    block, instruction, ValueOpcode::GetBuiltin,
 						    {Value(static_cast<uint32_t>(field.kind)), Value(0u)});
 						inst.SetArg(0, Value(&*input));
 						inst.SetArg(1, Value(offset.U32() - field.start));
-						lowered_ancillary.insert(source);
+						worklist.lowered_ancillary.insert(source);
 						return;
 					}
 				}
@@ -710,8 +763,7 @@ void FoldInstruction(Block& block, Block::iterator instruction,
 // ballot projection also survives scalar EXEC operations and loop-carried mask Phis.
 class LaneMaskProjection {
 public:
-	explicit LaneMaskProjection(std::unordered_set<Inst*>& lowered_ancillary)
+	explicit LaneMaskProjection(FoldWorklist& worklist): m_worklist(worklist) {}
-	    : m_lowered_ancillary(lowered_ancillary) {}
 
 	void Fold(Inst& inst, uint32_t wave_size) {
 		if (inst.GetOpcode() != ValueOpcode::INotEqual32 || !Immediate(Arg(inst, 1), 0u)) return;
@@ -778,19 +830,18 @@ private:
 		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(),
+		const auto where = m_worklist.locations.at(source).instruction;
-		                               [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)
+		auto result = op == ValueOpcode::Phi ? m_worklist.Prepend(*block, where, op)
 		            : op == ValueOpcode::LogicalNot
-		                ? block->PrependNewInst(where, op, {Value(false)})
+		                ? m_worklist.Prepend(*block, where, op, {Value(false)})
 		            : op == ValueOpcode::SelectU1
-		                ? block->PrependNewInst(where, op, {Arg(*source, 0), Value(false), Value(false)})
+		                ? m_worklist.Prepend(*block, where, op, {Arg(*source, 0), Value(false), Value(false)})
-		                : block->PrependNewInst(where, op, {Value(false), Value(false)});
+		                : m_worklist.Prepend(*block, 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) {
@@ -798,11 +849,11 @@ private:
 			if (op == ValueOpcode::Phi) result->AddPhiOperand(source->PhiBlock(arg), projected);
 			else result->SetArg(arg, projected);
 		}
-		FoldInstruction(*block, result, m_lowered_ancillary);
+		FoldInstruction(*block, result, m_worklist);
 		return Value(&*result).Resolve();
 	}
 
-	std::unordered_set<Inst*>& m_lowered_ancillary;
+	FoldWorklist& m_worklist;
 	std::unordered_set<const Inst*> m_visited;
 	std::unordered_map<const Inst*, Value> m_values;
 	bool m_grounded = false;
@@ -811,17 +862,43 @@ private:
 } // namespace
 
 void ConstantPropagationPass(const BlockList& blocks, uint32_t wave_size) {
-	std::unordered_set<Inst*> lowered_ancillary;
+	FoldWorklist worklist;
-	LaneMaskProjection mask_projection(lowered_ancillary);
+	LaneMaskProjection mask_projection(worklist);
-	for (auto* block: blocks) {
+	for (auto* block: blocks)
-		for (auto inst = block->begin(); inst != block->end(); ++inst) {
+		for (auto inst = block->begin(); inst != block->end(); ++inst) worklist.Register(inst);
+	std::vector<Use> users;
+	std::vector<Value> operands;
+	const auto drain = [&] {
+		for (size_t i = 0; i < worklist.pending.size(); ++i) {
+			auto* inst = worklist.pending[i];
+			auto& location = worklist.locations.at(inst);
+			location.queued = false;
+			const auto opcode = inst->GetOpcode();
+			if (opcode == ValueOpcode::Identity || opcode == ValueOpcode::Void) continue;
+			const auto where = location.instruction;
+			users.assign(inst->Uses().begin(), inst->Uses().end());
+			operands.clear();
+			for (size_t arg = 0; arg < inst->NumArgs(); ++arg) operands.push_back(inst->Arg(arg));
 			mask_projection.Fold(*inst, wave_size);
-			FoldInstruction(*block, inst, lowered_ancillary);
+			FoldInstruction(*inst->Parent(), where, worklist);
+			bool changed = opcode != inst->GetOpcode() || operands.size() != inst->NumArgs();
+			for (size_t arg = 0; !changed && arg < operands.size(); ++arg)
+				changed = operands[arg] != inst->Arg(arg);
+			if (changed) {
+				worklist.QueueUsers(users);
+				worklist.Queue(inst);
 			}
 		}
+		worklist.pending.clear();
+	};
+	drain();
 	// Normalize retained PHI/select values only after every supported field read has
 	// been lowered; direct raw consumers remain unsupported.
-	for (auto* source: lowered_ancillary) {
+	bool normalized;
+	do {
+		normalized = false;
+		for (auto* source: worklist.lowered_ancillary) {
+			if (source->GetOpcode() == ValueOpcode::Identity) continue;
 			const bool retained_only = std::ranges::all_of(source->Uses(), [](const Use& use) {
 				const auto& user = *use.user;
 				if (!user.HasUses() && !user.MayHaveSideEffects()) {
@@ -831,9 +908,13 @@ void ConstantPropagationPass(const BlockList& blocks, uint32_t wave_size) {
 				       (user.GetOpcode() == ValueOpcode::SelectU32 && use.operand == 2u);
 			});
 			if (retained_only) {
+				worklist.QueueUsers(source->Uses());
 				Replace(*source, Value(0u));
+				normalized = true;
 			}
 		}
+		drain();
+	} while (normalized);
 }
 
 } // namespace Libs::Graphics::ShaderRecompiler::IR

tests/shaderCfgTests.cpp

@@ -12534,6 +12534,86 @@ void TestLaneMaskProjection() {
   }
 }
 
+void TestLateConstantPropagation() {
+  using namespace ShaderRecompiler::IR;
+  for (const uint32_t wave_size : {32u, 64u}) {
+    for (const bool reverse : {false, true}) {
+      Program program;
+      program.stage = ShaderType::Compute;
+      program.wave_size = wave_size;
+      for (uint32_t id = 0; id < 4; ++id) {
+        program.block_storage.push_back(std::make_unique<Block>());
+        program.blocks.push_back(program.block_storage.back().get());
+        program.blocks.back()->id = id;
+      }
+      auto *entry = program.blocks[0];
+      auto *header = program.blocks[1];
+      auto *latch = program.blocks[2];
+      auto *exit = program.blocks[3];
+      entry->AddBranch(header);
+      header->AddBranch(latch);
+      header->AddBranch(exit);
+      latch->AddBranch(header);
+      entry->terminator.kind = ShaderRecompiler::CFG::TerminatorKind::Branch;
+      entry->terminator.true_block = header;
+      header->terminator.kind = ShaderRecompiler::CFG::TerminatorKind::ConditionalBranch;
+      header->terminator.true_block = latch;
+      header->terminator.false_block = exit;
+      latch->terminator.kind = ShaderRecompiler::CFG::TerminatorKind::Branch;
+      latch->terminator.true_block = header;
+      IREmitter e(entry);
+      const auto raw = e.GetUserData(static_cast<ScalarReg>(0));
+      const auto active = e.INotEqual(raw, U32(Value(0u)));
+      header->condition = active;
+      auto &predicate = header->AppendNewInst(ValueOpcode::Phi);
+      predicate.SetFlags(Type::U1);
+      auto &mask = header->AppendNewInst(ValueOpcode::Phi);
+      mask.SetFlags(Type::U32);
+      const auto thread_bit = [](IREmitter &ir, U32 word) {
+        const auto lane = ir.BitwiseAnd(U32(ir.Emit(ValueOpcode::LaneId)), U32(Value(31u)));
+        return ir.INotEqual(ir.BitwiseAnd(ir.ShiftRightLogical(word, lane), U32(Value(1u))),
+                           U32(Value(0u)));
+      };
+      IREmitter body(header);
+      const auto live = thread_bit(body, U32(Value(&mask)));
+      const auto selected = body.Select(live, U32(Value(42u)), raw);
+      const auto projected = body.Emit(ValueOpcode::SelectU1,
+                                      {Value(&predicate), live, Value(false)});
+      const auto absorbed = body.LogicalAnd(active,
+          body.LogicalOr(body.LogicalNot(active), U1(Value(&predicate))));
+      IREmitter tail(latch);
+      const auto full = thread_bit(tail, U32(Value(UINT32_MAX)));
+      const auto full_word = tail.Select(full, U32(Value(UINT32_MAX)), U32(Value(0u)));
+      predicate.AddPhiOperand(entry, Value(true));
+      predicate.AddPhiOperand(latch, full);
+      mask.AddPhiOperand(entry, Value(UINT32_MAX));
+      mask.AddPhiOperand(latch, full_word);
+      exit->condition = projected;
+      if (reverse) std::ranges::reverse(program.blocks);
+
+      ConstantPropagationPass(program.blocks, wave_size);
+      if (reverse) std::ranges::reverse(program.blocks);
+      Check(Value(&predicate).Resolve() == Value(true) &&
+                Value(&mask).Resolve() == Value(UINT32_MAX) && live.Resolve() == Value(true) &&
+                selected.Resolve() == Value(42u) && projected.Resolve() == Value(true) &&
+                absorbed.Resolve() == active.Resolve(),
+            "late constants did not reach Phi, nested mask patterns or new logical nodes");
+      ResolveControlFlowIdentities(program);
+      Check(exit->condition == Value(true), "late constant lost its control-flow use");
+      RemoveIdentities(program.blocks);
+      ValidateProgram(program, true);
+    }
+  }
+
+  Block cycle;
+  auto &phi = cycle.AppendNewInst(ValueOpcode::Phi);
+  phi.SetFlags(Type::U32);
+  phi.AddPhiOperand(&cycle, Value(&phi));
+  ConstantPropagationPass({&cycle});
+  Check(phi.GetOpcode() == ValueOpcode::Phi && phi.Arg(0).TryInstruction() == &phi,
+        "an ungrounded Phi cycle acquired a constant");
+}
+
 void TestU64ShiftConstantPropagation() {
   using namespace ShaderRecompiler::IR;
 
@@ -15722,6 +15802,7 @@ int main() {
 #endif
   TestValuePhiValidation();
   TestLaneMaskProjection();
+  TestLateConstantPropagation();
   TestU64ShiftConstantPropagation();
   TestU64RegisterPairRoundTrip();
 #if KYTY_PLATFORM != KYTY_PLATFORM_WINDOWS