feat(shader): resolve user-data s_swappc_b64 call targets

core/libs/prx/libSceAgcDriver/CMakeLists.txt

@@ -619,6 +619,7 @@ target_link_libraries(agc_driver_interpolation_modes_tests PRIVATE shader_recomp
 target_include_directories(agc_driver_prepared_shaders_tests PRIVATE
         $<TARGET_PROPERTY:shader_recompiler_agc_driver,INTERFACE_INCLUDE_DIRECTORIES>
 )
+target_include_directories(agc_driver_user_data_swappc_tests PRIVATE $<TARGET_PROPERTY:shader_recompiler_agc_driver,INTERFACE_INCLUDE_DIRECTORIES>)
 
 if(BUILD_TESTING)
     add_test(NAME agc_shader_bindless COMMAND agc_shader_memory_tests --bindless)
@@ -954,6 +955,16 @@ if(BUILD_TESTING)
         set_tests_properties(scalar_swappc PROPERTIES ENVIRONMENT_MODIFICATION
             "PATH=path_list_prepend:$<TARGET_FILE_DIR:libc>")
     endif()
+    add_test_executable(user_data_swappc_tests ${LIBS_INCLUDE_DIR}/tests/UserDataSwappc.cpp)
+    target_link_libraries(user_data_swappc_tests PRIVATE shader_recompiler_agc_driver libc)
+    configure_windows_unwind(user_data_swappc_tests)
+    add_test(NAME user_data_swappc COMMAND user_data_swappc_tests)
+    set_tests_properties(user_data_swappc PROPERTIES TIMEOUT 30)
+    set_tests_properties(agc_driver_user_data_swappc PROPERTIES TIMEOUT 30)
+    if(WIN32)
+        set_tests_properties(user_data_swappc PROPERTIES ENVIRONMENT_MODIFICATION
+            "PATH=path_list_prepend:$<TARGET_FILE_DIR:libc>")
+    endif()
     add_test_executable(scalar_bitcount64_tests ${LIBS_INCLUDE_DIR}/tests/ScalarBitCount64.cpp)
     target_link_libraries(scalar_bitcount64_tests PRIVATE shader_recompiler_agc_driver libc)
     configure_windows_unwind(scalar_bitcount64_tests)

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

@@ -79,7 +79,7 @@ void ResolvePreparedGraphics(const ShaderSnapshot& front, const std::shared_ptr<
 ShaderRecompiler::RectListShaders PreparedRectangle(const ShaderSnapshot& snapshot, std::uint64_t vertexId, std::uint64_t fragmentId);
 
 std::shared_ptr<const ShaderRecompiler::SourceHandle> SourceHandleFor(const ShaderSnapshot& snapshot, std::size_t codeOffset, const ShaderRecompiler::RecompileRequest& request);
-ShaderRecompiler::PreparedShaderInvocation InvocationFor(const ShaderSnapshot& snapshot, std::size_t codeOffset, const ShaderRecompiler::RecompileRequest& request);
+ShaderRecompiler::PreparedShaderInvocation InvocationFor(const ShaderSnapshot& snapshot, std::size_t codeOffset, const ShaderRecompiler::RecompileRequest& request, ShaderMemory* memory = nullptr);
 
 }
 

core/libs/prx/libSceAgcDriver/Execution/include/ShaderMemory.hpp

@@ -14,6 +14,7 @@
 
 namespace ShaderRecompiler {
 struct SourceHandle;
+struct CapturedCallProgram;
 class PreparedShaderInvocation;
 }
 
@@ -63,6 +64,7 @@ class ShaderMemory {
     // resolution.
     std::shared_ptr<const ShaderRecompiler::ResourceCapture> Capture(const ShaderRecompiler::RecompileRequest& request, const ShaderRecompiler::SourceHandle* handle = nullptr);
     std::shared_ptr<const ShaderRecompiler::ResourceCapture> Capture(const ShaderRecompiler::PreparedShaderInvocation& invocation);
+    ShaderRecompiler::CapturedCallProgram ResolveCalls(const ShaderRecompiler::RecompileRequest& request);
     [[nodiscard]] std::vector<ShaderRecompiler::MemoryRegion> Regions() const;
     // The page regions read since the previous call (or construction), a word read again
     // included, the initial regions excluded: one stage's own reads on the draw path's shared
@@ -91,6 +93,7 @@ class ShaderMemory {
     // the caller keeps them alive for as long as the capture is used.
     std::map<std::uint64_t, std::span<const std::byte>> initial;
     std::map<std::uint64_t, Page> pages;
+    std::vector<std::uint64_t> callCodeReads;
     PendingWriteQuery pendingWrite = nullptr;
     PendingWriteObserver observe = nullptr;
     HookWaitCounter hookWaits = nullptr;

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Dispatch/DataHitVerification.cpp

@@ -13,7 +13,7 @@ void Driver::verifyDataHit(const ShaderSnapshot& snapshot, std::size_t codeOffse
         std::abort();
     };
     auto verifyMemory = std::make_shared<ShaderMemory>(memory, &queryPendingWrite, &observePendingWrite, hookWaitCounter());
-    const auto invocation = InvocationFor(snapshot, codeOffset, request);
+    const auto invocation = InvocationFor(snapshot, codeOffset, request, verifyMemory.get());
     const auto capture = verifyMemory->Capture(invocation);
     const auto regions = verifyMemory->Regions();
     request.context.memory = regions;

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Dispatch/Dispatch.cpp

@@ -193,11 +193,12 @@ void Driver::dispatch(QueueState& queue, std::span<const std::uint32_t> packet,
 
         static const bool dumpShaders = std::getenv("APS5_DUMP_SHADERS") != nullptr;
         try {
-            const auto invocation = InvocationFor(snapshot, codeOffset, request);
-            timing.Mark("prepared_invocation");
-
             const auto waitedBefore = traceCapSync() ? Graphics::Recorder::ThreadWaitedMs() : 0.0;
             forgetAtCapture = GuestMemory::ForgetSerial();
+            const auto invocation = InvocationFor(snapshot, codeOffset, request, shaderMemory.get());
+            request.shader = invocation.Request().shader;
+            timing.Mark("prepared_invocation");
+
             capture = [&] {
                 const SampledReadScope sampling(evidenceReads);
                 return shaderMemory->Capture(invocation);

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

@@ -1,5 +1,6 @@
 #include "prx/libc/include/general/LogMacros.hpp"
 #include "ControlFlow/GraphBuilder.hpp"
+#include "ControlFlow/UserDataCalls.hpp"
 #include "ControlFlow/Structurizer.hpp"
 #include "prx/libSceAgcDriver/Execution/include/PerformanceTimer.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Driver/Driver.hpp"
@@ -175,9 +176,11 @@ std::shared_ptr<const ShaderRecompiler::SourceHandle> SourceHandleFor(const Shad
     for (const auto& entry : snapshot.prepared->entries) {
         if (entry.codeOffset == codeOffset && ShaderRecompiler::MatchesPreparedShader(request, *entry.handle, key)) return entry.handle;
     }
-    if (snapshot.header.empty()) {
-        if (snapshot.type != 0 || request.shader.stage != ShaderRecompiler::ShaderStage::Compute) throw std::runtime_error("AGC driver: unregistered program is not a compute shader");
-        APS5_LOG_ERR("Compute shader 0x%llx was not registered; preparing its artifact at dispatch", static_cast<unsigned long long>(snapshot.codeAddress));
+    if (snapshot.header.empty() || !request.shader.capturedCalls.empty()) {
+        if (snapshot.header.empty()) {
+            if (snapshot.type != 0 || request.shader.stage != ShaderRecompiler::ShaderStage::Compute) throw std::runtime_error("AGC driver: unregistered program is not a compute shader");
+            APS5_LOG_ERR("Compute shader 0x%llx was not registered; preparing its artifact at dispatch", static_cast<unsigned long long>(snapshot.codeAddress));
+        }
         auto handle = ShaderRecompiler::PrepareShader(request);
         snapshot.prepared->entries.push_back({codeOffset, handle});
         return handle;
@@ -202,26 +205,33 @@ ShaderRecompiler::RectListShaders PreparedRectangle(const ShaderSnapshot& snapsh
     throw std::runtime_error("AGC driver: prepared rectangle artifacts are missing");
 }
 
-ShaderRecompiler::PreparedShaderInvocation InvocationFor(const ShaderSnapshot& snapshot, std::size_t codeOffset, const ShaderRecompiler::RecompileRequest& request) {
+ShaderRecompiler::PreparedShaderInvocation InvocationFor(const ShaderSnapshot& snapshot, std::size_t codeOffset, const ShaderRecompiler::RecompileRequest& request, ShaderMemory* memory) {
     require(codeOffset < snapshot.code.size(), "prepared shader code offset is outside the snapshot");
     const auto code = std::span(snapshot.code).subspan(codeOffset);
     require(request.shader.code.data() == code.data() && request.shader.code.size() == code.size(), "prepared invocation does not refer to registered code");
     auto invocationRequest = request;
+    ShaderRecompiler::CapturedCallProgram captured;
     struct PreparedKeyStorage {};
     auto& key = HostThreadLocal<std::vector<std::uint64_t>, PreparedKeyStorage>();
-    ShaderRecompiler::BuildPreparedShaderKey(request, key);
     std::lock_guard lock(snapshot.prepared->mutex);
+    const bool captureCalls = std::ranges::none_of(snapshot.prepared->entries, [&](const auto& entry) { return entry.codeOffset == codeOffset; }) || std::ranges::any_of(snapshot.prepared->entries, [&](const auto& entry) {
+        return entry.codeOffset == codeOffset && !ShaderRecompiler::GetPreparedCalls(*entry.handle).empty();
+    });
+    if (memory != nullptr && captureCalls && request.shader.stage == ShaderRecompiler::ShaderStage::Compute) {
+        captured = memory->ResolveCalls(request);
+        if (!captured.calls.empty()) {
+            invocationRequest.shader.code = captured.code;
+            invocationRequest.shader.capturedCalls = captured.calls;
+        }
+    }
+    ShaderRecompiler::BuildPreparedShaderKey(invocationRequest, key);
     for (const auto& entry : snapshot.prepared->entries) {
         if (entry.codeOffset != codeOffset) continue;
-        invocationRequest.shader.code = ShaderRecompiler::GetPreparedCode(*entry.handle);
         if (auto invocation = ShaderRecompiler::PreparedShaderInvocation::TryCreate(invocationRequest, entry.handle, key)) return std::move(*invocation);
     }
-    if (snapshot.header.empty()) {
-        if (snapshot.type != 0 || request.shader.stage != ShaderRecompiler::ShaderStage::Compute) throw std::runtime_error("AGC driver: unregistered program is not a compute shader");
-        APS5_LOG_ERR("Compute shader 0x%llx was not registered; preparing its artifact at dispatch", static_cast<unsigned long long>(snapshot.codeAddress));
-        auto handle = ShaderRecompiler::PrepareShader(request);
-        invocationRequest = request;
-        invocationRequest.shader.code = ShaderRecompiler::GetPreparedCode(*handle);
+    if (snapshot.header.empty() || !captured.calls.empty()) {
+        if (snapshot.header.empty() && (snapshot.type != 0 || request.shader.stage != ShaderRecompiler::ShaderStage::Compute)) throw std::runtime_error("AGC driver: unregistered program is not a compute shader");
+        auto handle = ShaderRecompiler::PrepareShader(invocationRequest);
         auto invocation = ShaderRecompiler::PreparedShaderInvocation::TryCreate(invocationRequest, handle, key);
         if (!invocation.has_value()) throw std::runtime_error("AGC driver: raw compute artifact does not match its invocation");
         snapshot.prepared->entries.push_back({codeOffset, std::move(handle)});
@@ -366,6 +376,7 @@ std::vector<PreparedShaders::Entry> PrepareRegistered(const ShaderSnapshot& snap
     std::vector<std::uint32_t> userData(userCount);
     if (stage != Stage::Compute && stage != Stage::Fragment && snapshot.type != 6) vertex = Graphics::DecodeVertexStageInfo(snapshot.header, snapshot.headerAddress, userData, nullptr, true);
     const ShaderRecompiler::SwappcInfo swappc{vertex.has_value(), firstUser, userCount};
+    if (stage == Stage::Compute && !ShaderRecompiler::AnalyzeUserDataCalls(decoded, swappc).empty()) return {};
     auto graph = ShaderRecompiler::GraphBuilder{}.Build(decoded, &swappc);
     ShaderRecompiler::Structurizer{}.Structurize(graph);
     const std::array<ShaderRecompiler::MemoryRegion, 2> memory{{{snapshot.codeAddress, std::as_bytes(std::span(snapshot.code))}, {snapshot.headerAddress, snapshot.header}}};

core/libs/prx/libSceAgcDriver/Execution/src/ShaderMemory.cpp

@@ -4,6 +4,7 @@
 #include "Optimization/RequestMemoryView.hpp"
 #include "Optimization/ResourceMaterializer.hpp"
 #include "Optimization/ResourceProgram.hpp"
+#include "ControlFlow/UserDataCalls.hpp"
 #include <algorithm>
 #include <atomic>
 #include <chrono>
@@ -219,6 +220,21 @@ std::shared_ptr<const ShaderRecompiler::ResourceCapture> ShaderMemory::Capture(c
     return capture(invocation.Request(), nullptr, &invocation);
 }
 
+ShaderRecompiler::CapturedCallProgram ShaderMemory::ResolveCalls(const ShaderRecompiler::RecompileRequest& request) {
+    const GuestMemory::ReadSiteScope site(GuestMemory::ReadSite::Capture);
+    const auto reader = [](void* context, std::uint64_t address, std::uint32_t* value) {
+        auto& memory = *static_cast<ShaderMemory*>(context);
+        const auto next = memory.initial.upper_bound(address);
+        if (next != memory.initial.begin()) {
+            const auto previous = std::prev(next);
+            if (address - previous->first < previous->second.size()) throw std::runtime_error("AGC driver: captured callee code aliases a registered shader snapshot");
+        }
+        memory.callCodeReads.push_back(address);
+        return read(context, address, value);
+    };
+    return ShaderRecompiler::ResolveUserDataCalls(request, reader, this);
+}
+
 std::shared_ptr<const ShaderRecompiler::ResourceCapture> ShaderMemory::capture(const ShaderRecompiler::RecompileRequest& request, const ShaderRecompiler::SourceHandle* handle, const ShaderRecompiler::PreparedShaderInvocation* invocation) {
     // The capture's word and page reads (through `read`) are attributed to it ([hooksync], [guestmem]).
     const GuestMemory::ReadSiteScope site(GuestMemory::ReadSite::Capture);
@@ -234,6 +250,14 @@ std::shared_ptr<const ShaderRecompiler::ResourceCapture> ShaderMemory::capture(c
     runtime.readMemory = &read;
     runtime.readSpecializationMemory = &read;
     auto capture = invocation != nullptr ? invocation->Capture(runtime) : handle != nullptr ? ShaderRecompiler::CaptureResources(request, runtime, *handle) : ShaderRecompiler::CaptureResources(request, runtime);
+    if (!callCodeReads.empty()) {
+        auto traced = std::make_shared<ShaderRecompiler::ResourceCapture>(*capture);
+        auto& reads = traced->readTrace.otherReads;
+        reads.insert(reads.end(), callCodeReads.begin(), callCodeReads.end());
+        std::sort(reads.begin(), reads.end());
+        reads.erase(std::unique(reads.begin(), reads.end()), reads.end());
+        capture = std::move(traced);
+    }
     if (profile) {
         totals.captureNanoseconds += NanosecondsSince(started);
         totals.resolveNanoseconds += capture->sourceNanoseconds;

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

@@ -916,12 +916,12 @@ void verifyPixelRequestSerialization() {
     minimal.context.waveSize = 64;
     minimal.context.pixel = ShaderPixelStageInfo{};
     const auto encoded = serializer.Serialize(minimal);
-    require(requestPrefix(encoded, 8u) == "NVNQQQ0AAAA=", "new requests did not use serialization version 13");
+    require(requestPrefix(encoded, 8u) == "NVNQQQ4AAAA=", "new requests did not use serialization version 14");
     constexpr std::size_t mappingOffset = 8u + 37u + 18u + 163u;
     for (std::size_t bytes = 0; bytes < 8u; ++bytes) {
         expectFailure([&] { static_cast<void>(serializer.Deserialize(requestPrefix(encoded, mappingOffset + bytes))); }, "truncated data", "a truncated version-12 pixel mapping was accepted");
     }
-    for (const auto unsupported : {"NVNQQQAAAAA=", "NVNQQQ4AAAA="}) {
+    for (const auto unsupported : {"NVNQQQAAAAA=", "NVNQQQ8AAAA="}) {
         expectFailure([&] { static_cast<void>(serializer.Deserialize(unsupported)); }, "serialization version", "an unsupported request version was accepted");
     }
 }