fix(agc): read unregistered graphics programs beside registered shaders

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

@@ -87,15 +87,17 @@ struct ShaderSnapshot {
     std::shared_ptr<const RegisteredShaderState> registeredState;
     std::shared_ptr<std::atomic<std::uint8_t>> workgroupUse = std::make_shared<std::atomic<std::uint8_t>>(0u);
     std::shared_ptr<std::atomic<std::uint8_t>> positionVectorUse = std::make_shared<std::atomic<std::uint8_t>>(0u);
+    std::shared_ptr<std::atomic<bool>> rawGraphicsNoted = std::make_shared<std::atomic<bool>>(false);
     std::optional<std::uint32_t> userDataTailRegister;
 };
 
 
 std::shared_ptr<const ShaderSnapshot> ReadRawShader(std::uint64_t address);
 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();
+void NoteRawGraphicsDraw(const ShaderSnapshot& snapshot);
 bool RawGraphicsDrawn();
+bool RawGraphicsDrawnIn(std::uint64_t begin, std::uint64_t end);
 bool ProgramUsesWorkgroup(const ShaderSnapshot& snapshot, std::size_t codeOffset);
 // The PA_CL_VS_OUT_CNTL position export bits a vertex program is compiled with: none for a program
 // without pos1..pos3 exports, so its artifact does not depend on them.

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

