fix(agc): consume the write-watch faults of the driver's own CPU stores

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

@@ -74,6 +74,7 @@ bool UnchangedSinceAll(std::span<const UnchangedQuery> queries);
 // Returns the generation the blocks were stamped with (0 when the arena is not write-watched):
 // UnchangedSince(range, it) holds until the next store over the range.
 std::uint64_t MarkWritten(std::uint64_t address, std::size_t bytes);
+std::uint64_t StoreOwnBytes(std::uint64_t address, std::size_t bytes, const std::function<void()>& store);
 bool StoredOver(std::uint64_t address, std::size_t bytes, std::uint64_t generation);
 // Collect epoch: within one epoch a range already collected is not walked again, CollectWrites
 // returns the current generation instead. A guest write landing between two collects of the same

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

@@ -1093,6 +1093,40 @@ bool UnchangedSinceAll(std::span<const UnchangedQuery> queries) {
     return true;
 }
 
+std::uint64_t storeOwn(std::uint64_t address, std::size_t bytes, const std::function<std::pair<std::uint64_t, std::uint64_t>()>& store) {
+    if (bytes == 0) return 0;
+    const auto stampStored = [](WriteTracker& tracker, std::pair<std::uint64_t, std::uint64_t> stored) -> std::uint64_t {
+        if (stored.second <= stored.first || !tracker.watched || !tracker.covers(stored.first, static_cast<std::size_t>(stored.second - stored.first))) return 0;
+        ++tracker.generation;
+        for (auto block = tracker.blockOf(stored.first); block <= tracker.blockOf(stored.second - 1); ++block) {
+            tracker.stamp(block, tracker.generation, StampKind::Driver);
+            tracker.noteDriverStore(block, stored.first, stored.second, tracker.generation);
+        }
+        return tracker.generation;
+    };
+    auto& tracker = Tracker();
+    const auto lock = lockTracker(tracker);
+    tracker.initialize();
+    if (!tracker.watched || bytes > std::numeric_limits<std::uint64_t>::max() - address || !tracker.covers(address, bytes)) return stampStored(tracker, store());
+    constexpr std::uint64_t page = 4096;
+    const auto first = address & ~(page - 1);
+    const auto stop = (address + bytes + page - 1) & ~(page - 1);
+#ifdef _WIN32
+    if (GuestArena::GuestArenaHostRegionOverlaps_nid_postfix(first, static_cast<std::size_t>(stop - first))) return stampStored(tracker, store());
+#endif
+    if (!walkWrites(tracker, first, stop, StampKind::Cpu)) return stampStored(tracker, store());
+    const auto stored = store();
+    walkWrites(tracker, first, stop, StampKind::Driver);
+    return stampStored(tracker, stored);
+}
+
+std::uint64_t StoreOwnBytes(std::uint64_t address, std::size_t bytes, const std::function<void()>& store) {
+    return storeOwn(address, bytes, [&] {
+        store();
+        return std::pair<std::uint64_t, std::uint64_t>{address, address + bytes};
+    });
+}
+
 std::uint64_t MarkWritten(std::uint64_t address, std::size_t bytes) {
     auto& tracker = Tracker();
     const auto lock = lockTracker(tracker);
@@ -1641,23 +1675,25 @@ void WriteChanged(std::uint64_t address, std::span<const std::byte> current, std
         const auto length = std::min(block, size - at);
         return std::memcmp(current.data() + at, original.data() + at, length) != 0;
     };
-    std::size_t firstChanged = size;
-    std::size_t lastChanged = 0;
-    for (std::size_t at = 0; at < size; at += block) {
-        if (!differs(at)) continue;
-        const auto blockEnd = std::min(at + block, size);
-        for (std::size_t run = at; run < blockEnd;) {
-            if (current[run] == original[run]) { ++run; continue; }
-            auto runEnd = run + 1;
-            while (runEnd < blockEnd && current[runEnd] != original[runEnd]) ++runEnd;
-            std::memcpy(destination + run, current.data() + run, runEnd - run);
-            firstChanged = std::min(firstChanged, run);
-            lastChanged = std::max(lastChanged, runEnd);
-            run = runEnd;
-        }
-    }
     // Stamped like a GPU write: a collect memoized for this packet would not see the page fault.
-    if (firstChanged < lastChanged) MarkWritten(address + firstChanged, lastChanged - firstChanged);
+    storeOwn(address, size, [&] {
+        std::size_t firstChanged = size;
+        std::size_t lastChanged = 0;
+        for (std::size_t at = 0; at < size; at += block) {
+            if (!differs(at)) continue;
+            const auto blockEnd = std::min(at + block, size);
+            for (std::size_t run = at; run < blockEnd;) {
+                if (current[run] == original[run]) { ++run; continue; }
+                auto runEnd = run + 1;
+                while (runEnd < blockEnd && current[runEnd] != original[runEnd]) ++runEnd;
+                std::memcpy(destination + run, current.data() + run, runEnd - run);
+                firstChanged = std::min(firstChanged, run);
+                lastChanged = std::max(lastChanged, runEnd);
+                run = runEnd;
+            }
+        }
+        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};
+    });
 }
 
 void Write(std::uint64_t address, std::span<const std::byte> source, std::size_t alignment) {
@@ -1669,9 +1705,8 @@ void Write(std::uint64_t address, std::span<const std::byte> source, std::size_t
     CountWriteCaller(__builtin_return_address(0));
     auto* destination = reinterpret_cast<void*>(address);
     CheckRange(destination, source.size(), alignment, true);
-    std::memcpy(destination, source.data(), source.size());
     // Stamped like a GPU write: a collect memoized for this packet would not see the page fault.
-    MarkWritten(address, source.size());
+    StoreOwnBytes(address, source.size(), [&] { std::memcpy(destination, source.data(), source.size()); });
 }
 
 }

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

