fix(agc): reuse validated raw compute snapshots

This commit is contained in:
Nicolas Abosch 2026-10-05 22:52:03 -03:00
parent 41cd420cc6
commit e5191ebf13
4 changed files with 79 additions and 7 deletions

View File

@ -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<std::uintptr_t>(state), alignof(std::uint32_t), __func__);
Agc::Command::Require(version < 14, __func__, "unsupported AGC initialization version");
return 0;

View File

@ -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); });

View File

@ -6,6 +6,7 @@
#include "Optimization/ResourceProgram.hpp"
#include <cstdlib>
#include <cstring>
#include <list>
#include <stdexcept>
#include "RdnaDecoder/RdnaInstructionDecoder.hpp"
@ -13,6 +14,25 @@ namespace AgcDriver::DriverDetail {
std::shared_ptr<const ShaderSnapshot> ReadRawComputeShader(std::uint64_t address) {
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;
static std::size_t cacheBytes = 0;
std::shared_ptr<const ShaderSnapshot> 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<const ShaderSnapshot> 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<const ShaderSnapshot>(std::move(snapshot));
auto result = std::make_shared<const ShaderSnapshot>(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;
}

View File

@ -375,12 +375,35 @@ int main() {
const auto rawAddress = reinterpret_cast<std::uintptr_t>(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<std::shared_ptr<const AgcDriver::DriverDetail::ShaderSnapshot>, 8> concurrent;
std::vector<std::thread> 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<std::array<std::uint32_t, 64>, 65> programs{};
for (auto& program : programs) program[0] = 0xbf810000;
const auto firstAddress = reinterpret_cast<std::uintptr_t>(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<std::uintptr_t>(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