fix(agc): preserve GPU results in CPU-stamped full image blocks

core/libs/prx/libSceAgcDriver/CMakeLists.txt

@@ -463,6 +463,31 @@ target_link_libraries(agc_driver_recorder_tests PRIVATE
 set_target_properties(agc_driver_recorder_tests PROPERTIES CXX_EXTENSIONS OFF)
 configure_windows_unwind(agc_driver_recorder_tests)
 
+add_test_executable(agc_storage_edge_writeback_tests EXCLUDE_FROM_ALL tests/StorageEdgeWriteBack.cpp)
+target_include_directories(agc_storage_edge_writeback_tests PRIVATE
+        ${LIBS_INCLUDE_DIR}
+        ${CMAKE_SOURCE_DIR}/core/shader/recompiler
+        ${CMAKE_SOURCE_DIR}/3rdparty/Vulkan-Headers/include
+)
+target_link_libraries(agc_storage_edge_writeback_tests PRIVATE libSceAgcDriver SDL2-static libc)
+set_target_properties(agc_storage_edge_writeback_tests PROPERTIES CXX_EXTENSIONS OFF)
+configure_windows_unwind(agc_storage_edge_writeback_tests)
+if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
+    target_sources(agc_storage_edge_writeback_tests PRIVATE tests/StorageEdgeWriteWatch.cpp)
+    target_compile_definitions(agc_storage_edge_writeback_tests PRIVATE ANYPS5_TEST_WRITE_WATCH=1)
+    set_target_properties(agc_storage_edge_writeback_tests PROPERTIES ENABLE_EXPORTS ON)
+endif()
+if(BUILD_TESTING)
+    set_target_properties(agc_storage_edge_writeback_tests PROPERTIES EXCLUDE_FROM_ALL OFF)
+    add_test(NAME agc_storage_edge_writeback COMMAND agc_storage_edge_writeback_tests)
+    add_test(NAME agc_storage_full_block_writeback COMMAND agc_storage_edge_writeback_tests --full-blocks)
+    set_tests_properties(agc_storage_edge_writeback agc_storage_full_block_writeback PROPERTIES SKIP_RETURN_CODE 77)
+    if(WIN32)
+        set_tests_properties(agc_storage_edge_writeback agc_storage_full_block_writeback PROPERTIES ENVIRONMENT_MODIFICATION
+                "PATH=path_list_prepend:$<TARGET_FILE_DIR:libc>")
+    endif()
+endif()
+
 add_test_executable(agc_recorder_kept_bytes_tests EXCLUDE_FROM_ALL tests/RecorderKeptBytes.cpp)
 target_include_directories(agc_recorder_kept_bytes_tests PRIVATE ${LIBS_INCLUDE_DIR} ${CMAKE_SOURCE_DIR}/core/shader/recompiler ${CMAKE_SOURCE_DIR}/3rdparty/Vulkan-Headers/include)
 target_link_libraries(agc_recorder_kept_bytes_tests PRIVATE libSceAgcDriver libc)

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

@@ -45,6 +45,7 @@ Compare CompareMapped(std::uint64_t address, std::span<const std::byte> bytes);
 // re-read of a differing run).
 Compare CopyMapped(std::uint64_t address, std::span<std::byte> out);
 void WriteChangedCommitted(std::uint64_t address, std::span<const std::byte> current, std::span<const std::byte> original);
+bool MergeChangedCommitted(std::uint64_t address, std::span<const std::byte> current, std::span<const std::byte> original, std::uint64_t generation);
 void Write(std::uint64_t address, std::span<const std::byte> source, std::size_t alignment = 1);
 // 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.
@@ -76,6 +77,7 @@ bool UnchangedSinceAll(std::span<const UnchangedQuery> queries);
 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);
