fix(shader): align prepared graphics with draw stage decoding

core/libs/prx/libSceAgc/Shader/src/PrimState.cpp

@@ -1,6 +1,7 @@
 #include "prx/libSceAgc/Shader/include/PrimState.hpp"
 
 #include <cstdio>
+#include <array>
 #include <stdexcept>
 #include <prx/libc/include/General.hpp>
 
@@ -62,8 +63,8 @@ int APS5_VABI sceAgcCreatePrimState(ShaderRegister* cx_regs, ShaderRegister* uc_
     }
 
     if (cx_regs != nullptr && uc_regs != nullptr) {
-        AgcDriverResolveShaderAbi_nid_postfix(gs, {cx_regs, 2}, {uc_regs, 3});
-        if (hs != nullptr) AgcDriverResolveShaderAbi_nid_postfix(hs, {cx_regs, 2}, {uc_regs, 3});
+        const std::array<const Shader*, 2> stages{hs, gs};
+        AgcDriverResolveGraphicsStagesAbi_nid_postfix(std::span(stages).subspan(hs == nullptr ? 1u : 0u), {cx_regs, 2}, {uc_regs, 3});
     }
     return 0;
 }

core/libs/prx/libSceAgcDriver/CMakeLists.txt

@@ -606,6 +606,8 @@ foreach(agcExecutionTestSource IN LISTS agcExecutionTestSources)
     list(APPEND agcExecutionTests ${agcExecutionTest})
 endforeach()
 
+target_include_directories(agc_driver_prepared_graphics_tests PRIVATE $<TARGET_PROPERTY:shader_recompiler_agc_driver,INTERFACE_INCLUDE_DIRECTORIES>)
+
 target_include_directories(agc_driver_prepared_shaders_tests PRIVATE
         $<TARGET_PROPERTY:shader_recompiler_agc_driver,INTERFACE_INCLUDE_DIRECTORIES>
 )

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

@@ -18,6 +18,7 @@ extern "C" void AgcDriverShutdown_nid_postfix();
 extern "C" void AgcDriverSuspendPoint_nid_postfix();
 extern "C" void AgcDriverRegisterShader_nid_postfix(const Shader* shader);
 extern "C" void AgcDriverResolveShaderAbi_nid_postfix(const Shader* shader, std::span<const ShaderRegister> context, std::span<const ShaderRegister> primitive);
+extern "C" void AgcDriverResolveGraphicsStagesAbi_nid_postfix(std::span<const Shader* const> stages, std::span<const ShaderRegister> context, std::span<const ShaderRegister> primitive);
 extern "C" void AgcDriverResolveGraphicsAbi_nid_postfix(const Shader* vertex, const Shader* pixel, std::uint32_t primitiveType);
 
 #endif

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

@@ -32,6 +32,15 @@ struct DrawDecode {
     std::vector<ShaderRecompiler::ProgramRole> roles;
 };
 
