Merge pull request #627 from oneandonlydean/perf/agc-large-label-stores

perf(agc): keep label stores larger than 64 bytes out of the label table

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Memory/WriteEvidence.cpp

@@ -165,6 +165,10 @@ ShaderMemory::PendingWrite Driver::classifyPendingWrite(std::uint64_t address, s
     }
     const auto first = address & ~3ull;
     const auto limit = address + bytes;
+    if (Graphics::Recorder::WideLabelIn(first, static_cast<std::size_t>(limit - first))) {
+        reason = &ValidateCounters::syncedLabel;
+        return Policy::Sync;
+    }
     for (std::uint64_t dword = first; dword < limit; dword += 4) {
         if (Graphics::Recorder::LookupLabelValue(dword, 4, 0).has_value()) {
             reason = &ValidateCounters::syncedLabel;

core/libs/prx/libSceAgcDriver/Graphics/include/Recorder.hpp

@@ -235,6 +235,7 @@ public:
     // Entries leave the table with their batch (a memory bound only). The note also enters the
     // label's range as a pending write of the open batch (the flush hook syncs CPU reads), marked
     // as the label's own so it does not count as an overwrite of the entry.
+    static constexpr std::size_t LabelTableBytes = 64;
     void NoteLabel(std::uint64_t address, std::span<const std::byte> bytes, std::uint64_t stamp, std::uint32_t queue);
     // Late rule: an entry whose stamp is not newer than the wait's submission is still the value
     // memory will hold, unless a CPU store touched the dword since it was recorded (the table
@@ -285,6 +286,7 @@ public:
     static std::optional<LabelHit> LookupLabel(std::uint64_t address, std::size_t bytes, std::uint64_t afterStamp, LabelRefusal* refusal = nullptr);
     // The value alone, for callers asking whether any label is pending in a dword (afterStamp 0).
     static std::optional<std::uint64_t> LookupLabelValue(std::uint64_t address, std::size_t bytes, std::uint64_t afterStamp);
+    static bool WideLabelIn(std::uint64_t address, std::size_t bytes);
     // A label a queue worker decoded but has not recorded yet (Driver.cpp DeferredLabels): it
     // enters the table with no batch, and only a lookup made on the noting thread (a WAIT_REG_MEM
     // of the same queue) takes its value, since that queue's later work follows the label in queue
@@ -520,6 +522,7 @@ private:
         std::uint64_t readGeneration = 0;
         // Dword addresses this batch noted in the label table (removed when it finishes).
         std::vector<std::uint64_t> labelDwords;
+        std::vector<std::multimap<std::uint64_t, std::uint64_t>::iterator> wideLabels;
         // Labels stored by this batch's completion actions (see AfterCompletions); subtracted from
         // the lock-free pending count when the batch finishes, whether or not its completions ran.
         std::uint32_t completionLabelCount = 0;
@@ -709,6 +712,9 @@ private:
     // The queued entries (batch always null), by dword; a lookup on the noting queue's thread
     // prefers them (its newest store in program order), every other lookup sees `labels` alone.
     std::map<std::uint64_t, LabelEntry> queuedLabels;
+    std::multimap<std::uint64_t, std::uint64_t> wideLabels;
+    std::uint64_t wideLabelBytes = 0;
+    bool wideLabelInLocked(std::uint64_t address, std::size_t bytes) const;
     // labels.size(), readable without the table mutex: a noted write skips the mutex while the
     // table is empty (most of the time between label groups).
     std::atomic<std::size_t> recordedLabels{0};

core/libs/prx/libSceAgcDriver/Graphics/src/Recorder.cpp

@@ -1208,6 +1208,20 @@ std::optional<Recorder::LabelHit> Recorder::LookupLabel(std::uint64_t address, s
     return labelTableOwner->lookupLabel(address, bytes, afterStamp, refusal);
 }
 
+bool Recorder::WideLabelIn(std::uint64_t address, std::size_t bytes) {
+    std::lock_guard tableLock(labelTableMutex);
+    return labelTableOwner != nullptr && labelTableOwner->wideLabelInLocked(address, bytes);
+}
+
+bool Recorder::wideLabelInLocked(std::uint64_t address, std::size_t bytes) const {
+    if (wideLabels.empty() || bytes == 0) return false;
+    const auto end = address + bytes;
+    for (auto it = wideLabels.lower_bound(address >= wideLabelBytes ? address - wideLabelBytes + 1 : 0); it != wideLabels.end() && it->first < end; ++it) {
+        if (it->second > address) return true;
+    }
+    return false;
+}
+
 std::optional<std::uint64_t> Recorder::LookupLabelValue(std::uint64_t address, std::size_t bytes, std::uint64_t afterStamp) {
     const auto hit = LookupLabel(address, bytes, afterStamp);
     if (!hit.has_value()) return std::nullopt;
@@ -2501,7 +2515,16 @@ void Recorder::NoteLabel(std::uint64_t address, std::span<const std::byte> bytes
     ensureOpen();
     // The table entry before the write note: the note bumps the write generation a poller
     // watches, and a poller that sees the bump then finds the label without the GPU mutex.
-    noteLabelOn(*open, address, bytes, stamp, queue);
+    const bool tabled = bytes.size() <= LabelTableBytes;
+    if (tabled) {
+        noteLabelOn(*open, address, bytes, stamp, queue);
+    } else {
+        std::lock_guard tableLock(labelTableMutex);
+        labels.erase(labels.lower_bound(address), labels.lower_bound(address + bytes.size()));
+        recordedLabels.store(labels.size(), std::memory_order_relaxed);
+        open->wideLabels.push_back(wideLabels.emplace(address, address + bytes.size()));
+        wideLabelBytes = std::max<std::uint64_t>(wideLabelBytes, bytes.size());
+    }
     if (activeRecorder == this && pendingLabelSince.load(std::memory_order_relaxed) == NoPendingLabel) {
         pendingLabelSince.store(std::chrono::steady_clock::now().time_since_epoch().count(), std::memory_order_release);
     }
@@ -2524,6 +2547,7 @@ bool Recorder::PendingLabelIn(std::uint64_t address, std::size_t bytes) const {
     if (bytes == 0) return false;
     const auto end = address + bytes;
     if (const auto first = labels.lower_bound(address); first != labels.end() && first->first < end) return true;
+    if (wideLabelInLocked(address, bytes)) return true;
     // A queued label of another queue is unordered against the caller on hardware (nothing that
     // queue recorded could have satisfied a wait yet); it is still reported, since the table mutex
     // is held anyway and a caller deciding a CPU store wants the conservative answer.
@@ -3256,15 +3280,17 @@ void Recorder::finish(std::unique_ptr<Batch> batch, bool wait, int source) {
     }
     // The batch's label entries leave the table (a later label to the same dword already replaced
     // its entry and belongs to another batch). Correctness never depended on this removal.
-    if (!batch->labelDwords.empty()) {
+    if (!batch->labelDwords.empty() || !batch->wideLabels.empty()) {
         std::lock_guard tableLock(labelTableMutex);
         for (const auto dword : batch->labelDwords) {
             const auto found = labels.find(dword);
             if (found != labels.end() && found->second.batch == batch.get()) labels.erase(found);
         }
+        for (const auto entry : batch->wideLabels) wideLabels.erase(entry);
         recordedLabels.store(labels.size(), std::memory_order_relaxed);
     }
     batch->labelDwords.clear();
+    batch->wideLabels.clear();
     release(*batch);
 }
 

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

@@ -429,6 +429,38 @@ void lateLabelTests(Recorder& recorder) {
     Require(!Recorder::LookupLabel(a, 4, 5).has_value() && !Recorder::LookupLabel(b, 8, 5).has_value() && !recorder.PendingLabelIn(0x70000, 0x400), "(7) entries outlived their batch");
 }
 
+void largeLabelTests(Recorder& recorder) {
+    constexpr std::uint64_t slot = 0x71000, small = slot + 0x10, outside = slot - 4, wide = 0x72000;
+    const std::array<std::byte, 4> one{std::byte{1}, std::byte{0}, std::byte{0}, std::byte{0}};
+    std::array<std::byte, 256> clear{};
+    std::array<std::byte, Recorder::LabelTableBytes> limit{};
+    limit.fill(std::byte{3});
+    recorder.Sync();
+    Require(recorder.PendingLabels() == 0 && !Recorder::PendingLabelSince().has_value() && !Recorder::WideLabelIn(0, ~0ull), "the table is not empty before the large label tests");
+    recorder.NoteLabel(small, one, 100, 0);
+    recorder.NoteLabel(outside, one, 100, 0);
+    Recorder::CloseLabelGroup(AgcDriver::GuestMemory::TrackerGeneration());
+    Require(recorder.PendingLabels() == 2 && Recorder::LookupLabel(small, 4, 50).has_value(), "4-byte labels did not enter the table");
+    recorder.NoteLabel(slot, clear, 101, 0);
+    Recorder::CloseLabelGroup(AgcDriver::GuestMemory::TrackerGeneration());
+    Require(recorder.PendingLabels() == 1 && Recorder::LookupLabel(outside, 4, 50).has_value(), "a label larger than the table bound entered dwords or removed one outside its range");
+    Require(!Recorder::LookupLabel(small, 4, 50).has_value() && !Recorder::LookupLabel(slot, 8, 0).has_value(), "an older label under a larger one still composes the value");
+    Require(Recorder::WideLabelIn(slot, 4) && Recorder::WideLabelIn(slot + clear.size() - 4, 4) && Recorder::WideLabelIn(slot - 0x100, 0x104) && !Recorder::WideLabelIn(slot - 0x100, 0x100) && !Recorder::WideLabelIn(slot + clear.size(), 4), "a large label's range is wrong");
+    Require(recorder.PendingLabelIn(slot + 0x20, 0x20) && recorder.PendingWriteOverlaps(slot, clear.size()), "a large label is not pending");
+    Require(Recorder::PendingLabelSince().has_value(), "a large label did not start the label flush deadline");
+    recorder.NoteLabel(small, one, 102, 0);
+    Recorder::CloseLabelGroup(AgcDriver::GuestMemory::TrackerGeneration());
+    const auto later = Recorder::LookupLabel(small, 4, 101);
+    Require(later.has_value() && later->value == 1 && later->stamp == 102, "a label after a large one is not served");
+    recorder.NoteLabel(wide, limit, 103, 0);
+    Recorder::CloseLabelGroup(AgcDriver::GuestMemory::TrackerGeneration());
+    Require(recorder.PendingLabels() == 2 + Recorder::LabelTableBytes / 4 && !Recorder::WideLabelIn(wide, limit.size()), "a label at the table bound did not enter every dword");
+    const auto bound = Recorder::LookupLabel(wide + Recorder::LabelTableBytes - 8, 8, 0);
+    Require(bound.has_value() && bound->value == 0x0303030303030303ull && bound->stamp == 103, "a label at the table bound is not served");
+    recorder.Sync();
+    Require(recorder.PendingLabels() == 0 && !Recorder::WideLabelIn(0, ~0ull) && !recorder.PendingLabelIn(slot, clear.size()) && !recorder.PendingWriteOverlaps(slot, clear.size()), "large label tests left entries, ranges or writes behind");
+}
+
 // (6) GuestMemory::UnchangedSinceCollected: false outside the watched arena, and around a
 // MarkWritten or a CPU write of the block inside it.
 void unchangedSinceTests() {
@@ -2666,6 +2698,7 @@ int main() {
         batchStampTests(recorder);
         labelTests(recorder);
         lateLabelTests(recorder);
+        largeLabelTests(recorder);
         unchangedSinceTests();
         closeRaceTests(device, recorder);
         keyProofTests(device, recorder);