fix(agc): hold a label wait while its producing chain advances

core/libs/prx/libSceAgcDriver/Execution/include/Driver/Driver.hpp

@@ -235,6 +235,7 @@ class Driver {
     void noteLabelPass(std::span<const std::uint32_t> packet, std::uint32_t queue);
     void noteLabelIssue(std::span<const std::uint32_t> packet, std::uint32_t queue);
     void retireLabelTakes(std::uint32_t queue, std::uint64_t received);
+    bool findWaitProducer(std::uint64_t address, std::size_t bytes, std::uint64_t& token);
     bool storedSince(std::span<const std::uint32_t> packet, std::uint64_t address, std::size_t bytes, std::uint64_t received, std::uint32_t queue, std::uint32_t* writer = nullptr);
     static void traceLabel(std::span<const std::uint32_t> packet, std::uint32_t queue);
     static std::array<WriteRecord, 16384>& writeHistory();
@@ -399,6 +400,7 @@ class Driver {
     static std::optional<std::uint32_t> releaseInterrupt(std::span<const std::uint32_t> packet, std::uint32_t opcode, bool carriesFlip);
     void markCompleted(std::uint64_t serial);
     static const std::atomic<std::uint64_t>*& workerQueued();
+    static std::atomic<std::uint64_t>*& workerPacket();
     void reapCompletionLabels();
     void run(std::uint32_t id) noexcept;
 

core/libs/prx/libSceAgcDriver/Execution/include/Driver/Queues/Submission.hpp

@@ -55,6 +55,7 @@ struct QueueWorker {
     std::unordered_map<std::uint64_t, std::uint32_t> unfinishedWrites;
     std::uint64_t executingWaitFree = 0;
     bool executingWaits = false;
+    alignas(64) std::atomic<std::uint64_t> executingPacket{0};
     std::thread thread;
 };
 

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Packets/PacketExecution.cpp

@@ -277,6 +277,7 @@ void Driver::execute(const Submission& submission) {
             if (queue0) DrawFinisher::Get().Settle();
         }
     } settleFinisher{submission.queue == 0};
+    auto* const executingPacket = workerPacket();
     for (std::size_t cursor = 0; cursor < submission.commands.size();) {
         if (packetEpoch()) bumpEpoch(&EpochBumps::packets);
         CheckFailure();
@@ -285,6 +286,7 @@ void Driver::execute(const Submission& submission) {
         const auto count = Pm4::PacketWords(header);
         const auto packet = std::span(submission.commands).subspan(cursor, count);
         const auto opcode = (header >> 8u) & 0xffu;
+        if (executingPacket != nullptr) executingPacket->store((static_cast<std::uint64_t>(cursor + 1) << 8u) | opcode, std::memory_order_relaxed);
         auto nextCursor = cursor + count;
         const bool predicated = Pm4::Predicated(header) && queue.predication.operation != 0;
         if (submission.queue == 0 && DrawFinisher::Busy() && !keepsDrawFinisher(header, opcode)) {

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Queues/QueueWorker.cpp

@@ -134,6 +134,11 @@ const std::atomic<std::uint64_t>*& Driver::workerQueued() {
     return queued;
 }
 
+std::atomic<std::uint64_t>*& Driver::workerPacket() {
+    static thread_local std::atomic<std::uint64_t>* packet = nullptr;
+    return packet;
+}
+
 void Driver::reapCompletionLabels() {
     std::unique_lock gpuLock(GuestMemory::GpuMutex(), std::try_to_lock);
     if (!gpuLock.owns_lock()) {
@@ -190,6 +195,7 @@ void Driver::run(std::uint32_t id) noexcept {
                 auto& worker = workers.at(id);
                 auto& pending = worker.pending;
                 workerQueued() = &worker.queued;
+                workerPacket() = &worker.executingPacket;
                 if (traceGpu && pending.empty()) std::fprintf(stderr, "[gpu] %.1f idle queue=0x%x\n", TraceMs(), id);
                 const auto ready = [&] { return failure || stopping || !pending.empty(); };
                 const auto allQueuesIdle = [&] {
@@ -254,6 +260,7 @@ void Driver::run(std::uint32_t id) noexcept {
                 }
                 submission = std::move(pending.front());
                 pending.pop_front();
+                worker.executingPacket.store(0, std::memory_order_relaxed);
                 if (APS5_ENABLE_TIMING_LOG) submission.dequeuedAt = std::chrono::steady_clock::now();
                 if (id == 0) queue0Executing = submission.suspend ? 0 : submission.received;
                 struct OrderHold {

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Synchronization/LabelHistory.cpp

@@ -63,6 +63,44 @@ void Driver::noteLabelIssue(std::span<const std::uint32_t> packet, std::uint32_t
     for (const auto& store : stores) take(store.address, store.bytes);
 }
 
+bool Driver::findWaitProducer(std::uint64_t address, std::size_t bytes, std::uint64_t& token) {
+    std::lock_guard lock(mutex);
+    token = 0;
+    std::vector<std::uint64_t> frontier{address};
+    std::unordered_set<std::uint32_t> visited;
+    bool found = false;
+    for (std::size_t depth = 0; depth < 8 && !frontier.empty(); ++depth) {
+        const std::size_t span = depth == 0 ? bytes : 8;
+        std::vector<std::uint64_t> next;
+        for (const auto current : frontier) {
+            for (const auto& [queue, worker] : workers) {
+                bool holds = false;
+                for (std::size_t offset = 0; offset < span; offset += 4) {
+                    if (worker.unfinishedWrites.contains(current + offset)) {
+                        holds = true;
+                        break;
+                    }
+                }
+                if (!holds || !visited.insert(queue).second) continue;
+                found = true;
+                std::uint64_t motion = worker.executingPacket.load(std::memory_order_relaxed) & 0xffffffffull;
+                motion |= static_cast<std::uint64_t>(worker.pending.size() & 0xffu) << 32;
+                motion |= static_cast<std::uint64_t>(queue & 0xffu) << 40;
+                motion *= 0x9e3779b97f4a7c15ull + 2ull * queue;
+                const auto shift = depth * 8;
+                token ^= shift == 0 ? motion : ((motion << shift) | (motion >> (64 - shift)));
+                if (queue != 0 && (queue < 0x20 || queue > SubmissionQueueEnd)) continue;
+                const auto slot = queue == 0 ? 0 : queue - 0x20 + 1;
+                if (const auto awaited = awaitedSlots[slot].address.load(std::memory_order_acquire); awaited != 0) next.push_back(awaited);
+            }
+        }
+        std::sort(next.begin(), next.end());
+        next.erase(std::unique(next.begin(), next.end()), next.end());
+        frontier = std::move(next);
+    }
+    return found;
+}
+
 bool Driver::storedSince(std::span<const std::uint32_t> packet, std::uint64_t address, std::size_t bytes, std::uint64_t received, std::uint32_t queue, std::uint32_t* writer) {
     std::lock_guard lock(labelStoresMutex);
     const auto low = labelStores.find(address);

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Synchronization/WaitMemory.cpp

@@ -356,6 +356,8 @@ void Driver::waitMemory(std::span<const std::uint32_t> packet, std::uint32_t que
     bool spinning = pauseSpin;
     std::uint32_t polls = 0;
     StallWatch stall(queue);
+    std::uint64_t producerToken = 0;
+    auto lastProducerScan = std::chrono::steady_clock::now() - std::chrono::milliseconds(100);
     while (!memorySatisfied()) {
         if (std::uint32_t writer = queue; !requireMemory && storedSince(packet, awaited, awaitedBytes, received, queue, &writer)) {
             if (writer != queue) notePendingWait(packet, queue, received);
@@ -386,6 +388,13 @@ void Driver::waitMemory(std::span<const std::uint32_t> packet, std::uint32_t que
             lastDone = done;
             start = std::chrono::steady_clock::now();
             stall.Progressed();
+        } else if (const auto now = std::chrono::steady_clock::now(); now - lastProducerScan >= std::chrono::milliseconds(100)) {
+            lastProducerScan = now;
+            if (std::uint64_t token = 0; findWaitProducer(awaited, awaitedBytes, token) && token != producerToken) {
+                producerToken = token;
+                start = now;
+                stall.Progressed();
+            }
         }
         stall.Check(((packet[0] >> 8u) & 0xffu) == 0x93u ? "WAIT_REG_MEM_64 at" : "WAIT_REG_MEM at", awaited);
 

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

@@ -48,6 +48,14 @@ struct DriverTestAccess {
         std::lock_guard lock(driver.labelStoresMutex);
         return driver.labelTakes.Taken(queue, reinterpret_cast<std::uintptr_t>(address));
     }
+    static bool FindWaitProducer(std::uint64_t address, std::size_t bytes, std::uint64_t& token) {
+        return Driver::Get().findWaitProducer(address, bytes, token);
+    }
+    static std::uint64_t AwaitedAddress(std::uint32_t queue) {
+        const auto& driver = Driver::Get();
+        const auto slot = queue == 0 ? 0 : queue - 0x20 + 1;
+        return driver.awaitedSlots[slot].address.load(std::memory_order_acquire);
+    }
 };
 
 }
@@ -1100,6 +1108,63 @@ void testWorkerFailure() {
     check(expectFailure([&] { sceAgcDriverSubmitAcb(0x20, &packet); }) == messages[0], "subsequent ACB lost worker failure");
 }
 
+void testWaitProducerChain() {
+    using AgcDriver::DriverDetail::DriverTestAccess;
+    alignas(64) static volatile std::uint32_t gateFrozen = 0, gateA = 0, gateB = 0, gateC = 0;
+    alignas(64) static volatile std::uint32_t label = 0, scratch1 = 0, scratch2 = 0;
+
+    submit(0x20, commands(waitEqual(&gateFrozen, 1), writeData(&label, 0xAA)));
+    submit(0x21, commands(waitEqual(&gateA, 1), writeData(&scratch1, 1), waitEqual(&gateB, 1), writeData(&label, 0xAA)));
+    submit(0x22, commands(waitEqual(&gateC, 1), writeData(&scratch2, 1), writeData(&gateB, 1)));
+
+    const auto awaitedIs = [](volatile std::uint32_t* expected, std::uint32_t queue, const char* message) {
+        const auto address = reinterpret_cast<std::uintptr_t>(expected);
+        const auto start = std::chrono::steady_clock::now();
+        while (DriverTestAccess::AwaitedAddress(queue) != address) {
+            check(std::chrono::steady_clock::now() - start < std::chrono::seconds(10), message);
+            std::this_thread::sleep_for(std::chrono::milliseconds(1));
+        }
+    };
+    awaitedIs(&gateFrozen, 0x20, "queue 0x20 never reached the frozen gate wait");
+    awaitedIs(&gateA, 0x21, "queue 0x21 never reached the first producer gate");
+    awaitedIs(&gateC, 0x22, "queue 0x22 never reached its own gate");
+
+    const auto labelAddress = reinterpret_cast<std::uintptr_t>(&label);
+    std::uint64_t token1 = 0;
+    check(DriverTestAccess::FindWaitProducer(labelAddress, 4, token1) && token1 != 0, "no producer seen with two queues holding the awaited label");
+
+    gateA = 1;
+    awaitedIs(&gateB, 0x21, "queue 0x21 never advanced past the first gate");
+
+    std::uint64_t token2 = 0;
+    check(DriverTestAccess::FindWaitProducer(labelAddress, 4, token2), "producer lookup lost the advancing queue");
+    check(token2 != token1, "progress of queue 0x21 was invisible while the same label was frozen on queue 0x20");
+
+    gateC = 1;
+    waitFor(&scratch2, 1, "queue 0x22 never advanced behind queue 0x21's blocked label");
+    waitFor(&label, 0xAA, "the producer chain never wrote the awaited label");
+
+    alignas(64) static volatile std::uint32_t gateStep1 = 0, gateStep2 = 0, labelPair = 0;
+    submit(0x23, commands(waitEqual(&gateStep1, 1), waitEqual(&gateStep2, 1), writeData(&labelPair, 0xBB)));
+    submit(0x24, commands(waitEqual(&gateStep1, 1), waitEqual(&gateStep2, 1), writeData(&labelPair, 0xBB)));
+    awaitedIs(&gateStep1, 0x23, "queue 0x23 never reached the shared first gate");
+    awaitedIs(&gateStep1, 0x24, "queue 0x24 never reached the shared first gate");
+    const auto pairAddress = reinterpret_cast<std::uintptr_t>(&labelPair);
+    std::uint64_t tokenA = 0;
+    check(DriverTestAccess::FindWaitProducer(pairAddress, 4, tokenA) && tokenA != 0, "no producers seen for the paired wait");
+    gateStep1 = 1;
+    awaitedIs(&gateStep2, 0x23, "queue 0x23 never reached the shared second gate");
+    awaitedIs(&gateStep2, 0x24, "queue 0x24 never reached the shared second gate");
+    std::uint64_t tokenB = 0;
+    check(DriverTestAccess::FindWaitProducer(pairAddress, 4, tokenB), "paired producers were lost from the lookup");
+    check(tokenB != tokenA, "identical in-step progress on queues 0x23 and 0x24 was cancelled out of the motion token");
+    gateStep2 = 1;
+    waitFor(&labelPair, 0xBB, "the paired producers never wrote their awaited label");
+
+    gateFrozen = 1;
+    AgcDriverWaitIdle_nid_postfix();
+}
+
 }
 
 int main(int argc, char** argv) {
@@ -1120,6 +1185,12 @@ int main(int argc, char** argv) {
             std::puts("Read-only resummarize submission passed");
             return 0;
         }
+        if (argc == 2 && std::string_view(argv[1]) == "--wait-producer-chain") {
+            testWaitProducerChain();
+            LibcRunShutdown_nid_postfix();
+            std::puts("AGC wait producer chain test passed");
+            return 0;
+        }
         const auto device = OpenVulkanTestDevice();
         if (!device) return VulkanTestSkipped;
         if (argc == 2 && std::string_view(argv[1]) == "--small-copy-sync") {

core/libs/prx/libSceAgcDriver/tests/cmake/agc_driver_tests.cmake

@@ -23,8 +23,10 @@ if(BUILD_TESTING)
     set_tests_properties(agc_driver_readonly_resummarize PROPERTIES TIMEOUT 10)
     add_test(NAME agc_driver_eop_interrupts COMMAND agc_driver_tests --eop-interrupts)
     set_tests_properties(agc_driver_eop_interrupts PROPERTIES TIMEOUT 60)
+    add_test(NAME agc_driver_wait_producer_chain COMMAND agc_driver_tests --wait-producer-chain)
+    set_tests_properties(agc_driver_wait_producer_chain PROPERTIES TIMEOUT 30)
     if(WIN32)
-        set_tests_properties(agc_driver_readonly_resummarize agc_driver_eop_interrupts PROPERTIES ENVIRONMENT_MODIFICATION
+        set_tests_properties(agc_driver_readonly_resummarize agc_driver_eop_interrupts agc_driver_wait_producer_chain PROPERTIES ENVIRONMENT_MODIFICATION
             "PATH=path_list_prepend:$<TARGET_FILE_DIR:libc>")
     endif()
     add_test(NAME agc_driver_small_copy_sync COMMAND agc_driver_tests --small-copy-sync)