fix(agc): stamp a table before its pending-write query and note GPU writes before stamping

A table column's stamp vouches that nothing wrote its memory after the stamp's generation and
nothing was pending at the time. stampWrites asked for pending writes before it took the
generation, and two GPU write paths stamped their blocks before they noted the write pending:
a label or store kept behind completions, and a DMA copy. A capture on another queue worker that
ran between those two steps saw no pending write and took a generation newer than the write's
stamp, so once the batch was reaped the column vouched for the bytes from before the write.

stampWrites now flushes, takes the layout and generation, and asks for pending writes last. The
label path notes the completion before it stamps the store. The DMA copy still stamps first, so a
newer copy's ring entry never meets an older one's generation, then notes the write pending and
stamps again, and its ring entry carries the second generation. A store noted after the pending
query is now always stamped after the generation the column keeps.

The new two-thread test stamps a GPU write from another thread while the capture asks for
pending writes; the next capture must compare the table instead of reusing it.

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

@@ -301,11 +301,11 @@ ShaderRecompiler::SrtWriteStamp ShaderMemory::stampWrites(void* context, std::ui
     if (bytes == 0 || address > std::numeric_limits<std::uint64_t>::max() - bytes || self.overlapsInitial(address, address + bytes)) return {};
     const auto fetched = self.pages.lower_bound(address & ~static_cast<std::uint64_t>(PageBytes - 1));
     if (fetched != self.pages.end() && fetched->first < address + bytes) return {};
-    if (self.pendingWrite != nullptr && self.pendingWrite(address, static_cast<std::size_t>(bytes), {}) != PendingWrite::None) return {};
     GuestMemory::FlushGpuWrites(address, static_cast<std::size_t>(bytes));
     const auto layout = GuestAllocations::GuestAllocationsGeneration_nid_postfix();
     const auto generation = GuestMemory::CollectWrites(address, static_cast<std::size_t>(bytes));
     if (generation == 0) return {};
+    if (self.pendingWrite != nullptr && self.pendingWrite(address, static_cast<std::size_t>(bytes), {}) != PendingWrite::None) return {};
     return {generation, layout};
 }
 

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

@@ -2026,10 +2026,10 @@ int VulkanDevice::WriteLabelOnGpu(std::uint64_t address, std::span<const std::by
         // a title misbehaves with this path.
         if (drain) return 2;
         recordStore();
-        GuestMemory::MarkWritten(address, bytes.size());
         // Commands() opened the batch the store went into, so the completion is appended to that
         // batch and runs at its reap, after the write-backs of every batch before it.
         recorder.AfterCompletions(address, bytes, stamp, queue, true);
+        GuestMemory::MarkWritten(address, bytes.size());
         return 5;
     }
     recordStore();
