perf(agc): allow shader preparation on first use

core/libs/prx/libSceAgcDriver/CMakeLists.txt

@@ -1730,14 +1730,16 @@ if(BUILD_TESTING)
             endif()
         endforeach()
     endforeach()
-    foreach(preparedMode IN ITEMS indirect uncached debug debug_failure deferred deferred_undecodable registered_workgroup registered_workgroup_sync)
+    foreach(preparedMode IN ITEMS indirect uncached debug debug_failure deferred deferred_undecodable registered_workgroup registered_workgroup_sync lazy lazy_sync)
         set(preparedTest "agc_driver_prepared_shaders_${preparedMode}")
         if(preparedMode STREQUAL "indirect")
             add_test(NAME ${preparedTest} COMMAND agc_driver_prepared_shaders_tests --indirect)
         elseif(preparedMode STREQUAL "deferred")
             add_test(NAME ${preparedTest} COMMAND agc_driver_prepared_shaders_tests --deferred)
         elseif(preparedMode STREQUAL "deferred_undecodable")
             add_test(NAME ${preparedTest} COMMAND agc_driver_prepared_shaders_tests --deferred-undecodable)
+        elseif(preparedMode STREQUAL "lazy" OR preparedMode STREQUAL "lazy_sync")
+            add_test(NAME ${preparedTest} COMMAND agc_driver_prepared_shaders_tests --lazy)
         elseif(preparedMode STREQUAL "registered_workgroup" OR preparedMode STREQUAL "registered_workgroup_sync")
             add_test(NAME ${preparedTest} COMMAND agc_driver_prepared_shaders_tests --registered-workgroup)
         elseif(preparedMode STREQUAL "debug_failure")
@@ -1751,6 +1753,11 @@ if(BUILD_TESTING)
             set_property(TEST ${preparedTest} APPEND PROPERTY ENVIRONMENT "APS5_NO_DISPATCH_CACHE=1")
         elseif(preparedMode STREQUAL "registered_workgroup_sync")
             set_property(TEST ${preparedTest} APPEND PROPERTY ENVIRONMENT "APS5_SYNC_SHADER_PREPARE=1")
+        elseif(preparedMode STREQUAL "lazy" OR preparedMode STREQUAL "lazy_sync")
+            set_property(TEST ${preparedTest} APPEND PROPERTY ENVIRONMENT "APS5_LAZY_SHADER_PREPARE=1" "APS5_SUSPEND_DRAIN=1")
+            if(preparedMode STREQUAL "lazy_sync")
+                set_property(TEST ${preparedTest} APPEND PROPERTY ENVIRONMENT "APS5_SYNC_SHADER_PREPARE=1")
+            endif()
         elseif(preparedMode STREQUAL "debug" OR preparedMode STREQUAL "debug_failure")
             set_property(TEST ${preparedTest} APPEND PROPERTY ENVIRONMENT "APS5_PROFILE_DRAW=1" "APS5_NO_UNLOCKED_DEVICE=1")
         endif()

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

@@ -75,6 +75,7 @@ class Driver {
     };
     DispatchCacheFootprint DispatchCacheSize();
     static std::uint64_t DispatchCacheBytes();
+    std::shared_ptr<const ShaderRegistry> RegisteredShaders();
 
 private:
     friend class SampledReadScope;

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

@@ -1462,6 +1462,11 @@ ShaderSnapshot PrepareNullPixelProgram(const VulkanDevice& device) {
     return null;
 }
 
