fix(agc): keep guest stores over freed storage-image memory

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

@@ -50,6 +50,8 @@ void Write(std::uint64_t address, std::span<const std::byte> source, std::size_t
 // Stores the parts of `current` that differ from `original` (the guest bytes the GPU started from), so
 // guest writes made meanwhile to untouched bytes survive. Compares in 256-byte blocks.
 void WriteChanged(std::uint64_t address, std::span<const std::byte> current, std::span<const std::byte> original);
+bool WriteChangedOverOriginal(std::uint64_t address, std::span<const std::byte> current, std::span<const std::byte> original);
+bool WriteChangedOverOriginalCommitted(std::uint64_t address, std::span<const std::byte> current, std::span<const std::byte> original);
 
 // Write tracking. The guest arena is reserved with page write watching, so CollectWrites learns which
 // pages of a range the CPU wrote since they were last collected, stamps their 64 KiB blocks with a new

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

@@ -50,6 +50,7 @@ bool Driver::preparePacketMemory(const Submission& submission, QueueState& queue
                 countLabelOutcome(reason);
                 ++immediateLabels;
             } else {
+                if (localDevice != nullptr) Graphics::StorageTexture::StoreBeforeGuestCompletion("interrupt");
                 ++noOpLabels;
             }
             if (reason == 0 || reason == 5) {

core/libs/prx/libSceAgcDriver/Execution/src/GuestMemory.cpp

@@ -2225,6 +2225,52 @@ void WriteChanged(std::uint64_t address, std::span<const std::byte> current, std
     });
 }
 