@@ -1468,8 +1468,7 @@ VulkanDevice::CopyOutcome VulkanDevice::CopyBuffer(std::uint64_t destination, st
         // outstanding on the range) and nothing recorded.
         std::vector<std::byte> expected;
         if (verify) expected.assign(reinterpret_cast<const std::byte*>(source), reinterpret_cast<const std::byte*>(source) + bytes);
-        std::memcpy(reinterpret_cast<void*>(destination), reinterpret_cast<const void*>(source), bytes);
-        GuestMemory::MarkWritten(destination, bytes);
+        GuestMemory::StoreOwnBytes(destination, bytes, [&] { std::memcpy(reinterpret_cast<void*>(destination), reinterpret_cast<const void*>(source), bytes); });
         outcome.path = 0;
         outcome.sourceSettledBySignal = settledBySignal;
         if (verify) {

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

@@ -2377,8 +2377,7 @@ void Recorder::AfterCompletions(std::uint64_t address, std::span<const std::byte
         }
         completionStoresRun.fetch_add(1, std::memory_order_relaxed);
         GuestMemory::CheckRange(reinterpret_cast<const void*>(address), copy.size(), 4, true);
-        std::memcpy(reinterpret_cast<void*>(address), copy.data(), copy.size());
-        GuestMemory::MarkWritten(address, copy.size());
+        GuestMemory::StoreOwnBytes(address, copy.size(), [&] { std::memcpy(reinterpret_cast<void*>(address), copy.data(), copy.size()); });
     });
     // Counted per batch and counted down when the batch finishes, on every exit path of finish()
     // (a label whose range became unmapped is reported there, and a lost device throws before the

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

@@ -90,6 +90,33 @@ void CheckSharedBlock() {
     if (!Watched(outside, unwatched.size())) Require(StoredOver(outside, unwatched.size(), TrackerGeneration()), "an unwatched range reads as not stored over");
 }
 
