From e5191ebf130610e4605a3e4ce2c9e441a66ea7f6 Mon Sep 17 00:00:00 2001 From: Nicolas Abosch <88494836+nicolasabosch@users.noreply.github.com> Date: Mon, 5 Oct 2026 22:52:03 -0300 Subject: [PATCH] fix(agc): reuse validated raw compute snapshots --- core/libs/prx/libSceAgc/DcbState/src/Init.cpp | 4 +- core/libs/prx/libSceAgc/tests/Commands.cpp | 9 ++-- .../src/Driver/Shaders/ShaderRegistry.cpp | 41 ++++++++++++++++++- .../libs/prx/libSceAgcDriver/tests/Submit.cpp | 32 +++++++++++++++ 4 files changed, 79 insertions(+), 7 deletions(-) diff --git a/core/libs/prx/libSceAgc/DcbState/src/Init.cpp b/core/libs/prx/libSceAgc/DcbState/src/Init.cpp index 7c8283eaa..1b63094d1 100644 --- a/core/libs/prx/libSceAgc/DcbState/src/Init.cpp +++ b/core/libs/prx/libSceAgc/DcbState/src/Init.cpp @@ -9,8 +9,8 @@ extern "C" { -APS5_EXPORT("23LRUSvYu1M", sceAgcInitState_nid_no_patch); -int APS5_VABI sceAgcInitState_nid_no_patch(std::uint32_t* state, std::uint32_t version) { +APS5_EXPORT("23LRUSvYu1M", sceAgcUnknownInitState); +int APS5_VABI sceAgcUnknownInitState(std::uint32_t* state, std::uint32_t version) { Agc::Command::CheckAddress(reinterpret_cast(state), alignof(std::uint32_t), __func__); Agc::Command::Require(version < 14, __func__, "unsupported AGC initialization version"); return 0; diff --git a/core/libs/prx/libSceAgc/tests/Commands.cpp b/core/libs/prx/libSceAgc/tests/Commands.cpp index feefffe0b..4e11a5bed 100644 --- a/core/libs/prx/libSceAgc/tests/Commands.cpp +++ b/core/libs/prx/libSceAgc/tests/Commands.cpp @@ -14,7 +14,7 @@ extern "C" std::uint32_t* APS5_VABI sceAgcDcbResetQueue(CommandBuffer* buf, std: extern "C" std::uint32_t* APS5_VABI sceAgcDcbSetFlip(CommandBuffer* buf, std::uint32_t handle, std::int32_t index, std::uint32_t mode, std::int64_t argument); extern "C" int APS5_VABI sceAgcSuspendPoint(); extern "C" int APS5_VABI sceAgcInit(std::uint32_t version); -extern "C" int APS5_VABI sceAgcInitState_nid_no_patch(std::uint32_t* state, std::uint32_t version); +extern "C" int APS5_VABI sceAgcUnknownInitState(std::uint32_t* state, std::uint32_t version); extern "C" std::uint32_t* APS5_VABI sceAgcDcbDrawIndexAuto(CommandBuffer* buf, std::uint32_t indexCount, std::uint64_t modifier); extern "C" std::uint32_t* APS5_VABI sceAgcDcbDrawIndexIndirect(CommandBuffer* buf, std::uint32_t dataOffsetInBytes, std::uint64_t modifier); extern "C" std::uint32_t* APS5_VABI sceAgcDcbDrawIndexIndirectMulti(CommandBuffer* buf, std::uint32_t dataOffsetInBytes, std::uint32_t countIndirect, std::uint32_t maxCountOrCount, const volatile void* countAddress, std::uint32_t strideInBytes, std::uint64_t modifier); @@ -337,9 +337,10 @@ void testMemory() { void testDefaults() { std::uint32_t state = 0x12345678; - check(sceAgcInitState_nid_no_patch(&state, 8) == 0 && state == 0x12345678, "AGC initialization failed or modified caller state"); - expectFailure([] { sceAgcInitState_nid_no_patch(nullptr, 8); }); - expectFailure([&] { sceAgcInitState_nid_no_patch(&state, 14); }); + check(sceAgcUnknownInitState(&state, 8) == 0 && state == 0x12345678, "AGC initialization failed or modified caller state"); + check(sceAgcUnknownInitState(&state, 13) == 0 && state == 0x12345678, "AGC version 13 initialization changed caller state"); + expectFailure([] { sceAgcUnknownInitState(nullptr, 8); }); + expectFailure([&] { sceAgcUnknownInitState(&state, 14); }); check(sceAgcInit(8) == 0, "AGC version initialization failed"); expectFailure([] { sceAgcInit(14); }); expectFailure([] { sceAgcInit(0xffffffffu); }); diff --git a/core/libs/prx/libSceAgcDriver/Execution/src/Driver/Shaders/ShaderRegistry.cpp b/core/libs/prx/libSceAgcDriver/Execution/src/Driver/Shaders/ShaderRegistry.cpp index d40f7e4d4..8ba86476a 100644 --- a/core/libs/prx/libSceAgcDriver/Execution/src/Driver/Shaders/ShaderRegistry.cpp +++ b/core/libs/prx/libSceAgcDriver/Execution/src/Driver/Shaders/ShaderRegistry.cpp @@ -6,6 +6,7 @@ #include "Optimization/ResourceProgram.hpp" #include #include +#include #include #include "RdnaDecoder/RdnaInstructionDecoder.hpp" @@ -13,6 +14,25 @@ namespace AgcDriver::DriverDetail { std::shared_ptr ReadRawComputeShader(std::uint64_t address) { GuestMemory::CheckRange(reinterpret_cast(address), sizeof(std::uint32_t), 256); + static std::mutex cacheMutex; + static std::list> cache; + static std::size_t cacheBytes = 0; + std::shared_ptr cached; + { + 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; + } + if (cached) { + const auto code = std::as_bytes(std::span(cached->code)); + GuestMemory::FlushGpuWrites(address, code.size()); + if (GuestMemory::CompareMapped(address, code) == GuestMemory::Compare::Equal) { + std::lock_guard lock(cacheMutex); + const auto found = std::find(cache.begin(), cache.end(), cached); + if (found != cache.end()) cache.splice(cache.begin(), cache, found); + return cached; + } + } constexpr std::size_t limit = 1024 * 1024; const auto ranges = GuestMemory::CommittedRanges(address, limit); std::uint64_t end = address; @@ -31,7 +51,26 @@ std::shared_ptr ReadRawComputeShader(std::uint64_t address const auto decoded = ShaderRecompiler::RdnaInstructionDecoder{}.Decode(snapshot.code); const auto& last = decoded.instructions.back(); snapshot.code.resize(last.programCounter / sizeof(std::uint32_t) + last.wordCount); - return std::make_shared(std::move(snapshot)); + auto result = std::make_shared(std::move(snapshot)); + 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()) { + if ((*found)->code == result->code) { + result = *found; + cache.splice(cache.begin(), cache, found); + return result; + } + cacheBytes -= (*found)->code.size() * sizeof(std::uint32_t); + cache.erase(found); + } + const auto bytes = result->code.size() * sizeof(std::uint32_t); + while (!cache.empty() && (cache.size() >= 64 || cacheBytes + bytes > 8 * 1024 * 1024)) { + cacheBytes -= cache.back()->code.size() * sizeof(std::uint32_t); + cache.pop_back(); + } + cache.push_front(result); + cacheBytes += bytes; + return result; } catch (const std::out_of_range&) { if (snapshot.code.size() == available) break; } diff --git a/core/libs/prx/libSceAgcDriver/tests/Submit.cpp b/core/libs/prx/libSceAgcDriver/tests/Submit.cpp index a37a89f14..6888c50b6 100644 --- a/core/libs/prx/libSceAgcDriver/tests/Submit.cpp +++ b/core/libs/prx/libSceAgcDriver/tests/Submit.cpp @@ -375,12 +375,35 @@ int main() { const auto rawAddress = reinterpret_cast(rawCode.data()); const auto literal = AgcDriver::DriverDetail::ReadRawComputeShader(rawAddress); check(literal->code.size() == 3 && literal->header.empty(), "raw compute stopped at an instruction literal"); + check(AgcDriver::DriverDetail::ReadRawComputeShader(rawAddress) == literal, "unchanged raw compute code lost its snapshot identity"); rawCode[0] = 0xbf820002; rawCode[1] = 0xbf810000; rawCode[2] = 0xbf800000; rawCode[3] = 0xbf810000; const auto branched = AgcDriver::DriverDetail::ReadRawComputeShader(rawAddress); check(branched->code.size() == 4 && literal->code[0] == 0xbe8003ff, "raw compute lost branch targets or modified an earlier snapshot"); + check(branched != literal, "changed raw compute code reused a stale snapshot"); + std::array, 8> concurrent; + std::vector readers; + for (auto& snapshot : concurrent) readers.emplace_back([&snapshot, rawAddress] { + snapshot = AgcDriver::DriverDetail::ReadRawComputeShader(rawAddress); + }); + for (auto& reader : readers) reader.join(); + for (const auto& snapshot : concurrent) check(snapshot == branched, "concurrent raw compute readers lost snapshot reuse"); + alignas(256) std::array, 65> programs{}; + for (auto& program : programs) program[0] = 0xbf810000; + const auto firstAddress = reinterpret_cast(programs.front().data()); + readers.clear(); + for (auto& snapshot : concurrent) readers.emplace_back([&snapshot, firstAddress] { + snapshot = AgcDriver::DriverDetail::ReadRawComputeShader(firstAddress); + }); + for (auto& reader : readers) reader.join(); + for (const auto& snapshot : concurrent) check(snapshot == concurrent.front(), "concurrent raw compute misses duplicated snapshots"); + const auto evicted = AgcDriver::DriverDetail::ReadRawComputeShader(firstAddress); + for (std::size_t i = 1; i < programs.size(); ++i) + AgcDriver::DriverDetail::ReadRawComputeShader(reinterpret_cast(programs[i].data())); + check(AgcDriver::DriverDetail::ReadRawComputeShader(firstAddress) != evicted && evicted->code[0] == 0xbf810000, + "raw compute cache eviction lost snapshot lifetime or exceeded its entry limit"); check(!expectFailure([&] { AgcDriver::DriverDetail::ReadRawComputeShader(rawAddress + 4); }).empty(), "raw compute accepted a misaligned entry"); check(!expectFailure([] { AgcDriver::DriverDetail::ReadRawComputeShader(0); }).empty(), "raw compute accepted an unmapped entry"); #ifdef _WIN32 @@ -394,6 +417,15 @@ int main() { const auto unterminated = expectFailure([&] { AgcDriver::DriverDetail::ReadRawComputeShader(boundedAddress); }); boundedCode[63] = 0xbf810000; const auto bounded = AgcDriver::DriverDetail::ReadRawComputeShader(boundedAddress); + check(VirtualProtect(mapping + 1024, 4096, PAGE_READWRITE, &protection) != 0, "cannot extend raw compute mapping"); + boundedCode[63] = 0xbf800000; + boundedCode[64] = 0xbf810000; + const auto extended = AgcDriver::DriverDetail::ReadRawComputeShader(boundedAddress); + check(extended != bounded && extended->code.size() == 65, "raw compute reused code before its end changed"); + check(VirtualProtect(mapping + 1024, 4096, PAGE_NOACCESS, &protection) != 0, "cannot revoke cached raw compute tail"); + check(!expectFailure([&] { AgcDriver::DriverDetail::ReadRawComputeShader(boundedAddress); }).empty(), "raw compute reused an inaccessible cached tail"); + check(VirtualProtect(mapping, 4096, PAGE_NOACCESS, &protection) != 0, "cannot revoke cached raw compute code"); + check(!expectFailure([&] { AgcDriver::DriverDetail::ReadRawComputeShader(boundedAddress); }).empty(), "raw compute reused inaccessible cached code"); check(VirtualFree(mapping, 0, MEM_RELEASE) != 0, "cannot release raw compute boundary test"); check(!unterminated.empty() && bounded->code.size() == 64, "raw compute crossed inaccessible memory or missed its last instruction"); #endif