shader: handle reduced guards for image loops

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

@@ -1080,33 +1080,42 @@ private:
 		}
 		if (!induction) return {};
 
-		const auto contains = [&](auto&& self, Value value, const Inst* comparison) -> bool {
+		// Native branches advance the whole guest wave, so the induction and runtime
-			value = value.Resolve();
+		// bound are uniform. A reduced condition witnesses a lane with its edge
-			if (value.TryInstruction() == comparison) return true;
+		// polarity; implication of the comparison therefore bounds the whole wave.
-			const auto* inst = value.TryInstruction();
+		const auto implies = [](auto&& self, Value value, bool positive,
-			return inst != nullptr && inst->GetOpcode() == ValueOpcode::LogicalAnd &&
+		                        const Inst* comparison) -> bool {
-			       (self(self, inst->Arg(0), comparison) || self(self, inst->Arg(1), comparison));
+			const auto* inst = value.Resolve().TryInstruction();
+			if (inst == comparison) return positive;
+			if (inst == nullptr) return false;
+			if (inst->GetOpcode() == ValueOpcode::ConditionRef)
+				return self(self, inst->Arg(0), positive, comparison);
+			if (inst->GetOpcode() == ValueOpcode::LogicalNot)
+				return self(self, inst->Arg(0), !positive, comparison);
+			return positive && inst->GetOpcode() == ValueOpcode::LogicalAnd &&
+			       (self(self, inst->Arg(0), true, comparison) ||
+			        self(self, inst->Arg(1), true, comparison));
 		};
 		for (const auto& use_of_key: phi->Uses()) {
 			const auto* compare = use_of_key.user;
 			if (compare->GetOpcode() != ValueOpcode::SLessThan32 || use_of_key.operand != 0u ||
 			    !ValidateRuntimeValue(m_program, compare->Arg(1))) continue;
-			for (uint32_t i = 0; i < m_program.block_info.size(); ++i) {
+			const auto position = std::ranges::find(m_program.blocks, compare->Parent());
-				const auto& info = m_program.block_info[i];
+			if (position == m_program.blocks.end()) continue;
-				const auto* negated = info.condition.Resolve().TryInstruction();
+			const auto* block = *position;
-				if (info.terminator.kind != CFG::TerminatorKind::ConditionalBranch ||
+			const auto& info = m_program.block_info[position - m_program.blocks.begin()];
-				    compare->Parent() != m_program.blocks[i] || negated == nullptr ||
+			if (info.terminator.kind != CFG::TerminatorKind::ConditionalBranch || use == block ||
-				    negated->GetOpcode() != ValueOpcode::LogicalNot ||
+			    CanReach(m_program.blocks.front(), use, block) ||
-				    !contains(contains, negated->Arg(0), compare) ||
+			    CanReach(phi->Parent(), use, block)) continue;
-				    use == m_program.blocks[i] ||
+			const auto* yes = FindBlock(info.terminator.true_block);
-				    CanReach(m_program.blocks.front(), use, m_program.blocks[i]) ||
+			const auto* no = FindBlock(info.terminator.false_block);
-				    CanReach(phi->Parent(), use, m_program.blocks[i]) ||
+			const bool yes_reaches = CanReach(yes, use, phi->Parent());
-				    CanReach(FindBlock(info.terminator.true_block), use, phi->Parent()) ||
+			const bool no_reaches = CanReach(no, use, phi->Parent());
-				    !CanReach(FindBlock(info.terminator.false_block), use, phi->Parent())) {
+			if (yes_reaches == no_reaches) continue;
-					continue;
+			Value condition;
-				}
+			bool positive = false;
-				return compare->Arg(1);
+			if (ConditionalEdge(block, yes_reaches ? yes : no, condition, positive) &&
-			}
+			    implies(implies, condition, positive, compare)) return compare->Arg(1);
 		}
 		return {};
 	}

tests/ResourceTrackingTests.cpp

