shader: prove narrowing loop masks through shared implication analysis

Co-authored-by: fxpw <fxpw@outlook.com>

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

@@ -1264,7 +1264,7 @@ private:
 		return (invariant.IsEmpty() ? guard : invariant).Resolve();
 	}
 
-	bool Implies(Value guard, Value required) const {
+	bool Implies(Value guard, Value required, bool allow_phi = true) const {
 		guard = SimplifyGuard(guard);
 		required = SimplifyGuard(required);
 		if (required == Value(true) || EquivalentValue(m_program, guard, required)) return true;
@@ -1277,9 +1277,20 @@ private:
 		const auto* inst = guard.TryInstruction();
 		if (inst == nullptr) return false;
 		if (inst->GetOpcode() == ValueOpcode::LogicalAnd)
-			return Implies(inst->Arg(0), required) || Implies(inst->Arg(1), required);
+			return Implies(inst->Arg(0), required, allow_phi) || Implies(inst->Arg(1), required, allow_phi);
-		return inst->GetOpcode() == ValueOpcode::LogicalOr &&
+		if (inst->GetOpcode() == ValueOpcode::LogicalOr)
-		       Implies(inst->Arg(0), required) && Implies(inst->Arg(1), required);
+			return Implies(inst->Arg(0), required, allow_phi) && Implies(inst->Arg(1), required, allow_phi);
+		if (!allow_phi || inst->GetOpcode() != ValueOpcode::Phi || inst->GetType() != Type::U1 ||
+		    test == nullptr || test->Parent() == inst->Parent() || !AvailableBefore(*test, *inst)) return false;
+		// The required predicate must be invariant across this loop. Every entry
+		// establishes it; carried masks may only remove lanes from this same Phi.
+		bool entry = false;
+		for (size_t arm = 0; arm < inst->NumArgs(); ++arm) {
+			const bool carried = Dominates(inst->Parent(), inst->PhiBlock(arm));
+			if (!Implies(inst->Arg(arm), carried ? guard : required, false)) return false;
+			entry |= !carried;
+		}
+		return entry;
 	}
 
 	Value EqualLocalValue(Value guard, const auto& matches) const {
@@ -1502,17 +1513,6 @@ private:
 			if (executing.IsEmpty() ? !HasActiveLane(guard, first->Parent())
 			                        : !Implies(executing, guard)) return nullptr;
 			local = first->Arg(0).Resolve();
-			const auto* phi = guard.TryInstruction();
-			if (phi != nullptr && phi->GetOpcode() == ValueOpcode::Phi &&
-			    phi->GetType() == Type::U1 && phi->NumArgs() == 2u) {
-				for (uint32_t initial = 0; initial < 2u; ++initial) {
-					if (Implies(phi->Arg(initial ^ 1u), guard)) {
-						// The backedge only removes lanes from the initial mask.
-						guard = phi->Arg(initial).Resolve();
-						break;
-					}
-				}
-			}
 		} else {
 			guard = PositiveLaneWitness(use.Parent());
 			if (guard.IsEmpty()) return nullptr;
@@ -1763,12 +1763,11 @@ private:
 		return false;
 	}
 
-	bool AvailableBefore(const Inst& initial, const Inst& use, Value witness = {}) const {
+	bool Dominates(const Block* initial, const Block* use, Value witness = {}) const {
-		if (initial.Parent() == use.Parent()) return BeforeInBlock(initial, use);
+		std::vector<const Block*> pending {use};
-		std::vector<const Block*> pending {use.Parent()};
 		for (size_t index = 0; index < pending.size(); ++index) {
 			const auto* current = pending[index];
-			if (current == initial.Parent()) continue;
+			if (current == initial) continue;
 			if (current->ImmPredecessors().empty()) return false;
 			for (const auto* previous: current->ImmPredecessors()) {
 				if (!witness.IsEmpty()) {
@@ -1783,6 +1782,11 @@ private:
 		return true;
 	}
 
+	bool AvailableBefore(const Inst& initial, const Inst& use, Value witness = {}) const {
+		return initial.Parent() == use.Parent() ? BeforeInBlock(initial, use)
+		                                      : Dominates(initial.Parent(), use.Parent(), witness);
+	}
+
 	std::optional<std::array<uint64_t, 5>> LocalIndex(Value value, Value active,
 	                                                Value& induction, bool allow_induction) const {
 		return AffineIndex<5>(value, [&](Value at) { return ActiveValue(at, active); },

tests/ResourceTrackingTests.cpp

@@ -5387,7 +5387,8 @@ void TestGuardedWaterfallBufferRecords() {
   using Libs::Graphics::DstSel;
   enum class Guard {
     Exact, Reversed, Greater, GreaterEqual, WrongGreater, SignedGreater,
-    NoRange, NoEquality, OtherLocal, UnguardedUse, SeparateVotes
+    NoRange, NoEquality, OtherLocal, UnguardedUse, SeparateVotes,
+    Loop, LoopReversed, LoopEntries, LoopBadEntry, LoopWiden, LoopUnanchored, LoopVarying
   };
   const auto make = [](Guard guard) {
     auto f = std::make_unique<Fixture>();
@@ -5398,7 +5399,6 @@ void TestGuardedWaterfallBufferRecords() {
     const auto local = f->Emit(ValueOpcode::LoadBufferU32,
         {input, lane, Value(0u), Value(0u), Value(true)},
         f->AddMemory({.kind = ResourceKind::Buffer, .idxen = true}, 0x200u));
-    const auto key = f->Emit(ValueOpcode::ReadLane, {local, Value(0u)});
     const auto initial = f->Emit(ValueOpcode::ULessThan32, {lane, Value(32u)});
     const auto range_local = guard == Guard::OtherLocal
         ? f->Emit(ValueOpcode::IAdd32, {local, Value(1u)}) : local;
@@ -5413,11 +5413,48 @@ void TestGuardedWaterfallBufferRecords() {
       range = f->Emit(ValueOpcode::LogicalAnd, {initial, compare});
     } else {
       // Match ANDN1_SAVEEXEC, including its initial-EXEC witness.
-      const auto excluded = f->Emit(ValueOpcode::LogicalAnd,
+      auto executing = initial;
-          {initial, f->Emit(ValueOpcode::ULessThanEqual32, {count, range_local})});
+      auto comparison_mask = initial;
-      range = f->Emit(ValueOpcode::LogicalAnd,
+      if (guard >= Guard::Loop) {
-          {initial, f->Emit(ValueOpcode::LogicalNot, {excluded})});
+        auto* entry = f->block;
+        auto* header = f->AddBlock();
+        f->block = header;
+        auto& mask = header->AppendNewInst(ValueOpcode::Phi, {}, uint64_t(Type::U1));
+        if (guard != Guard::LoopUnanchored) entry->AddBranch(header);
+        header->AddBranch(header);
+        if (guard == Guard::LoopVarying) {
+          auto& predicate = header->AppendNewInst(ValueOpcode::Phi, {}, uint64_t(Type::U1));
+          predicate.AddPhiOperand(entry, initial);
+          predicate.AddPhiOperand(header, f->Emit(ValueOpcode::LogicalNot, {Value(&predicate)}));
+          comparison_mask = Value(&predicate);
         }
+        const auto next = guard == Guard::LoopWiden
+            ? f->Emit(ValueOpcode::LogicalNot, {initial})
+            : guard == Guard::LoopVarying ? comparison_mask
+            : f->Emit(ValueOpcode::ULessThan32, {lane, Value(16u)});
+        const auto carried = f->Emit(guard == Guard::LoopWiden
+            ? ValueOpcode::LogicalOr : ValueOpcode::LogicalAnd, {Value(&mask), next});
+        if (guard == Guard::LoopReversed) mask.AddPhiOperand(header, carried);
+        if (guard != Guard::LoopUnanchored) mask.AddPhiOperand(entry, initial);
+        if (guard != Guard::LoopReversed) mask.AddPhiOperand(header, carried);
+        if (guard == Guard::LoopEntries || guard == Guard::LoopBadEntry ||
+            guard == Guard::LoopUnanchored) {
+          auto* alternate = f->AddBlock();
+          (guard == Guard::LoopUnanchored ? header : entry)->AddBranch(alternate);
+          alternate->AddBranch(header);
+          const auto incoming = guard == Guard::LoopUnanchored
+              ? f->Emit(ValueOpcode::LogicalAnd, {Value(&mask), initial}, 0, alternate)
+              : guard == Guard::LoopBadEntry ? Value(true) : initial;
+          mask.AddPhiOperand(alternate, incoming);
+        }
+        executing = Value(&mask);
+      }
+      const auto excluded = f->Emit(ValueOpcode::LogicalAnd,
+          {comparison_mask, f->Emit(ValueOpcode::ULessThanEqual32, {count, range_local})});
+      range = f->Emit(ValueOpcode::LogicalAnd,
+          {executing, f->Emit(ValueOpcode::LogicalNot, {excluded})});
+    }
+    const auto key = f->Emit(ValueOpcode::ReadLane, {local, Value(0u)});
     const auto equal = f->Emit(ValueOpcode::IEqual32,
         guard == Guard::Reversed ? std::initializer_list<Value>{local, key}
                                  : std::initializer_list<Value>{key, local});
@@ -5466,14 +5503,17 @@ void TestGuardedWaterfallBufferRecords() {
   ResourceSpecialization specialization;
   for (const auto guard : {Guard::Exact, Guard::Reversed, Guard::Greater, Guard::GreaterEqual,
                           Guard::WrongGreater, Guard::SignedGreater, Guard::NoRange, Guard::NoEquality,
-                          Guard::OtherLocal, Guard::UnguardedUse, Guard::SeparateVotes}) {
+                          Guard::OtherLocal, Guard::UnguardedUse, Guard::SeparateVotes,
+                          Guard::Loop, Guard::LoopReversed, Guard::LoopEntries, Guard::LoopBadEntry,
+                          Guard::LoopWiden, Guard::LoopUnanchored, Guard::LoopVarying}) {
     auto fixture = make(guard);
     const auto plan = ExtractResourcePlan(fixture->program);
     const auto root = std::ranges::find_if(plan.info.buffers, [&](const auto& buffer) {
       return plan.descriptor_sources[buffer.source].indirect_descriptor.has_value();
     });
     Check(root != plan.info.buffers.end(), "waterfall buffer lost its indirect resource");
-    const bool proven = guard == Guard::Exact || guard == Guard::Reversed || guard == Guard::Greater;
+    const bool proven = guard == Guard::Exact || guard == Guard::Reversed || guard == Guard::Greater ||
+        guard == Guard::Loop || guard == Guard::LoopReversed || guard == Guard::LoopEntries;
     const auto& domain = *plan.descriptor_sources[root->source].indirect_descriptor;
     Check(domain.key_count.IsEmpty() != proven && !domain.selector &&
               domain.table_stride == 488u && domain.table_immediate == 216u,

tests/ShaderRecompilerComputeTests.cpp

@@ -37987,16 +37987,19 @@ void CheckGuardedWaterfallBufferRecords(VulkanHarness &vulkan) {
                             .limits.minStorageBufferOffsetAlignment;
   const u32 slot = (4096u + alignment - 1u) / alignment * alignment;
   constexpr u32 selectors = 256u, table = 512u;
-  for (const bool direct_gt : {false, true})
+  enum class Guard { AndNot, DirectGt, NarrowingLoop };
+  for (const auto guard : {Guard::AndNot, Guard::DirectGt, Guard::NarrowingLoop})
   for (const auto [wave, upper_only] : {std::pair{32u, false}, std::pair{64u, false},
                                       std::pair{64u, true}}) {
+    const bool direct_gt = guard == Guard::DirectGt;
+    const bool outer_loop = guard == Guard::NarrowingLoop;
+    const u32 iterations = outer_loop ? 2u : 1u;
+    const std::string name = std::string(direct_gt ? "GuardedWaterfallGtRecords" : outer_loop
+        ? "GuardedWaterfallLoopRecords" : "GuardedWaterfallRecords") +
+        (upper_only ? "Upper64" : std::to_string(wave));
     TestCase test;
-    test.name = direct_gt
+    test.name = name.c_str();
-        ? upper_only ? "GuardedWaterfallGtRecordsUpper64" : wave == 32u
+    test.initial.assign((slot * 3u + wave * iterations * 16u + 16u) / 4u, 0xdeadbeefu);
-            ? "GuardedWaterfallGtRecords32" : "GuardedWaterfallGtRecords64"
-        : upper_only ? "GuardedWaterfallRecordsUpper64" : wave == 32u
-            ? "GuardedWaterfallRecords32" : "GuardedWaterfallRecords64";
-    test.initial.assign((slot * 3u + wave * 16u + 16u) / 4u, 0xdeadbeefu);
     std::fill_n(test.initial.begin() + table / 4u, 976u / 4u, 0u);
     constexpr std::array poison{0x012b6890u, 0x501502f9u, 0x493ba30eu, 0x3ea2f352u};
     std::copy(poison.begin(), poison.end(), test.initial.begin() + (table + 480u) / 4u);
@@ -38004,7 +38007,7 @@ void CheckGuardedWaterfallBufferRecords(VulkanHarness &vulkan) {
     const std::array<u32, 8> descriptors{table, 488u << 16u, 2u, 0x5204u,
                                         selectors, 4u << 16u, wave, 0x5204u};
     std::copy(descriptors.begin(), descriptors.end(), test.user_data.begin());
-    const std::array output{slot * 3u, 16u << 16u, wave, 0x5204u};
+    const std::array output{slot * 3u, 16u << 16u, wave * iterations, 0x5204u};
     std::copy(output.begin(), output.end(), test.user_data.begin() + 16u);
     for (u32 row = 0; row < 2u; ++row) {
       const std::array descriptor{slot * (row + 1u), 12u << 16u, wave,
@@ -38027,9 +38030,20 @@ void CheckGuardedWaterfallBufferRecords(VulkanHarness &vulkan) {
       AppendSMovLiteral(&code, 126, 0u);
       AppendSMovLiteral(&code, 127, UINT32_MAX);
     }
+    if (outer_loop) code.insert(code.end(), {EncodeVop2(0x1b, 10, InlineU32(8), 30),
+                                           EncodeVopc(0xd2, InlineU32(0), 10)});
     code.insert(code.end(), {EncodeSop1(0x04, 20, 126),
         EncodeMubuf0(0x0c, 0, true, false), EncodeMubuf1(5, 1, 30), EncodeSopp(0x0c, 0)});
+    if (outer_loop) {
+      // Save the initial-EXEC comparison before the loop, as in native 9d4c.
+      code.insert(code.end(), {EncodeVopc(0xc3, 2, 5), EncodeSop1(0x04, 52, 106)});
+      AppendSMovLiteral(&code, 36, 0u);
+    }
+    const auto outer = code.size();
     if (direct_gt) code.push_back(EncodeVopc(0xd4, 2, 5)); // Keep N > local key in EXEC.
+    else if (outer_loop) code.insert(code.end(), {
+        EncodeSop2(0x1e, 37, 36, InlineU32(std::countr_zero(wave))),
+        EncodeVop2(0x25, 31, 37, 30), EncodeSop1(0x37, 50, 52)});
     else code.insert(code.end(), {EncodeVopc(0xc3, 2, 5),
                                  EncodeSop1(0x37, 52, 106)}); // Exclude N <= local key.
     const auto empty = code.size();
@@ -38044,16 +38058,27 @@ void CheckGuardedWaterfallBufferRecords(VulkanHarness &vulkan) {
     code.insert(code.end(), {EncodeSop2(0x26, 33, 32, 255), 488u,
         EncodeSmem0(0x0a, 12, 0), EncodeSmem1(216u, 33), EncodeSopp(0x0c, 0),
         EncodeMubuf0(0x02, 0, true, false, true), EncodeMubuf1(6, 3, 30),
-        EncodeSopp(0x0c, 0), EncodeMubuf0(0x1f, 0, true, false), EncodeMubuf1(6, 4, 30)});
+        EncodeSopp(0x0c, 0), EncodeMubuf0(0x1f, 0, true, false),
+        EncodeMubuf1(6, 4, outer_loop ? 31u : 30u)});
     code[skip] = EncodeSopp(0x08, code.size() - skip - 1u);
     code.insert(code.end(), {EncodeSop2(0x15, 54, 54, 106), EncodeSop1(0x04, 126, 58)});
     code.push_back(EncodeSopp(0x05, static_cast<int16_t>(loop - code.size() - 1u)));
     code[empty] = EncodeSopp(0x08, code.size() - empty - 1u);
+    if (outer_loop) {
+      code.push_back(EncodeSop1(0x04, 126, 50));
+      AppendVMovLiteral(&code, 9, 0x13579bdfu);
+      code.insert(code.end(), {EncodeMubuf0(0x1c, 12u, true, false), EncodeMubuf1(9, 4, 31),
+          EncodeVop2(0x1b, 10, InlineU32(4), 30), EncodeVopc(0xd2, InlineU32(0), 10),
+          EncodeSop2(0x00, 36, 36, InlineU32(1)), EncodeSopc(0x0a, 36, InlineU32(iterations))});
+      code.push_back(EncodeSopp(0x05, static_cast<int16_t>(outer - code.size() - 1u)));
+    }
     code.push_back(EncodeSop1(0x04, 126, 20));
     AppendVMovLiteral(&code, 9, 0x76543210u);
     code.insert(code.end(), {EncodeMubuf0(0x1c, 12u, true, false), EncodeMubuf1(9, 4, 30)});
     AppendEnd(&code);
     test.decoded_counts = {{direct_gt ? "V_CMPX_GT_U32" : "V_CMP_LE_U32", 1u}};
+    if (outer_loop) test.decoded_counts.insert(test.decoded_counts.end(),
+        {{"V_CMPX_EQ_U32", 2u}, {"S_CBRANCH_SCC1", 2u}});
     test.required_spirv = {"OpLoopMerge", "OpGroupNonUniformShuffle", "Volatile"};
     test.forbidden_spirv = {"load_indirect_buffer_", "indirect_buffer_format_geometry"};
     const auto compiled = CompileCase(test, vulkan.SubgroupSize());
@@ -38080,18 +38105,23 @@ void CheckGuardedWaterfallBufferRecords(VulkanHarness &vulkan) {
       AppendBufferStoreDword(&producer.code, 1, 2);
       test.expected[selectors / 4u + lane] = key;
       if (upper_only && lane < 32u) continue;
-      const u32 out = slot * 3u / 4u + lane * 4u;
+      if (outer_loop && (lane & 8u) != 0u) continue;
+      for (u32 iteration = 0; iteration < iterations; ++iteration) {
+        if (iteration != 0u && (lane & 4u) != 0u) continue;
+        const u32 out = slot * 3u / 4u + (iteration * wave + lane) * 4u;
         if (key < 2u)
           std::copy_n(test.initial.begin() + slot * (key + 1u) / 4u + lane * 3u, 3u,
                       test.expected.begin() + out);
-      test.expected[out + 3u] = 0x76543210u;
+        if (outer_loop) test.expected[out + 3u] = 0x13579bdfu;
+      }
+      test.expected[slot * 3u / 4u + lane * 4u + 3u] = 0x76543210u;
     }
     AppendEnd(&producer.code);
     const auto producer_shader = CompileCase(producer, vulkan.SubgroupSize());
     auto buffer = vulkan.CreateStorageBuffer(test.name, test.initial, test.initial.size());
     vulkan.Dispatch(producer, producer_shader, buffer);
     vulkan.Dispatch(test, compiled, buffer);
-    CompareWords(test, "same-lane bound, key == N, wrapped poison, EXEC restoration and canaries",
+    CompareWords(test, "same-lane bound, narrowing loop, key == N, poison, EXEC restoration and canaries",
                  test.expected, vulkan.ReadBuffer(test.name, buffer, test.expected.size()));
     vulkan.DestroyBuffer(&buffer);
   }