+bool WriteChangedOverOriginal(std::uint64_t address, std::span<const std::byte> current, std::span<const std::byte> original) {
+    require(current.size() == original.size(), "write-back snapshot sizes differ");
+    if (current.empty()) return true;
+    const ReadSiteScope site(ReadSite::Store);
+    FlushGpuWrites(address, current.size());
+    const TimedAccess timed(CounterChangedWrite, current.size());
+    auto* destination = reinterpret_cast<std::byte*>(address);
+    CheckRange(destination, current.size(), 1, true);
+    const auto size = current.size();
+    bool complete = true;
+    storeOwn(address, size, [&] {
+        std::size_t firstChanged = size;
+        std::size_t lastChanged = 0;
+        const auto exchange = [&]<typename TPiece>(std::size_t at) {
+            TPiece expected, desired;
+            std::memcpy(&expected, original.data() + at, sizeof(TPiece));
+            std::memcpy(&desired, current.data() + at, sizeof(TPiece));
+            if (expected == desired) return;
+            if (!std::atomic_ref(*reinterpret_cast<TPiece*>(destination + at)).compare_exchange_strong(expected, desired, std::memory_order_relaxed)) {
+                complete = false;
+                return;
+            }
+            firstChanged = std::min(firstChanged, at);
+            lastChanged = std::max(lastChanged, at + sizeof(TPiece));
+        };
+        std::size_t at = 0;
+        for (; at < size && (address + at) % sizeof(std::uint64_t) != 0; ++at) exchange.template operator()<std::uint8_t>(at);
+        for (; at + sizeof(std::uint64_t) <= size; at += sizeof(std::uint64_t)) exchange.template operator()<std::uint64_t>(at);
+        for (; at < size; ++at) exchange.template operator()<std::uint8_t>(at);
+        return firstChanged < lastChanged ? std::pair<std::uint64_t, std::uint64_t>{address + firstChanged, address + lastChanged} : std::pair<std::uint64_t, std::uint64_t>{0, 0};
+    });
+    return complete;
+}
+
+bool WriteChangedOverOriginalCommitted(std::uint64_t address, std::span<const std::byte> current, std::span<const std::byte> original) {
+    require(current.size() == original.size(), "write-back snapshot sizes differ");
+    if (Accessible(reinterpret_cast<const void*>(address), current.size(), true)) return WriteChangedOverOriginal(address, current, original);
+    bool complete = true;
+    for (const auto& [begin, end] : CommittedRanges(address, current.size(), true)) {
+        const auto offset = static_cast<std::size_t>(begin - address);
+        const auto length = static_cast<std::size_t>(end - begin);
+        complete = WriteChangedOverOriginal(begin, current.subspan(offset, length), original.subspan(offset, length)) && complete;
+    }
+    return complete;
+}
+
 void Write(std::uint64_t address, std::span<const std::byte> source, std::size_t alignment) {
     if (source.empty()) return;
     // A store, as in WriteChanged.

core/libs/prx/libSceAgcDriver/Execution/src/VulkanDevice.cpp

@@ -2160,6 +2160,7 @@ bool VulkanDevice::StoresPendingOver(std::uint64_t address, std::size_t bytes) c
 int VulkanDevice::WriteLabelOnGpu(std::uint64_t address, std::span<const std::byte> bytes, std::uint64_t stamp, std::uint32_t queue, bool reapFirst) {
     static const bool storeAtLabel = Graphics::StorageTexture::StoreAtLabelRequested(std::getenv("APS5_STORE_AT_LABEL"));
     if (storeAtLabel) Graphics::StorageTexture::FlushAllPending("label");
+    Graphics::StorageTexture::StoreBeforeGuestCompletion("label");
     if (!state->recorder || bytes.empty() || bytes.size() > 65536 || address % 4 != 0 || bytes.size() % 4 != 0) return 4;
     auto& recorder = *state->recorder;
     // Batches that already finished are retired before the checks: their completions ran, so they

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

@@ -195,6 +195,7 @@ class StorageTexture : public std::enable_shared_from_this<StorageTexture> {
     // overlapping pending image, as before).
     static bool FlushPending(std::uint64_t address, std::size_t bytes, const StorageTexture* except = nullptr, const char* reason = "memory access", PublishScope scope = PublishScope::Whole, bool* published = nullptr);
     static void FlushAllPending(const char* reason);
+    static void StoreBeforeGuestCompletion(const char* reason);
     static void SaveBeforeUnwatch(std::uint64_t address, std::size_t bytes);
     static constexpr std::uint64_t SmallStoreBytes = 64 * 1024;
     static std::size_t StoreSmallPending(const char* reason);
@@ -349,6 +350,7 @@ class StorageTexture : public std::enable_shared_from_this<StorageTexture> {
     static std::uint64_t SinglePassMoves();
     static std::uint64_t RefreshesProved();
     static std::uint64_t UntrackedCompares();
+    static std::uint64_t CpuWonStores();
     std::uint64_t GuestBytes() const;
     std::size_t OriginalSnapshotBytes() const { return original.capacity(); }
     VkDeviceSize AllocationBytes() const { return memoryBytes; }
@@ -590,6 +592,7 @@ class StorageTexture : public std::enable_shared_from_this<StorageTexture> {
     VkFormat storageFormat = VK_FORMAT_UNDEFINED;
     // Results are on the GPU only (guarded by the pending-write registry lock).
     bool dirty = false;
+    std::uint64_t pendingSinceCompletion = 0;
     std::uint64_t version = 0;
     // `original` holds the guest bytes; false after a GPU-side clear, which never read them.
     std::vector<bool> layerOriginal;

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

@@ -87,6 +87,8 @@ std::atomic<std::uint64_t> partialUploads{0}, partialUploadBytes{0}, partialWrit
 // writeBackLayers), and pending units a DCC clear -> uncompressed key flip kept as the texels
 // (Refresh).
 std::atomic<std::uint64_t> pretestSkipped{0}, unregisteredDropped{0}, keyFlipKept{0};
+std::atomic<std::uint64_t> cpuWonStores{0};
+std::atomic<std::uint64_t> guestCompletions{0};
 // Images stored by a FlushPending after a hook skip of theirs (AccessKeptByCpu), of which by the
 // hook for the read site of the last skip; images evictStale dropped (the [hooksync] and [storage] lines).
 std::atomic<std::uint64_t> flushedAfterSkip{0}, flushedAfterSkipSameSite{0}, staleEvicted{0};
@@ -2054,6 +2056,10 @@ std::uint64_t StorageTexture::UntrackedCompares() {
     return untrackedCompares.load(std::memory_order_relaxed);
 }
 
+std::uint64_t StorageTexture::CpuWonStores() {
+    return cpuWonStores.load(std::memory_order_relaxed);
+}
+
 bool StorageTexture::ServesKeysAt(std::uint64_t dccAddress) const {
     const bool locked = GuestMemory::GpuMutex().HeldByThisThread();
     const auto readFollowed = [&] {
@@ -3123,6 +3129,7 @@ void StorageTexture::markLayersPending(std::uint32_t first, std::uint32_t count)
         anyUnsaved = anyUnsaved || unsaved[layer];
     }
     if (anyUnsaved && !GuestMemory::Watched(descriptor.baseAddress, static_cast<std::size_t>(guestBytes))) captureGuestBytes(&unsaved);
+    if (!anyLayerPending()) pendingSinceCompletion = guestCompletions.load(std::memory_order_acquire);
     for (std::uint32_t layer = first; layer < first + count; ++layer) layerPending[layer] = true;
     // Results of this image now cover the alias's results borrowed into these units: this image
     // stores them, the alias no longer has to (its results, unchanged since the borrow, are in
@@ -3485,6 +3492,22 @@ void StorageTexture::FlushAllPending(const char* reason) {
     static_cast<void>(FlushPending(0, std::numeric_limits<std::size_t>::max(), nullptr, reason));
 }
 
+void StorageTexture::StoreBeforeGuestCompletion(const char* reason) {
+    static const bool deferred = std::getenv("APS5_DEFER_COMPARED_RESULTS") != nullptr;
+    if (!deferred) {
+        std::vector<std::pair<std::uint64_t, std::size_t>> ranges;
+        {
+            auto& pending = Pending();
+            std::lock_guard lock(pending.mutex);
+            for (const auto* texture : pending.textures) ranges.emplace_back(texture->descriptor.baseAddress, static_cast<std::size_t>(texture->guestBytes));
+        }
+        for (const auto& [address, bytes] : ranges) {
+            if (!GuestMemory::Watched(address, bytes)) static_cast<void>(FlushPending(address, bytes, nullptr, reason, PublishScope::None));
+        }
+    }
+    guestCompletions.fetch_add(1, std::memory_order_acq_rel);
+}
+
 std::size_t StorageTexture::StoreSmallPending(const char* reason) {
     std::vector<std::shared_ptr<StorageTexture>> small;
     {
@@ -4564,6 +4587,18 @@ void StorageTexture::writeBackLayers(const std::vector<bool>& layers) {
     constexpr std::uint64_t block = 65536;
     const auto spanBegin = firstStored & ~(block - 1);
     const auto spanBlocks = static_cast<std::size_t>((lastStored - spanBegin + block - 1) / block);
+    static const bool recordedWriteBack = std::getenv("APS5_RECORDED_WRITEBACK") != nullptr;
+    std::vector<ImportSegment> segments;
+    if (recordedWriteBack || pendingSinceCompletion == guestCompletions.load(std::memory_order_acquire)) {
+        if (const auto* import = HostImportFor(context, descriptor.baseAddress, static_cast<std::size_t>(guestBytes))) segments.push_back({import, descriptor.baseAddress, descriptor.baseAddress + guestBytes});
+        else segments = importSegments();
+    }
+    if (segments.empty()) {
+        original.resize(static_cast<std::size_t>(guestBytes));
+        std::vector<bool> unsaved = layerOriginal;
+        unsaved.flip();
+        for (const auto& [begin, end] : unitRuns(unsaved)) GuestMemory::ReadCommitted(descriptor.baseAddress + begin, std::span(original).subspan(static_cast<std::size_t>(begin), static_cast<std::size_t>(end - begin)));
+    }
     // The keep-blocks decision must see every CPU write to the stored span up to now, so it
     // bypasses the per-packet collect memo (a Refresh collect earlier in the same packet would
     // satisfy it); the unselected layers keep their own generations and are collected by their
@@ -4643,7 +4678,6 @@ void StorageTexture::writeBackLayers(const std::vector<bool>& layers) {
     // own stamps, so every later import writer stamps newer.
     std::vector<ShadowedRange> shadowed;
     std::vector<std::pair<std::uint64_t, std::uint64_t>> imported;
-    std::vector<ImportSegment> segments;
     const auto markShadowed = [&](std::uint64_t generation) {
         for (const auto& segment : segments) {
             std::vector<ShadowedRange> ranges;
@@ -4679,8 +4713,6 @@ void StorageTexture::writeBackLayers(const std::vector<bool>& layers) {
         settle(true);
         return;
     }
-    if (const auto* import = HostImportFor(context, descriptor.baseAddress, static_cast<std::size_t>(guestBytes))) segments.push_back({import, descriptor.baseAddress, descriptor.baseAddress + guestBytes});
-    else segments = importSegments();
     if (!segments.empty()) {
         Require(!unsavedOutsideWatch, "storing storage image results pending over memory outside the write watch without the guest bytes saved before them is not implemented");
         // Outside the write watch the stored bytes are what the next store compares with.
@@ -4854,15 +4886,7 @@ void StorageTexture::writeBackLayers(const std::vector<bool>& layers) {
         }
         return;
     }
-    original.resize(static_cast<std::size_t>(guestBytes));
-    // The store below compares against the guest bytes, read now when `original` does not hold
-    // them (a GPU clear, a dropped or kept unit). Reading them says nothing about the image: a
-    // stale unit (dropped, or kept for the CPU) still differs from them, so `original` vouches for
-    // the content again only where it did before, or once every unit is stored.
     Require(!unsavedOutsideWatch, "storing storage image results pending over memory outside the write watch without the guest bytes saved before them is not implemented");
-    std::vector<bool> unsaved = layerOriginal;
-    unsaved.flip();
-    for (const auto& [begin, end] : unitRuns(unsaved)) GuestMemory::ReadCommitted(descriptor.baseAddress + begin, std::span(original).subspan(static_cast<std::size_t>(begin), static_cast<std::size_t>(end - begin)));
     DeviceBuffer linear(context, static_cast<std::size_t>(sliceLinearBytes * arrayLayers), VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT);
     DeviceBuffer tiled(context, original.size(), VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT);
     Buffer host(context, original.size(), VK_BUFFER_USAGE_TRANSFER_SRC_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT);
@@ -4936,7 +4960,15 @@ void StorageTexture::writeBackLayers(const std::vector<bool>& layers) {
     for (const auto& [from, to] : keep) {
         const auto offset = static_cast<std::size_t>(from - descriptor.baseAddress);
         const auto length = static_cast<std::size_t>(to - from);
-        GuestMemory::WriteChangedCommitted(from, current.subspan(offset, length), std::span<const std::byte>(original).subspan(offset, length));
+        if (!GuestMemory::WriteChangedOverOriginalCommitted(from, current.subspan(offset, length), std::span<const std::byte>(original).subspan(offset, length))) {
+            for (std::uint32_t layer = 0; layer < trackedLayers; ++layer) {
+                if (!layers[layer] || layerBegin(layer) >= to || layerBegin(layer) + layerBytes(layer) <= from) continue;
+                skippedLayer[layer] = true;
+            }
+            cpuWonStores.fetch_add(1, std::memory_order_relaxed);
+            static std::atomic<int> reports{0};
+            if (reports.fetch_add(1) < 8) std::fprintf(stderr, "[gpu] storage image 0x%llx: the CPU stored into 0x%llx+0x%zx during its write-back; keeping the CPU's bytes\n", static_cast<unsigned long long>(descriptor.baseAddress), static_cast<unsigned long long>(from), length);
+        }
         std::memcpy(original.data() + offset, current.data() + offset, length);
     }
     // The texels now hold the whole image, so later reads must see them rather than a fast clear