@@ -741,10 +741,12 @@ void TestGuardedDirectImageTable() {
 void TestBoundedComputeImageLoop() {
   namespace CFG = Libs::Graphics::ShaderRecompiler::CFG;
   enum class Variant {
-    Bounded, WrongGuard, EntryBypass, ExitBypass, GuardBlock, WrongStep
+    Bounded, Plain, Nonzero, TrueEdge, WrongGuard, EntryBypass, ExitBypass, GuardBlock,
+    WrongStep, DivergentBound, DivergentKey, Disjunction, WrongPolarity
   };
   const auto make_plan = [](Variant variant) {
     Fixture fixture;
+    fixture.program.wave_size = 64u;
     auto *entry = fixture.block;
     auto *header = fixture.AddBlock();
     auto *body = fixture.AddBlock();
@@ -763,7 +765,8 @@ void TestBoundedComputeImageLoop() {
         .true_block = 1u, .false_block = 2u};
     fixture.program.block_info[1].terminator = {
         .kind = CFG::TerminatorKind::ConditionalBranch,
-        .true_block = 4u, .false_block = 2u};
+        .true_block = variant == Variant::TrueEdge ? 2u : 4u,
+        .false_block = variant == Variant::TrueEdge ? 4u : 2u};
     fixture.program.block_info[2].terminator = {
         .kind = CFG::TerminatorKind::Branch, .true_block = 3u};
     fixture.program.block_info[3].terminator = {
@@ -776,19 +779,34 @@ void TestBoundedComputeImageLoop() {
     auto &phi = header->AppendNewInst(ValueOpcode::Phi, {},
                                       static_cast<uint64_t>(Type::U32));
     const auto key = Value(&phi);
-    const auto count = fixture.UserData(2);
+    const auto local = fixture.Emit(
+        ValueOpcode::GetBuiltin,
+        {Value(static_cast<uint32_t>(StageInputKind::LocalInvocationId)), Value(0u)});
+    const auto count = variant == Variant::DivergentBound ? local : fixture.UserData(2);
     const auto in_range = fixture.Emit(ValueOpcode::SLessThan32,
                                        {variant == Variant::WrongGuard
                                             ? Value(0u) : key,
                                         count}, 0, header);
-    const auto allowed = fixture.Emit(ValueOpcode::LogicalAnd,
+    const auto active = fixture.Emit(ValueOpcode::INotEqual32,
-                                      {in_range, Value(true)}, 0, header);
+                                     {local, Value(0u)}, 0, header);
-    fixture.program.block_info[1].condition =
+    const auto allowed = fixture.Emit(variant == Variant::Disjunction
-        fixture.Emit(ValueOpcode::LogicalNot, {allowed}, 0, header);
+                                         ? ValueOpcode::LogicalOr : ValueOpcode::LogicalAnd,
+                                      {in_range, active}, 0, header);
+    const bool nonzero = variant == Variant::Nonzero || variant == Variant::TrueEdge;
+    auto condition = allowed;
+    if (!nonzero)
+      condition = fixture.Emit(ValueOpcode::LogicalNot, {condition}, 0, header);
+    if (variant != Variant::Plain)
+      condition = fixture.Emit(ValueOpcode::ConditionRef, {condition},
+          nonzero ? CFG::BranchCondition::ExecNonZero
+                  : CFG::BranchCondition::ExecZero, header);
+    if (variant == Variant::Nonzero || variant == Variant::WrongPolarity)
+      condition = fixture.Emit(ValueOpcode::LogicalNot, {condition}, 0, header);
+    fixture.program.block_info[1].condition = condition;
     const auto step = fixture.Emit(ValueOpcode::IAdd32,
                                    {key, Value(variant == Variant::WrongStep ? 2u : 1u)},
                                    0, latch);
-    phi.AddPhiOperand(entry, Value(0u));
+    phi.AddPhiOperand(entry, variant == Variant::DivergentKey ? local : Value(0u));
     phi.AddPhiOperand(latch, step);
 
     fixture.block = variant == Variant::GuardBlock ? header : body;
@@ -825,6 +843,9 @@ void TestBoundedComputeImageLoop() {
   };
 
   auto plan = make_plan(Variant::Bounded);
+  make_plan(Variant::Plain);
+  make_plan(Variant::Nonzero);
+  make_plan(Variant::TrueEdge);
   CheckFatal([&] { make_plan(Variant::WrongGuard); },
              "not a valid runtime value",
              "compute image loop accepted an unrelated guard");
@@ -840,6 +861,11 @@ void TestBoundedComputeImageLoop() {
   CheckFatal([&] { make_plan(Variant::WrongStep); },
              "not a valid runtime value",
              "compute image loop accepted a two-step induction");
+  for (const auto variant : {Variant::DivergentBound, Variant::DivergentKey,
+                             Variant::Disjunction, Variant::WrongPolarity}) {
+    CheckFatal([&] { make_plan(variant); }, "not a valid runtime value",
+               "compute image loop accepted a guard without a uniform bound");
+  }
 
   LinearTestMemory memory;
   const auto table = 0x1800u + 0x6b0u;