@@ -65,7 +65,7 @@ void DecodeGraphicsPrograms(DrawDecode& decoded, const QueueState& queue, const
         const auto& snapshot = *program;
         require((address - snapshot.codeAddress) % sizeof(std::uint32_t) == 0, "graphics entry point is not dword aligned");
         require(snapshot.header.empty() || snapshot.type == type, "graphics program refers to an incompatible shader binary type");
-        if (!nullPixel && snapshot.header.empty()) NoteRawGraphicsDraw();
+        if (!nullPixel && snapshot.header.empty()) NoteRawGraphicsDraw(snapshot);
         const auto codeOffset = static_cast<std::size_t>((address - snapshot.codeAddress) / sizeof(std::uint32_t));
         DrawProgram result{
             {stage, address, std::span(snapshot.code).subspan(codeOffset), snapshot.headerAddress, snapshot.header},

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

@@ -24,7 +24,10 @@
 #include <deque>
 #include <exception>
 #include <functional>
+#include <iterator>
 #include <list>
+#include <map>
+#include <mutex>
 #include <atomic>
 #include <pthread.h>
 #ifdef __GLIBC__
@@ -142,22 +145,49 @@ std::shared_ptr<const ShaderSnapshot> ReadUserDataTailContinuation(std::span<con
 
 std::shared_ptr<const ShaderSnapshot> ProgramSnapshot(const ShaderRegistry& shaders, std::uint64_t address, std::initializer_list<std::uint8_t> types) {
     if (auto registered = RegisteredProgram(shaders, address, types)) return registered;
-    require(shaders.empty(), "graphics program does not belong to a compatible registered shader");
     return ReadRawShader(address);
 }
 
 namespace {
 std::atomic<bool> rawGraphicsDrawn{false};
+std::mutex rawGraphicsMutex;
+std::map<std::uint64_t, std::uint64_t> rawGraphicsRanges;
 }
 
-void NoteRawGraphicsDraw() {
+void NoteRawGraphicsDraw(const ShaderSnapshot& snapshot) {
+    if (snapshot.rawGraphicsNoted->load(std::memory_order_acquire)) return;
+    auto begin = snapshot.codeAddress;
+    auto end = snapshot.codeAddress + snapshot.code.size() * sizeof(std::uint32_t);
+    {
+        std::lock_guard lock(rawGraphicsMutex);
+        auto next = rawGraphicsRanges.upper_bound(begin);
+        if (next != rawGraphicsRanges.begin() && std::prev(next)->second >= begin) {
+            const auto previous = std::prev(next);
+            begin = previous->first;
+            end = std::max(end, previous->second);
+            rawGraphicsRanges.erase(previous);
+        }
+        while (next != rawGraphicsRanges.end() && next->first <= end) {
+            end = std::max(end, next->second);
+            next = rawGraphicsRanges.erase(next);
+        }
+        rawGraphicsRanges.emplace(begin, end);
+    }
+    snapshot.rawGraphicsNoted->store(true, std::memory_order_release);
     rawGraphicsDrawn.store(true, std::memory_order_relaxed);
 }
 
 bool RawGraphicsDrawn() {
     return rawGraphicsDrawn.load(std::memory_order_relaxed);
 }
 
+bool RawGraphicsDrawnIn(std::uint64_t begin, std::uint64_t end) {
+    if (begin >= end) return false;
+    std::lock_guard lock(rawGraphicsMutex);
+    const auto next = rawGraphicsRanges.lower_bound(end);
+    return next != rawGraphicsRanges.begin() && std::prev(next)->second > begin;
+}
+
 namespace {
 
 void AwaitRegisteredPreparation(PreparedShaders& prepared, std::unique_lock<std::mutex>& lock) {
@@ -1434,7 +1464,6 @@ void Driver::RegisterShader(const Shader* shader) {
     PerformanceTimer timing("Shader.Register");
     ShaderPreparationTransaction transaction;
     CheckFailure();
-    require(!RawGraphicsDrawn(), "a shader was registered after graphics programs were drawn without registration");
     GuestMemory::CheckRange(shader, sizeof(Shader), 1);
     Shader fields;
     std::memcpy(&fields, static_cast<const void*>(shader), sizeof(Shader));
@@ -1444,6 +1473,7 @@ void Driver::RegisterShader(const Shader* shader) {
     GuestMemory::CheckRange(shader, fields.header_size, 1);
     const auto* code = const_cast<const void*>(fields.code);
     GuestMemory::CheckRange(code, fields.shader_size, 256);
+    require(!RawGraphicsDrawnIn(reinterpret_cast<std::uintptr_t>(code), reinterpret_cast<std::uintptr_t>(code) + fields.shader_size), "a shader was registered over a graphics program that was drawn without registration");
     ShaderSnapshot snapshot{reinterpret_cast<std::uintptr_t>(code), reinterpret_cast<std::uintptr_t>(shader), fields.type, {}, {}};
     snapshot.code.resize(fields.shader_size / sizeof(std::uint32_t));
     std::memcpy(snapshot.code.data(), code, fields.shader_size);

core/libs/prx/libSceAgcDriver/tests/Submit.cpp

@@ -671,9 +671,10 @@ void testProgramSnapshots() {
     shaders[registeredAddress] = {entry};
     check(AgcDriver::DriverDetail::ProgramSnapshot(shaders, registeredAddress, {2}) == entry, "a registered program was not resolved to its registration");
     check(AgcDriver::DriverDetail::ProgramSnapshot(shaders, registeredAddress + 4 * 8, {2}) == entry, "an entry inside registered code was not resolved to its registration");
-    check(!expectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::ProgramSnapshot(shaders, rawAddress, {2})); }).empty(), "an unregistered program was read raw while shaders are registered");
-    const auto unregistered = AgcDriver::DriverDetail::ProgramSnapshot({}, rawAddress, {2});
-    check(unregistered != entry && unregistered->codeAddress == rawAddress && unregistered->header.empty() && unregistered->code.size() == 2, "an empty registry did not read the program as raw code");
+    const auto unregistered = AgcDriver::DriverDetail::ProgramSnapshot(shaders, rawAddress, {2});
+    check(unregistered != entry && unregistered->codeAddress == rawAddress && unregistered->header.empty() && unregistered->code.size() == 2, "an unregistered program beside registered shaders was not read as raw code");
+    check(AgcDriver::DriverDetail::ProgramSnapshot({}, rawAddress, {2}) == unregistered, "an empty registry did not read the program as raw code");
+    check(!expectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::ProgramSnapshot(shaders, registeredAddress, {1})); }).empty(), "a registered program was accepted as another shader type");
     check(AgcDriver::DriverDetail::ProgramSnapshot({}, rawAddress, {2}) == unregistered, "unchanged raw code lost its snapshot identity");
     check(!expectFailure([] { static_cast<void>(AgcDriver::DriverDetail::ProgramSnapshot({}, 0, {2})); }).empty(), "an unmapped unregistered program was accepted");
     AgcDriver::DriverDetail::DrawDecode decode{};
@@ -690,6 +691,27 @@ void testProgramSnapshots() {
     check(AgcDriver::DriverDetail::DecodeProgramsCurrent(decode, shaders), "a decode with the rewritten program was not current");
 }
 
+void testRawGraphicsRanges() {
+    using AgcDriver::DriverDetail::RawGraphicsDrawnIn;
+    constexpr std::uint64_t base = 0x7e0000000000ull;
+    const auto note = [](std::uint64_t address, std::size_t dwords) {
+        AgcDriver::DriverDetail::ShaderSnapshot snapshot{address, 0, 0, std::vector<std::uint32_t>(dwords, 0xbf810000u), {}};
+        AgcDriver::DriverDetail::NoteRawGraphicsDraw(snapshot);
+        check(snapshot.rawGraphicsNoted->load(), "a raw graphics program was not marked as noted");
+        AgcDriver::DriverDetail::NoteRawGraphicsDraw(snapshot);
+    };
+    note(base + 0x100, 0x10);
+    note(base, 0x100);
+    note(base + 0x3f0, 0x8);
+    note(base + 0x800, 0x1);
+    check(RawGraphicsDrawnIn(base + 0x200, base + 0x210), "code inside a raw program that contains a smaller one was not reported");
+    check(RawGraphicsDrawnIn(base + 0x40c, base + 0x500), "the end of a raw program that overlaps another was not reported");
+    check(!RawGraphicsDrawnIn(base + 0x410, base + 0x800), "code between raw programs was reported");
+    check(!RawGraphicsDrawnIn(base - 4, base), "code before a raw program was reported");
+    check(RawGraphicsDrawnIn(base + 0x700, base + 0x900), "a range covering a raw program was not reported");
+    check(!RawGraphicsDrawnIn(base + 0x804, base + 0x808) && !RawGraphicsDrawnIn(base + 0x200, base + 0x200), "code after a raw program or an empty range was reported");
+}
+
 void testProgramWorkgroupUse() {
     const auto snapshot = [](std::vector<std::uint32_t> code) {
         AgcDriver::DriverDetail::ShaderSnapshot result{0x1000, 0, 2, std::move(code), {}};
@@ -860,6 +882,7 @@ int main(int argc, char** argv) {
         check(!unterminated.empty() && bounded->code.size() == 64, "raw compute crossed inaccessible memory or missed its last instruction");
 #endif
         testProgramSnapshots();
+        testRawGraphicsRanges();
         testProgramWorkgroupUse();
         testVertexTailSnapshots();
         testRegisteredVertexWorkgroupRouting();

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

@@ -304,7 +304,13 @@ void Check(AgcDriver::VulkanDevice& device, AgcDriver::Graphics::ShaderPath path
         registry.erase(late.snapshot->codeAddress);
     }
     ShaderRegistry invalidRegistry;
-    invalidRegistry[1] = {std::make_shared<const ShaderSnapshot>(ShaderSnapshot{1, 0, 2, {}, {}})};
+    for (const auto& [address, snapshots] : registry) {
+        for (const auto& snapshot : snapshots) {
+            auto compute = *snapshot;
+            compute.type = 0;
+            invalidRegistry[address].push_back(std::make_shared<const ShaderSnapshot>(std::move(compute)));
+        }
+    }
     DrawDecode invalid{};
     invalid.state = prepared.state;
     Reject([&] { DecodeGraphicsPrograms(invalid, queue, invalidRegistry, true, true); }, "registered");

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

@@ -51,13 +51,13 @@ void Check(AgcDriver::VulkanDevice& device) {
     Require(draw.programs.size() == 2 && draw.roles.back() == ProgramRole::Fragment, "unregistered graphics programs did not decode as a vertex and a fragment stage");
     for (const auto& program : draw.programs) Require(program.snapshot != nullptr && program.snapshot->header.empty() && program.snapshot->prepared->entries.empty(), "an unregistered graphics program was not read as raw code");
     Require(RawGraphicsDrawn(), "a draw of unregistered graphics programs was not noted");
-    bool refused = false;
-    try {
-        Driver::Get().RegisterShader(nullptr);
-    } catch (const std::exception& error) {
-        refused = std::string(error.what()).find("registered after graphics programs were drawn") != std::string::npos;
+    for (const auto& program : draw.programs) {
+        const auto begin = program.snapshot->codeAddress;
+        Require(RawGraphicsDrawnIn(begin, begin + 4u) && RawGraphicsDrawnIn(begin - 4u, begin + 1u), "a raw-drawn program's code is not recorded");
+        Require(!RawGraphicsDrawnIn(begin - 4u, begin), "the raw-drawn range starts before the program");
     }
-    Require(refused, "a shader registration after unregistered graphics draws was accepted");
+    alignas(256) static std::array<std::uint32_t, 4> Elsewhere{0xbf810000u, 0xbf810000u, 0xbf810000u, 0xbf810000u};
+    Require(!RawGraphicsDrawnIn(reinterpret_cast<std::uintptr_t>(Elsewhere.data()), reinterpret_cast<std::uintptr_t>(Elsewhere.data()) + sizeof(Elsewhere)), "unrelated code reads as drawn raw");
     const auto target = device.Target();
     std::vector<LinkedProgram> linked;
     std::vector<MemoryRegion> memory;