@@ -2483,6 +2483,8 @@ VulkanDevice::CopyOutcome VulkanDevice::CopyBuffer(std::uint64_t destination, st
     recorder.EndGpuTiming(classTiming, bytes);
     recorder.MarkCovered(copiedAccess);
     recorder.NotePendingRead(source, bytes, Graphics::Recorder::ReadKind::CopySource);
+    GuestMemory::MarkWritten(destination, bytes);
+    recorder.NotePendingWrite(destination, bytes);
     // The destination's bytes after the transfer are the source's now when nothing recorded,
     // pending or still to land can change the source before the transfer reads it (a CPU store
     // to it in between is a hardware race too): the caller's ring entry carries them, stamped

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

@@ -41,6 +41,8 @@
 #include <new>
 #include <stdexcept>
 #include <span>
+#include <atomic>
+#include <thread>
 #include <string>
 #include <string_view>
 #include <tuple>
@@ -1649,6 +1651,66 @@ void verifyTableSnapshotReuse() {
     std::cout << "table snapshot reuse " << (watched ? "with" : "without") << " write watching passed\n";
 }
 
+std::atomic<int> tableRacePhase{0};
+std::uint64_t tableRaceBegin = 0;
+std::uint64_t tableRaceBytes = 0;
+
+AgcDriver::ShaderMemory::PendingWrite RacingTableWrite(std::uint64_t address, std::size_t bytes, std::span<std::byte>) {
+    if (tableRaceBytes != 0u && address == tableRaceBegin && bytes >= tableRaceBytes) {
+        tableRaceBytes = 0;
+        tableRacePhase.store(1, std::memory_order_release);
+        while (tableRacePhase.load(std::memory_order_acquire) != 2) std::this_thread::yield();
+    }
+    return AgcDriver::ShaderMemory::PendingWrite::None;
+}
+
+void verifyTableStampRace() {
+    using namespace ShaderRecompiler;
+    constexpr std::uint32_t Records = 4096;
+    constexpr std::uint32_t Stride = 48;
+    constexpr std::size_t Bytes = static_cast<std::size_t>(Records) * Stride;
+    constexpr std::size_t RecordWords = Stride / 4u;
+    const auto texture = [](std::uint32_t index) { return TableT(reinterpret_cast<const void*>(static_cast<std::uintptr_t>(0x100000000ull + index * 256ull)), TableFormat8888UNorm, TableType2D, 1u + index % 5u); };
+    auto* heap = AllocateWatchedTable(Bytes);
+    for (std::uint32_t record = 0; record < Records; record++) {
+        const auto words = texture(record % 1024u);
+        std::copy(words.begin(), words.end(), heap + record * RecordWords);
+        std::fill_n(heap + record * RecordWords + 8u, 4u, 0x5a5a5a5au);
+    }
+    const auto address = static_cast<std::uint64_t>(reinterpret_cast<std::uintptr_t>(heap));
+    if (AgcDriver::GuestMemory::CollectWrites(address, Bytes) == 0) {
+        std::cout << "table stamp race without write watching skipped\n";
+        return;
+    }
+    alignas(4096) static TableGuest guest;
+    FillTableSrt(guest, heap, Stride, Records);
+    const auto userData = TableUserData(guest);
+    auto wideCode = TableWholeCode;
+    std::replace(wideCode.begin(), wideCode.end(), 0x8f108510u, 0x9310b010u);
+    const auto request = TableRequest(wideCode, userData);
+    const auto capture = [&](AgcDriver::ShaderMemory::PendingWriteQuery pending) {
+        const auto before = ResourceMaterializer::TableReuse();
+        AgcDriver::ShaderMemory memory({}, pending);
+        static_cast<void>(memory.Capture(request));
+        const auto after = ResourceMaterializer::TableReuse();
+        return ImageTableReuseCounts{after.unwritten - before.unwritten, after.compared - before.compared, after.wordsRead - before.wordsRead, after.shared - before.shared, after.columns - before.columns};
+    };
+    tableRacePhase.store(0, std::memory_order_relaxed);
+    tableRaceBegin = address;
+    tableRaceBytes = Bytes - 16u;
+    std::thread writer([&] {
+        while (tableRacePhase.load(std::memory_order_acquire) != 1) std::this_thread::yield();
+        AgcDriver::GuestMemory::MarkWritten(address + 5u * Stride, 32u);
+        tableRacePhase.store(2, std::memory_order_release);
+    });
+    const auto first = capture(&RacingTableWrite);
+    writer.join();
+    require(tableRacePhase.load(std::memory_order_relaxed) == 2 && first.columns == 1u, "table stamp race: the racing store did not run while the table was stamped");
+    const auto next = capture(nullptr);
+    require(next.unwritten == 0u && next.compared == 1u, "table stamp race: a GPU write stamped while the table was stamped, before it was noted pending, was hidden by the table's stamp");
+    std::cout << "table stamp race with write watching passed\n";
+}
+
 void verifyImageTables() {
     using namespace ShaderRecompiler;
     namespace Abi = ImageTableAbi;
@@ -5660,6 +5722,7 @@ int main(int argc, char** argv) {
         }
         if (argc == 2 && std::string_view(argv[1]) == "--table-snapshot-reuse") {
             verifyTableSnapshotReuse();
+            verifyTableStampRace();
             return 0;
         }
         if (argc == 2 && std::string_view(argv[1]) == "--srt-read-patching") {