fix(recompiler): exempt proven control conditions and own-bit ballots in stencil replay

MaskReplaySafety counted every use reaching a conditional branch condition as a quad observation, including EXEC/VCC guards that control() itself already proves replay-safe, and every Ballot regardless of how its bits are read. Nested optional regions (a merged phi re-used as the predicate of an inner execz) and ballots consumed only through the own-lane funnel (word extract, lane<32 select, lane&31 shift, &1 isolate) were refused even though neither crosses lane boundaries; the same reasoning exempts explicit-lod samples, which do not compute implicit quad derivatives.

The check now runs conditional branches through a memoized controlProven() with a visiting set, exempts those condition uses in quadObserved, walks past own-bit ballots into the isolated per-lane bools, and treats explicit-lod samples as lane-local in both quadObserved and helperObserved. optionalBody now uses a local region set so nested control proofs do not corrupt the outer walk. Anything not matching these shapes is refused as before.

Generated with [Devin](https://devin.ai)

Co-Authored-By: Devin <158243242+devin-ai-integration[bot]@users.noreply.github.com>

core/libs/prx/libSceAgcDriver/tests/MaskedSelect.cpp

@@ -1295,6 +1295,114 @@ int main() {
     passed &= run("an optional wave body spanning several blocks preserves inactive values through its own phis", [&] { return optionalPixelRegion(false); });
     passed &= run("an optional wave body spanning several blocks cannot write an inactive lane on any of its paths", [&] { return !optionalPixelRegion(true); });
 
+    const auto nestedOptionalGuard = [&](bool unsafeInner, bool ballotEscape) {
+        Builder b;
+        b.program.Resources().stage = IrShaderStage::Pixel;
+        auto& entry = *b.block;
+        auto& predicate = pixelPredicate(b);
+        auto& condition = b.emit(IrOpcode::LogicalNot, IrType::Bool, {&predicate});
+        auto& base = helperPredicate(b, false);
+        auto& body = b.newBlock();
+        auto& merge = b.newBlock();
+        auto& inner = b.newBlock();
+        auto& join = b.newBlock();
+        entry.AddBranch(&body);
+        entry.AddBranch(&merge);
+        body.AddBranch(&merge);
+        merge.AddBranch(&inner);
+        merge.AddBranch(&join);
+        inner.AddBranch(&join);
+        auto& entryInfo = b.program.Metadata().blockInfo[0u];
+        entryInfo.terminator.kind = TerminatorKind::ConditionalBranch;
+        entryInfo.terminator.condition = BranchCondition::ExecZero;
+        entryInfo.terminator.trueBlock = merge.Id();
+        entryInfo.terminator.falseBlock = body.Id();
+        entryInfo.condition = &condition;
+        b.program.Metadata().blockInfo[1u].terminator.kind = TerminatorKind::Branch;
+        b.program.Metadata().blockInfo[1u].terminator.trueBlock = merge.Id();
+        b.block = &body;
+        auto& taken = b.emit(IrOpcode::SelectU1, IrType::Bool, {&predicate, &b.truth(), &base});
+        b.block = &merge;
+        auto& merged = b.phi(IrType::Bool, entry, base);
+        merged.AddPhiOperand(&body, &taken);
+        auto& innerCondition = b.emit(IrOpcode::LogicalNot, IrType::Bool, {&merged});
+        auto& mergeInfo = b.program.Metadata().blockInfo[2u];
+        mergeInfo.terminator.kind = TerminatorKind::ConditionalBranch;
+        mergeInfo.terminator.condition = BranchCondition::ExecZero;
+        mergeInfo.terminator.trueBlock = join.Id();
+        mergeInfo.terminator.falseBlock = inner.Id();
+        mergeInfo.condition = &innerCondition;
+        b.program.Metadata().blockInfo[3u].terminator.kind = TerminatorKind::Branch;
+        b.program.Metadata().blockInfo[3u].terminator.trueBlock = join.Id();
+        b.block = &inner;
+        auto& written = b.select(merged, b.constant(42u), b.constant(0u));
+        if (unsafeInner) b.emit(IrOpcode::Discard, IrType::Void, {});
+        b.block = &join;
+        auto& joined = b.phi(IrType::U32, merge, b.constant(0u));
+        joined.AddPhiOperand(&inner, &written);
+        if (ballotEscape) {
+            const auto [low, high] = b.ballotWords(merged);
+            b.keep(b.threadBit(*low, *high));
+            b.keep(b.add(*low, 1u));
+        }
+        b.keep(joined);
+        exportUnder(b, b.truth(), true, true);
+        return StencilReplaySafe(b.program);
+    };
+
+    passed &= run("an optional wave body merge may guard another optional wave body", [&] { return nestedOptionalGuard(false, false); });
+    passed &= run("an optional wave body merge cannot guard a body with side effects", [&] { return !nestedOptionalGuard(true, false); });
+    passed &= run("an optional wave body merge cannot feed a ballot escaping the lane's own bit", [&] { return !nestedOptionalGuard(false, true); });
+
+    passed &= run("an optional wave body merge may feed an own-lane ballot read", [&] {
+        Builder b;
+        b.program.Resources().stage = IrShaderStage::Pixel;
+        auto& entry = *b.block;
+        auto& predicate = pixelPredicate(b);
+        auto& condition = b.emit(IrOpcode::LogicalNot, IrType::Bool, {&predicate});
+        auto& base = helperPredicate(b, false);
+        auto& body = b.newBlock();
+        auto& merge = b.newBlock();
+        entry.AddBranch(&body);
+        entry.AddBranch(&merge);
+        body.AddBranch(&merge);
+        auto& entryInfo = b.program.Metadata().blockInfo[0u];
+        entryInfo.terminator.kind = TerminatorKind::ConditionalBranch;
+        entryInfo.terminator.condition = BranchCondition::ExecZero;
+        entryInfo.terminator.trueBlock = merge.Id();
+        entryInfo.terminator.falseBlock = body.Id();
+        entryInfo.condition = &condition;
+        b.program.Metadata().blockInfo[1u].terminator.kind = TerminatorKind::Branch;
+        b.program.Metadata().blockInfo[1u].terminator.trueBlock = merge.Id();
+        b.block = &body;
+        auto& taken = b.emit(IrOpcode::SelectU1, IrType::Bool, {&predicate, &b.truth(), &base});
+        b.block = &merge;
+        auto& merged = b.phi(IrType::Bool, entry, base);
+        merged.AddPhiOperand(&body, &taken);
+        b.keep(b.ballotBit(merged));
+        exportUnder(b, b.truth(), true, true);
+        return StencilReplaySafe(b.program);
+    });
+
+    passed &= run("an active mask may flag a sample with an explicit lod", [&] {
+        Builder b;
+        b.program.Resources().stage = IrShaderStage::Pixel;
+        MemoryInfo memory{};
+        memory.kind = ResourceKind::Image;
+        memory.imageSampleFlags = RdnaImageSampleFlagLevelZero;
+        b.program.Resources().memoryInfo.push_back(memory);
+        auto& active = b.ballotBit(b.truth());
+        auto& address = b.emit(IrOpcode::MakeImageAddress, IrType::ImageAddress, {&b.constant(0u)});
+        auto& image = b.emit(IrOpcode::GetImageResource, IrType::ImageResource, {});
+        auto& sampler = b.emit(IrOpcode::GetSamplerResource, IrType::SamplerResource, {});
+        MemoryFlags flags{0u, 0u};
+        std::uint64_t bits = 0;
+        std::memcpy(&bits, &flags, sizeof(flags));
+        b.emit(IrOpcode::ImageSampleRaw, IrType::U32x4, {&image, &sampler, &address, &active}, bits);
+        exportUnder(b, b.truth(), true, true);
+        return StencilReplaySafe(b.program);
+    });
+
     passed &= run("varying attribute sources cannot be treated as uniform any-active predicates", [&] {
         for (const auto opcode : {IrOpcode::GetAttribute, IrOpcode::GetInterpolationParameter, IrOpcode::GetInterpolationParameterF16, IrOpcode::GetTessellationAttribute}) {
             Builder b;

core/shader/recompiler/StencilReplay.cpp

@@ -108,7 +108,7 @@ public:
         }
         for (const auto& info : program.Metadata().blockInfo) {
             if (info.terminator.kind != TerminatorKind::ConditionalBranch) continue;
-            if (!control(info)) return false;
+            if (!controlProven(info)) return false;
         }
         for (const auto& [value, mask] : ownMasks) if (helperObserved(value)) return false;
         for (const auto* block : program.BlockOrder()) {
@@ -493,27 +493,88 @@ private:
         return left == right || (left->HasImmediate() && right->HasImmediate() && left->Type() == IrType::U32 && right->Type() == IrType::U32 && left->ImmediateU32() == right->ImmediateU32());
     }
 
-    bool inactiveEquals(const IrValue* value, const IrValue* other, const IrValue* predicate, unsigned depth = 0) {
+    bool inactiveEquals(const IrValue* value, const IrValue* other, const IrValue* predicate, const std::set<const IrBlock*>& region, unsigned depth = 0) {
         if (depth > 32u) return false;
         value = value->Resolve();
         if (sameValue(value, other)) return true;
-        if ((value->Opcode() == IrOpcode::SelectU32 || value->Opcode() == IrOpcode::SelectF32 || value->Opcode() == IrOpcode::SelectU1) && value->ArgumentCount() == 3u && (value->Argument(0)->Resolve() == predicate->Resolve() || GuardImplies(program, value->Argument(0), predicate))) return inactiveEquals(value->Argument(2), other, predicate, depth + 1u);
+        if ((value->Opcode() == IrOpcode::SelectU32 || value->Opcode() == IrOpcode::SelectF32 || value->Opcode() == IrOpcode::SelectU1) && value->ArgumentCount() == 3u && (value->Argument(0)->Resolve() == predicate->Resolve() || GuardImplies(program, value->Argument(0), predicate))) return inactiveEquals(value->Argument(2), other, predicate, region, depth + 1u);
         if (value->IsPhi() && region.contains(value->Parent()) && value->ArgumentCount() != 0u && value->ArgumentCount() == value->PhiBlockCount()) {
-            return std::all_of(value->Arguments().begin(), value->Arguments().end(), [&](const IrValue* argument) { return inactiveEquals(argument, other, predicate, depth + 1u); });
+            return std::all_of(value->Arguments().begin(), value->Arguments().end(), [&](const IrValue* argument) { return inactiveEquals(argument, other, predicate, region, depth + 1u); });
         }
         return false;
     }
 
-    bool quadObserved(const IrValue* value) const {
+    bool explicitLodSample(const IrValue* value) const {
+        if (value->Opcode() == IrOpcode::ImageSampleExplicitLod) return true;
+        if (value->Opcode() != IrOpcode::ImageSampleRaw && value->Opcode() != IrOpcode::ImageGatherRaw) return false;
+        const auto index = value->Flags<MemoryFlags>().index;
+        return index < program.Resources().memoryInfo.size() && ImageSampleExplicitLod(program.Resources().memoryInfo[index].imageSampleFlags, program.Resources().stage);
+    }
+
+    bool ownBitShifted(const IrValue* user, std::size_t operand, std::vector<const IrValue*>& terminals) const {
+        if (user->Opcode() != IrOpcode::ShiftRightLogical32 || user->ArgumentCount() != 2u || operand != 0u) return false;
+        const IrValue* amount = user->Argument(1)->Resolve();
+        if (amount->Opcode() != IrOpcode::BitwiseAnd32) return false;
+        const IrValue* lane = Other(amount, 31u);
+        if (lane == nullptr || lane->Opcode() != IrOpcode::LaneId) return false;
+        for (const auto& use : user->OperandUses()) {
+            const IrValue* isolated = use.user;
+            if (isolated->Opcode() != IrOpcode::BitwiseAnd32 || Other(isolated, 1u) != user) return false;
+            terminals.push_back(isolated);
+        }
+        return true;
+    }
+
+    bool ownBitSelected(const IrValue* user, std::size_t operand, unsigned half, std::vector<const IrValue*>& terminals) const {
+        if (user->Opcode() != IrOpcode::SelectU32 || user->ArgumentCount() != 3u || operand != half + 1u) return false;
+        const IrValue* condition = user->Argument(0)->Resolve();
+        if (condition->Opcode() != IrOpcode::ULessThan32 || condition->ArgumentCount() != 2u || condition->Argument(0)->Resolve()->Opcode() != IrOpcode::LaneId || !Immediate(condition->Argument(1), 32u)) return false;
+        for (const auto& use : user->OperandUses()) if (!ownBitShifted(use.user, use.operand, terminals)) return false;
+        return true;
+    }
+
+    bool ownBitBallot(const IrValue* ballot, std::vector<const IrValue*>& terminals) const {
+        if (ballot->Opcode() != IrOpcode::Ballot || ballot->ArgumentCount() != 1u) return false;
+        for (const auto& use : ballot->OperandUses()) {
+            const IrValue* extract = use.user;
+            if (extract->Opcode() != IrOpcode::CompositeExtractU32x4 || extract->ArgumentCount() != 2u || extract->Argument(0)->Resolve() != ballot) return false;
+            const IrValue* index = extract->Argument(1)->Resolve();
+            if (!index->HasImmediate() || index->ImmediateU32() >= (waveSize == 32u ? 1u : 2u)) return false;
+            const auto half = index->ImmediateU32();
+            for (const auto& euse : extract->OperandUses()) {
+                const IrValue* word = euse.user;
+                if (word->Opcode() == IrOpcode::BitwiseAnd32) {
+                    for (const auto& muse : word->OperandUses()) {
+                        const bool read = waveSize == 64u ? ownBitSelected(muse.user, muse.operand, half, terminals) : ownBitShifted(muse.user, muse.operand, terminals);
+                        if (!read) return false;
+                    }
+                } else {
+                    const bool read = waveSize == 64u ? ownBitSelected(word, euse.operand, half, terminals) : ownBitShifted(word, euse.operand, terminals);
+                    if (!read) return false;
+                }
+            }
+        }
+        return true;
+    }
+
+    bool quadObserved(const IrValue* value) {
         std::set<const IrValue*> visited;
         std::vector<const IrValue*> pending{value};
         while (!pending.empty()) {
             const IrValue* current = pending.back();
             pending.pop_back();
             if (!visited.insert(current).second) continue;
-            for (const auto& info : program.Metadata().blockInfo) if (info.condition && info.condition->Resolve() == current) return true;
+            for (const auto& info : program.Metadata().blockInfo) if (info.condition && info.condition->Resolve() == current && !controlProven(info)) return true;
             for (const auto& use : current->OperandUses()) {
                 const auto opcode = use.user->Opcode();
+                if (opcode == IrOpcode::Ballot) {
+                    std::vector<const IrValue*> terminals;
+                    if (ownBitBallot(use.user, terminals)) {
+                        pending.insert(pending.end(), terminals.begin(), terminals.end());
+                        continue;
+                    }
+                }
+                if (ImageOpcodeInfoOf(opcode).needsSampler && explicitLodSample(use.user)) { pending.push_back(use.user); continue; }
                 if (ImageOpcodeInfoOf(opcode).needsSampler || opcode == IrOpcode::Ballot || opcode == IrOpcode::WqmU64 || opcode == IrOpcode::ReadLane || opcode == IrOpcode::ReadFirstLane || opcode == IrOpcode::DppMoveU32 || opcode == IrOpcode::DppUpdateU32 || opcode == IrOpcode::SwizzleU32 || opcode == IrOpcode::PermuteU32 || opcode == IrOpcode::BpermuteU32 || opcode == IrOpcode::Permlane16U32 || opcode == IrOpcode::WriteLane || SharedAccessOf(opcode) != SharedAccess::None) return true;
                 if (opcode == IrOpcode::SetAttribute || opcode == IrOpcode::Reference || opcode == IrOpcode::ReferenceU32) continue;
                 pending.push_back(use.user);
@@ -536,6 +597,7 @@ private:
             }
             for (const auto& use : current->OperandUses()) {
                 const auto opcode = use.user->Opcode();
+                if (ImageOpcodeInfoOf(opcode).needsSampler && explicitLodSample(use.user)) { pending.push_back(use.user); continue; }
                 if (ImageOpcodeInfoOf(opcode).needsSampler) return true;
                 if (restricted.contains(use.user) || opcode == IrOpcode::SetAttribute || opcode == IrOpcode::Reference || opcode == IrOpcode::ReferenceU32) continue;
                 pending.push_back(use.user);
@@ -545,8 +607,8 @@ private:
     }
 
     bool optionalBody(const IrBlock* branch, const IrBlock* body, const IrBlock* merge, const IrValue* predicate) {
-        region.clear();
         if (body == merge) return false;
+        std::set<const IrBlock*> region;
         std::vector<const IrBlock*> pending{body};
         while (!pending.empty()) {
             const IrBlock* block = pending.back();
@@ -571,8 +633,9 @@ private:
                 else if (region.contains(instruction->PhiBlock(index))) taken.push_back(instruction->Argument(index)->Resolve());
                 else return false;
             }
-            if (!skipped || taken.empty() || quadObserved(instruction)) return false;
-            for (const auto* value : taken) if (!inactiveEquals(value, skipped, predicate)) return false;
+            if (!skipped || taken.empty()) return false;
+            if (quadObserved(instruction)) return false;
+            for (const auto* value : taken) if (!inactiveEquals(value, skipped, predicate, region)) return false;
             merged.insert(instruction);
         }
         for (const auto* block : region) {
@@ -587,6 +650,15 @@ private:
         return !merged.empty();
     }
 
+    bool controlProven(const BlockInfo& info) {
+        if (const auto found = controlResults.find(info.id); found != controlResults.end()) return found->second;
+        if (!controlVisiting.insert(info.id).second) return false;
+        const bool result = control(info);
+        controlVisiting.erase(info.id);
+        controlResults[info.id] = result;
+        return result;
+    }
+
     bool control(const BlockInfo& info) {
         const auto* branch = program.BlockOrder()[indices.at(info.id)];
         const auto* trueBlock = program.BlockOrder()[indices.at(info.terminator.trueBlock)];
@@ -633,7 +705,8 @@ private:
     std::map<std::uint32_t, std::size_t> indices;
     std::map<const IrBlock*, std::size_t> blockIndices;
     std::vector<std::vector<const IrBlock*>> allPaths;
-    std::set<const IrBlock*> region;
+    std::map<std::uint32_t, bool> controlResults;
+    std::set<std::uint32_t> controlVisiting;
 };
 
 std::optional<std::set<std::uint32_t>> Values(const IrProgram& program, const IrValue* value, std::set<const IrValue*>& visiting, unsigned depth, std::optional<unsigned> component = {}, const IrValue* active = nullptr) {