+void CheckOwnStore() {
+    void* memory = AllocateWatched(2 * Block);
+    const auto base = reinterpret_cast<std::uint64_t>(memory);
+    const auto second = base + Block;
+    std::memset(memory, 0x11, 2 * Block);
+    const auto synced = CollectWrites(base, 2 * Block);
+    Require(synced != 0, "the own-store block is not collected");
+
+    std::array<std::byte, 8192> bytes{};
+    bytes.fill(std::byte{0x44});
+    Write(second, bytes);
+    const auto stored = TrackerGeneration();
+    Require(!UnchangedSince(second, bytes.size(), synced), "the driver store is not stamped");
+    CollectWritesUncached(base, 2 * Block);
+    Require(UnchangedSince(second, Block, stored), "a walk after a driver store reports the store again as a newer write");
+    Require(UnchangedSinceCollected(second, bytes.size(), synced), "the driver store's own page faults are stamped as a CPU write");
+
+    static_cast<volatile std::uint8_t*>(memory)[Block + 3 * 4096 + 8] = 0x22;
+    const std::array<std::byte, 64> label{};
+    Write(second + 3 * 4096 + 64, label);
+    Require(!UnchangedSinceCollected(second + 3 * 4096, 4096, stored), "a CPU write before a driver store in its page is not stamped as one");
+
+    const auto beforeCpu = TrackerGeneration();
+    static_cast<volatile std::uint8_t*>(memory)[16] = 0x33;
+    CollectWritesUncached(base, 2 * Block);
+    Require(!UnchangedSince(base, 64, beforeCpu), "a CPU write after the driver store is not seen");
+}
 }
 
 int main() {
@@ -99,6 +126,7 @@ int main() {
             return 77;
         }
         CheckSharedBlock();
+        CheckOwnStore();
     } catch (const std::exception& error) {
         std::cerr << "write tracking test failed: " << error.what() << "\n";
         return 1;

core/libs/prx/libc/include/GuestArena.hpp

@@ -24,6 +24,7 @@ void GuestArenaSetProtection_nid_postfix(std::uintptr_t address, std::size_t byt
 bool GuestArenaHandleWrite_nid_postfix(std::uintptr_t address);
 bool GuestArenaProtection_nid_postfix(std::uintptr_t address, std::uint32_t* protection);
 bool GuestArenaCollectWrites_nid_postfix(std::uintptr_t address, std::size_t bytes, void** pages, std::size_t* count, bool clear);
+bool GuestArenaHostRegionOverlaps_nid_postfix(std::uintptr_t address, std::size_t bytes);
 void GuestArenaCommit_nid_postfix(void* pointer, std::size_t bytes, std::uint32_t protection, std::size_t granule);
 void GuestArenaReset_nid_postfix(void* pointer, std::size_t bytes);
 void GuestArenaMap_nid_postfix(void* pointer, std::size_t bytes, void* section, std::uint64_t offset, std::uint32_t protection);

core/libs/prx/libc/src/GuestArena.cpp

@@ -115,6 +115,7 @@ class Arena {
             if (last < regionEnd) _holes.emplace_back(last, regionEnd);
             cursor = regionEnd;
         }
+        _hostRegions = _holes;
         _holes.emplace_back(SystemReservedStart, SystemReservedEnd);
         for (const auto& [holeStart, holeEnd] : _holes) _used.emplace(holeStart, holeEnd);
         _base = ArenaStart;
@@ -125,6 +126,7 @@ class Arena {
     std::mutex _lock;
     std::map<std::uintptr_t, std::uintptr_t> _used;
     std::vector<std::pair<std::uintptr_t, std::uintptr_t>> _holes;
+    std::vector<std::pair<std::uintptr_t, std::uintptr_t>> _hostRegions;
     std::uintptr_t _base = 0;
     std::uintptr_t _end = 0;
     bool _writeWatched = false;
@@ -133,6 +135,12 @@ class Arena {
     std::uintptr_t Base() const { return _base; }
     std::size_t Size() const { return _end - _base; }
     bool WriteWatched() const { return _writeWatched; }
+    bool OverlapsHostRegion(std::uintptr_t start, std::size_t bytes) const {
+        for (const auto& [regionStart, regionEnd] : _hostRegions) {
+            if (regionStart < start + bytes && start < regionEnd) return true;
+        }
+        return false;
+    }
 };
 
 const bool g_reserved = (Arena::Get(), true);
@@ -185,6 +193,10 @@ bool GuestArenaCollectWrites_nid_postfix(std::uintptr_t address, std::size_t byt
     return WindowsMappings::Get().Collect(address, bytes, pages, count, clear);
 }
 
+bool GuestArenaHostRegionOverlaps_nid_postfix(std::uintptr_t address, std::size_t bytes) {
+    return Arena::Get().OverlapsHostRegion(address, bytes);
+}
+
 void GuestArenaCommit_nid_postfix(void* pointer, std::size_t bytes, std::uint32_t protection, std::size_t granule) {
     if (!Arena::Get().Contains(pointer, bytes)) throw std::invalid_argument("commit outside the guest arena");
     WindowsMappings::Get().Commit(pointer, bytes, protection, granule, Arena::Get().WriteWatched());