perf(agc): prune and gate the label take records

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

@@ -13,6 +13,7 @@
 #include "prx/libSceAgcDriver/Execution/include/Driver/Dispatch/IndirectDispatch.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Driver/Memory/WriteEvidence.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Driver/Packets/PacketTiming.hpp"
+#include "prx/libSceAgcDriver/Execution/include/Driver/Synchronization/LabelTakes.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Driver/Synchronization/WaitMemory.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Driver.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Pm4.hpp"
@@ -220,6 +221,7 @@ class Driver {
     void noteLabelStore(std::uint64_t address, std::span<const std::byte> bytes, std::uint64_t stamp, std::uint32_t queue);
     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 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();
@@ -461,8 +463,7 @@ class Driver {
 
     AgcDriver::Mutex labelStoresMutex;
     std::unordered_map<std::uint64_t, std::array<LabelStore, LabelStoreHistory>> labelStores;
-    std::unordered_map<std::uint64_t, std::uint64_t> labelPasses;
-    std::unordered_map<std::uint64_t, std::uint64_t> labelTaken;
+    LabelTakes labelTakes;
 
     std::vector<std::shared_ptr<VulkanDevice>> replacedDevices;
     std::stop_token shutdownToken = LibcShutdownToken_nid_postfix();

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

@@ -0,0 +1,88 @@
+#ifndef CORE_LIBS_PRX_LIBSCEAGCDRIVER_EXECUTION_INCLUDE_DRIVER_SYNCHRONIZATION_LABELTAKES_HPP
+#define CORE_LIBS_PRX_LIBSCEAGCDRIVER_EXECUTION_INCLUDE_DRIVER_SYNCHRONIZATION_LABELTAKES_HPP
+
+#include <array>
+#include <atomic>
+#include <cstddef>
+#include <cstdint>
+#include <deque>
+#include <unordered_map>
+#include <utility>
+
+namespace AgcDriver::DriverDetail {
+
+class LabelTakes {
+public:
+    void Pass(std::uint32_t queue, std::uint64_t dword, std::uint64_t serial) {
+        const auto key = queueDword(queue, dword);
+        auto& entry = entries[key];
+        if (entry.passed == 0) passCounts[slot(key)].fetch_add(1, std::memory_order_relaxed);
+        entry.passed = serial;
+        passes[queue].emplace_back(serial, key);
+    }
+
+    bool MayHavePass(std::uint32_t queue, std::uint64_t dword) const {
+        return passCounts[slot(queueDword(queue, dword))].load(std::memory_order_relaxed) != 0;
+    }
+
+    void Take(std::uint32_t queue, std::uint64_t dword) {
+        const auto key = queueDword(queue, dword);
+        const auto found = entries.find(key);
+        if (found == entries.end() || found->second.passed == 0) return;
+        found->second.taken = found->second.passed;
+        found->second.passed = 0;
+        passCounts[slot(key)].fetch_sub(1, std::memory_order_relaxed);
+    }
+
+    std::uint64_t Taken(std::uint32_t queue, std::uint64_t dword) const {
+        const auto found = entries.find(queueDword(queue, dword));
+        return found == entries.end() ? 0 : found->second.taken;
+    }
+
+    void Retire(std::uint32_t queue, std::uint64_t received) {
+        const auto queued = passes.find(queue);
+        if (queued == passes.end()) return;
+        auto& order = queued->second;
+        while (!order.empty() && order.front().first <= received) {
+            const auto [serial, key] = order.front();
+            order.pop_front();
+            const auto found = entries.find(key);
+            if (found == entries.end()) continue;
+            if (found->second.passed == serial) {
+                found->second.passed = 0;
+                passCounts[slot(key)].fetch_sub(1, std::memory_order_relaxed);
+            }
+            if (found->second.taken == serial) found->second.taken = 0;
+            if (found->second.passed == 0 && found->second.taken == 0) entries.erase(found);
+        }
+        if (order.empty()) passes.erase(queued);
+    }
+
+    bool Empty() const {
+        return entries.empty() && passes.empty();
+    }
+
+private:
+    struct Entry {
+        std::uint64_t passed = 0;
+        std::uint64_t taken = 0;
+    };
+
+    static constexpr std::size_t PassSlotBits = 12;
+
+    static std::uint64_t queueDword(std::uint32_t queue, std::uint64_t dword) {
+        return dword | (static_cast<std::uint64_t>(queue) << 56u);
+    }
+
+    static std::size_t slot(std::uint64_t key) {
+        return static_cast<std::size_t>(((key >> 2u) * 0x9e3779b97f4a7c15ull) >> (64u - PassSlotBits));
+    }
+
+    std::unordered_map<std::uint64_t, Entry> entries;
+    std::unordered_map<std::uint32_t, std::deque<std::pair<std::uint64_t, std::uint64_t>>> passes;
+    std::array<std::atomic<std::uint32_t>, std::size_t{1} << PassSlotBits> passCounts{};
+};
+
+}
+
+#endif

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

@@ -278,6 +278,7 @@ void Driver::run(std::uint32_t id) noexcept {
                 running.Start();
                 if (profile && submission.enqueuedAt != std::chrono::steady_clock::time_point{}) costs.dequeueNs += static_cast<std::uint64_t>(std::chrono::duration_cast<std::chrono::nanoseconds>(std::chrono::steady_clock::now() - submission.enqueuedAt).count());
             }
+            retireLabelTakes(id, submission.received);
             execute(submission);
             if (submission.progress && submission.rewindTail == nullptr && !completionDeferred()) submitOpenWork(submission.progress);
             if (traceGpu) std::fprintf(stderr, "[gpu] %.1f done serial=%llu queue=0x%x\n", TraceMs(), static_cast<unsigned long long>(submission.serial), id);

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

@@ -362,6 +362,7 @@ void Driver::executeRewindTail(const Submission& stalled) {
     if (APS5_ENABLE_TIMING_LOG) tail.roomReadyAt = std::chrono::steady_clock::now();
     auto* const capture = SubmissionCapture::Active();
     auto captured = capture != nullptr ? SubmissionCapture::Prepare(tail.queue, SubmissionKind::Rewind, tail.commands, tail.conditionalEnds) : CapturedSubmission{};
+    std::uint64_t retiredTakes = 0;
     {
         std::lock_guard lock(mutex);
         rethrowFailure();
@@ -370,12 +371,15 @@ void Driver::executeRewindTail(const Submission& stalled) {
         tail.shaders = shaders;
         tail.serial = stalled.serial;
         tail.received = ++eventSerial;
+        const auto& waiting = workers[tail.queue].pending;
+        retiredTakes = waiting.empty() ? tail.received : std::min(tail.received, waiting.front().received);
         if (capture != nullptr) capture->Take(captured, tail.serial, captureTime(), tail.heldAtSubmit);
         if (ring) RecordSubmission(tail.queue, tail.serial, tail.commands, stalled.rewindTail, true);
     }
     if (capture != nullptr) writeCaptured(*this, *capture, captured);
     if (APS5_ENABLE_TIMING_LOG) tail.enqueuedAt = tail.dequeuedAt = tail.orderedAt = std::chrono::steady_clock::now();
     timing.Finish();
+    retireLabelTakes(tail.queue, retiredTakes);
     execute(tail);
 }
 

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

@@ -20,35 +20,35 @@ void Driver::noteLabelStore(std::uint64_t address, std::span<const std::byte> by
     }
 }
 
-namespace {
-
-std::uint64_t queueDword(std::uint32_t queue, std::uint64_t dword) {
-    return dword | (static_cast<std::uint64_t>(queue) << 56u);
-}
-
-}
-
 void Driver::noteLabelPass(std::span<const std::uint32_t> packet, std::uint32_t queue) {
     const auto opcode = (packet[0] >> 8u) & 0xffu;
     if ((opcode != 0x3c && opcode != 0x93) || packet.size() < 7 || ((packet[1] >> 4u) & 3u) != 1u) return;
     const auto bytes = Pm4::WaitAwaitedBytes(packet);
     const std::uint64_t address = (packet[2] & (bytes == 8 ? ~7u : ~3u)) | (static_cast<std::uint64_t>(packet[3]) << 32u);
     const auto passed = eventSerial.load(std::memory_order_acquire);
     std::lock_guard lock(labelStoresMutex);
-    for (std::size_t offset = 0; offset < bytes; offset += 4) labelPasses[queueDword(queue, address + offset)] = passed;
+    for (std::size_t offset = 0; offset < bytes; offset += 4) labelTakes.Pass(queue, address + offset, passed);
+}
+
+void Driver::retireLabelTakes(std::uint32_t queue, std::uint64_t received) {
+    std::lock_guard lock(labelStoresMutex);
+    labelTakes.Retire(queue, received);
 }
 
 void Driver::noteLabelIssue(std::span<const std::uint32_t> packet, std::uint32_t queue) {
-    const auto take = [&](std::uint64_t address, std::uint64_t bytes) {
+    const auto mayTake = [&](std::uint64_t address, std::uint64_t bytes) {
         for (std::uint64_t offset = 0; offset < bytes; offset += 4) {
-            const auto key = queueDword(queue, address + offset);
-            if (const auto pass = labelPasses.find(key); pass != labelPasses.end()) labelTaken[key] = pass->second;
+            if (labelTakes.MayHavePass(queue, address + offset)) return true;
         }
+        return false;
+    };
+    const auto take = [&](std::uint64_t address, std::uint64_t bytes) {
+        for (std::uint64_t offset = 0; offset < bytes; offset += 4) labelTakes.Take(queue, address + offset);
     };
     switch ((packet[0] >> 8u) & 0xffu) {
         case 0x37: case 0x49: {
             const auto label = Pm4::DecodeLabelWrite(packet);
-            if (!label.has_value() || !KeepsLabelStore(label->address, label->Bytes().size())) return;
+            if (!label.has_value() || !KeepsLabelStore(label->address, label->Bytes().size()) || !mayTake(label->address, label->Bytes().size())) return;
             std::lock_guard lock(labelStoresMutex);
             take(label->address, label->Bytes().size());
             return;
@@ -57,7 +57,7 @@ void Driver::noteLabelIssue(std::span<const std::uint32_t> packet, std::uint32_t
         default: return;
     }
     auto stores = Pm4::DecodeStores(packet);
-    std::erase_if(stores, [](const Pm4::PacketStore& store) { return store.conditional || !KeepsLabelStore(store.address, static_cast<std::size_t>(store.bytes)); });
+    std::erase_if(stores, [&](const Pm4::PacketStore& store) { return store.conditional || !KeepsLabelStore(store.address, static_cast<std::size_t>(store.bytes)) || !mayTake(store.address, store.bytes); });
     if (stores.empty()) return;
     std::lock_guard lock(labelStoresMutex);
     for (const auto& store : stores) take(store.address, store.bytes);
@@ -70,9 +70,7 @@ bool Driver::storedSince(std::span<const std::uint32_t> packet, std::uint64_t ad
     const auto high = bytes == 8 ? labelStores.find(address + 4) : labelStores.end();
     if (bytes == 8 && high == labelStores.end()) return false;
     std::uint64_t taken = 0;
-    for (std::size_t offset = 0; offset < bytes; offset += 4) {
-        if (const auto found = labelTaken.find(queueDword(queue, address + offset)); found != labelTaken.end()) taken = std::max(taken, found->second);
-    }
+    for (std::size_t offset = 0; offset < bytes; offset += 4) taken = std::max(taken, labelTakes.Taken(queue, address + offset));
     for (const auto& store : low->second) {
         if (store.stamp <= received || (store.queue != queue && store.stamp <= taken)) continue;
         std::uint64_t value = store.value;

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

@@ -39,6 +39,18 @@ static_assert(offsetof(Packet, addr) == 0);
 static_assert(offsetof(Packet, dw_num) == 8);
 static_assert(offsetof(Packet, flags) == 12);
 
+namespace AgcDriver::DriverDetail {
+
+struct DriverTestAccess {
+    static std::uint64_t LabelTaken(std::uint32_t queue, const volatile std::uint32_t* address) {
+        auto& driver = Driver::Get();
+        std::lock_guard lock(driver.labelStoresMutex);
+        return driver.labelTakes.Taken(queue, reinterpret_cast<std::uintptr_t>(address));
+    }
+};
+
+}
+
 namespace {
 
 void check(bool condition, const char* message) {
@@ -451,6 +463,56 @@ void testLabelTakenAndResetByEarlierSubmission() {
     AgcDriverWaitIdle_nid_postfix();
 }
 
+void testLabelTakesRetire() {
+    using AgcDriver::DriverDetail::LabelTakes;
+    auto takes = std::make_unique<LabelTakes>();
+    constexpr std::uint64_t label = 0x1000, other = 0x2000;
+    takes->Pass(0x20, label, 10);
+    check(takes->MayHavePass(0x20, label) && takes->Taken(0x20, label) == 0, "a passed wait was not kept for the store that takes it");
+    takes->Take(0x20, label);
+    check(takes->Taken(0x20, label) == 10 && !takes->MayHavePass(0x20, label), "the queue's own store did not turn its pass into a take");
+    takes->Retire(0x20, 9);
+    check(takes->Taken(0x20, label) == 10, "a take was dropped while a submission of the queue entered before it could still wait");
+    takes->Retire(0x28, 20);
+    check(takes->Taken(0x20, label) == 10, "another queue's submission retired a take");
+    takes->Retire(0x20, 10);
+    check(takes->Taken(0x20, label) == 0 && takes->Empty(), "a take was kept after every submission of the queue that could wait entered after it");
+    takes->Pass(0x20, label, 30);
+    takes->Pass(0x20, label, 40);
+    takes->Retire(0x20, 30);
+    check(takes->MayHavePass(0x20, label), "an older pass on the dword retired its newer pass");
+    takes->Take(0x20, label);
+    takes->Retire(0x20, 39);
+    check(takes->Taken(0x20, label) == 40, "a take was retired by a submission entered before it");
+    takes->Pass(0x20, other, 50);
+    takes->Pass(0x28, label, 50);
+    takes->Retire(0x20, 50);
+    check(takes->Taken(0x20, label) == 0 && !takes->MayHavePass(0x20, other), "a submission of the queue entered after its passes and take did not retire them");
+    check(takes->MayHavePass(0x28, label) && !takes->Empty(), "a submission of one queue retired another queue's pass");
+    takes->Retire(0x28, 50);
+    check(takes->Empty(), "a pass was kept after every submission of its queue entered after it");
+}
+
+void testLabelTakeRetiredByLaterSubmission() {
+    using AgcDriver::DriverDetail::DriverTestAccess;
+    alignas(64) static volatile std::uint32_t gate = 0, label = 0, first = 0, second = 0, third = 0;
+    submit(0x20, commands(waitEqual(&gate, 1), waitEqual(&label, 1), writeData(&label, 0), writeData(&first, 1)));
+    submit(0x20, commands(waitEqual(&label, 1), writeData(&second, 1)));
+    submit(0, commands(writeData(&label, 1)));
+    waitFor(&label, 1, "producer label never landed");
+    gate = 1;
+    waitFor(&first, 1, "consumer never passed its first label wait");
+    std::this_thread::sleep_for(std::chrono::milliseconds(1200));
+    check(second == 0, "a label store an earlier submission of the queue had taken and reset satisfied a later submission's wait, or the wait expired");
+    check(DriverTestAccess::LabelTaken(0x20, &label) != 0, "a take was dropped while a submission of the queue entered before it was waiting");
+    submit(0, commands(writeData(&label, 1)));
+    waitFor(&second, 1, "the later submission did not resume after a new producer label");
+    submit(0x20, commands(writeData(&third, 1)));
+    waitFor(&third, 1, "a submission entered after the take never ran");
+    check(DriverTestAccess::LabelTaken(0x20, &label) == 0, "a take was kept after a submission of the queue entered after it started");
+    AgcDriverWaitIdle_nid_postfix();
+}
+
 void testWaitFreeSubmissionAfterEarlierQueue0Work() {
     alignas(64) static volatile std::uint32_t first = 0, second = 0;
     constexpr std::uint32_t writes = 20000;
@@ -1103,6 +1165,8 @@ int main(int argc, char** argv) {
         testLabelHeldAtSubmission();
         testLabelTakenAndResetInOneSubmission();
         testLabelTakenAndResetByEarlierSubmission();
+        testLabelTakesRetire();
+        testLabelTakeRetiredByLaterSubmission();
         testWideLabelStoredSinceSubmission();
         testWaitFreeSubmissionAfterEarlierQueue0Work();
         testWaitFreeSubmissionQueue0WaitsOn();