fix(agc): order raw shader reads after queued code writes

core/libs/prx/libSceAgcDriver/CMakeLists.txt

@@ -1875,6 +1875,7 @@ if(BUILD_TESTING)
         set_tests_properties(agc_driver_draw_pipeline_trace_disabled PROPERTIES ENVIRONMENT_MODIFICATION
                 "PATH=path_list_prepend:$<TARGET_FILE_DIR:libc>")
     endif()
+    set_tests_properties(agc_driver_queued_raw_shader PROPERTIES ENVIRONMENT "ANYPS5_NO_SHADER_CACHE=1;APS5_PIPELINED_DRAWS=8" TIMEOUT 30)
     add_test(NAME agc_driver_partially_resident_pending_stores COMMAND agc_driver_partially_resident_snapshot_tests --pending-stores)
     set_tests_properties(agc_driver_partially_resident_pending_stores PROPERTIES ENVIRONMENT "ANYPS5_NO_SHADER_CACHE=1" SKIP_RETURN_CODE 77 WORKING_DIRECTORY "${CMAKE_SOURCE_DIR}" TIMEOUT 30)
     if(WIN32)

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

@@ -2,6 +2,7 @@
 #include "ControlFlow/GraphBuilder.hpp"
 #include "prx/libSceAgcDriver/Execution/include/PerformanceTimer.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Driver/Driver.hpp"
+#include "prx/libSceAgcDriver/Execution/include/Driver/Draw/DrawPipeline.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Driver/Shaders/ShaderRegistry.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Driver/Diagnostics.hpp"
 #include "prx/libSceAgcDriver/Execution/include/GuestMemory.hpp"
@@ -69,6 +70,11 @@ std::shared_ptr<const ShaderRecompiler::SourceHandle> PrepareShaderWithDiagnosti
 }
 
 std::shared_ptr<const ShaderSnapshot> ReadRawShader(std::uint64_t address) {
+    const auto drain = [](std::uint64_t first, std::size_t bytes) {
+        if (!DrawPipeline::Active()) return;
+        auto& pipeline = DrawPipeline::Queue0();
+        if (pipeline.Busy() && pipeline.Overlaps(first, bytes)) pipeline.Drain(DrawPipeline::DrainReason::Capture);
+    };
     GuestMemory::CheckRange(reinterpret_cast<const void*>(address), sizeof(std::uint32_t), 256);
     static std::mutex cacheMutex;
     static std::list<std::shared_ptr<const ShaderSnapshot>> cache;
@@ -81,6 +87,7 @@ std::shared_ptr<const ShaderSnapshot> ReadRawShader(std::uint64_t address) {
     }
     if (cached) {
         const auto code = std::as_bytes(std::span(cached->code));
+        drain(address, code.size());
         GuestMemory::FlushGpuWrites(address, code.size());
         if (GuestMemory::CompareMapped(address, code) == GuestMemory::Compare::Equal) {
             std::lock_guard lock(cacheMutex);
@@ -101,6 +108,7 @@ std::shared_ptr<const ShaderSnapshot> ReadRawShader(std::uint64_t address) {
     while (snapshot.code.size() < available) {
         const auto previous = snapshot.code.size();
         snapshot.code.resize(std::min(available, std::max<std::size_t>(64, previous * 2)));
+        drain(address + previous * sizeof(std::uint32_t), (snapshot.code.size() - previous) * sizeof(std::uint32_t));
         GuestMemory::Read(address + previous * sizeof(std::uint32_t),
             std::as_writable_bytes(std::span(snapshot.code).subspan(previous)), alignof(std::uint32_t));
         try {

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

@@ -0,0 +1,66 @@
+#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"
+#include <array>
+#include <atomic>
+#include <chrono>
+#include <cstdint>
+#include <exception>
+#include <iostream>
+#include <span>
+#include <stdexcept>
+#include <string>
+#include <thread>
+
+namespace {
+
+void QueuedCode(bool cached) {
+    using namespace AgcDriver::DriverDetail;
+    alignas(256) static std::array<std::array<std::uint32_t, 64>, 2> programs{};
+    auto& code = programs[cached ? 0 : 1];
+    code.fill(0xbf810000u);
+    code[0] = 0x7e000200u;
+    const auto address = reinterpret_cast<std::uintptr_t>(code.data());
+    const auto before = cached ? ReadRawShader(address) : nullptr;
+    std::atomic<bool> finished{false};
+    std::atomic<bool> started{false};
+    struct Restore {
+        std::atomic<bool>& finished;
+        bool active = DrawPipeline::Active();
+        ~Restore() {
+            finished.store(true, std::memory_order_release);
+            DrawPipeline::Active() = active;
+            DrawPipeline::Queue0().Drain(DrawPipeline::DrainReason::Submission);
+        }
+    } restore{finished};
+    DrawPipeline::Queue0().Enqueue([&, address] {
+        const auto deadline = std::chrono::steady_clock::now() + std::chrono::milliseconds(200);
+        started.store(true, std::memory_order_release);
+        while (!finished.load(std::memory_order_acquire) && std::chrono::steady_clock::now() < deadline) std::this_thread::yield();
+        const std::uint32_t replacement = 0x7e0002f2u;
+        AgcDriver::GuestMemory::Write(address, std::as_bytes(std::span(&replacement, 1)), alignof(std::uint32_t));
+    }, {{address, address + sizeof(std::uint32_t)}});
+    while (!started.load(std::memory_order_acquire)) std::this_thread::yield();
+    DrawPipeline::Active() = true;
+    const auto after = ReadRawShader(address);
+    finished.store(true, std::memory_order_release);
+    DrawPipeline::Queue0().Drain(DrawPipeline::DrainReason::Submission);
+    if (after->code.front() != 0x7e0002f2u || (before != nullptr && (before->code.front() != 0x7e000200u || before == after))) {
+        throw std::runtime_error(std::string(cached ? "cached" : "fresh") + " raw shader read ignored a queued write to its code");
+    }
+}
+
+}
+
+int main(int argc, char** argv) {
+    try {
+        if (AgcDriver::DriverDetail::DrawPipeline::Depth() == 0) return 77;
+        if (argc == 1 || std::string(argv[1]) == "cached") QueuedCode(true);
+        if (argc == 1 || std::string(argv[1]) == "fresh") QueuedCode(false);
+        std::cout << "Raw shader reads observed queued code writes\n";
+        return 0;
+    } catch (const std::exception& error) {
+        std::cerr << error.what() << '\n';
+        return 1;
+    }
+}