+std::shared_ptr<const ShaderRegistry> Driver::RegisteredShaders() {
+    std::lock_guard lock(mutex);
+    return shaders;
+}
+
 void Driver::RegisterShader(const Shader* shader) {
     PerformanceContext timingContext(FrameTiming::Preparation());
     PerformanceTimer timing("Shader.Register");
@@ -1514,7 +1519,11 @@ void Driver::RegisterShader(const Shader* shader) {
     registered.context = snapshot.registeredState->context;
     registered.userConfig = snapshot.registeredState->userConfig;
     std::shared_ptr<RegisteredPreparation> preparation;
-    if (AsyncRegistrationPrepare()) {
+    static const bool lazy = std::getenv("APS5_LAZY_SHADER_PREPARE") != nullptr;
+    if (lazy) {
+        static_cast<void>(PlanRegistered(snapshot, *localDevice, registered, true, &snapshot.prepared->deferred));
+        snapshot.prepared->deferred = true;
+    } else if (AsyncRegistrationPrepare()) {
         preparation = PlanRegistered(snapshot, *localDevice, registered, true, &snapshot.prepared->deferred);
         if (preparation != nullptr) preparation->prebuildDevice = localDevice;
         if (preparation != nullptr && (snapshot.type == 0 || snapshot.type == 1)) BuildRegisteredAbiKey(registered, *localDevice, preparation->abiKey);

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

@@ -4,7 +4,10 @@
 #include "prx/libSceAgcDriver/Execution/include/Driver/Shaders/ShaderRegistry.hpp"
 #include "prx/libSceAgcDriver/Graphics/include/ShaderInputState.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Driver.hpp"
+#include "prx/libSceAgcDriver/Execution/include/Driver/Driver.hpp"
 #include "prx/libSceAgcDriver/Execution/include/ShaderPreparation.hpp"
+#include "../NarrowConstantStoreFixture.hpp"
+#include "prx/libc/include/GuestAllocations.hpp"
 #include "prx/libSceAgcDriver/Submit/include/Acb.hpp"
 #include "prx/libSceAgcDriver/Submit/include/Dcb.hpp"
 #include <spirv/unified1/spirv.hpp>
@@ -27,6 +30,12 @@
 #include <fstream>
 #include <sstream>
 #include <thread>
+#ifdef _WIN32
+#ifndef NOMINMAX
+#define NOMINMAX
+#endif
+#include <windows.h>
+#endif
 
 namespace {
 
@@ -600,7 +609,7 @@ void RegisteredVertexWorkgroup(const ShaderRecompiler::SpirvTarget& target, bool
         header.specials.vgt_shader_stages_en.value = 0;
         ExpectFailure([&] { AgcDriverRegisterShader_nid_postfix(&header.shader); }, "missing AGC user-data header");
         header.shader.user_data = &header.userData;
-        if (target.subgroupSize < 64u && std::getenv("APS5_SYNC_SHADER_PREPARE") != nullptr) {
+        if (target.subgroupSize < 64u && std::getenv("APS5_SYNC_SHADER_PREPARE") != nullptr && std::getenv("APS5_LAZY_SHADER_PREPARE") == nullptr) {
             ExpectFailure([&] { AgcDriverRegisterShader_nid_postfix(&header.shader); }, "v_readlane_b32 of lane 63 is outside");
         }
         std::cout << "registered NGG preparation deferral passed; subgroup=" << target.subgroupSize << '\n';
@@ -612,6 +621,102 @@ void RegisteredVertexWorkgroup(const ShaderRecompiler::SpirvTarget& target, bool
     std::cout << "registered NGG wave64 preparation passed; subgroup=" << target.subgroupSize << '\n';
 }
 
+void LazyRegistration() {
+    using namespace NarrowConstantStoreFixture;
+    Require(std::getenv("APS5_LAZY_SHADER_PREPARE") != nullptr, "lazy preparation is not enabled");
+    struct Header {
+        Shader shader{};
+        std::array<ShaderRegister, 7> registers{};
+        ShaderSpecialRegs specials{};
+    } header;
+    const auto address = reinterpret_cast<std::uintptr_t>(Code.data());
+    header.shader.file_header = 0x34333231u;
+    header.shader.version = 0x18;
+    header.shader.header_size = sizeof(header);
+    header.shader.shader_size = sizeof(Code);
+    header.shader.code = Code.data();
+    header.shader.sh_registers = header.registers.data();
+    header.shader.num_sh_registers = header.registers.size();
+    header.shader.specials = &header.specials;
+    header.specials.dispatch_modifier = 0x8000;
+    header.registers = {{{0x20c, static_cast<std::uint32_t>(address >> 8u)}, {0x20d, static_cast<std::uint32_t>(address >> 40u)}, {0x207, Threads}, {0x208, 1}, {0x209, 1}, {0x212, 0}, {0x213, 4}}};
+    AgcDriverRegisterShader_nid_postfix(&header.shader);
+    ++header.registers[0].value;
+    ExpectFailure([&] { AgcDriverRegisterShader_nid_postfix(&header.shader); }, "entry point is outside shader code");
+    --header.registers[0].value;
+    header.registers[2].value = 0;
+    ExpectFailure([&] { AgcDriverRegisterShader_nid_postfix(&header.shader); }, "must be nonzero");
+    header.registers[2].value = Threads;
+    header.shader.num_sh_registers = 255;
+    ExpectFailure([&] { AgcDriverRegisterShader_nid_postfix(&header.shader); }, "truncated shader metadata");
+    header.shader.num_sh_registers = header.registers.size();
+    const auto registered = [&] {
+        const auto registry = AgcDriver::DriverDetail::Driver::Get().RegisteredShaders();
+        Require(registry != nullptr, "lazy registration published no shader registry");
+        auto snapshot = AgcDriver::DriverDetail::RegisteredHeader(*registry, &header.shader);
+        Require(snapshot != nullptr, "lazy registration published no snapshot");
+        return snapshot;
+    };
+    const auto lazySnapshot = registered();
+    {
+        AgcDriver::DriverDetail::ShaderPreparationTransaction transaction;
+        const auto& state = transaction.Read(*lazySnapshot);
+        Require(state.deferred && state.entries.empty(), "lazy registration prepared the shader before its first dispatch");
+    }
+    constexpr std::size_t outputBytes = 65536;
+#ifdef _WIN32
+    auto* allocation = static_cast<std::uint8_t*>(VirtualAlloc(nullptr, outputBytes, MEM_RESERVE | MEM_COMMIT, PAGE_READWRITE));
+#else
+    auto* allocation = static_cast<std::uint8_t*>(std::aligned_alloc(outputBytes, outputBytes));
+#endif
+    Require(allocation != nullptr, "cannot allocate lazy shader test output");
+    const auto release = [](std::uint8_t* block) {
+        GuestAllocations::Mutation().Remove(block);
+#ifdef _WIN32
+        VirtualFree(block, 0, MEM_RELEASE);
+#else
+        std::free(block);
+#endif
+    };
+    const std::unique_ptr<std::uint8_t, decltype(release)> ownedOutput(allocation, release);
+    const std::span output(allocation, outputBytes);
+    GuestAllocations::Mutation().Add(output.data(), output.size(), true, true, true);
+    const auto outputAddress = reinterpret_cast<std::uintptr_t>(output.data());
+    std::vector<std::uint32_t> commands;
+    for (const auto reg : header.registers) commands.insert(commands.end(), {0xc0017600u, reg.offset, reg.value});
+    commands.insert(commands.end(), {0xc0027600u, 0x240, static_cast<std::uint32_t>(outputAddress), static_cast<std::uint32_t>(outputAddress >> 32u), 0xc0031500u, 1, 1, 1, 0x8041});
+    Packet packet{commands.data(), static_cast<std::uint32_t>(commands.size()), 0, {}};
+    for (unsigned iteration = 0; iteration < 2; ++iteration) {
+        std::ranges::fill(output, Fill);
+        sceAgcDriverSubmitAcb(0x20, &packet);
+        AgcDriverWaitIdle_nid_postfix();
+        AgcDriverSuspendPoint_nid_postfix();
+        AgcDriverWaitIdle_nid_postfix();
+        for (std::size_t offset = 0; offset < output.size(); ++offset) {
+            const auto expected = offset < Threads * LaneBytes ? ExpectedLane[offset % LaneBytes] : Fill;
+            if (output[offset] != expected) throw std::runtime_error("lazy dispatch " + std::to_string(iteration) + " byte " + std::to_string(offset) + " is " + std::to_string(output[offset]) + ", expected " + std::to_string(expected));
+        }
+    }
+    {
+        AgcDriver::DriverDetail::ShaderPreparationTransaction transaction;
+        Require(!transaction.Read(*lazySnapshot).entries.empty(), "the first dispatch did not prepare the lazily registered shader");
+    }
+    AgcDriverResolveShaderAbi_nid_postfix(&header.shader, {}, {});
+    alignas(256) static const std::array<std::uint32_t, 1> invalid{0xbf9f0000u};
+    const auto invalidAddress = reinterpret_cast<std::uintptr_t>(invalid.data());
+    header.shader.code = invalid.data();
+    header.shader.shader_size = sizeof(invalid);
+    header.registers[0].value = static_cast<std::uint32_t>(invalidAddress >> 8u);
+    header.registers[1].value = static_cast<std::uint32_t>(invalidAddress >> 40u);
+    AgcDriverRegisterShader_nid_postfix(&header.shader);
+    ExpectFailure([&] { AgcDriverResolveShaderAbi_nid_postfix(&header.shader, {}, {}); }, "s_code_end");
+    commands[2] = header.registers[0].value;
+    commands[5] = header.registers[1].value;
+    sceAgcDriverSubmitAcb(0x20, &packet);
+    ExpectFailure([] { AgcDriverWaitIdle_nid_postfix(); }, "s_code_end");
+    ExpectFailure([] { AgcDriverShutdown_nid_postfix(); }, "s_code_end");
+}
+
 void DeferredRegistration() {
     alignas(256) std::array<std::uint32_t, 2> code{0xbe842104u, 0xbf810000u};
     struct Header {
@@ -672,7 +777,13 @@ void DeferredUndecodableRegistration() {
 
 int main(int argc, char** argv) {
     try {
-        Require(argc == 1 || (argc == 2 && (std::string_view(argv[1]) == "--indirect" || std::string_view(argv[1]) == "--fail-before-registration" || std::string_view(argv[1]) == "--deferred" || std::string_view(argv[1]) == "--deferred-undecodable" || std::string_view(argv[1]) == "--registered-workgroup")), "invalid test arguments");
+        Require(argc == 1 || (argc == 2 && (std::string_view(argv[1]) == "--indirect" || std::string_view(argv[1]) == "--fail-before-registration" || std::string_view(argv[1]) == "--deferred" || std::string_view(argv[1]) == "--deferred-undecodable" || std::string_view(argv[1]) == "--registered-workgroup" || std::string_view(argv[1]) == "--lazy")), "invalid test arguments");
+        if (argc == 2 && std::string_view(argv[1]) == "--lazy") {
+            if (!OpenVulkanTestDevice()) return VulkanTestSkipped;
+            LazyRegistration();
+            std::cout << "lazy shader preparation tests passed\n";
+            return 0;
+        }
         if (argc == 2 && std::string_view(argv[1]) == "--registered-workgroup") {
             const auto device = OpenVulkanTestDevice();
             if (!device) return VulkanTestSkipped;

docs/dev/ARCHITECTURE.md

@@ -100,6 +100,7 @@ flowchart LR
 - Queue 0's packet worker decodes, captures and prepares a draw and hands its recording (`Graphics::Draw` under the GPU lock) to a committer thread in [`DrawPipeline.cpp`](../../core/libs/prx/libSceAgcDriver/Execution/src/Driver/Draw/DrawPipeline.cpp), so the next packet is decoded while the previous draw is recorded. Each queued item carries the guest ranges it will write (`DrawWriteRanges`: written buffers and storage images, color targets with their DCC and CMASK, depth and stencil). A later packet that reads or writes one of those ranges, or that the pipeline can't order, drains it first; labels, submission ends, buffer fills and (with `APS5_PIPELINE_DMA=1`) DMA packets are queued as items instead. An exception thrown by a commit is rethrown on the worker at its next enqueue or drain, as an inline draw error would be. On by default; `APS5_PIPELINED_DRAWS=0` records every draw inline, and `APS5_EXACT_DRAW_WRITES=0` widens the color and depth ranges to estimates.
 - The shader registry maps a code address to every header registered for that code, in registration order. Static, graphics and rectangle ABI resolution pick the header at the bound address (or one equal to it, `SameHeader`, which also takes a hull header fused by `sceAgcFuseShaderHalves` (type 3) as the registered back half (type 7) it copies, the fused registers aside; graphics ABI resolution then uses the fused header's shader registers); draws and dispatches pick the program by the stage's binary type (`RegisteredProgram`), and a front entry matching both type 2 and type 4 throws. Within a type the last registration wins; re-registering an identical earlier header makes its snapshot current again without preparing it again, and changed code drops the older headers.
 - The snapshot is published as pending. `SourceHandleFor`, `InvocationFor` and `ShaderPreparationTransaction::Edit` / `Read` wait for the pending job before reading `PreparedShaders`, and rethrow its failure. `APS5_SYNC_SHADER_PREPARE=1` prepares on the registering thread instead.
+- `APS5_LAZY_SHADER_PREPARE=1` retains registration validation and decoding but discards the temporary plan without queueing preparation. The snapshot is deferred, not pending; draws and dispatches prepare and cache it on first use. Explicit ABI resolution still prepares its requested stages. This opt-in mode takes precedence over synchronous registration preparation and trades background CPU/memory use for possible first-use stalls; disk caches and shader output are unchanged.
 - On devices with graphics pipeline libraries (`SpirvTarget::registrationModules`), a vertex or pixel artifact is emitted in registration form: the buffer V# words, image modes, swizzles and mip ranges, depth-compare state, export mappings and vertex component selectors it used to specialize on per draw are read at run time (image `ResourceMetadata` and a few state dwords in the shader data), descriptor heaps keep every element in place (an unused one gets a null or placeholder descriptor), and the pixel stage's bindings live in descriptor set 1 so each stage's library depends only on its own set layout (`INDEPENDENT_SETS`). An artifact above `registrationWordLimit` SPIR-V words is re-emitted specialized, and `CompiledShaderArtifact::registrationModule` records which form it has. For artifacts in registration form, `RegistrationResult` materializes the module once, and `PrebuildStageLibrary` builds its pre-rasterization or fragment-shader library on a separate one-thread pool (`APS5_PREBUILD_THREADS`) with a fixed default state, pairing it with a stub stage. A later draw with that state finds both libraries in the store and only links.
 - A draw or dispatch whose state matches no artifact prepared at registration (a title that builds its static ABI state itself, without `sceAgcDriverResolve*Abi`) gets one prepared at use by `SourceHandleFor` or `InvocationFor`, logged once per artifact and kept in the snapshot for the next use.
 - A compute shader whose `s_swappc_b64` calls a user-data pointer cannot be linked at registration, because the callee's address is only known at dispatch. `InvocationFor` in [`ShaderRegistry.cpp`](../../core/libs/prx/libSceAgcDriver/Execution/src/Driver/Shaders/ShaderRegistry.cpp) reads each callee from guest memory, appends one copy per call after the caller (`LinkUserDataCalls`) and prepares that linked code, so a different callee gets its own prepared entry.