+void DecodeGraphicsPrograms(DrawDecode& decoded, const QueueState& queue, const ShaderRegistry& registry, bool staticAbi, bool includeFragment);
+
+struct PreparedGraphicsStage {
+    std::shared_ptr<const ShaderSnapshot> snapshot;
+    PreparedShaders::Entry entry;
+};
+
+std::vector<PreparedGraphicsStage> PrepareGraphicsStages(const DrawDecode& decoded, const ShaderRecompiler::SpirvTarget& target);
+
 struct DrawRecipeRecord {
     std::vector<std::weak_ptr<const DispatchVariant>> stages;
     std::shared_ptr<const DrawRecipe> recipe;

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

@@ -57,6 +57,7 @@ public:
     void ReleaseWindow(void* window);
     void RegisterShader(const Shader* shader);
     void ResolveShaderAbi(const Shader* shader, std::span<const ShaderRegister> context, std::span<const ShaderRegister> primitive);
+    void ResolveGraphicsStagesAbi(std::span<const Shader* const> stages, std::span<const ShaderRegister> context, std::span<const ShaderRegister> primitive);
     void ResolveGraphicsAbi(const Shader* vertex, const Shader* pixel, std::uint32_t primitiveType);
 
 private:
@@ -96,7 +97,7 @@ private:
     void lookupDispatch(std::uint64_t address, const Submission& submission, std::uint64_t key, bool noDispatchCache, bool traceCache, bool profile, std::span<const ShaderRecompiler::MemoryRegion> memory, DispatchPhaseTiming& phaseTiming, std::array<double, DriverPhaseCount>& phaseMs, std::shared_ptr<const ShaderRecompiler::RecompileResult>& compiledResult, std::shared_ptr<DispatchVariant>& keepVariant, std::vector<ShaderRecompiler::MemoryRegion>& captured, std::vector<std::uint32_t>& liveWords, bool& dataHit, bool& cached, bool& validated, std::shared_ptr<DispatchEntry>& missedEntry, bool& missedDiffering);
     void insertDispatch(std::uint64_t address, std::uint64_t key, bool noDispatchCache, bool profile, const std::shared_ptr<const ShaderSnapshot>& registeredShader, std::uint64_t forgetAtCapture, std::span<const ShaderRecompiler::MemoryRegion> memory, const std::shared_ptr<ShaderMemory>& shaderMemory, const std::vector<ShaderRecompiler::MemoryRegion>& captured, const std::shared_ptr<const ShaderRecompiler::ResourceCapture>& capture, const std::shared_ptr<const ShaderRecompiler::RecompileResult>& compiledResult, const std::shared_ptr<DispatchEntry>& missedEntry, bool missedDiffering, std::shared_ptr<DispatchVariant>& attachVariant, DispatchPhaseTiming& phaseTiming);
     void dispatch(QueueState& queue, std::span<const std::uint32_t> packet, const Submission& submission, std::uint64_t indirectArguments = 0);
-    void verifyDataHit(const ShaderSnapshot& snapshot, std::size_t codeOffset, std::uint64_t deviceSerial, ShaderRecompiler::RecompileRequest request, std::span<const ShaderRecompiler::MemoryRegion> memory, std::uint64_t address, const DispatchVariant& variant, std::span<const std::uint32_t> liveWords, const ShaderRecompiler::RecompileResult& patched);
+    void verifyDataHit(const ShaderSnapshot& snapshot, std::size_t codeOffset, ShaderRecompiler::RecompileRequest request, std::span<const ShaderRecompiler::MemoryRegion> memory, std::uint64_t address, const DispatchVariant& variant, std::span<const std::uint32_t> liveWords, const ShaderRecompiler::RecompileResult& patched);
     static bool drawEntries();
     static bool verifyDrawEntries();
     static bool registerKeyEnabled();

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

@@ -17,7 +17,6 @@ namespace AgcDriver::DriverDetail {
 struct PreparedShaders {
     struct Entry {
         std::size_t codeOffset;
-        std::uint64_t deviceSerial;
         std::shared_ptr<const ShaderRecompiler::SourceHandle> handle;
     };
     struct Rectangle {
@@ -45,8 +44,8 @@ std::shared_ptr<const ShaderSnapshot> ReadRawComputeShader(std::uint64_t address
 std::uint64_t NullPixelProgramAddress();
 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, std::uint64_t deviceSerial, const ShaderRecompiler::RecompileRequest& request);
-ShaderRecompiler::PreparedShaderInvocation InvocationFor(const ShaderSnapshot& snapshot, std::size_t codeOffset, std::uint64_t deviceSerial, const ShaderRecompiler::RecompileRequest& request);
+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);
 
 }
 

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

@@ -6,14 +6,14 @@
 
 namespace AgcDriver::DriverDetail {
 
-void Driver::verifyDataHit(const ShaderSnapshot& snapshot, std::size_t codeOffset, std::uint64_t deviceSerial, ShaderRecompiler::RecompileRequest request, std::span<const ShaderRecompiler::MemoryRegion> memory, std::uint64_t address, const DispatchVariant& variant, std::span<const std::uint32_t> liveWords, const ShaderRecompiler::RecompileResult& patched) {
+void Driver::verifyDataHit(const ShaderSnapshot& snapshot, std::size_t codeOffset, ShaderRecompiler::RecompileRequest request, std::span<const ShaderRecompiler::MemoryRegion> memory, std::uint64_t address, const DispatchVariant& variant, std::span<const std::uint32_t> liveWords, const ShaderRecompiler::RecompileResult& patched) {
     const auto fail = [&](const char* what, std::size_t position, std::size_t slot) {
         std::fprintf(stderr, "[dispatch-cache] APS5_VERIFY_DATA_HITS: %s disagrees (program 0x%llx, position %zu, slot %zu)\n", what, static_cast<unsigned long long>(address), position, slot);
         std::fflush(stderr);
         std::abort();
     };
     auto verifyMemory = std::make_shared<ShaderMemory>(memory, &queryPendingWrite, &observePendingWrite, hookWaitCounter());
-    const auto invocation = InvocationFor(snapshot, codeOffset, deviceSerial, request);
+    const auto invocation = InvocationFor(snapshot, codeOffset, request);
     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

@@ -95,7 +95,7 @@ void Driver::dispatch(QueueState& queue, std::span<const std::uint32_t> packet,
         explicit ProbeScope(bool active) : active(active) { if (active) ShaderRecompiler::SetDebugProbeActive(true); }
         ~ProbeScope() { if (active) ShaderRecompiler::SetDebugProbeActive(false); }
     } probeScope{probeThis};
-    const auto invocation = InvocationFor(snapshot, codeOffset, localDevice->Serial(), request);
+    const auto invocation = InvocationFor(snapshot, codeOffset, request);
     if (fillBuffer(queue, submission.queue, packet, std::span(snapshot.code).subspan(codeOffset), userData, compute, localDevice)) {
         pendingDispatchPhases().outcome = DispatchOutcome::FillHle;
         return;
@@ -225,7 +225,7 @@ void Driver::dispatch(QueueState& queue, std::span<const std::uint32_t> packet,
         phaseTiming.Phase(PhaseRecompile);
         insertDispatch(address, key, noDispatchCache, profile, registeredShader, forgetAtCapture, memory, shaderMemory, captured, capture, compiledResult, missedEntry, missedDiffering, attachVariant, phaseTiming);
     }
-    if (verifyDataHits() && dataHit) verifyDataHit(snapshot, codeOffset, localDevice->Serial(), request, memory, address, *keepVariant, liveWords, *compiledResult);
+    if (verifyDataHits() && dataHit) verifyDataHit(snapshot, codeOffset, request, memory, address, *keepVariant, liveWords, *compiledResult);
 
     if (recordQueuedLabelsAfterCapture(submission.queue, captured)) {
         dispatch(queue, packet, submission, indirectArguments);

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

@@ -22,7 +22,7 @@ ShaderRecompiler::RecompileResult Driver::materializeDrawStage(std::size_t i, st
         ShaderRecompiler::GraphicsCompileContext{program.firstUserSgpr, linked, graphics.stages.mesh, graphics.stages.tessellation, {drawParameters.indexAddress, drawParameters.indexCount, drawParameters.indexSize, drawParameters.instanceCount}}
     };
     const auto waitedBefore = traceCapSync() || profile ? Graphics::Recorder::ThreadWaitedMs() : 0.0;
-    const auto invocation = InvocationFor(*program.snapshot, program.codeOffset, localDevice->Serial(), request);
+    const auto invocation = InvocationFor(*program.snapshot, program.codeOffset, request);
     auto& stageCapture = stageCaptures[i];
     stageCapture.forgetSerial = GuestMemory::ForgetSerial();
     stageCapture.pushOffset = pushOffset;

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

@@ -5,28 +5,38 @@
 
 namespace AgcDriver::DriverDetail {
 
-std::shared_ptr<DrawDecode> Driver::decodeDraw(const QueueState& queue, const Submission& submission) {
+namespace {
+
+std::uint32_t ReadGraphicsRegister(const Registers& registers, std::uint32_t offset) {
+    const auto found = registers.find(offset);
+    return found == registers.end() ? 0u : found->second;
+}
+
+}
+
+void DecodeGraphicsPrograms(DrawDecode& decoded, const QueueState& queue, const ShaderRegistry& registry, bool staticAbi, bool includeFragment) {
     using Stage = ShaderRecompiler::ShaderStage;
     using Role = ShaderRecompiler::ProgramRole;
     const auto programAddress = [&](std::uint32_t base) {
         Graphics::NoteRegisterRead(Graphics::RegisterBank::Shader, base);
         Graphics::NoteRegisterRead(Graphics::RegisterBank::Shader, base + 1);
-        const auto high = readRegister(queue.shader, base + 1);
+        const auto high = ReadGraphicsRegister(queue.shader, base + 1);
         require((high & ~0xffu) == 0, "reserved graphics program address bits are set");
-        return (static_cast<std::uint64_t>(readRegister(queue.shader, base)) << 8u) | (static_cast<std::uint64_t>(high) << 40u);
+        return (static_cast<std::uint64_t>(ReadGraphicsRegister(queue.shader, base)) << 8u) | (static_cast<std::uint64_t>(high) << 40u);
     };
     const bool pixelSkipped = Graphics::PixelProgramSkipped(queue);
     const auto prepare = [&](std::uint64_t address, std::uint8_t type, Stage stage, std::uint32_t rsrc2, std::uint32_t userDataBase) {
         const bool nullPixel = stage == Stage::Fragment && (address == 0 || pixelSkipped);
         if (nullPixel) address = NullPixelProgramAddress();
-        auto it = submission.shaders->upper_bound(address);
-        require(it != submission.shaders->begin(), "graphics program does not belong to a registered shader");
+        auto it = registry.upper_bound(address);
+        require(it != registry.begin(), "graphics program does not belong to a registered shader");
         --it;
         const auto& snapshot = *it->second;
         require(address - snapshot.codeAddress < snapshot.code.size() * sizeof(std::uint32_t), "graphics program is outside registered shader code");
+        require((address - snapshot.codeAddress) % sizeof(std::uint32_t) == 0, "graphics entry point is not dword aligned");
         require(snapshot.type == type, "graphics program refers to an incompatible shader binary type");
         Graphics::NoteRegisterRead(Graphics::RegisterBank::Shader, rsrc2);
-        const auto resources = nullPixel && !queue.shader.contains(rsrc2) ? 0u : readRegister(queue.shader, rsrc2);
+        const auto resources = nullPixel && !queue.shader.contains(rsrc2) ? 0u : ReadGraphicsRegister(queue.shader, rsrc2);
         const auto userCount = ((resources >> 1u) & 0x1fu) | (((resources >> 27u) & 1u) << 5u);
         require(userCount <= 32, "graphics user SGPR count exceeds the register bank");
         const auto codeOffset = static_cast<std::size_t>((address - snapshot.codeAddress) / sizeof(std::uint32_t));
@@ -41,13 +51,12 @@ std::shared_ptr<DrawDecode> Driver::decodeDraw(const QueueState& queue, const Su
         };
         for (std::uint32_t i = 0; i < userCount; ++i) {
             Graphics::NoteRegisterRead(Graphics::RegisterBank::Shader, userDataBase + i);
-            result.userData.push_back(readUserData(queue.shader, userDataBase + i));
+            result.userData.push_back((staticInterface ? ReadGraphicsRegister(queue.shader, userDataBase + i) : readUserData(queue.shader, userDataBase + i)));
         }
         return result;
     };
     {
-        auto product = std::make_shared<DrawDecode>();
-        product->state = Graphics::DecodeState(queue);
+        auto* product = &decoded;
         auto& programs = product->programs;
         auto& roles = product->roles;
         programs.reserve(5);
@@ -61,9 +70,10 @@ std::shared_ptr<DrawDecode> Driver::decodeDraw(const QueueState& queue, const Su
             program.userData.insert(program.userData.begin(), 8, 0);
             Graphics::NoteRegisterRead(Graphics::RegisterBank::Shader, pointerBase);
             Graphics::NoteRegisterRead(Graphics::RegisterBank::Shader, pointerBase + 1);
+            if (staticAbi) return;
             if (!pointerRequired && !queue.shader.contains(pointerBase) && !queue.shader.contains(pointerBase + 1)) return;
-            const auto low = readRegister(queue.shader, pointerBase);
-            const auto high = readRegister(queue.shader, pointerBase + 1);
+            const auto low = ReadGraphicsRegister(queue.shader, pointerBase);
+            const auto high = ReadGraphicsRegister(queue.shader, pointerBase + 1);
             const auto address = static_cast<std::uint64_t>(low) | (static_cast<std::uint64_t>(high) << 32u);
             require(address != 0 || !pointerRequired, "merged shader user-data address is null");
             if (address == 0) return;
@@ -79,8 +89,8 @@ std::shared_ptr<DrawDecode> Driver::decodeDraw(const QueueState& queue, const Su
             append(0x0c8, 2, Stage::TessellationEvaluation, 0x08b, 0x08c, Role::Domain);
         } else if (graphics.stages.path == Graphics::ShaderPath::Geometry) {
             const auto frontAddress = programAddress(0xc8);
-            auto snapshot = submission.shaders->upper_bound(frontAddress);
-            require(snapshot != submission.shaders->begin(), "geometry front program is not registered");
+            auto snapshot = registry.upper_bound(frontAddress);
+            require(snapshot != registry.begin(), "geometry front program is not registered");
             --snapshot;
             const auto type = snapshot->second->type;
             require(type == 2 || type == 4, "invalid geometry front binary type");
@@ -90,18 +100,25 @@ std::shared_ptr<DrawDecode> Driver::decodeDraw(const QueueState& queue, const Su
         } else {
             append(0xc8, 2, Stage::Vertex, 0x8b, 0x8c, Role::Main);
         }
-        const bool nullPixel = pixelSkipped;
-        if (nullPixel) {
-            const auto rejection = Graphics::NullPixelProgramRejection(queue);
-            require(rejection.empty(), rejection.c_str());
+        if (includeFragment) {
+            if (pixelSkipped) {
+                const auto rejection = Graphics::NullPixelProgramRejection(queue);
+                require(rejection.empty(), rejection.c_str());
+            }
+            append(0x008, 1, Stage::Fragment, 0x00b, 0x00c, Role::Fragment);
+            programs.back().firstUserSgpr = 0;
         }
-        append(0x008, 1, Stage::Fragment, 0x00b, 0x00c, Role::Fragment);
-        programs.back().firstUserSgpr = 0;
-        product->pixel = Graphics::DecodePixelStageInfo(queue.context, Graphics::ExportMappings(graphics), nullPixel);
-        return product;
     }
 }
 
+std::shared_ptr<DrawDecode> Driver::decodeDraw(const QueueState& queue, const Submission& submission) {
+    auto product = std::make_shared<DrawDecode>();
+    product->state = Graphics::DecodeState(queue);
+    DecodeGraphicsPrograms(*product, queue, *submission.shaders, false, true);
+    product->pixel = Graphics::DecodePixelStageInfo(queue.context, Graphics::ExportMappings(product->state), Graphics::PixelProgramUnset(queue));
+    return product;
+}
+
 void Driver::resolveDrawDecode(const QueueState& queue, const Submission& submission, std::shared_ptr<const DrawDecode>& decode, bool registerKey, std::uint64_t drawKey, bool profile) {
     if (decode == nullptr || verifyDrawRecipe()) {
         std::vector<Graphics::RegisterRead> readLog;

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

@@ -116,4 +116,7 @@ extern "C" void AgcDriverResolveShaderAbi_nid_postfix(const Shader* shader, std:
 
 extern "C" void AgcDriverResolveGraphicsAbi_nid_postfix(const Shader* vertex, const Shader* pixel, std::uint32_t primitiveType) {
     AgcDriver::DriverDetail::Driver::Get().ResolveGraphicsAbi(vertex, pixel, primitiveType);
-}
+}
+extern "C" void AgcDriverResolveGraphicsStagesAbi_nid_postfix(std::span<const Shader* const> stages, std::span<const ShaderRegister> context, std::span<const ShaderRegister> primitive) {
+    AgcDriver::DriverDetail::Driver::Get().ResolveGraphicsStagesAbi(stages, context, primitive);
+}

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

@@ -6,6 +6,7 @@
 #include "prx/libSceAgcDriver/Execution/include/Driver/Diagnostics.hpp"
 #include "prx/libSceAgcDriver/Execution/include/GuestMemory.hpp"
 #include "CacheKey.hpp"
+#include "CompiledVariant.hpp"
 #include "Optimization/ResourceProgram.hpp"
 #include <cstdlib>
 #include <cstring>
@@ -83,16 +84,16 @@ std::shared_ptr<const ShaderSnapshot> ReadRawComputeShader(std::uint64_t address
 
 
 
-std::shared_ptr<const ShaderRecompiler::SourceHandle> SourceHandleFor(const ShaderSnapshot& snapshot, std::size_t codeOffset, std::uint64_t deviceSerial, const ShaderRecompiler::RecompileRequest& request) {
+std::shared_ptr<const ShaderRecompiler::SourceHandle> SourceHandleFor(const ShaderSnapshot& snapshot, std::size_t codeOffset, const ShaderRecompiler::RecompileRequest& request) {
     std::lock_guard lock(snapshot.prepared->mutex);
     for (const auto& entry : snapshot.prepared->entries) {
-        if (entry.codeOffset == codeOffset && entry.deviceSerial == deviceSerial && ShaderRecompiler::MatchesPreparedShader(request, *entry.handle)) return entry.handle;
+        if (entry.codeOffset == codeOffset && ShaderRecompiler::MatchesPreparedShader(request, *entry.handle)) 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));
         auto handle = ShaderRecompiler::PrepareShader(request);
-        snapshot.prepared->entries.push_back({codeOffset, deviceSerial, handle});
+        snapshot.prepared->entries.push_back({codeOffset, handle});
         return handle;
     }
     throw std::runtime_error("AGC driver: prepared shader artifact is missing for the requested static ABI");
@@ -106,14 +107,14 @@ 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, std::uint64_t deviceSerial, const ShaderRecompiler::RecompileRequest& request) {
+ShaderRecompiler::PreparedShaderInvocation InvocationFor(const ShaderSnapshot& snapshot, std::size_t codeOffset, const ShaderRecompiler::RecompileRequest& request) {
     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;
     std::lock_guard lock(snapshot.prepared->mutex);
     for (const auto& entry : snapshot.prepared->entries) {
-        if (entry.codeOffset != codeOffset || entry.deviceSerial != deviceSerial) continue;
+        if (entry.codeOffset != codeOffset) continue;
         invocationRequest.shader.code = ShaderRecompiler::GetPreparedCode(*entry.handle);
         if (auto invocation = ShaderRecompiler::PreparedShaderInvocation::TryCreate(invocationRequest, entry.handle)) return std::move(*invocation);
     }
@@ -230,7 +231,7 @@ std::vector<PreparedShaders::Entry> PrepareRegistered(const ShaderSnapshot& snap
     ShaderRecompiler::RecompileRequest request{{stage, address, code, snapshot.headerAddress, snapshot.header}, {wave, firstUser, userData, compute, pixel, vertex, memory}, stage == Stage::Compute ? device.ComputeTarget(wave) : device.Target(), {0, 0, 0, 128}, graphics};
     if (graphics && graphics->mesh) request.layout.pushConstantSizeBytes = ShaderRecompiler::MeshDrawPushOffsetBytes;
     std::vector<PreparedShaders::Entry> entries;
-    const auto append = [&] { entries.push_back({codeOffset, device.Serial(), ShaderRecompiler::PrepareShader(request)}); };
+    const auto append = [&] { entries.push_back({codeOffset, ShaderRecompiler::PrepareShader(request)}); };
     append();
     if (compute) {
         request.context.compute->partialThreads = {1, 1, 1};
@@ -244,9 +245,91 @@ std::vector<PreparedShaders::Entry> PrepareRegistered(const ShaderSnapshot& snap
 
 }
 
+std::vector<PreparedGraphicsStage> PrepareGraphicsStages(const DrawDecode& decoded, const ShaderRecompiler::SpirvTarget& target) {
+    require(decoded.programs.size() == decoded.roles.size(), "graphics ABI program roles are incomplete");
+    for (const auto& program : decoded.programs) require(program.snapshot != nullptr, "graphics ABI has no registered shader snapshot");
+    std::vector<ShaderRecompiler::ProgramRole> expected;
+    if (decoded.state.stages.path == Graphics::ShaderPath::Tessellation) expected = {ShaderRecompiler::ProgramRole::Local, ShaderRecompiler::ProgramRole::Hull, ShaderRecompiler::ProgramRole::Domain};
+    else {
+        expected.push_back(ShaderRecompiler::ProgramRole::Main);
+        if (decoded.state.stages.path == Graphics::ShaderPath::Geometry && !decoded.programs.empty() && decoded.programs.front().snapshot->type == 4) expected.push_back(ShaderRecompiler::ProgramRole::GeometryBack);
+    }
+    if (!decoded.roles.empty() && decoded.roles.back() == ShaderRecompiler::ProgramRole::Fragment) expected.push_back(ShaderRecompiler::ProgramRole::Fragment);
+    require(expected == decoded.roles, "graphics ABI linked programs do not match the stage routing");
+    std::vector<ShaderRecompiler::LinkedProgram> linked;
+    std::vector<ShaderRecompiler::MemoryRegion> memory;
+    for (std::size_t index = 0; index < decoded.programs.size(); ++index) {
+        const auto& program = decoded.programs[index];
+        require(program.snapshot != nullptr && program.codeOffset < program.snapshot->code.size(), "graphics ABI has an invalid registered entry point");
+        const auto code = std::span(program.snapshot->code).subspan(program.codeOffset);
+        require(program.binary.code.data() == code.data() && program.binary.code.size() == code.size() && program.binary.codeAddress == program.snapshot->codeAddress + program.codeOffset * sizeof(std::uint32_t), "graphics ABI entry point differs from its snapshot");
+        linked.push_back({decoded.roles[index], program.binary, program.userDataBase, program.firstUserSgpr, program.userData});
+        memory.insert(memory.end(), program.memory.begin(), program.memory.end());
+    }
+    std::vector<PreparedGraphicsStage> prepared;
+    std::uint32_t pushOffset = 0;
+    const auto capacity = decoded.state.stages.mesh ? ShaderRecompiler::MeshDrawPushOffsetBytes : Graphics::PipelinePushConstantBytes;
+    for (std::size_t index = 0; index < decoded.programs.size(); ++index) {
+        if (decoded.roles[index] == ShaderRecompiler::ProgramRole::GeometryBack) continue;
+        const auto& program = decoded.programs[index];
+        const bool fragment = program.binary.stage == ShaderRecompiler::ShaderStage::Fragment;
+        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);
+        ShaderRecompiler::RecompileRequest request{program.binary, {wave, program.firstUserSgpr, program.userData, {}, fragment ? std::optional(decoded.pixel) : std::nullopt, vertex, memory}, target, {0, 0, pushOffset, capacity - pushOffset}, ShaderRecompiler::GraphicsCompileContext{program.firstUserSgpr, linked, decoded.state.stages.mesh, decoded.state.stages.tessellation, {}}};
+        const auto handle = ShaderRecompiler::PrepareShader(request);
+        const auto bytes = handle->artifact->bindings.pushConstantSizeBytes;
+        require(bytes <= capacity - pushOffset, "prepared graphics stages exceed the push constant block");
+        prepared.push_back({program.snapshot, {program.codeOffset, handle}});
+        pushOffset += bytes;
+    }
+    return prepared;
+}
+
+void Driver::ResolveGraphicsStagesAbi(std::span<const Shader* const> stages, std::span<const ShaderRegister> context, std::span<const ShaderRegister> primitive) {
+    CheckFailure();
+    require(!stages.empty(), "graphics ABI has no shader headers");
+    QueueState state{};
+    std::shared_ptr<const ShaderRegistry> registry;
+    {
+        std::lock_guard lock(mutex);
+        registry = shaders;
+        for (const auto* shader : stages) {
+            GuestMemory::CheckRange(shader, sizeof(Shader), alignof(Shader));
+            const auto address = reinterpret_cast<std::uintptr_t>(const_cast<const void*>(shader->code));
+            require(registry != nullptr && registry->contains(address), "graphics ABI refers to an unregistered shader");
+            const auto& snapshot = *registry->at(address);
+            require(snapshot.headerAddress == reinterpret_cast<std::uintptr_t>(shader), "graphics ABI refers to a replaced shader header");
+            const auto registered = RegisteredState(snapshot);
+            for (const auto& [offset, value] : registered.shader) state.shader.insert_or_assign(offset, value);
+            for (const auto& [offset, value] : registered.context) state.context.insert_or_assign(offset, value);
+            for (const auto& [offset, value] : registered.userConfig) state.userConfig.insert_or_assign(offset, value);
+        }
+    }
+    for (const auto reg : context) state.context.insert_or_assign(reg.offset, reg.value);
+    for (const auto reg : primitive) state.userConfig.insert_or_assign(reg.offset, reg.value);
+    DrawDecode decoded{};
+    decoded.state.stages = Graphics::DecodeShaderStages(state);
+    DecodeGraphicsPrograms(decoded, state, *registry, true, false);
+    const auto localDevice = device.Load();
+    require(localDevice != nullptr, "graphics ABI device is missing");
+    auto prepared = PrepareGraphicsStages(decoded, localDevice->Target());
+    for (auto& stage : prepared) {
+        std::lock_guard lock(stage.snapshot->prepared->mutex);
+        auto& entries = stage.snapshot->prepared->entries;
+        const auto duplicate = std::ranges::any_of(entries, [&](const auto& existing) { return existing.codeOffset == stage.entry.codeOffset && existing.handle->artifact == stage.entry.handle->artifact; });
+        if (!duplicate) entries.push_back(std::move(stage.entry));
+    }
+}
+
 void Driver::ResolveShaderAbi(const Shader* shader, std::span<const ShaderRegister> context, std::span<const ShaderRegister> primitive) {
     CheckFailure();
     GuestMemory::CheckRange(shader, sizeof(Shader), alignof(Shader));
+    if (shader->type != 0 && shader->type != 1) {
+        const std::array<const Shader*, 1> stages{shader};
+        ResolveGraphicsStagesAbi(stages, context, primitive);
+        return;
+    }
     std::shared_ptr<const ShaderSnapshot> snapshot;
     {
         std::lock_guard lock(mutex);

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

@@ -0,0 +1,197 @@
+#include "VulkanTestDevice.hpp"
+#include "SceShaders.hpp"
+#include "prx/libSceAgcDriver/Execution/include/Driver/Draw/DrawCache.hpp"
+#include "prx/libSceAgcDriver/Graphics/include/ShaderInputState.hpp"
+#include "Optimization/ResourceProgram.hpp"
+#include <spirv/unified1/spirv.hpp>
+#include <array>
+#include <vector>
+#include <cstring>
+#include <filesystem>
+#include <fstream>
+#include <iostream>
+#include <stdexcept>
+#include <string>
+
+namespace {
+
+void Require(bool condition, const char* message) {
+    if (!condition) throw std::runtime_error(message);
+}
+
+template<typename TAction>
+void Reject(TAction action, const char* expected) {
+    try { action(); }
+    catch (const std::exception& error) {
+        Require(std::string(error.what()).find(expected) != std::string::npos, error.what());
+        return;
+    }
+    throw std::runtime_error("invalid graphics ABI was accepted");
+}
+
+struct Fixture {
+    struct Header {
+        Shader shader{};
+        ShaderUserData users{};
+    } header;
+    alignas(256) std::array<std::uint32_t, 65> code{};
+    std::shared_ptr<AgcDriver::DriverDetail::ShaderSnapshot> snapshot;
+
+    void Initialize(std::uint8_t type) {
+        header.shader.type = type;
+        header.shader.user_data = &header.users;
+        code.fill(0xffffffffu);
+        code.back() = 0xbf810000u;
+        const auto address = reinterpret_cast<std::uintptr_t>(code.data());
+        const auto headerAddress = reinterpret_cast<std::uintptr_t>(&header);
+        snapshot = std::make_shared<AgcDriver::DriverDetail::ShaderSnapshot>();
+        snapshot->type = type;
+        snapshot->codeAddress = address;
+        snapshot->headerAddress = headerAddress;
+        snapshot->code.assign(code.begin(), code.end());
+        snapshot->header.resize(sizeof(header));
+        std::memcpy(snapshot->header.data(), &header, sizeof(header));
+    }
+
+    void Bind(AgcDriver::QueueState& queue, std::uint32_t program, std::uint32_t resources) const {
+        const auto address = snapshot->codeAddress + 256u;
+        queue.shader[program] = static_cast<std::uint32_t>(address >> 8u);
+        queue.shader[program + 1u] = static_cast<std::uint32_t>(address >> 40u);
+        queue.shader[resources] = 16u << 1u;
+    }
+};
+
+void Check(AgcDriver::VulkanDevice& device, AgcDriver::Graphics::ShaderPath path, const std::filesystem::path& dump) {
+    using namespace AgcDriver::DriverDetail;
+    using namespace ShaderRecompiler;
+    Fixture front;
+    Fixture back;
+    Fixture domain;
+    Fixture fragment;
+    const bool tessellation = path == AgcDriver::Graphics::ShaderPath::Tessellation;
+    const bool mesh = path == AgcDriver::Graphics::ShaderPath::Geometry;
+    front.Initialize(tessellation ? 5u : mesh ? 4u : 2u);
+    back.Initialize(tessellation ? 7u : 6u);
+    domain.Initialize(2u);
+    fragment.Initialize(1u);
+    const std::array<std::uint32_t, 7> pixelCode{0xc8020002u, 0xc8060102u, 0xc80a0202u, 0xc80e0302u, 0xf800180fu, 0x03020100u, 0xbf810000u};
+    fragment.snapshot->code.resize(64);
+    fragment.snapshot->code.insert(fragment.snapshot->code.end(), pixelCode.begin(), pixelCode.end());
+    ShaderRegistry registry;
+    for (const auto* fixture : {&front, &back, &domain, &fragment}) registry.emplace(fixture->snapshot->codeAddress, fixture->snapshot);
+    AgcDriver::QueueState queue{};
+    front.Bind(queue, tessellation ? 0x148u : 0xc8u, tessellation ? 0x10bu : 0x8bu);
+    if (tessellation || mesh) back.Bind(queue, tessellation ? 0x108u : 0x88u, tessellation ? 0x10bu : 0x8bu);
+    if (tessellation) domain.Bind(queue, 0xc8u, 0x8bu);
+    fragment.Bind(queue, 0x8u, 0xbu);
+    alignas(8) const std::array<std::uint32_t, 2> merged{};
+    const auto mergedAddress = reinterpret_cast<std::uintptr_t>(merged.data());
+    const auto pointerBase = tessellation ? 0x102u : 0x82u;
+    queue.shader[pointerBase] = static_cast<std::uint32_t>(mergedAddress);
+    queue.shader[pointerBase + 1u] = static_cast<std::uint32_t>(mergedAddress >> 32u);
+    DrawDecode prepared{};
+    prepared.state.stages.path = path;
+    prepared.state.stages.vertexWaveSize = tessellation || mesh ? 64u : 32u;
+    prepared.state.stages.fragmentWaveSize = 32u;
+    if (mesh) prepared.state.stages.mesh = MeshConfiguration{4, 1, 3, 3, 1, 64, 128, 0, 4};
+    if (tessellation) prepared.state.stages.tessellation = TessellationConfiguration{3, 4, 1, 2, 2};
+    prepared.pixel.wave32 = true;
+    prepared.pixel.interpolatorCount = 2;
+    prepared.pixel.interpolatorSettings[0] = 0x403u;
+    prepared.pixel.interpolatorSettings[1] = 0x220u;
+    prepared.pixel.targetOutputMode[0] = 9;
+    prepared.pixel.targetExportMapping.fill(0xe4u);
+    DecodeGraphicsPrograms(prepared, queue, registry, true, true);
+    DrawDecode draw{};
+    draw.state = prepared.state;
+    draw.pixel = prepared.pixel;
+    DecodeGraphicsPrograms(draw, queue, registry, false, true);
+    Require(prepared.programs.size() == draw.programs.size() && prepared.roles == draw.roles, "prepared graphics programs differ from draw programs");
+    for (std::size_t index = 0; index < prepared.programs.size(); ++index) {
+        const auto& expected = prepared.programs[index];
+        const auto& actual = draw.programs[index];
+        Require(expected.codeOffset == 64u && expected.binary.code.size() == (prepared.roles[index] == ProgramRole::Fragment ? pixelCode.size() : 1u) && expected.binary.code[0] == (prepared.roles[index] == ProgramRole::Fragment ? pixelCode[0] : 0xbf810000u), "graphics entry point did not trim the code prefix");
+        Require(expected.binary.stage == actual.binary.stage && expected.firstUserSgpr == actual.firstUserSgpr && expected.userData.size() == actual.userData.size(), "graphics preparation changed the user SGPR ABI");
+    }
+    auto target = device.Target();
+    std::vector<std::uint32_t> capabilities(target.supportedCapabilities.begin(), target.supportedCapabilities.end());
+    std::vector<std::string_view> extensions(target.supportedExtensions.begin(), target.supportedExtensions.end());
+    if (mesh) {
+        capabilities.push_back(spv::CapabilityMeshShadingEXT);
+        extensions.push_back("SPV_EXT_mesh_shader");
+        target.supportedCapabilities = capabilities;
+        target.supportedExtensions = extensions;
+        target.mesh = MeshTargetLimits{{128, 1, 1}, 128, 32768, 256, 256, 128, 32768, 1, 1};
+    }
+    const auto stages = PrepareGraphicsStages(prepared, target);
+    Require(stages.size() == (tessellation ? 4u : 2u), "graphics preparation compiled the wrong stages");
+    for (const auto& stage : stages) stage.snapshot->prepared->entries.push_back(stage.entry);
+    const auto& pixelArtifact = GetPreparedArtifact(*stages.back().entry.handle);
+    Require(!pixelArtifact.fragmentParameters.empty() && pixelArtifact.fragmentParameters.front().sourceLocation == 3u, "prepared fragment lost interpolant mapping");
+    if (!dump.empty()) {
+        for (std::size_t index = 0; index < stages.size(); ++index) {
+            const auto& words = GetPreparedArtifact(*stages[index].entry.handle).spirv;
+            const auto name = dump / (std::to_string(static_cast<unsigned>(path)) + "-" + std::to_string(index) + ".spv");
+            std::ofstream output(name, std::ios::binary);
+            output.write(reinterpret_cast<const char*>(words.data()), static_cast<std::streamsize>(words.size() * sizeof(std::uint32_t)));
+            Require(static_cast<bool>(output), "cannot write prepared graphics SPIR-V");
+        }
+    }
+    auto incomplete = prepared;
+    incomplete.roles.pop_back();
+    Reject([&] { static_cast<void>(PrepareGraphicsStages(incomplete, target)); }, "program roles are incomplete");
+    auto wrongEntry = prepared;
+    wrongEntry.programs.front().binary.codeAddress += 4;
+    Reject([&] { static_cast<void>(PrepareGraphicsStages(wrongEntry, target)); }, "entry point differs");
+    std::vector<LinkedProgram> linked;
+    std::vector<MemoryRegion> memory;
+    for (std::size_t index = 0; index < draw.programs.size(); ++index) {
+        const auto& program = draw.programs[index];
+        linked.push_back({draw.roles[index], program.binary, program.userDataBase, program.firstUserSgpr, program.userData});
+        memory.insert(memory.end(), program.memory.begin(), program.memory.end());
+    }
+    for (std::size_t index = 0; index < draw.programs.size(); ++index) {
+        if (draw.roles[index] == ProgramRole::GeometryBack) continue;
+        const auto& program = draw.programs[index];
+        const bool pixel = program.binary.stage == ShaderStage::Fragment;
+        std::optional<ShaderVertexStageInfo> vertex;
+        if (!pixel) vertex = AgcDriver::Graphics::DecodeVertexStageInfo(program.binary.header, program.binary.headerAddress, program.userData, nullptr, true);
+        RecompileRequest request{program.binary, {pixel ? 32u : prepared.state.stages.vertexWaveSize, program.firstUserSgpr, program.userData, {}, pixel ? std::optional(draw.pixel) : std::nullopt, vertex, memory}, target, {0, 0, 0, mesh ? MeshDrawPushOffsetBytes : 128u}, GraphicsCompileContext{program.firstUserSgpr, linked, prepared.state.stages.mesh, prepared.state.stages.tessellation, {0, 3, 4, 1}}};
+        static_cast<void>(InvocationFor(*program.snapshot, program.codeOffset, request));
+        request.layout.pushConstantOffsetBytes = 4;
+        request.layout.pushConstantSizeBytes -= 4;
+        Reject([&] { static_cast<void>(InvocationFor(*program.snapshot, program.codeOffset, request)); }, "artifact is missing");
+        request.layout = {0, 0, 0, mesh ? MeshDrawPushOffsetBytes : 128u};
+        if (pixel) request.context.pixel->interpolatorSettings[0] ^= 1u;
+        else if (tessellation) ++request.graphics->tessellation->outputControlPoints;
+        else if (mesh) ++request.graphics->mesh->maxVertices;
+        else request.context.userDataBaseRegister = 0;
+        Reject([&] { static_cast<void>(InvocationFor(*program.snapshot, program.codeOffset, request)); }, "artifact is missing");
+    }
+    ShaderRegistry invalidRegistry;
+    DrawDecode invalid{};
+    invalid.state = prepared.state;
+    Reject([&] { DecodeGraphicsPrograms(invalid, queue, invalidRegistry, true, true); }, "registered");
+    const auto frontBase = tessellation ? 0x148u : 0xc8u;
+    queue.shader[frontBase + 1u] |= 0x100u;
+    Reject([&] { DecodeGraphicsPrograms(invalid, queue, registry, true, true); }, "reserved graphics program address");
+}
+
+}
+
+int main(int argc, char** argv) {
+    try {
+        auto device = OpenVulkanTestDevice();
+        if (!device) return VulkanTestSkipped;
+        Require(argc <= 2, "invalid test arguments");
+        const auto dump = argc == 2 ? std::filesystem::path(argv[1]) : std::filesystem::path{};
+        Check(*device, AgcDriver::Graphics::ShaderPath::Vertex, dump);
+        Check(*device, AgcDriver::Graphics::ShaderPath::Geometry, dump);
+        Check(*device, AgcDriver::Graphics::ShaderPath::Tessellation, dump);
+        std::cout << "prepared graphics ABI tests passed\n";
+        return 0;
+    } catch (const std::exception& error) {
+        std::cerr << error.what() << '\n';
+        return 1;
+    }
+}

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

@@ -34,16 +34,15 @@ void Run(AgcDriver::VulkanDevice& device) {
     const ShaderRecompiler::ShaderComputeStageInfo compute{{1, 1, 1}, 0, {false, false, false}, false, 1, {}};
     ShaderRecompiler::RecompileRequest request{{ShaderRecompiler::ShaderStage::Compute, reinterpret_cast<std::uintptr_t>(code.data()), code, 0, {}}, {32, 0, users, compute, {}, {}, {}}, device.ComputeTarget(32), {0, 0, 0, 128}};
     AgcDriver::DriverDetail::ShaderSnapshot snapshot{request.shader.codeAddress, 0, 0, {code.begin(), code.end()}, {}};
-    ExpectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::SourceHandleFor(snapshot, 0, device.Serial(), request)); }, "artifact is missing");
+    ExpectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::SourceHandleFor(snapshot, 0, request)); }, "artifact is missing");
     const auto handle = ShaderRecompiler::PrepareShader(request);
-    snapshot.prepared->entries.push_back({0, device.Serial(), handle});
-    ExpectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::SourceHandleFor(snapshot, 1, device.Serial(), request)); }, "artifact is missing");
-    ExpectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::SourceHandleFor(snapshot, 0, device.Serial() + 1, request)); }, "artifact is missing");
+    snapshot.prepared->entries.push_back({0, handle});
+    ExpectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::SourceHandleFor(snapshot, 1, request)); }, "artifact is missing");
     const auto& artifact = ShaderRecompiler::GetPreparedArtifact(*handle);
     for (std::uint32_t value = 0; value < 8; ++value) {
         users.fill(value);
         request.useCache = value % 2 == 0;
-        const auto ready = AgcDriver::DriverDetail::SourceHandleFor(snapshot, 0, device.Serial(), request);
+        const auto ready = AgcDriver::DriverDetail::SourceHandleFor(snapshot, 0, request);
         ShaderRecompiler::SrtRuntime runtime{};
         runtime.userData = users;
         const auto capture = ShaderRecompiler::CaptureResources(request, runtime, *ready);
@@ -55,12 +54,11 @@ void Run(AgcDriver::VulkanDevice& device) {
     auto registeredRequest = request;
     registeredRequest.shader.code = snapshot.code;
     registeredRequest.useCache = true;
-    const auto invocation = AgcDriver::DriverDetail::InvocationFor(snapshot, 0, device.Serial(), registeredRequest);
-    ExpectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::InvocationFor(snapshot, 0, device.Serial(), request)); }, "does not refer to registered code");
-    ExpectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::InvocationFor(snapshot, 1, device.Serial(), registeredRequest)); }, "outside the snapshot");
-    ExpectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::InvocationFor(snapshot, 0, device.Serial() + 1, registeredRequest)); }, "artifact is missing");
+    const auto invocation = AgcDriver::DriverDetail::InvocationFor(snapshot, 0, registeredRequest);
+    ExpectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::InvocationFor(snapshot, 0, request)); }, "does not refer to registered code");
+    ExpectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::InvocationFor(snapshot, 1, registeredRequest)); }, "outside the snapshot");
     registeredRequest.context.compute->numThreads[0] = 2;
