fix(agc): retain raw draw decodes across snapshot eviction

core/libs/prx/libSceAgcDriver/Execution/include/Driver/Shaders/ShaderRegistry.hpp

@@ -303,7 +303,7 @@ struct ShaderSnapshot {
 };
 
 
-std::shared_ptr<const ShaderSnapshot> ReadRawShader(std::uint64_t address);
+std::shared_ptr<const ShaderSnapshot> ReadRawShader(std::uint64_t address, const std::shared_ptr<const ShaderSnapshot>& retained = {});
 std::shared_ptr<const ShaderSnapshot> ReadUserDataTailContinuation(std::span<const std::uint32_t> code, std::uint32_t userDataBase, std::span<const std::uint32_t> userData);
 std::shared_ptr<const ShaderSnapshot> ProgramSnapshot(const ShaderRegistry& shaders, std::uint64_t address, std::initializer_list<std::uint8_t> types);
 void NoteRawGraphicsDraw(const ShaderSnapshot& snapshot);

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Draw/DrawCache.cpp

@@ -227,8 +227,8 @@ void Driver::eraseDrawEntry(std::uint64_t key, const std::shared_ptr<DrawEntry>&
 bool DecodeProgramsCurrent(const DrawDecode& decode, const ShaderRegistry&) {
     return std::all_of(decode.programs.begin(), decode.programs.end(), [&](const DrawProgram& program) {
         if (!program.snapshot->header.empty()) return true;
-        return ReadRawShader(program.snapshot->codeAddress + program.codeOffset * sizeof(std::uint32_t)) == program.snapshot &&
-            (program.tailContinuation == nullptr || ReadRawShader(program.tailContinuation->codeAddress) == program.tailContinuation);
+        return ReadRawShader(program.snapshot->codeAddress + program.codeOffset * sizeof(std::uint32_t), program.snapshot) == program.snapshot &&
+            (program.tailContinuation == nullptr || ReadRawShader(program.tailContinuation->codeAddress, program.tailContinuation) == program.tailContinuation);
     });
 }
 

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Shaders/ShaderRegistry.cpp

@@ -61,7 +61,7 @@
 namespace AgcDriver::DriverDetail {
 
 
-std::shared_ptr<const ShaderSnapshot> ReadRawShader(std::uint64_t address) {
+std::shared_ptr<const ShaderSnapshot> ReadRawShader(std::uint64_t address, const std::shared_ptr<const ShaderSnapshot>& retained) {
     const auto drain = [](std::uint64_t first, std::size_t bytes) {
         if (!DrawPipeline::Active()) return;
         auto& pipeline = DrawPipeline::Queue0();
@@ -70,8 +70,8 @@ std::shared_ptr<const ShaderSnapshot> ReadRawShader(std::uint64_t address) {
     static std::mutex cacheMutex;
     static std::list<std::shared_ptr<const ShaderSnapshot>> cache;
     static std::size_t cacheBytes = 0;
-    std::shared_ptr<const ShaderSnapshot> cached;
-    {
+    std::shared_ptr<const ShaderSnapshot> cached = retained != nullptr && retained->codeAddress == address && retained->header.empty() && !retained->code.empty() ? retained : nullptr;
+    if (cached == nullptr) {
         std::lock_guard lock(cacheMutex);
         const auto found = std::find_if(cache.begin(), cache.end(), [address](const auto& entry) { return entry->codeAddress == address; });
         if (found != cache.end()) cached = *found;

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

@@ -10,6 +10,7 @@
 #include <iostream>
 #include <stdexcept>
 #include <string>
+#include <utility>
 #include <vector>
 
 namespace {
@@ -28,7 +29,42 @@ void Bind(AgcDriver::QueueState& queue, const void* code, std::uint32_t program,
     queue.shader[resources] = 16u << 1u;
 }
 
+void CheckRetainedDecode() {
+    using namespace AgcDriver::DriverDetail;
+    alignas(256) static std::array<std::uint32_t, 64> code{0xbf810000u};
+    alignas(256) static std::array<std::uint32_t, 64> tail{0xbf810000u};
+    alignas(256) static std::array<std::array<std::uint32_t, 64>, 4097> others{};
+    const auto address = reinterpret_cast<std::uintptr_t>(code.data());
+    DrawProgram program{};
+    program.snapshot = ReadRawShader(address);
+    program.tailContinuation = ReadRawShader(reinterpret_cast<std::uintptr_t>(tail.data()));
+    DrawDecode decode{};
+    decode.programs.push_back(std::move(program));
+    const ShaderRegistry registry;
+    Require(DecodeProgramsCurrent(decode, registry), "a newly captured raw draw decode is not current");
+    for (auto& other : others) {
+        other.fill(0xbf810000u);
+        static_cast<void>(ReadRawShader(reinterpret_cast<std::uintptr_t>(other.data())));
+    }
+    const auto reread = ReadRawShader(address);
+    Require(reread != decode.programs.front().snapshot, "raw snapshot churn did not evict the retained snapshot");
+    Require(reread->code == decode.programs.front().snapshot->code, "raw snapshot churn changed the shader contents");
+    std::cout << "retained raw snapshot evicted=" << (reread != decode.programs.front().snapshot) << '\n';
+    Require(DecodeProgramsCurrent(decode, registry), "an unchanged retained raw draw decode was invalidated by snapshot cache eviction");
+    code[0] = 0xbf800000u;
+    code[1] = 0xbf810000u;
+    Require(!DecodeProgramsCurrent(decode, registry), "a retained raw draw decode ignored changed main code");
+    code[0] = 0xbf810000u;
+    Require(DecodeProgramsCurrent(decode, registry), "restoring raw main code did not restore its retained decode");
+    tail[0] = 0xbf800000u;
+    tail[1] = 0xbf810000u;
+    Require(!DecodeProgramsCurrent(decode, registry), "a retained raw draw decode ignored changed tail code");
+    tail[0] = 0xbf810000u;
+    Require(DecodeProgramsCurrent(decode, registry), "restoring raw tail code did not restore its retained decode");
+}
+
 void Check(AgcDriver::VulkanDevice& device) {
+    CheckRetainedDecode();
     using namespace AgcDriver::DriverDetail;
     using namespace ShaderRecompiler;
     AgcDriver::QueueState queue{};

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

@@ -1,3 +1,4 @@
+#include "prx/libSceAgcDriver/Execution/include/Driver/Draw/DrawCache.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Driver/Draw/DrawPipeline.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Driver/Shaders/ShaderRegistry.hpp"
 #include "prx/libSceAgcDriver/Execution/include/GuestMemory.hpp"
@@ -11,6 +12,7 @@
 #include <stdexcept>
 #include <string>
 #include <thread>
+#include <utility>
 
 namespace {
 
@@ -22,6 +24,12 @@ void QueuedCode(bool cached) {
     code[0] = 0x7e000200u;
     const auto address = reinterpret_cast<std::uintptr_t>(code.data());
     const auto before = cached ? ReadRawShader(address) : nullptr;
+    DrawDecode decode{};
+    if (before != nullptr) {
+        DrawProgram program{};
+        program.snapshot = before;
+        decode.programs.push_back(std::move(program));
+    }
     std::atomic<bool> finished{false};
     std::atomic<bool> started{false};
     struct Restore {
@@ -42,6 +50,9 @@ void QueuedCode(bool cached) {
     }, {{address, address + sizeof(std::uint32_t)}});
     while (!started.load(std::memory_order_acquire)) std::this_thread::yield();
     DrawPipeline::Active() = true;
+    if (before != nullptr && DecodeProgramsCurrent(decode, ShaderRegistry{})) {
+        throw std::runtime_error("a retained raw draw decode ignored a queued write to its code");
+    }
     const auto after = ReadRawShader(address);
     finished.store(true, std::memory_order_release);
     DrawPipeline::Queue0().Drain(DrawPipeline::DrainReason::Submission);