fix(agc): do not wait for a storage image's own stamped writes before flushing it

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

@@ -150,12 +150,15 @@ class Recorder {
     void NotePendingWordFill(std::uint64_t address, std::size_t bytes, std::uint32_t word);
     // Notes several [begin, end) ranges and publishes the snapshot once (a dispatch writes many buffers).
     void NotePendingWrites(std::span<const std::pair<std::uint64_t, std::uint64_t>> ranges);
+    void NoteStampedWrites(std::span<const std::pair<std::uint64_t, std::uint64_t>> ranges);
     bool PendingWriteOverlaps(std::uint64_t address, std::size_t bytes) const;
     // Whether no unfinished batch writes the range any more: no open-batch overlap, and every
     // overlapping in-flight batch's fence has signaled (its in-place GPU stores are final in host
     // memory; its completion stores are the caller's business, see VulkanDevice::CopyBuffer).
     bool PendingWriteSettled(std::uint64_t address, std::size_t bytes) const;
     std::uint64_t LastWriteNote(std::uint64_t address, std::size_t bytes) const;
+    std::uint64_t WriteNoteCount() const { return writeNoteCount; }
+    std::vector<std::uint64_t> StampedNotesSince(std::uint64_t note) const;
     std::uint64_t NewestWriteNote(std::uint64_t address, std::size_t bytes) const;
     struct PendingWrite {
         std::uint64_t begin;
@@ -610,6 +613,7 @@ class Recorder {
         std::vector<std::function<void()>> completions;
         std::vector<std::pair<std::uint64_t, std::uint64_t>> writes;
         std::vector<std::uint64_t> writeNotes;
+        std::vector<std::uint64_t> stampedNotes;
         std::vector<std::int16_t> writeValues;
         std::vector<std::int64_t> writeWords;
         // In-place reads (see NotePendingRead), dying with the batch: a finished batch's reads are done.

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

@@ -830,6 +830,7 @@ class StorageTexture : public std::enable_shared_from_this<StorageTexture> {
     bool pendingFromOthers = false;
     std::vector<DrawnRect> drawnRects;
     std::vector<std::uint64_t> storeNotes;
+    std::vector<std::uint64_t> flushNotes;
     std::uint64_t writtenFlip = 0;
     std::atomic<bool> writtenOccasionally{false};
     std::uint64_t version = 0;

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

@@ -3337,6 +3337,13 @@ void Recorder::NotePendingWrites(std::span<const std::pair<std::uint64_t, std::u
     std::atomic_thread_fence(std::memory_order_seq_cst);
 }
 
+void Recorder::NoteStampedWrites(std::span<const std::pair<std::uint64_t, std::uint64_t>> ranges) {
+    const auto first = writeNoteCount + 1;
+    NotePendingWrites(ranges);
+    if (open == nullptr) return;
+    for (auto note = first; note <= writeNoteCount; ++note) open->stampedNotes.push_back(note);
+}
+
 void Recorder::publishPendingWrites(std::span<const std::pair<std::uint64_t, std::uint64_t>> additions) const {
     if (activeRecorder != this) {
         publishedWriteRangesValid = false;
@@ -3467,6 +3474,18 @@ bool Recorder::PendingWriteSettled(std::uint64_t address, std::size_t bytes) con
     return true;
 }
 
+std::vector<std::uint64_t> Recorder::StampedNotesSince(std::uint64_t note) const {
+    std::vector<std::uint64_t> notes;
+    const auto collect = [&](const Batch& batch) {
+        for (const auto stamped : batch.stampedNotes) {
+            if (stamped > note) notes.push_back(stamped);
+        }
+    };
+    for (const auto& batch : inFlight) collect(*batch);
+    if (open != nullptr) collect(*open);
+    return notes;
+}
+
 std::uint64_t Recorder::LastWriteNote(std::uint64_t address, std::size_t bytes) const {
     if (open == nullptr || open->writes.empty() || open->writeNotes.size() != open->writes.size()) return 0;
     if (open->writes.back() != std::pair<std::uint64_t, std::uint64_t>{address, address + bytes}) return 0;

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

@@ -3958,7 +3958,7 @@ std::uint64_t StorageTexture::writeBackWindows(std::span<const ImportSegment> se
         recorder->MarkCovered(storedAccess);
         // Only the pieces written into the import are stores a CPU reader must wait for; the
         // shadowed ones reach it through a publish, which notes its own.
-        if (!imported.empty()) recorder->NotePendingWrites(imported);
+        if (!imported.empty()) recorder->NoteStampedWrites(imported);
     }
     partialWriteBacks.fetch_add(1, std::memory_order_relaxed);
     partialWriteBackBytes.fetch_add(storedBytes, std::memory_order_relaxed);
@@ -4633,6 +4633,7 @@ bool StorageTexture::FlushPending(std::uint64_t address, std::size_t bytes, cons
         if (!flush.empty()) BumpPendingSerial();
         for (const auto& texture : flush) {
             if (texture->drawnListed) ownNotes.insert(ownNotes.end(), texture->storeNotes.begin(), texture->storeNotes.end());
+            ownNotes.insert(ownNotes.end(), texture->flushNotes.begin(), texture->flushNotes.end());
         }
     }
     if (flush.empty()) {
@@ -4713,6 +4714,8 @@ bool StorageTexture::FlushPending(std::uint64_t address, std::size_t bytes, cons
             throw;
         }
     }
+    const auto* noting = Recorder::Active();
+    const auto firstNote = noting != nullptr ? noting->WriteNoteCount() : 0;
     std::exception_ptr failure;
     const auto previousReason = std::exchange(flushReason, reason);
     for (const auto& texture : flush) {
@@ -4733,6 +4736,15 @@ bool StorageTexture::FlushPending(std::uint64_t address, std::size_t bytes, cons
         if (units != 0) waitForIdlePublish(reason);
         if (published != nullptr) *published = units != 0;
     }
+    if (flush.size() == 1 && noting != nullptr && noting == Recorder::Active()) {
+        if (const auto notes = noting->StampedNotesSince(firstNote); !notes.empty()) {
+            auto& pending = Pending();
+            std::lock_guard lock(pending.mutex);
+            auto& own = flush.front()->flushNotes;
+            own.insert(own.end(), notes.begin(), notes.end());
+            if (own.size() > 32) own.erase(own.begin(), own.end() - 32);
+        }
+    }
     return true;
 }
 
@@ -6935,7 +6947,8 @@ void StorageTexture::writeBackLayers(const std::vector<bool>& layers) {
         } else {
             recorder->EndGpuTiming(timing, storedBytes);
             recorder->MarkShaderReadsCovered();
-            recorder->NotePendingWrite(firstStored, static_cast<std::size_t>(lastStored - firstStored));
+            const std::array stamped{std::pair{firstStored, lastStored}};
+            recorder->NoteStampedWrites(stamped);
             if (drawn.drawn) noteStore(recorder->LastWriteNote(firstStored, static_cast<std::size_t>(lastStored - firstStored)));
         }
         countStorageWriteBack(storedBytes, true, drawn.drawn);

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

@@ -320,7 +320,7 @@ void recordCopies(const UnitShadow& shadow, const std::vector<SlabCopies>& copie
     recorder->MarkCovered(publishedAccess);
     for (const auto& entry : copies) recorder->Keep(entry.slab);
     // CPU readers reach the import's bytes only after the hook's SyncThrough waits by this note.
-    recorder->NotePendingWrites(ranges);
+    recorder->NoteStampedWrites(ranges);
 }
 
 // APS5_SHADOW_VERIFY=1: the published bytes read back from the slabs equal the guest bytes.

core/libs/prx/libSceAgcDriver/tests/cmake/agc_driver_storage_flush_ordering_tests.cmake

@@ -0,0 +1,8 @@
+if(BUILD_TESTING)
+    add_test(NAME agc_driver_storage_flush_ordering_no_import COMMAND agc_driver_storage_flush_ordering_tests --stored-again)
+    set_tests_properties(agc_driver_storage_flush_ordering_no_import PROPERTIES ENVIRONMENT "APS5_NO_HOST_IMPORT=1" SKIP_RETURN_CODE 77 TIMEOUT 60)
+    if(WIN32)
+        set_tests_properties(agc_driver_storage_flush_ordering_no_import PROPERTIES ENVIRONMENT_MODIFICATION
+            "PATH=path_list_prepend:$<TARGET_FILE_DIR:libc>")
+    endif()
+endif()

core/libs/prx/libSceAgcDriver/tests/execution/StorageFlushOrdering.cpp

@@ -1,5 +1,6 @@
 #include "prx/libSceAgcDriver/Execution/include/GuestMemory.hpp"
 #include "prx/libSceAgcDriver/Execution/include/VulkanDevice.hpp"
+#include "prx/libSceAgcDriver/Graphics/include/GuestBufferMemory.hpp"
 #include "prx/libSceAgcDriver/Graphics/include/Recorder.hpp"
 #include "prx/libSceAgcDriver/Graphics/include/Texture.hpp"
 #include "prx/libc/include/GuestAllocations.hpp"
@@ -25,6 +26,14 @@
 #include <string>
 #include <vector>
 
+namespace AgcDriver {
+
+struct VulkanDeviceTestAccess {
+    static Graphics::Context ContextOf(const VulkanDevice& device) { return device.graphicsContext(); }
+};
+
+}
+
 namespace {
 
 using AgcDriver::Graphics::Require;
@@ -39,6 +48,7 @@ constexpr std::uint32_t Type2D = 9;
 constexpr std::uint32_t Initial = 0x3f800000u;
 constexpr std::uint32_t Stored = 0x11111111u;
 constexpr std::uint32_t Overwritten = 0x22222222u;
+constexpr std::uint32_t Restored = 0x33333333u;
 constexpr std::size_t OverwrittenDword = 0x100u / 4u;
 constexpr std::size_t FarDword = 100u * Width;
 constexpr std::uint32_t Filled = 0x33333333u;
@@ -48,6 +58,7 @@ constexpr std::size_t ShadowUnitBytes = 0x10000u;
 constexpr std::uint32_t Standard256B = 1u;
 
 alignas(256) constexpr std::array<std::uint32_t, 11> StoreImageCode{0x7e000280u, 0x7e020280u, 0x7e0402ffu, Stored, 0xf0200108u, 0x00010200u, 0x7e0202ffu, 0x00000064u, 0xf0200108u, 0x00010200u, 0xbf810000u};
+alignas(256) constexpr std::array<std::uint32_t, 11> RestoreImageCode{0x7e000280u, 0x7e020280u, 0x7e0402ffu, Restored, 0xf0200108u, 0x00010200u, 0x7e0202ffu, 0x00000064u, 0xf0200108u, 0x00010200u, 0xbf810000u};
 alignas(256) constexpr std::array<std::uint32_t, 7> StoreGlobalCode{0x7e000208u, 0x7e020209u, 0x7e0402ffu, Overwritten, 0xdc708100u, 0x007d0200u, 0xbf810000u};
 
 class GuestBlock {
@@ -226,18 +237,91 @@ void Run(AgcDriver::VulkanDevice& device, const GuestBlock& guest, std::size_t o
     Require(words[OverwrittenDword] == Overwritten, what + "the pending image's store put its older texels over a newer address-based store");
 }
 
+std::uint64_t OpenSerial() {
+    std::lock_guard gpu(AgcDriver::GuestMemory::GpuMutex());
+    const auto* recorder = AgcDriver::Graphics::Recorder::Active();
+    Require(recorder != nullptr, "storage flush ordering: no active recorder");
+    return recorder->OpenSerial();
 }
 
-int main() {
+void RunStoredAgain(AgcDriver::VulkanDevice& device, const GuestBlock& guest, std::size_t offset, std::size_t images, bool unrelated) {
+    const auto first = guest.Address() + offset;
+    const auto target = device.Target();
+    std::vector<ShaderRecompiler::RecompileResult> stores, restores;
+    for (std::size_t index = 0; index < images; ++index) {
+        stores.push_back(Compile(StoreImageCode, ImageUserData(first + index * ImageBytes), target));
+        restores.push_back(Compile(RestoreImageCode, ImageUserData(first + index * ImageBytes), target));
+    }
+    for (const auto& store : stores) device.Dispatch(store, 1u, 1u, 1u, {}, reinterpret_cast<std::uintptr_t>(StoreImageCode.data()));
+    const std::string subject = "storage flush ordering, " + std::to_string(images) + (images == 1 ? " image" : " images") + (unrelated ? " under another stamped write" : "");
+    Require(AgcDriver::Graphics::StorageTexture::FlushPending(first, ImageBytes * images, nullptr, "test"), subject + " stored again in the open batch: the first store left no results pending");
+    const auto& context = AgcDriver::VulkanDeviceTestAccess::ContextOf(device);
+    const bool imported = AgcDriver::Graphics::ExistingHostImportFor(context, first, ImageBytes * images) != nullptr;
+    const bool compared = imported && AgcDriver::Graphics::ComparedHostImportCovers(context, first, ImageBytes * images);
+    const bool skips = images == 1 && imported && !unrelated;
+    const std::string what = subject + (imported ? compared ? " over a compared host import" : "" : " without a host import") + " stored again in the open batch: ";
+    for (const auto& restore : restores) device.Dispatch(restore, 1u, 1u, 1u, {}, reinterpret_cast<std::uintptr_t>(RestoreImageCode.data()));
+    if (unrelated) {
+        std::lock_guard gpu(AgcDriver::GuestMemory::GpuMutex());
+        const std::array stamped{std::pair{first, first + ImageBytes}};
+        AgcDriver::Graphics::Recorder::Active()->NoteStampedWrites(stamped);
+    }
+    const auto open = OpenSerial();
+    Require(AgcDriver::Graphics::StorageTexture::FlushPending(first, ImageBytes * images, nullptr, "test"), what + "the second store left no results pending");
+    if (!compared || unrelated) {
+        if (skips) Require(OpenSerial() == open, what + "the second store waited for the GPU to finish the image's own earlier store");
+        else if (unrelated) Require(OpenSerial() != open, what + "the second store did not wait for a stamped write that is not the image's own");
+        else Require(OpenSerial() != open, what + "the second store did not wait for the earlier store, which only a single image with a host import may skip");
+    }
+    device.WaitIdle();
+    for (std::size_t index = 0; index < images; ++index) {
+        const auto* words = guest.Words() + (offset + index * ImageBytes) / 4u;
+        Require(words[0] == Restored && words[FarDword] == Restored, what + "the newer texels did not land over the earlier store");
+    }
+}
+
+void RunStoredAloneAfterJoint(AgcDriver::VulkanDevice& device, const GuestBlock& guest, std::size_t offset) {
+    const auto first = guest.Address() + offset;
+    const auto target = device.Target();
+    const std::array stores{Compile(StoreImageCode, ImageUserData(first), target), Compile(StoreImageCode, ImageUserData(first + ImageBytes), target)};
+    const auto restore = Compile(RestoreImageCode, ImageUserData(first), target);
+    for (const auto& store : stores) device.Dispatch(store, 1u, 1u, 1u, {}, reinterpret_cast<std::uintptr_t>(StoreImageCode.data()));
+    const std::string what = "storage flush ordering, 1 image stored alone after a joint store of 2 images in the open batch: ";
+    Require(AgcDriver::Graphics::StorageTexture::FlushPending(first, ImageBytes * 2u, nullptr, "test"), what + "the joint store left no results pending");
+    const bool compared = AgcDriver::Graphics::ComparedHostImportCovers(AgcDriver::VulkanDeviceTestAccess::ContextOf(device), first, ImageBytes * 2u);
+    device.Dispatch(restore, 1u, 1u, 1u, {}, reinterpret_cast<std::uintptr_t>(RestoreImageCode.data()));
+    const auto open = OpenSerial();
+    Require(AgcDriver::Graphics::StorageTexture::FlushPending(first, ImageBytes, nullptr, "test"), what + "the second store left no results pending");
+    Require(compared || OpenSerial() != open, what + "the second store did not wait for the joint store, whose writes are not the image's own");
+    device.WaitIdle();
+    const auto* words = guest.Words() + offset / 4u;
+    Require(words[0] == Restored && words[FarDword] == Restored, what + "the newer texels did not land over the joint store");
+}
+
+}
+
+int main(int argc, char** argv) {
     try {
         const auto device = OpenVulkanTestDevice();
         if (!device) return VulkanTestSkipped;
         const bool watched = AgcDriver::GuestMemory::WriteWatched();
         GuestBlock guest(watched);
+        if (argc > 1 && std::string(argv[1]) == "--stored-again") {
+            RunStoredAgain(*device, guest, BlockBytes / 4u, 1, false);
+            RunStoredAgain(*device, guest, BlockBytes / 8u * 5u, 2, false);
+            RunStoredAgain(*device, guest, BlockBytes / 8u * 3u, 1, true);
+            RunStoredAloneAfterJoint(*device, guest, BlockBytes / 8u * 5u);
+            std::puts("storage flush ordering stored-again tests passed");
+            return 0;
+        }
         if (!watched) std::puts("guest memory has no write watch: the writer-in-flight case is skipped");
         else if (!device->CanWaitUnlocked()) std::puts("the recorder cannot wait for the GPU without running completions: the writer-in-flight case is skipped");
         else Run(*device, guest, 0, true, "writer in flight");
         Run(*device, guest, BlockBytes / 2u, false, "writer in the open batch");
+        RunStoredAgain(*device, guest, BlockBytes / 4u, 1, false);
+        RunStoredAgain(*device, guest, BlockBytes / 8u * 5u, 2, false);
+        RunStoredAgain(*device, guest, BlockBytes / 8u * 3u, 1, true);
+        RunStoredAloneAfterJoint(*device, guest, BlockBytes / 8u * 5u);
         RunWriterOutsideLayers(*device, guest, BlockBytes / 4u);
         RunWriterInsideLayer(*device, guest, 3u * BlockBytes / 4u);
         GuestBlock shadowed(watched);