-    ExpectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::InvocationFor(snapshot, 0, device.Serial(), registeredRequest)); }, "artifact is missing");
+    ExpectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::InvocationFor(snapshot, 0, registeredRequest)); }, "artifact is missing");
     ShaderRecompiler::SrtRuntime preparedRuntime{};
     preparedRuntime.userData = users;
     const auto preparedCapture = invocation.Capture(preparedRuntime);
@@ -73,11 +71,25 @@ void Run(AgcDriver::VulkanDevice& device) {
     Require(changedResult != firstResult && changedResult->spirv.data() == firstResult->spirv.data(), "prepared invocation did not distinguish changed runtime data");
     users[0] -= 1;
     Require(invocation.Materialize(*preparedCapture) == firstResult, "prepared invocation evicted the previous runtime data");
+    auto replacement = OpenVulkanTestDevice();
+    Require(replacement != nullptr && replacement->Serial() != device.Serial(), "replacement device was not created");
+    auto replacementRequest = request;
+    replacementRequest.target = replacement->ComputeTarget(32);
+    Require(AgcDriver::DriverDetail::SourceHandleFor(snapshot, 0, replacementRequest) == handle, "device replacement discarded a compatible prepared artifact");
+    replacementRequest.shader.code = snapshot.code;
+    const auto replacementInvocation = AgcDriver::DriverDetail::InvocationFor(snapshot, 0, replacementRequest);
+    const auto replacementCapture = replacementInvocation.Capture(preparedRuntime);
+    const auto replacementResult = replacementInvocation.Materialize(*replacementCapture);
+    Require(replacementResult->spirv.data() == artifact.spirv.data(), "device replacement recompiled the prepared shader");
+    replacement->Dispatch(*replacementResult, 1, 1, 1);
+    replacement->WaitIdle();
+    replacementRequest.target.nonConstantImageOffsets = !replacementRequest.target.nonConstantImageOffsets;
+    ExpectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::InvocationFor(snapshot, 0, replacementRequest)); }, "artifact is missing");
     ShaderRecompiler::SrtRuntime runtime{};
     runtime.userData = users;
     const auto capture = ShaderRecompiler::CaptureResources(request, runtime, *handle);
     request.context.compute->numThreads[0] = 2;
