feat(shader): link raw vertex prologs through user-data tails

core/libs/prx/libSceAgcDriver/Execution/include/Driver/Draw/DrawCache.hpp

@@ -23,6 +23,7 @@ struct DrawProgram {
 
     std::shared_ptr<const ShaderSnapshot> snapshot;
     std::size_t codeOffset = 0;
+    std::shared_ptr<const ShaderSnapshot> tailContinuation;
 };
 
 struct DrawDecode {

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

@@ -87,17 +87,19 @@ 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::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();
 bool RawGraphicsDrawn();
 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.
-std::uint32_t PositionExportLayout(const ShaderSnapshot* snapshot, std::size_t codeOffset, std::span<const std::uint32_t> code, std::uint32_t paClVsOutCntl);
+std::uint32_t PositionExportLayout(const ShaderSnapshot* snapshot, std::size_t codeOffset, std::span<const std::uint32_t> code, std::uint32_t paClVsOutCntl, const ShaderSnapshot* continuation = nullptr);
 bool VertexWorkgroupRequired(const QueueState& queue, const ShaderRegistry& registry);
 
 std::shared_ptr<const ShaderRecompiler::SourceHandle> PrepareShaderWithDiagnostics(const ShaderRecompiler::RecompileRequest& request, ShaderRecompiler::ShaderPreparationContext* preparation = nullptr);

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

@@ -145,7 +145,7 @@ DrawVerdict Driver::draw(QueueState& queue, std::span<const std::uint32_t> packe
         if (program.binary.stage == Stage::Fragment || roles[i] == Role::GeometryBack) return;
         decodeReads[i].clear();
         vertexInfos[i] = Graphics::DecodeVertexStageInfo(program.binary.header, program.binary.headerAddress, program.userData, &decodeReads[i], false, 8u - program.firstUserSgpr);
-        if (vertexInfos[i]) vertexInfos[i]->paClVsOutCntl = PositionExportLayout(program.snapshot.get(), program.codeOffset, program.binary.code, readRegister(queue.context, 0x207)) | graphics.layerExports;
+        if (vertexInfos[i]) vertexInfos[i]->paClVsOutCntl = PositionExportLayout(program.snapshot.get(), program.codeOffset, program.binary.code, readRegister(queue.context, 0x207), program.tailContinuation.get()) | graphics.layerExports;
     };
     if (!registerKey) {
         for (std::size_t i = 0; i < programs.size(); ++i) decodeVertexInfo(i);

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

@@ -120,7 +120,9 @@ void Driver::eraseDrawEntry(std::uint64_t key, const std::shared_ptr<DrawEntry>&
 
 bool DecodeProgramsCurrent(const DrawDecode& decode, const ShaderRegistry&) {
     return std::all_of(decode.programs.begin(), decode.programs.end(), [&](const DrawProgram& program) {
-        return !program.snapshot->header.empty() || ReadRawShader(program.snapshot->codeAddress + program.codeOffset * sizeof(std::uint32_t)) == program.snapshot;
+        if (!program.snapshot->header.empty()) return true;
+        return ReadRawShader(program.snapshot->codeAddress + program.codeOffset * sizeof(std::uint32_t)) == program.snapshot &&
+            (program.tailContinuation == nullptr || ReadRawShader(program.tailContinuation->codeAddress) == program.tailContinuation);
     });
 }
 

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

@@ -51,6 +51,7 @@ void DecodeGraphicsPrograms(DrawDecode& decoded, const QueueState& queue, const
             Graphics::NoteRegisterRead(Graphics::RegisterBank::Shader, userDataBase + i);
             result.userData.push_back((staticAbi ? ReadGraphicsRegister(queue.shader, userDataBase + i) : readUserData(queue.shader, userDataBase + i)));
         }
+        if (stage == Stage::Vertex && snapshot.header.empty() && snapshot.userDataTailRegister) result.tailContinuation = ReadUserDataTailContinuation(result.binary.code, result.firstUserSgpr, result.userData);
         return result;
     };
     {

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

@@ -137,7 +137,7 @@ bool Driver::sameDecode(const DrawDecode& a, const DrawDecode& b) {
     for (std::size_t i = 0; i < a.programs.size(); ++i) {
         const auto& x = a.programs[i];
         const auto& y = b.programs[i];
-        if (x.binary.stage != y.binary.stage || x.binary.codeAddress != y.binary.codeAddress || x.userDataBase != y.userDataBase || x.firstUserSgpr != y.firstUserSgpr || x.userData != y.userData || x.snapshot != y.snapshot || x.codeOffset != y.codeOffset) return false;
+        if (x.binary.stage != y.binary.stage || x.binary.codeAddress != y.binary.codeAddress || x.userDataBase != y.userDataBase || x.firstUserSgpr != y.firstUserSgpr || x.userData != y.userData || x.snapshot != y.snapshot || x.codeOffset != y.codeOffset || x.tailContinuation != y.tailContinuation) return false;
     }
     if (p.targetExportPacking != q.targetExportPacking || p.dualSourceBlend != q.dualSourceBlend || p.alphaToCoverage != q.alphaToCoverage || p.rasterPrimitive != q.rasterPrimitive) return false;
     return true;

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

@@ -91,6 +91,7 @@ std::shared_ptr<const ShaderSnapshot> ReadRawShader(std::uint64_t address) {
             std::as_writable_bytes(std::span(snapshot.code).subspan(previous)), alignof(std::uint32_t));
         try {
             const auto decoded = ShaderRecompiler::RdnaInstructionDecoder{}.Decode(snapshot.code);
+            if (const auto tail = ShaderRecompiler::FindUserDataTail(decoded, 0u, 106u)) snapshot.userDataTailRegister = tail->userDataRegister;
             const auto& last = decoded.instructions.back();
             snapshot.code.resize(last.programCounter / sizeof(std::uint32_t) + last.wordCount);
             auto result = std::make_shared<const ShaderSnapshot>(std::move(snapshot));
@@ -120,6 +121,16 @@ std::shared_ptr<const ShaderSnapshot> ReadRawShader(std::uint64_t address) {
     throw std::runtime_error("AGC driver: raw shader program has no reachable end within mapped code or the size limit");
 }
 
+std::shared_ptr<const ShaderSnapshot> ReadUserDataTailContinuation(std::span<const std::uint32_t> code, std::uint32_t userDataBase, std::span<const std::uint32_t> userData) {
+    const auto program = ShaderRecompiler::RdnaInstructionDecoder{}.Decode(code);
+    const auto tail = ShaderRecompiler::FindUserDataTail(program, userDataBase, static_cast<std::uint32_t>(userData.size()));
+    if (!tail) return nullptr;
+    const auto slot = tail->userDataRegister - userDataBase;
+    const auto address = userData[slot] | (static_cast<std::uint64_t>(userData[slot + 1u]) << 32u);
+    require(address != 0 && (address & 3u) == 0u, "user-data tail continuation address is null or not dword aligned");
+    return ReadRawShader(address);
+}
+
 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");
@@ -168,10 +179,11 @@ bool ProgramUsesWorkgroup(const ShaderSnapshot& snapshot, std::size_t codeOffset
     return used;
 }
 
-std::uint32_t PositionExportLayout(const ShaderSnapshot* snapshot, std::size_t codeOffset, std::span<const std::uint32_t> code, std::uint32_t paClVsOutCntl) {
+std::uint32_t PositionExportLayout(const ShaderSnapshot* snapshot, std::size_t codeOffset, std::span<const std::uint32_t> code, std::uint32_t paClVsOutCntl, const ShaderSnapshot* continuation) {
     constexpr std::uint8_t Unknown = 0, Unused = 1, Used = 2;
     const auto bits = paClVsOutCntl & Graphics::PositionExportLayoutBits;
     if (bits == 0) return 0;
+    if (continuation != nullptr && PositionExportLayout(continuation, 0, continuation->code, bits) != 0) return bits;
     const bool cached = snapshot != nullptr && codeOffset == 0;
     if (cached) {
         const auto use = snapshot->positionVectorUse->load(std::memory_order_acquire);
@@ -199,7 +211,22 @@ bool VertexWorkgroupRequired(const QueueState& queue, const ShaderRegistry& regi
     const auto address = (static_cast<std::uint64_t>(low->second) << 8u) | (static_cast<std::uint64_t>(high->second) << 40u);
     const auto snapshot = ProgramSnapshot(registry, address, {2});
     require((address - snapshot->codeAddress) % sizeof(std::uint32_t) == 0, "graphics entry point is not dword aligned");
-    return (snapshot->header.empty() || snapshot->type == 2) && ProgramUsesWorkgroup(*snapshot, (address - snapshot->codeAddress) / sizeof(std::uint32_t));
+    const auto offset = (address - snapshot->codeAddress) / sizeof(std::uint32_t);
+    if (!snapshot->header.empty()) return snapshot->type == 2 && ProgramUsesWorkgroup(*snapshot, offset);
+    if (!snapshot->userDataTailRegister) return ProgramUsesWorkgroup(*snapshot, offset);
+    Graphics::NoteRegisterRead(Graphics::RegisterBank::Shader, 0x8b);
+    const auto resources = queue.shader.find(0x8b);
+    const auto rsrc2 = resources == queue.shader.end() ? 0u : resources->second;
+    const auto count = ((rsrc2 >> 1u) & 0x1fu) | (((rsrc2 >> 27u) & 1u) << 5u);
+    require(count <= 32, "graphics user SGPR count exceeds the register bank");
+    std::vector<std::uint32_t> userData;
+    for (std::uint32_t index = 0; index < count; ++index) {
+        Graphics::NoteRegisterRead(Graphics::RegisterBank::Shader, 0x8c + index);
+        const auto found = queue.shader.find(0x8c + index);
+        userData.push_back(found == queue.shader.end() ? 0u : found->second);
+    }
+    const auto continuation = ReadUserDataTailContinuation(std::span(snapshot->code).subspan(offset), 8u, userData);
+    return ProgramUsesWorkgroup(*snapshot, offset) || (continuation != nullptr && ProgramUsesWorkgroup(*continuation, 0u));
 }
 
 struct ShaderPreparationTransaction::State {
@@ -458,6 +485,15 @@ ShaderRecompiler::RectListShaders DrawRectangle(const ShaderSnapshot& front, con
 
 namespace {
 
+std::vector<std::uint32_t> LinkVertexUserDataTail(const ShaderRecompiler::RecompileRequest& request, const ShaderSnapshot& snapshot) {
+    if (request.shader.stage != ShaderRecompiler::ShaderStage::Vertex || !request.shader.header.empty() || !snapshot.userDataTailRegister) return {};
+    const auto program = ShaderRecompiler::RdnaInstructionDecoder{}.Decode(request.shader.code);
+    const auto tail = ShaderRecompiler::FindUserDataTail(program, request.context.userDataBaseRegister, static_cast<std::uint32_t>(request.context.userData.size()));
+    if (!tail) return {};
+    const auto continuation = ReadUserDataTailContinuation(request.shader.code, request.context.userDataBaseRegister, request.context.userData);
+    return ShaderRecompiler::LinkUserDataTail(program, *tail, ShaderRecompiler::RdnaInstructionDecoder{}.Decode(continuation->code));
+}
+
 std::string hexText(std::uint64_t value) {
     char text[24];
     std::snprintf(text, sizeof(text), "0x%llx", static_cast<unsigned long long>(value));
@@ -512,7 +548,8 @@ ShaderRecompiler::PreparedShaderInvocation InvocationFor(const ShaderSnapshot& s
     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");
-    const auto linked = LinkUserDataCallees(request);
+    auto linked = LinkUserDataCallees(request);
+    if (linked.empty()) linked = LinkVertexUserDataTail(request, snapshot);
     auto source = request;
     if (!linked.empty()) source.shader.code = linked;
     auto invocationRequest = source;
@@ -970,7 +1007,7 @@ std::vector<PreparedGraphicsStage> PrepareGraphicsStages(const DrawDecode& decod
         const auto wave = fragment ? decoded.state.stages.fragmentWaveSize : decoded.state.stages.vertexWaveSize;
         std::optional<ShaderRecompiler::ShaderVertexStageInfo> vertex;
         if (!fragment) vertex = Graphics::DecodeVertexStageInfo(program.binary.header, program.binary.headerAddress, program.userData, nullptr, true, 8u - program.firstUserSgpr);
-        if (vertex) vertex->paClVsOutCntl = PositionExportLayout(program.snapshot.get(), program.codeOffset, program.binary.code, decoded.state.paClVsOutCntl) | decoded.state.layerExports;
+        if (vertex) vertex->paClVsOutCntl = PositionExportLayout(program.snapshot.get(), program.codeOffset, program.binary.code, decoded.state.paClVsOutCntl, program.tailContinuation.get()) | decoded.state.layerExports;
         pushOffset = Graphics::StagePushOffset(pushOffset, program.binary.stage, fixedSlots);
         const auto offset = pushOffset;
         const auto room = (fixedSlots && fragment ? Graphics::PipelinePushSlotBytes : capacity) - pushOffset % Graphics::PipelinePushSlotBytes;

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

@@ -480,6 +480,38 @@ void verifyNullSwappc() {
 
 }
 
+void VerifyUserDataTail() {
+    const RdnaInstructionDecoder decoder;
+    const std::vector<std::uint32_t> code{0x7e000280u, 0xbe802008u, 0xbf9f0000u};
+    const auto decoded = decoder.Decode(code);
+    const auto tail = FindUserDataTail(decoded, 8u, 2u);
+    require(tail && tail->programCounter == 4u && tail->userDataRegister == 8u, "unchanged user-data tail was not recognized");
+    require(!FindUserDataTail(decoded, 9u, 2u) && !FindUserDataTail(decoded, 8u, 1u), "tail outside user data was accepted");
+    const std::vector<std::vector<std::uint32_t>> refused{
+        {0xbe880380u, 0xbe802008u, 0xbf9f0000u},
+        {0xbe890380u, 0xbe802008u, 0xbf9f0000u},
+        {0xbe880480u, 0xbe802008u, 0xbf9f0000u},
+        {0xf4080200u, 0xfa000000u, 0xbe802008u, 0xbf9f0000u},
+        {0xbf820000u, 0xbe802008u, 0xbf9f0000u},
+        {0xbe802009u, 0xbf9f0000u},
+        {0xbe802008u, 0xbf810000u},
+        {0xbe802008u, 0xbf800000u, 0xbf9f0000u},
+    };
+    for (const auto& rejected : refused) require(!FindUserDataTail(decoder.Decode(rejected), 8u, 4u), "a modified, branching or malformed tail was accepted");
+    const std::vector<std::uint32_t> continuation{0xbf820001u, 0x7e000281u, 0xbf810000u};
+    const auto linked = LinkUserDataTail(decoded, *tail, decoder.Decode(continuation));
+    require(linked == std::vector<std::uint32_t>({0x7e000280u, 0xbf820000u, 0xbf820001u, 0x7e000281u, 0xbf810000u}), "linking did not preserve relative branches and register operations");
+    static_cast<void>(GraphBuilder{}.Build(decoder.Decode(linked)));
+    const std::vector<std::uint32_t> readsPc{0xbe8c1f00u, 0xbf810000u};
+    bool refusedPc = false;
+    try {
+        static_cast<void>(LinkUserDataTail(decoded, *tail, decoder.Decode(readsPc)));
+    } catch (const std::invalid_argument& error) {
+        refusedPc = std::string_view(error.what()).find("reads its program counter") != std::string_view::npos;
+    }
+    require(refusedPc, "a continuation reading its relocated PC was accepted");
+}
+
 int main(int argc, char** argv) {
     try {
         verifyLoopMergeDoesNotEnterNestedSelection();
@@ -1279,6 +1311,7 @@ int main(int argc, char** argv) {
     }
     try {
         verifyNullSwappc();
+        VerifyUserDataTail();
         verifyLargeLoopExitFallback();
         verifyNestedExitForwarders();
     } catch (const std::exception& error) {

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

@@ -658,6 +658,39 @@ void testProgramWorkgroupUse() {
     check(!AgcDriver::DriverDetail::ProgramUsesWorkgroup(entry, 2) && AgcDriver::DriverDetail::ProgramUsesWorkgroup(entry, 0), "the workgroup check did not start at the program's entry");
 }
 
+void testVertexTailSnapshots() {
+    using namespace AgcDriver::DriverDetail;
+    alignas(256) static std::array<std::uint32_t, 64> prolog{0xbe802008u, 0xbf9f0000u};
+    alignas(256) static std::array<std::uint32_t, 64> continuation{0xbf800000u, 0xbf810000u};
+    const auto address = reinterpret_cast<std::uintptr_t>(prolog.data());
+    const auto next = reinterpret_cast<std::uintptr_t>(continuation.data());
+    const std::array<std::uint32_t, 2> data{static_cast<std::uint32_t>(next), static_cast<std::uint32_t>(next >> 32u)};
+    const auto source = ReadRawShader(address);
+    check(source->userDataTailRegister == 8u, "raw tail metadata is missing");
+    const auto target = ReadUserDataTailContinuation(source->code, 8u, data);
+    check(target && target->codeAddress == next, "tail did not read its user-data continuation");
+    DrawDecode decode{};
+    decode.programs.resize(1);
+    decode.programs[0].snapshot = source;
+    decode.programs[0].tailContinuation = target;
+    check(DecodeProgramsCurrent(decode, {}), "unchanged tail snapshot became stale");
+    continuation[0] = 0xbf8a0000u;
+    check(!DecodeProgramsCurrent(decode, {}), "rewriting the continuation did not invalidate the draw cache");
+    AgcDriver::QueueState queue{};
+    queue.context[0x2d5] = 0x2000u;
+    queue.shader = {{0xc8, static_cast<std::uint32_t>(address >> 8u)}, {0xc9, static_cast<std::uint32_t>(address >> 40u)},
+        {0x8b, 4u}, {0x8c, data[0]}, {0x8d, data[1]}};
+    check(VertexWorkgroupRequired(queue, {}), "a barrier in the continuation was hidden from workgroup routing");
+    queue.context[0x2d5] |= 0x02000000u;
+    check(!VertexWorkgroupRequired(queue, {}), "a passthrough program was rerouted to a workgroup");
+    const std::array<std::uint32_t, 2> nullAddress{};
+    check(expectFailure([&] { static_cast<void>(ReadUserDataTailContinuation(source->code, 8u, nullAddress)); }).find("address is null") != std::string::npos, "null tail address was accepted");
+    const std::array<std::uint32_t, 2> unaligned{data[0] + 1u, data[1]};
+    check(expectFailure([&] { static_cast<void>(ReadUserDataTailContinuation(source->code, 8u, unaligned)); }).find("dword aligned") != std::string::npos, "unaligned tail address was accepted");
+    const std::array<std::uint32_t, 2> unmapped{0x1000u, 0u};
+    check(!expectFailure([&] { static_cast<void>(ReadUserDataTailContinuation(source->code, 8u, unmapped)); }).empty(), "unmapped tail address was accepted");
+}
+
 void testRegisteredVertexWorkgroupRouting() {
     using namespace AgcDriver::DriverDetail;
     ShaderSnapshot snapshot{0x20000, 0x30000, 2, {0xd8340000u, 0x00000100u, 0xbf810000u, 0xbf810000u}, {}};
@@ -781,6 +814,7 @@ int main(int argc, char** argv) {
 #endif
         testProgramSnapshots();
         testProgramWorkgroupUse();
+        testVertexTailSnapshots();
         testRegisteredVertexWorkgroupRouting();
         testEvents();
         testValidation();