perf(agc): reuse heap mirror reference buffers between builds

core/libs/prx/libSceAgcDriver/CMakeLists.txt

@@ -345,6 +345,7 @@ set(agcDriverSources
         ${CMAKE_CURRENT_SOURCE_DIR}/Graphics/src/ShaderAddressResources.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/Graphics/src/GuestBufferMemory.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/Graphics/src/ResidentMirror.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/Graphics/src/ReferenceStorage.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/Graphics/src/BdaResources.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/Graphics/src/BufferAddress.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/Graphics/src/ShaderValidation.cpp

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

@@ -520,6 +520,7 @@ class GuestBufferMemory {
     // the barriers that order them after earlier recorded writes and before the shaders reading them.
     void recordGpuCopies(std::span<Region* const> copies, bool addressable);
     void takeHeapReferences();
+    void releaseHeapReferences() noexcept;
     Context context;
     bool stagingAllowed = false;
     bool addressStaging = true;

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

@@ -0,0 +1,46 @@
+#ifndef CORE_LIBS_PRX_LIBSCEAGCDRIVER_GRAPHICS_INCLUDE_REFERENCESTORAGE_HPP
+#define CORE_LIBS_PRX_LIBSCEAGCDRIVER_GRAPHICS_INCLUDE_REFERENCESTORAGE_HPP
+
+#include <array>
+#include <chrono>
+#include <cstddef>
+#include <cstdint>
+#include <mutex>
+#include <vector>
+
+namespace AgcDriver::Graphics {
+
+class ReferenceStorage {
+public:
+    static constexpr std::size_t Slots = 8;
+
+    explicit ReferenceStorage(std::size_t capacityBytes) noexcept;
+    ReferenceStorage(const ReferenceStorage&) = delete;
+    ReferenceStorage& operator=(const ReferenceStorage&) = delete;
+
+    using Clock = std::chrono::steady_clock;
+    static constexpr std::chrono::milliseconds IdleLimit{2000};
+
+    std::vector<std::byte> Take(std::size_t bytes, Clock::time_point now = Clock::now()) noexcept;
+    void Return(std::vector<std::byte>&& storage, Clock::time_point now = Clock::now()) noexcept;
+    std::size_t ReleaseIdle(Clock::time_point now = Clock::now()) noexcept;
+    std::size_t IdleBytes() const noexcept;
+    std::size_t IdleCount() const noexcept;
+
+    static ReferenceStorage& Shared() noexcept;
+
+private:
+    std::size_t releaseIdle(Clock::time_point now, std::array<std::vector<std::byte>, Slots>& released) noexcept;
+
+    mutable std::mutex mutex;
+    std::array<std::vector<std::byte>, Slots> idle;
+    std::array<std::uint64_t, Slots> returned{};
+    std::array<Clock::time_point, Slots> returnedAt{};
+    std::uint64_t returns = 0;
+    std::size_t bytes = 0;
+    std::size_t capacityBytes;
+};
+
+}
+
+#endif

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

@@ -10,6 +10,7 @@
 #include "prx/libSceAgcDriver/Graphics/include/Texture.hpp"
 #include "prx/libSceAgcDriver/Graphics/include/Recorder.hpp"
 #include "prx/libSceAgcDriver/Graphics/include/ResidentMirror.hpp"
+#include "prx/libSceAgcDriver/Graphics/include/ReferenceStorage.hpp"
 #include "prx/libSceAgcDriver/Graphics/include/ShaderResources.hpp"
 #include "prx/libSceAgcDriver/Graphics/include/UnitShadow.hpp"
 #include "prx/libSceAgcDriver/Execution/include/GuestMemory.hpp"
@@ -3328,6 +3329,7 @@ std::shared_ptr<ImageMirror> findMirror(const Context& context, std::uint64_t be
 }
 
 void sweepMirrors(std::uint64_t neededBytes, std::uint64_t protectAddress) {
+    ReferenceStorage::Shared().ReleaseIdle();
     const bool released = releaseExpiredJoinedHeaps();
     auto space = Spaces().current.load();
     auto& state = Mirrors();
@@ -5351,6 +5353,7 @@ void ShutdownGuestBufferWorkers() {
 }
 
 GuestBufferMemory::~GuestBufferMemory() {
+    releaseHeapReferences();
     static const bool profile = std::getenv("APS5_PROFILE_DRAW") != nullptr;
     if (!profile) return;
     // A staged region whose last use recorded its copy-in but no copy-back (the use threw before
@@ -6465,9 +6468,14 @@ std::vector<std::pair<std::uint64_t, std::uint64_t>> mergeWrites(std::vector<std
 
 }
 
-void GuestBufferMemory::takeHeapReferences() {
+void GuestBufferMemory::releaseHeapReferences() noexcept {
+    for (auto& reference : heapReferences) ReferenceStorage::Shared().Return(std::move(reference.second));
     heapReferences.clear();
     heapReadbacks.clear();
+}
+
+void GuestBufferMemory::takeHeapReferences() {
+    releaseHeapReferences();
     for (const auto& [begin, end] : mergeWrites(writes)) {
         const auto* found = owner(begin, end);
         if (found == nullptr || found->mirror == nullptr || !found->mirror->heap || !found->mirror->writable) continue;
@@ -6476,10 +6484,16 @@ void GuestBufferMemory::takeHeapReferences() {
         const auto served = std::min(end, servedEnd(*found));
         if (begin >= served) continue;
         growResident(mirror, {{begin, served}});
-        std::vector<std::byte> bytes(static_cast<std::size_t>(served - begin));
-        if (mirror.resident != nullptr) mirror.resident->Read(begin, bytes);
-        else std::memcpy(bytes.data(), mirror.buffer->Bytes().data() + (begin - mirror.base), bytes.size());
-        heapReferences.emplace_back(begin, std::move(bytes));
+        const auto size = static_cast<std::size_t>(served - begin);
+        auto storage = ReferenceStorage::Shared().Take(size);
+        if (mirror.resident != nullptr) {
+            storage.resize(size);
+            mirror.resident->Read(begin, storage);
+        } else {
+            const auto bytes = mirror.buffer->Bytes().subspan(static_cast<std::size_t>(begin - mirror.base), size);
+            storage.assign(bytes.begin(), bytes.end());
+        }
+        heapReferences.emplace_back(begin, std::move(storage));
     }
 }
 
@@ -6580,8 +6594,7 @@ void GuestBufferMemory::WriteBack() {
     committed = true;
     lease.clear();
     space.reset();
-    heapReferences.clear();
-    heapReadbacks.clear();
+    releaseHeapReferences();
 }
 
 std::optional<std::vector<std::pair<std::uint64_t, std::uint64_t>>> GuestBufferMemory::DirectRegions() const {

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

@@ -4,6 +4,7 @@
 #include "prx/libSceAgcDriver/Execution/include/Mutex.hpp"
 #include "prx/libc/include/HostThreadLocal.hpp"
 #include "prx/libSceAgcDriver/Graphics/include/Recorder.hpp"
+#include "prx/libSceAgcDriver/Graphics/include/ReferenceStorage.hpp"
 #include "prx/libSceAgcDriver/Graphics/include/DepthSurface.hpp"
 #include "prx/libkernel/Time/include/SleepWake.hpp"
 #include "prx/libSceAgcDriver/Graphics/include/Resources.hpp"
@@ -2024,6 +2025,7 @@ void Recorder::NoteAccess(CommandClass which, const Access& access) {
 void Recorder::CountPresent() {
     timingPresents.fetch_add(1, std::memory_order_relaxed);
     presentSerial.fetch_add(1, std::memory_order_relaxed);
+    ReferenceStorage::Shared().ReleaseIdle();
 }
 
 std::uint64_t Recorder::Presents() {

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

@@ -0,0 +1,89 @@
+#include "prx/libSceAgcDriver/Graphics/include/ReferenceStorage.hpp"
+
+#include <utility>
+
+namespace AgcDriver::Graphics {
+
+ReferenceStorage::ReferenceStorage(std::size_t capacityBytes) noexcept : capacityBytes(capacityBytes) {}
+
+std::size_t ReferenceStorage::releaseIdle(Clock::time_point now, std::array<std::vector<std::byte>, Slots>& released) noexcept {
+    std::size_t freed = 0;
+    for (std::size_t slot = 0; slot < Slots; ++slot) {
+        if (idle[slot].capacity() == 0 || now - returnedAt[slot] < IdleLimit) continue;
+        freed += idle[slot].capacity();
+        bytes -= idle[slot].capacity();
+        released[slot] = std::move(idle[slot]);
+        idle[slot] = std::vector<std::byte>();
+    }
+    return freed;
+}
+
+std::size_t ReferenceStorage::ReleaseIdle(Clock::time_point now) noexcept {
+    std::array<std::vector<std::byte>, Slots> released;
+    std::lock_guard lock(mutex);
+    return releaseIdle(now, released);
+}
+
+std::vector<std::byte> ReferenceStorage::Take(std::size_t requested, Clock::time_point now) noexcept {
+    std::array<std::vector<std::byte>, Slots> released;
+    std::lock_guard lock(mutex);
+    releaseIdle(now, released);
+    std::size_t best = Slots;
+    for (std::size_t slot = 0; slot < Slots; ++slot) {
+        const auto capacity = idle[slot].capacity();
+        if (capacity == 0 || capacity < requested) continue;
+        if (best == Slots || capacity < idle[best].capacity()) best = slot;
+    }
+    if (best == Slots) return {};
+    auto storage = std::move(idle[best]);
+    idle[best] = std::vector<std::byte>();
+    bytes -= storage.capacity();
+    return storage;
+}
+
+void ReferenceStorage::Return(std::vector<std::byte>&& storage, Clock::time_point now) noexcept {
+    auto kept = std::move(storage);
+    const auto capacity = kept.capacity();
+    if (capacity == 0 || capacity > capacityBytes) return;
+    std::array<std::vector<std::byte>, Slots> evicted;
+    std::array<std::vector<std::byte>, Slots> released;
+    std::lock_guard lock(mutex);
+    releaseIdle(now, released);
+    std::size_t free = Slots;
+    for (std::size_t slot = 0; slot < Slots; ++slot) {
+        if (idle[slot].capacity() == 0) free = slot;
+    }
+    while (free == Slots || bytes + capacity > capacityBytes) {
+        std::size_t oldest = Slots;
+        for (std::size_t slot = 0; slot < Slots; ++slot) {
+            if (idle[slot].capacity() != 0 && (oldest == Slots || returned[slot] < returned[oldest])) oldest = slot;
+        }
+        bytes -= idle[oldest].capacity();
+        evicted[oldest] = std::move(idle[oldest]);
+        idle[oldest] = std::vector<std::byte>();
+        free = oldest;
+    }
+    bytes += capacity;
+    returned[free] = ++returns;
+    returnedAt[free] = now;
+    idle[free] = std::move(kept);
+}
+
+std::size_t ReferenceStorage::IdleBytes() const noexcept {
+    std::lock_guard lock(mutex);
+    return bytes;
+}
+
+std::size_t ReferenceStorage::IdleCount() const noexcept {
+    std::lock_guard lock(mutex);
+    std::size_t count = 0;
+    for (const auto& storage : idle) count += storage.capacity() != 0 ? 1 : 0;
+    return count;
+}
+
+ReferenceStorage& ReferenceStorage::Shared() noexcept {
+    static ReferenceStorage storage(std::size_t{2} << 30u);
+    return storage;
+}
+
+}

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

@@ -3585,6 +3585,67 @@ void RunSparseHeapMirrorBudgetTests(const Context& context, const BdaTestAccess&
     std::cout << "a sparse heap whose mirror would pass the heap mirror budget is copied\n";
 }
 
+void RunHeapMirrorReferenceSizeTests(const Context& context, const BdaTestAccess& access) {
+    constexpr std::size_t block = 65536;
+    constexpr std::size_t heapBytes = 4 * block;
+#ifdef _WIN32
+    if (!GuestArena::GuestArenaAvailable_nid_postfix() || !GuestArena::GuestArenaWriteWatched_nid_postfix()) {
+        std::cout << "guest arena unavailable or not write-watched: heap mirror reference sizes not tested\n";
+        return;
+    }
+    auto* guest = static_cast<std::uint8_t*>(GuestArena::GuestArenaAllocate_nid_postfix(heapBytes, block));
+    Require(guest != nullptr, "cannot allocate the heap in the guest arena");
+    GuestArena::GuestArenaCommit_nid_postfix(guest, heapBytes, PAGE_READWRITE, heapBytes);
+#else
+    if (!GuestWriteWatch::GuestWriteWatchAvailable_nid_postfix()) {
+        std::cout << "write watch unavailable: heap mirror reference sizes not tested\n";
+        return;
+    }
+    void* mapping = mmap(nullptr, heapBytes + block, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+    Require(mapping != MAP_FAILED, "cannot map the heap");
+    auto* guest = reinterpret_cast<std::uint8_t*>((reinterpret_cast<std::uintptr_t>(mapping) + block - 1) & ~std::uintptr_t{block - 1});
+    GuestWriteWatch::GuestWriteWatchRegister_nid_postfix(guest, heapBytes);
+#endif
+    std::memset(guest, 0x11, heapBytes);
+    const auto address = reinterpret_cast<std::uintptr_t>(guest);
+    const auto registry = [&](bool add) {
+        auto* mutation = GuestAllocations::GuestAllocationsBegin_nid_postfix();
+        if (add) GuestAllocations::GuestAllocationsAdd_nid_postfix(mutation, guest, heapBytes, true, true, true);
+        else GuestAllocations::GuestAllocationsRemove_nid_postfix(mutation, guest);
+        GuestAllocations::GuestAllocationsEnd_nid_postfix(mutation);
+    };
+    registry(true);
+    const auto before = MirrorCounters();
+    const auto store = [&](std::uint64_t written, std::size_t bytes, std::size_t at, std::byte value) {
+        GuestBufferMemory leased(context);
+        leased.AcquireRegistered();
+        leased.AddWritable(written, bytes);
+        leased.Upload(true);
+        std::uint32_t adjustment = 0;
+        const auto view = leased.Descriptor(written, static_cast<std::uint32_t>(bytes), adjustment);
+        access.bytes(view.buffer)[view.offset + adjustment + at] = value;
+        leased.WriteBack();
+    };
+    store(address + 4096, 8192, 8000, std::byte{0x7a});
+    Require(MirrorCounters().heapMirrors == before.heapMirrors + 1, "the writable heap was not mirrored");
+    Require(guest[4096 + 8000] == 0x7a, "the larger write-back lost the GPU's store");
+    store(address + 16384, 2048, 2000, std::byte{0x7b});
+    Require(guest[16384 + 2000] == 0x7b && guest[4096 + 8000] == 0x7a, "a smaller write-back after a larger one lost the GPU's store");
+    store(address + 32768, 8192, 100, std::byte{0x7c});
+    Require(guest[32768 + 100] == 0x7c && guest[16384 + 2000] == 0x7b, "a write-back after a smaller one lost the GPU's store");
+    registry(false);
+    ClearImageMirrors(context.device);
+    ClearSnapshotCopies(context.device);
+#ifdef _WIN32
+    GuestArena::GuestArenaReset_nid_postfix(guest, heapBytes);
+    GuestArena::GuestArenaRelease_nid_postfix(guest, heapBytes);
+#else
+    GuestWriteWatch::GuestWriteWatchUnregister_nid_postfix(guest, heapBytes);
+    munmap(mapping, heapBytes + block);
+#endif
+    std::cout << "a writable heap mirror writes back a smaller range after a larger one\n";
+}
+
 void RunHeapMirrorRunUnwatchTests(const Context& context, const BdaTestAccess& access) {
     constexpr std::size_t block = 65536;
     constexpr std::size_t heapBytes = 32 * block;

core/libs/prx/libSceAgcDriver/tests/BdaTests.hpp

@@ -22,6 +22,7 @@ void RunSparseHeapMirrorUnwatchTests(const AgcDriver::Graphics::Context& context
 void RunSparseHeapMirrorBudgetTests(const AgcDriver::Graphics::Context& context, const BdaTestAccess& access);
 void RunHeapMirrorRunUnwatchTests(const AgcDriver::Graphics::Context& context, const BdaTestAccess& access);
 void RunAddressCopyCapTests(const AgcDriver::Graphics::Context& context, const BdaTestAccess& access, bool remembered, int frames);
+void RunHeapMirrorReferenceSizeTests(const AgcDriver::Graphics::Context& context, const BdaTestAccess& access);
 void RunAddressCopyCapMirrorTests(const AgcDriver::Graphics::Context& context, const BdaTestAccess& access, int frames);
 void RunUnwatchedMirrorImportTests(const AgcDriver::Graphics::Context& context, const BdaTestAccess& access, int frames);
 void RunGuestAllocationTests();

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

@@ -12934,6 +12934,27 @@ int main(int argc, char** argv) {
             Require(mock.live == 0, "the sparse heap mirror budget test leaked Vulkan objects");
             return 0;
         }
+        if (argc == 2 && std::string_view(argv[1]) == "--heap-mirror-reference-size") {
+            {
+                mock = MockVulkan{};
+                auto context = mockContext();
+                context.bufferDeviceAddress = true;
+                context.limits.maxStorageBufferRange = 1u << 27;
+                RunHeapMirrorReferenceSizeTests(context, {
+                    [](VkBuffer buffer) -> std::span<std::byte> { return mockBufferMemory(buffer); },
+                    {},
+                    [](VkDeviceAddress address) {
+                        const auto offset = address - 0x100000000000ULL;
+                        const auto buffer = reinterpret_cast<VkBuffer>(offset / 0x10000);
+                        return mockBufferMemory(buffer).subspan(offset % 0x10000);
+                    },
+                    {},
+                    {}
+                });
+            }
+            Require(mock.live == 0, "the heap mirror reference size test leaked Vulkan objects");
+            return 0;
+        }
         if (argc == 2 && std::string_view(argv[1]) == "--heap-mirror-run-unwatch") {
             {
                 mock = MockVulkan{};