-    ExpectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::SourceHandleFor(snapshot, 0, device.Serial(), request)); }, "artifact is missing");
+    ExpectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::SourceHandleFor(snapshot, 0, request)); }, "artifact is missing");
     ExpectFailure([&] { static_cast<void>(ShaderRecompiler::CaptureResources(request, runtime, *handle)); }, "does not match the static ABI");
     ExpectFailure([&] { static_cast<void>(ShaderRecompiler::MaterializeShader(request, *capture, *handle)); }, "does not match the static ABI");
     const auto otherHandle = ShaderRecompiler::PrepareShader(request);
@@ -86,10 +98,10 @@ void Run(AgcDriver::VulkanDevice& device) {
     ExpectFailure([&] { static_cast<void>(ShaderRecompiler::MaterializeShader(request, *capture, *otherHandle)); }, "another prepared shader");
     request.context.compute->numThreads[0] = 1;
     request.layout.pushConstantSizeBytes = 124;
-    ExpectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::SourceHandleFor(snapshot, 0, device.Serial(), request)); }, "artifact is missing");
+    ExpectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::SourceHandleFor(snapshot, 0, request)); }, "artifact is missing");
     request.layout.pushConstantSizeBytes = 128;
     code[0] = 0xffffffffu;
-    ExpectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::SourceHandleFor(snapshot, 0, device.Serial(), request)); }, "artifact is missing");
+    ExpectFailure([&] { static_cast<void>(AgcDriver::DriverDetail::SourceHandleFor(snapshot, 0, request)); }, "artifact is missing");
 }
 
 void PrepareMultisampledStorage(AgcDriver::VulkanDevice& device) {
@@ -97,6 +109,8 @@ void PrepareMultisampledStorage(AgcDriver::VulkanDevice& device) {
     std::array<std::uint32_t, 13> code{0xd7460000u, 0x0401060cu, 0xd7460001u, 0x0405060du, 0x7e04020eu, 0x7e060280u, 0x7e080208u, 0x7e0a0209u, 0x7e0c020au, 0x7e0e020bu, 0xf0200f38u, 0x00000400u, 0xbf810000u};
     RecompileRequest request{};
     request.shader = {ShaderStage::Compute, 0x30000u, code, 0, {}};
+    const std::array<std::uint32_t, 16> users{};
+    request.context.userData = users;
     request.context.waveSize = 32;
     request.context.compute = ShaderComputeStageInfo{{8, 8, 1}, 0, {true, true, false}, false, 2, {}};
     request.target = device.ComputeTarget(32);

core/shader/recompiler/CacheKey.hpp

@@ -52,6 +52,7 @@ public:
         append(key, request.context.pixel);
         append(key, request.context.vertex);
         appendMesh(key, request);
+        appendTessellation(key, request);
         std::uint64_t hash = 0xcbf29ce484222325ull;
         for (const auto value : key) {
             hash ^= value;
@@ -82,6 +83,7 @@ private:
         append(key, request.context.pixel);
         append(key, request.context.vertex);
         appendMesh(key, request);
+        appendTessellation(key, request);
         append(key, request.target);
         append(key, DebugProbeActive());
         append(key, RayTracingStrict());
@@ -96,6 +98,14 @@ private:
         for (const auto value : {mesh->inputPrimitive, mesh->primitivesPerGroup, mesh->verticesPerGroup, mesh->maxVertices, mesh->maxPrimitives, mesh->threadsPerGroup, mesh->ldsSizeDwords, mesh->provokingVertex, mesh->esgsItemSize}) append(key, value);
     }
 
+    static void appendTessellation(std::vector<std::uint64_t>& key, const RecompileRequest& request) {
+        if (request.shader.stage != ShaderStage::Local && request.shader.stage != ShaderStage::TessellationControl && request.shader.stage != ShaderStage::TessellationEvaluation) return;
+        const auto* tessellation = request.graphics && request.graphics->tessellation ? &*request.graphics->tessellation : nullptr;
+        append(key, tessellation != nullptr);
+        if (tessellation == nullptr) return;
+        for (const auto value : {tessellation->inputControlPoints, tessellation->outputControlPoints, tessellation->domain, tessellation->partitioning, tessellation->outputTopology}) append(key, value);
+    }
+
     template<typename TValue>
     static void append(std::vector<std::uint64_t>& key, TValue value) requires (std::is_integral_v<TValue> || std::is_enum_v<TValue>) {
         key.push_back(static_cast<std::uint64_t>(value));

core/shader/recompiler/Recompiler.cpp

@@ -86,7 +86,8 @@ std::uint32_t HostSubgroupSize(const RecompileRequest& request) {
 
 ShaderStageInputInfo RequestInputInfo(const RecompileRequest& request) {
     const auto* mesh = request.graphics && request.graphics->mesh ? &*request.graphics->mesh : nullptr;
-    return BuildShaderStageInputInfo(toShaderStageKind(request.shader.stage), request.context, HostSubgroupSize(request), mesh);
+    const auto* tessellation = request.graphics && request.graphics->tessellation ? &*request.graphics->tessellation : nullptr;
+    return BuildShaderStageInputInfo(toShaderStageKind(request.shader.stage), request.context, HostSubgroupSize(request), mesh, tessellation);
 }
 
 }

core/shader/recompiler/Translation/include/Translation/ShaderInputInfoBuilder.hpp

@@ -9,7 +9,7 @@ namespace ShaderRecompiler {
 
 // hostSubgroupSize is the device subgroup width; wave64 workgroups on 32-wide subgroups run each
 // invocation as two guest lanes so a guest wave stays one host subgroup.
-ShaderStageInputInfo BuildShaderStageInputInfo(ShaderStageKind stage, const GuestContext& context, std::uint32_t hostSubgroupSize, const MeshConfiguration* mesh = nullptr);
+ShaderStageInputInfo BuildShaderStageInputInfo(ShaderStageKind stage, const GuestContext& context, std::uint32_t hostSubgroupSize, const MeshConfiguration* mesh = nullptr, const TessellationConfiguration* tessellation = nullptr);
 
 }
 

core/shader/recompiler/Translation/src/ShaderInputInfoBuilder.cpp

@@ -63,7 +63,7 @@ void _detectVertexBuffers(ShaderVertexInputInfo& info) {
 
 }
 
-ShaderStageInputInfo BuildShaderStageInputInfo(ShaderStageKind stage, const GuestContext& context, std::uint32_t hostSubgroupSize, const MeshConfiguration* mesh) {
+ShaderStageInputInfo BuildShaderStageInputInfo(ShaderStageKind stage, const GuestContext& context, std::uint32_t hostSubgroupSize, const MeshConfiguration* mesh, const TessellationConfiguration* tessellation) {
     switch (stage) {
     case ShaderStageKind::Compute: {
         if (!context.compute.has_value()) {
@@ -171,6 +171,17 @@ ShaderStageInputInfo BuildShaderStageInputInfo(ShaderStageKind stage, const Gues
             vertexStorage.resourcesDst[i].fetchIndex = vertex.resourcesDst[i].fetchIndex;
         }
         _detectVertexBuffers(vertexStorage);
+        if (stage == ShaderStageKind::Local || stage == ShaderStageKind::TessellationControl || stage == ShaderStageKind::TessellationEvaluation) {
+            if (tessellation == nullptr) throw std::runtime_error("ShaderInputInfoBuilder: tessellation configuration is missing");
+            if (tessellation->inputControlPoints == 0u || tessellation->inputControlPoints > 32u || tessellation->outputControlPoints == 0u || tessellation->outputControlPoints > 32u) throw std::runtime_error("ShaderInputInfoBuilder: invalid tessellation control-point counts");
+            if (tessellation->domain != 1u || tessellation->partitioning != 2u || tessellation->outputTopology != 2u) throw std::runtime_error("ShaderInputInfoBuilder: unsupported tessellation configuration");
+            auto& target = vertexStorage.tess;
+            target.inputControlPoints = tessellation->inputControlPoints;
+            target.outputControlPoints = tessellation->outputControlPoints;
+            target.domain = tessellation->domain;
+            target.partitioning = tessellation->partitioning;
+            target.outputTopology = tessellation->outputTopology;
+        }
         if (stage == ShaderStageKind::Mesh) {
             if (mesh == nullptr) throw std::runtime_error("ShaderInputInfoBuilder: a mesh-stage program has no mesh configuration");
             auto& target = vertexStorage.mesh;