perf(agc): find a range's newest writer through a granule index instead of scanning the ring

core/libs/prx/libSceAgcDriver/CMakeLists.txt

@@ -277,6 +277,7 @@ set(agcDriverSources
         ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/Driver/Memory/CaptureValidation.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/Driver/Memory/WriteEvidence.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/Driver/Memory/WrittenBuffers.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/Driver/Memory/WrittenBufferIndex.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/Driver/Packets/CommandValidation.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/Driver/Packets/PacketExecution.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/Execution/src/Driver/Packets/PacketFlush.cpp
@@ -866,6 +867,24 @@ if(BUILD_TESTING AND NOT WIN32)
     set_tests_properties(agc_collect_memo_generation PROPERTIES TIMEOUT 30)
 endif()
 
+add_test_executable(agc_written_buffer_index_tests EXCLUDE_FROM_ALL
+        tests/WrittenBufferIndex.cpp
+        Execution/src/Driver/Memory/WrittenBufferIndex.cpp
+)
+target_include_directories(agc_written_buffer_index_tests PRIVATE ${LIBS_INCLUDE_DIR})
+target_link_libraries(agc_written_buffer_index_tests PRIVATE libc)
+set_target_properties(agc_written_buffer_index_tests PROPERTIES CXX_EXTENSIONS OFF)
+configure_windows_unwind(agc_written_buffer_index_tests)
+if(BUILD_TESTING)
+    set_target_properties(agc_written_buffer_index_tests PROPERTIES EXCLUDE_FROM_ALL OFF)
+    add_test(NAME agc_written_buffer_index COMMAND agc_written_buffer_index_tests)
+    set_tests_properties(agc_written_buffer_index PROPERTIES TIMEOUT 60)
+    if(WIN32)
+        set_tests_properties(agc_written_buffer_index PROPERTIES ENVIRONMENT_MODIFICATION
+            "PATH=path_list_prepend:$<TARGET_FILE_DIR:libc>")
+    endif()
+endif()
+
 add_test_executable(agc_thread_priority_tests EXCLUDE_FROM_ALL
         tests/ThreadPriority.cpp
         Execution/src/ThreadPriority.cpp

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

@@ -482,7 +482,7 @@ class Driver {
     std::atomic<void (*)()> aheadAdoptHook{nullptr};
 
     AgcDriver::Mutex writtenBuffersMutex;
-    std::deque<WrittenBuffer> writtenBuffers;
+    WrittenBufferIndex writtenBuffers;
     std::uint64_t writtenBufferSerial = 0;
     std::unordered_map<std::uint64_t, DwordEvidence> dwordEvidence;
     std::atomic<std::uint64_t> observedUnchanged{0};

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

@@ -1,6 +1,7 @@
 #ifndef CORE_LIBS_PRX_LIBSCEAGCDRIVER_EXECUTION_INCLUDE_DRIVER_WRITEEVIDENCE_HPP
 #define CORE_LIBS_PRX_LIBSCEAGCDRIVER_EXECUTION_INCLUDE_DRIVER_WRITEEVIDENCE_HPP
 
+#include "prx/libSceAgcDriver/Execution/include/Driver/Memory/WrittenBufferIndex.hpp"
 #include "prx/libSceAgcDriver/Graphics/include/Recorder.hpp"
 #include <atomic>
 #include <chrono>
@@ -11,21 +12,8 @@
 
 namespace AgcDriver::DriverDetail {
 
-inline constexpr std::size_t WrittenBufferRing = 4096;
-
 inline constexpr std::size_t DwordEvidenceEntries = 65536;
 
-struct WrittenBuffer {
-    std::uint64_t program;
-    std::uint64_t begin;
-    std::uint64_t end;
-    std::uint64_t serial;
-    std::uint32_t queue;
-    bool atomic;
-    std::shared_ptr<const std::vector<std::byte>> value;
-    std::uint64_t generation = 0;
-};
-
 struct DwordEvidence {
     std::uint32_t streak = 0;
     std::uint32_t changed = 0;

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

@@ -0,0 +1,54 @@
+#ifndef CORE_LIBS_PRX_LIBSCEAGCDRIVER_EXECUTION_INCLUDE_DRIVER_MEMORY_WRITTENBUFFERINDEX_HPP
+#define CORE_LIBS_PRX_LIBSCEAGCDRIVER_EXECUTION_INCLUDE_DRIVER_MEMORY_WRITTENBUFFERINDEX_HPP
+
+#include <cstddef>
+#include <cstdint>
+#include <deque>
+#include <memory>
+#include <optional>
+#include <unordered_map>
+#include <vector>
+
+namespace AgcDriver::DriverDetail {
+
+inline constexpr std::size_t WrittenBufferRing = 4096;
+
+struct WrittenBuffer {
+    std::uint64_t program;
+    std::uint64_t begin;
+    std::uint64_t end;
+    std::uint64_t serial;
+    std::uint32_t queue;
+    bool atomic;
+    std::shared_ptr<const std::vector<std::byte>> value;
+    std::uint64_t generation = 0;
+};
+
+class WrittenBufferIndex {
+public:
+    explicit WrittenBufferIndex(std::size_t capacity = WrittenBufferRing);
+
+    void Push(WrittenBuffer buffer);
+    std::optional<WrittenBuffer> Newest(std::uint64_t begin, std::uint64_t end) const;
+    const std::deque<WrittenBuffer>& Entries() const;
+
+private:
+    static constexpr unsigned GranuleBits = 16;
+    static constexpr std::uint64_t WideGranules = 16;
+
+    static std::uint64_t firstGranule(const WrittenBuffer& buffer);
+    static std::uint64_t lastGranule(const WrittenBuffer& buffer);
+    static bool wide(const WrittenBuffer& buffer);
+    std::optional<WrittenBuffer> scan(std::uint64_t begin, std::uint64_t end) const;
+    void evictOldest();
+
+    std::size_t capacity;
+    std::uint64_t evicted = 0;
+    std::deque<WrittenBuffer> entries;
+    std::deque<std::uint64_t> wideEntries;
+    std::unordered_map<std::uint64_t, std::deque<std::uint64_t>> granules;
+};
+
+}
+
+#endif

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

@@ -128,8 +128,7 @@ void Driver::noteCopyWriter(std::uint64_t program, std::uint64_t begin, std::uin
         knownValueEntries.fetch_add(1, std::memory_order_relaxed);
     }
     std::lock_guard lock(writtenBuffersMutex);
-    writtenBuffers.push_back({program, begin, end, ++writtenBufferSerial, queue, false, std::move(known), generation});
-    while (writtenBuffers.size() > WrittenBufferRing) writtenBuffers.pop_front();
+    writtenBuffers.Push({program, begin, end, ++writtenBufferSerial, queue, false, std::move(known), generation});
 }
 
 bool Driver::copyKnownValues() {

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

@@ -0,0 +1,78 @@
+#include "prx/libSceAgcDriver/Execution/include/Driver/Memory/WrittenBufferIndex.hpp"
+#include <algorithm>
+#include <utility>
+
+namespace AgcDriver::DriverDetail {
+
+WrittenBufferIndex::WrittenBufferIndex(std::size_t capacity) : capacity(std::max<std::size_t>(capacity, 1)) {}
+
+std::uint64_t WrittenBufferIndex::firstGranule(const WrittenBuffer& buffer) {
+    return buffer.begin >> GranuleBits;
+}
+
+std::uint64_t WrittenBufferIndex::lastGranule(const WrittenBuffer& buffer) {
+    return buffer.end > buffer.begin ? (buffer.end - 1) >> GranuleBits : buffer.begin >> GranuleBits;
+}
+
+bool WrittenBufferIndex::wide(const WrittenBuffer& buffer) {
+    return lastGranule(buffer) - firstGranule(buffer) >= WideGranules;
+}
+
+void WrittenBufferIndex::Push(WrittenBuffer buffer) {
+    const auto sequence = evicted + entries.size();
+    if (wide(buffer)) {
+        wideEntries.push_back(sequence);
+    } else {
+        for (auto granule = firstGranule(buffer); granule <= lastGranule(buffer); ++granule) granules[granule].push_back(sequence);
+    }
+    entries.push_back(std::move(buffer));
+    while (entries.size() > capacity) evictOldest();
+}
+
+void WrittenBufferIndex::evictOldest() {
+    const auto& oldest = entries.front();
+    if (wide(oldest)) {
+        wideEntries.pop_front();
+    } else {
+        for (auto granule = firstGranule(oldest); granule <= lastGranule(oldest); ++granule) {
+            const auto found = granules.find(granule);
+            found->second.pop_front();
+            if (found->second.empty()) granules.erase(found);
+        }
+    }
+    entries.pop_front();
+    ++evicted;
+}
+
+std::optional<WrittenBuffer> WrittenBufferIndex::scan(std::uint64_t begin, std::uint64_t end) const {
+    for (auto it = entries.rbegin(); it != entries.rend(); ++it) {
+        if (begin < it->end && it->begin < end) return *it;
+    }
+    return std::nullopt;
+}
+
+std::optional<WrittenBuffer> WrittenBufferIndex::Newest(std::uint64_t begin, std::uint64_t end) const {
+    if (end <= begin || ((end - 1) >> GranuleBits) - (begin >> GranuleBits) >= WideGranules) return scan(begin, end);
+    std::optional<std::uint64_t> newest;
+    const auto consider = [&](const std::deque<std::uint64_t>& sequences) {
+        for (auto it = sequences.rbegin(); it != sequences.rend() && (!newest || *it > *newest); ++it) {
+            const auto& entry = entries[static_cast<std::size_t>(*it - evicted)];
+            if (begin < entry.end && entry.begin < end) {
+                newest = *it;
+                return;
+            }
+        }
+    };
+    consider(wideEntries);
+    for (auto granule = begin >> GranuleBits; granule <= (end - 1) >> GranuleBits; ++granule) {
+        if (const auto found = granules.find(granule); found != granules.end()) consider(found->second);
+    }
+    if (!newest) return std::nullopt;
+    return entries[static_cast<std::size_t>(*newest - evicted)];
+}
+
+const std::deque<WrittenBuffer>& WrittenBufferIndex::Entries() const {
+    return entries;
+}
+
+}

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

@@ -48,9 +48,8 @@ void Driver::noteWrittenBuffers(std::uint64_t program, std::uint32_t queue, cons
     std::lock_guard lock(writtenBuffersMutex);
     const auto serial = ++writtenBufferSerial;
     forEachWrittenBuffer(compiled, [&](std::uint32_t, std::uint64_t begin, std::uint64_t end, bool atomic) {
-        writtenBuffers.push_back({program, begin, end, serial, queue, atomic});
+        writtenBuffers.Push({program, begin, end, serial, queue, atomic});
     });
-    while (writtenBuffers.size() > WrittenBufferRing) writtenBuffers.pop_front();
 }
 
 std::vector<std::pair<std::uint64_t, std::uint64_t>> DispatchWriteRanges(const ShaderRecompiler::RecompileResult& compiled, bool& bounded) {
@@ -82,8 +81,7 @@ std::vector<std::pair<std::uint64_t, std::uint64_t>> DispatchWriteRanges(const S
 void Driver::noteForeignWriter(std::uint64_t begin, std::uint64_t end, std::uint32_t queue) {
     if (!foreignWriters() || !(writeEvidenceEnabled() || traceCapSync()) || end <= begin) return;
     std::lock_guard lock(writtenBuffersMutex);
-    writtenBuffers.push_back({0, begin, end, ++writtenBufferSerial, queue, false});
-    while (writtenBuffers.size() > WrittenBufferRing) writtenBuffers.pop_front();
+    writtenBuffers.Push({0, begin, end, ++writtenBufferSerial, queue, false});
 }
 
 void Driver::noteDrawWriters(std::span<const Graphics::CompiledShader> stages, std::uint32_t queue) {
@@ -196,10 +194,7 @@ std::vector<std::pair<std::uint64_t, std::uint64_t>> QueuedDrawReadRanges(const
 }
 
 std::optional<WrittenBuffer> Driver::newestWriterLocked(std::uint64_t begin, std::uint64_t end) const {
-    for (auto it = writtenBuffers.rbegin(); it != writtenBuffers.rend(); ++it) {
-        if (begin < it->end && it->begin < end) return *it;
-    }
-    return std::nullopt;
+    return writtenBuffers.Newest(begin, end);
 }
 
 std::optional<WrittenBuffer> Driver::newestWriter(std::uint64_t begin, std::uint64_t end) {
@@ -211,7 +206,7 @@ std::string Driver::describeWriters(std::uint64_t begin, std::uint64_t end) {
     std::string text;
     std::lock_guard lock(writtenBuffersMutex);
     int shown = 0;
-    for (auto it = writtenBuffers.rbegin(); it != writtenBuffers.rend() && shown < 3; ++it) {
+    for (auto it = writtenBuffers.Entries().rbegin(); it != writtenBuffers.Entries().rend() && shown < 3; ++it) {
         if (begin >= it->end || it->begin >= end) continue;
         char item[128];
         if (it->program == 0) std::snprintf(item, sizeof(item), " [foreign q0x%x 0x%llx+0x%llx age %llu]", it->queue, static_cast<unsigned long long>(it->begin), static_cast<unsigned long long>(it->end - it->begin), static_cast<unsigned long long>(writtenBufferSerial - it->serial));

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

@@ -0,0 +1,95 @@
+#include "prx/libSceAgcDriver/Execution/include/Driver/Memory/WrittenBufferIndex.hpp"
+#include <cstdint>
+#include <iostream>
+#include <optional>
+#include <random>
+#include <sstream>
+#include <stdexcept>
+#include <string>
+
+namespace {
+
+using AgcDriver::DriverDetail::WrittenBuffer;
+using AgcDriver::DriverDetail::WrittenBufferIndex;
+
+constexpr std::uint64_t Base = std::uint64_t{1} << 32;
+
+void Require(bool condition, const std::string& message) {
+    if (!condition) throw std::runtime_error(message);
+}
+
+WrittenBuffer Make(std::uint64_t begin, std::uint64_t end, std::uint64_t serial) {
+    return WrittenBuffer{1, begin, end, serial, 0, false};
+}
+
+std::optional<WrittenBuffer> Scan(const WrittenBufferIndex& index, std::uint64_t begin, std::uint64_t end) {
+    for (auto it = index.Entries().rbegin(); it != index.Entries().rend(); ++it) {
+        if (begin < it->end && it->begin < end) return *it;
+    }
+    return std::nullopt;
+}
+
+void RequireSame(const WrittenBufferIndex& index, std::uint64_t begin, std::uint64_t end) {
+    const auto expected = Scan(index, begin, end);
+    const auto found = index.Newest(begin, end);
+    if (expected.has_value() == found.has_value() && (!found || found->serial == expected->serial)) return;
+    std::ostringstream text;
+    text << std::hex << "query 0x" << begin << "..0x" << end << " differs from the linear scan";
+    throw std::runtime_error(text.str());
+}
+
+std::uint64_t Length(std::mt19937_64& random) {
+    const auto kind = random() % 100;
+    if (kind < 3) return 0;
+    if (kind < 70) return 4 * (1 + random() % 1024);
+    if (kind < 95) return 4 * (1 + random() % 65536);
+    return 0x200000 + random() % 0x600000;
+}
+
+void NewestOverlapWins() {
+    WrittenBufferIndex index(8);
+    index.Push(Make(Base, Base + 16, 1));
+    index.Push(Make(Base + 8, Base + 32, 2));
+    Require(index.Newest(Base + 8, Base + 12)->serial == 2, "the newer of two overlapping writers was not returned");
+    Require(index.Newest(Base, Base + 4)->serial == 1, "the only writer of a range was not returned");
+    Require(!index.Newest(Base + 32, Base + 36), "a range nothing wrote returned a writer");
+}
+
+void OldestEvicted() {
+    WrittenBufferIndex index(2);
+    index.Push(Make(Base, Base + 4, 1));
+    index.Push(Make(Base + 0x100000, Base + 0x100004, 2));
+    index.Push(Make(Base + 0x200000, Base + 0x200004, 3));
+    Require(index.Entries().size() == 2, "the ring kept more entries than its capacity");
+    Require(!index.Newest(Base, Base + 4), "an evicted writer was returned");
+    Require(index.Newest(Base + 0x200000, Base + 0x200004)->serial == 3, "the newest writer was not returned after an eviction");
+}
+
+void MatchesLinearScan() {
+    std::mt19937_64 random(0x5eed);
+    WrittenBufferIndex index(64);
+    for (std::uint64_t serial = 1; serial <= 20000; ++serial) {
+        const auto begin = Base + 4 * (random() % (0x800000 / 4));
+        index.Push(Make(begin, begin + Length(random), serial));
+        Require(index.Entries().size() <= 64, "the ring kept more entries than its capacity");
+        for (int query = 0; query < 8; ++query) {
+            const auto at = Base + 4 * (random() % (0xa00000 / 4));
+            RequireSame(index, at, at + Length(random));
+        }
+    }
+}
+
+}
+
+int main() {
+    try {
+        NewestOverlapWins();
+        OldestEvicted();
+        MatchesLinearScan();
+    } catch (const std::exception& error) {
+        std::cerr << "written buffer index tests failed: " << error.what() << "\n";
+        return 1;
+    }
+    std::cout << "written buffer index tests passed\n";
+    return 0;
+}