+bool DriverStoredOver(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
 // epoch is seen by the next epoch, which is the ordering real hardware gives a CPU write made while a

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

@@ -768,13 +768,15 @@ struct WriteTracker {
     std::vector<std::uint32_t> cpuBlocks;
     std::vector<std::uint32_t> writtenBlocks;
     WriteWatchCoverage coverage;
+    std::vector<std::uint32_t> driverBlocks;
 #else
     static constexpr std::size_t LeafBlocks = std::size_t{1} << 16;
     static constexpr std::size_t LeafCount = std::size_t{1} << 15;
     struct Leaf {
         std::array<std::uint32_t, LeafBlocks> blocks{};
         std::array<std::uint32_t, LeafBlocks> cpuBlocks{};
         std::array<std::uint32_t, LeafBlocks> writtenBlocks{};
+        std::array<std::uint32_t, LeafBlocks> driverBlocks{};
     };
     std::vector<std::unique_ptr<Leaf>> leaves;
 #endif
@@ -812,6 +814,7 @@ struct WriteTracker {
         cpuBlocks.assign(size / WriteBlockBytes + 1, 0);
         writtenBlocks.assign(size / WriteBlockBytes + 1, 0);
         coverage.Initialize(base, size);
+        driverBlocks.assign(size / WriteBlockBytes + 1, 0);
         pages.resize(1u << 16);
         GuestArena::GuestArenaSetPrivateMappingObserver_nid_postfix(&watchPrivateMapping);
 #else
@@ -919,21 +922,53 @@ struct WriteTracker {
 #endif
     }
 
+    std::uint32_t driverStampOf(std::uint64_t block) const {
+#ifdef _WIN32
+        return driverBlocks[block];
+#else
+        const auto& leaf = leaves[block / LeafBlocks];
+        return leaf != nullptr ? leaf->driverBlocks[block % LeafBlocks] : 0;
+#endif
+    }
+
     void stamp(std::uint64_t block, std::uint32_t stampGeneration, StampKind kind) {
         const bool cpu = kind != StampKind::Driver;
         const bool written = kind != StampKind::ImportWindow;
 #ifdef _WIN32
         blocks[block] = stampGeneration;
         if (cpu) cpuBlocks[block] = stampGeneration;
         if (written) writtenBlocks[block] = stampGeneration;
+        if (kind == StampKind::Driver) driverBlocks[block] = stampGeneration;
 #else
         auto& leaf = leaves[block / LeafBlocks];
         if (leaf == nullptr) leaf = std::make_unique<Leaf>();
         leaf->blocks[block % LeafBlocks] = stampGeneration;
         if (cpu) leaf->cpuBlocks[block % LeafBlocks] = stampGeneration;
         if (written) leaf->writtenBlocks[block % LeafBlocks] = stampGeneration;
+        if (kind == StampKind::Driver) leaf->driverBlocks[block % LeafBlocks] = stampGeneration;
 #endif
     }
+    bool driverStoredOver(std::uint64_t address, std::size_t bytes, std::uint64_t generation) const {
+        if (bytes == 0) return false;
+        if (!watched || generation == 0 || bytes > std::numeric_limits<std::uint64_t>::max() - address || !covers(address, bytes)) return true;
+        const auto end = address + bytes;
+        const auto first = blockOf(address);
+        const auto last = blockOf(end - 1);
+        for (auto block = first; block <= last; ++block) {
+            if (driverStampOf(block) <= generation) continue;
+            const auto found = driverPieces.find(block);
+            if (found == driverPieces.end()) return true;
+            const auto& entry = found->second;
+            if (entry.whole > generation || entry.dropped > generation) return true;
+            const auto begin = blockBegin(block);
+            const auto from = std::max(address, begin) - begin;
+            const auto to = std::min<std::uint64_t>(end, begin + WriteBlockBytes) - begin;
+            for (const auto& piece : entry.pieces) {
+                if (piece.generation > generation && piece.begin < to && from < piece.end) return true;
+            }
+        }
+        return false;
+    }
 };
 
 WriteTracker& Tracker() {
@@ -1305,6 +1340,13 @@ bool StoredOver(std::uint64_t address, std::size_t bytes, std::uint64_t generati
     return false;
 }
 
+bool DriverStoredOver(std::uint64_t address, std::size_t bytes, std::uint64_t generation) {
+    auto& tracker = Tracker();
+    const auto lock = lockTracker(tracker);
+    tracker.initialize();
+    return tracker.driverStoredOver(address, bytes, generation);
+}
+
 std::uint64_t TrackerGeneration() {
     auto& tracker = Tracker();
     // Under the mutex: a collect bumps the generation before its walk, and a walk still running
@@ -1797,6 +1839,34 @@ void WriteChangedCommitted(std::uint64_t address, std::span<const std::byte> cur
     }
 }
 
+bool MergeChangedCommitted(std::uint64_t address, std::span<const std::byte> current, std::span<const std::byte> original, std::uint64_t generation) {
+    require(current.size() == original.size(), "write-back snapshot sizes differ");
+    if (current.empty()) return true;
+    const ReadSiteScope site(ReadSite::Store);
+    FlushGpuWrites(address, current.size());
+    auto* destination = reinterpret_cast<std::byte*>(address);
+    if (!Accessible(destination, current.size(), true)) return false;
+    const TimedAccess timed(CounterChangedWrite, current.size());
+    bool merged = false;
+    storeOwn(address, current.size(), [&] {
+        if (Tracker().driverStoredOver(address, current.size(), generation)) return std::pair<std::uint64_t, std::uint64_t>{0, 0};
+        CheckRange(destination, current.size(), 1, true);
+        merged = true;
+        std::size_t firstChanged = current.size();
+        std::size_t lastChanged = 0;
+        for (std::size_t byte = 0; byte < current.size(); ++byte) {
+            if (current[byte] == original[byte]) continue;
+            auto expected = original[byte];
+            if (std::atomic_ref(destination[byte]).compare_exchange_strong(expected, current[byte], std::memory_order_relaxed)) {
+                firstChanged = std::min(firstChanged, byte);
+                lastChanged = byte + 1;
+            }
+        }
+        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 merged;
+}
+
 void WriteChanged(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;

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

@@ -18,10 +18,15 @@
 #include <utility>
 #include <vector>
 
+namespace GuestAllocations {
+struct Range;
+}
+
 namespace AgcDriver::Graphics {
 
 class ResidentColor;
 class CommandBatch;
+class Buffer;
 class StorageTexture;
 struct HostImport;
 
@@ -304,6 +309,27 @@ class StorageTexture : public std::enable_shared_from_this<StorageTexture> {
     // APS5_BLOCK_WRITEBACK_EACH=1 stores the touched units only).
     void writeBack(std::uint64_t address, std::size_t bytes);
     void writeBackLayers(const std::vector<bool>& layers);
+    struct EdgeSnapshot {
+        std::uint64_t begin;
+        std::uint64_t end;
+        std::shared_ptr<Buffer> bytes;
+        std::weak_ptr<const void> allocation;
+        std::size_t offset = 0;
+    };
+    struct EdgeMerge {
+        EdgeSnapshot baseline;
+        std::shared_ptr<Buffer> result;
+        std::uint64_t generation;
+        std::size_t offset = 0;
+    };
+    const EdgeSnapshot* edgeSnapshot(std::uint64_t begin, std::uint64_t end) const;
+    void prepareEdgeCapture(std::uint64_t begin, std::uint64_t end);
+    void captureEdges(VkCommandBuffer commands, const HostImport& import, VkBuffer source, std::span<const VkBufferCopy> copies, std::span<const std::pair<std::uint64_t, std::uint64_t>> ranges);
+    void captureCpuEdges(std::span<const std::pair<std::uint64_t, std::uint64_t>> ranges);
+    void captureCpuEdges(std::span<const std::pair<std::uint64_t, std::uint64_t>> ranges, std::span<const std::byte> source, std::span<const std::shared_ptr<const GuestAllocations::Range>> allocations);
+    void invalidateEdges(std::uint64_t begin, std::uint64_t end);
+    void finishEdgeMerges(std::span<EdgeMerge> merges);
+    static std::vector<std::pair<std::uint64_t, std::uint64_t>> withoutEdgeMerges(std::span<const std::pair<std::uint64_t, std::uint64_t>> ranges, std::span<const EdgeMerge> merges);
     // Tracked units as 64 KiB write-stamp blocks (`blockUnits`: a thin tiled surface at a 64 KiB
     // aligned base; APS5_NO_BLOCK_TRACKING=1 tracks array layers as above instead): a fill of one
     // layer, a CPU write or another image's store then costs the blocks it touched, moved through
@@ -354,7 +380,7 @@ class StorageTexture : public std::enable_shared_from_this<StorageTexture> {
     // The write-back retiles into the import's unit shadow where a slab takes the piece
     // (`shadowed` receives those guest ranges, `imported` the pieces written to the import).
     std::uint64_t uploadWindows(const HostImport& import, std::span<const std::pair<std::uint64_t, std::uint64_t>> runs, bool discard = false);
-    std::uint64_t writeBackWindows(const HostImport& import, std::span<const std::pair<std::uint64_t, std::uint64_t>> keep, std::uint64_t firstStored, std::uint64_t lastStored, std::vector<ShadowedRange>& shadowed, std::vector<std::pair<std::uint64_t, std::uint64_t>>& imported);
+    std::uint64_t writeBackWindows(const HostImport& import, std::span<const std::pair<std::uint64_t, std::uint64_t>> keep, std::uint64_t firstStored, std::uint64_t lastStored, std::vector<ShadowedRange>& shadowed, std::vector<std::pair<std::uint64_t, std::uint64_t>>& imported, std::span<EdgeMerge> merges);
     // Per-layer validity (thin array surfaces whose layers are 64 KiB-aligned guest slices; every
     // other surface is one tracked layer): a layer's generation is the write generation its guest
     // bytes were last known to match the image at, and a pending layer holds results (a clear, a
@@ -425,6 +451,7 @@ class StorageTexture : public std::enable_shared_from_this<StorageTexture> {
     std::uint64_t sliceLinearBytes = 0;
     SurfaceGeometry geometry;
     std::vector<std::byte> original;
+    std::vector<EdgeSnapshot> edgeSnapshots;
     // DCC keys the image content was uploaded under: a fast-cleared surface starts as its clear value.
     DccKeys uploadedKeys = DccKeys::Uncompressed;
     mutable DccKeys filledKeys = DccKeys::Uncompressed;