fix(agc): prepare the null pixel program at draw when its variant was not prepared

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

@@ -18,6 +18,10 @@ namespace ShaderRecompiler {
 class ShaderPreparationContext;
 }
 
+namespace AgcDriver {
+class VulkanDevice;
+}
+
 namespace AgcDriver::DriverDetail {
 
 struct ShaderSnapshot;
@@ -78,6 +82,7 @@ std::shared_ptr<const ShaderSnapshot> ReadRawComputeShader(std::uint64_t address
 std::shared_ptr<const ShaderRecompiler::SourceHandle> PrepareShaderWithDiagnostics(const ShaderRecompiler::RecompileRequest& request, ShaderRecompiler::ShaderPreparationContext* preparation = nullptr);
 
 std::uint64_t NullPixelProgramAddress();
+ShaderSnapshot PrepareNullPixelProgram(const VulkanDevice& device);
 std::optional<ShaderRecompiler::ShaderFloatMode> RegisteredFloatMode(const ShaderSnapshot& snapshot);
 void PublishRegisteredShader(std::shared_ptr<ShaderRegistry>& registry, const std::shared_ptr<const ShaderSnapshot>& snapshot);
 

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

@@ -35,8 +35,8 @@ void DecodeGraphicsPrograms(DrawDecode& decoded, const QueueState& queue, const
         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 : ReadGraphicsRegister(queue.shader, rsrc2);
+        if (!nullPixel) Graphics::NoteRegisterRead(Graphics::RegisterBank::Shader, rsrc2);
+        const auto resources = nullPixel ? 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));

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

@@ -179,6 +179,21 @@ void PublishRegisteredShader(std::shared_ptr<ShaderRegistry>& registry, const st
     transaction.Commit();
 }
 
+namespace {
+
+bool PreparedAtUse(const ShaderSnapshot& snapshot, const ShaderRecompiler::RecompileRequest& request) {
+    if (snapshot.codeAddress == NullPixelProgramAddress() && request.shader.stage == ShaderRecompiler::ShaderStage::Fragment && request.context.waveSize == 32u) {
+        APS5_LOG_ERR("The null pixel program has no wave%u artifact for this draw; preparing it at draw", request.context.waveSize);
+        return true;
+    }
+    if (!snapshot.header.empty()) return snapshot.prepared->deferred;
+    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));
+    return true;
+}
+
+}
+
 std::shared_ptr<const ShaderRecompiler::SourceHandle> SourceHandleFor(const ShaderSnapshot& snapshot, std::size_t codeOffset, const ShaderRecompiler::RecompileRequest& request) {
     struct PreparedKeyStorage {};
     auto& key = HostThreadLocal<std::vector<std::uint64_t>, PreparedKeyStorage>();
@@ -187,11 +202,7 @@ std::shared_ptr<const ShaderRecompiler::SourceHandle> SourceHandleFor(const Shad
     for (const auto& entry : snapshot.prepared->entries) {
         if (entry.codeOffset == codeOffset && ShaderRecompiler::MatchesPreparedShader(request, *entry.handle, key)) return entry.handle;
     }
-    if (snapshot.header.empty() || snapshot.prepared->deferred) {
-        if (snapshot.header.empty()) {
-            if (snapshot.type != 0 || request.shader.stage != ShaderRecompiler::ShaderStage::Compute) throw std::runtime_error("AGC driver: unregistered program is not a compute shader");
-            APS5_LOG_ERR("Compute shader 0x%llx was not registered; preparing its artifact at dispatch", static_cast<unsigned long long>(snapshot.codeAddress));
-        }
+    if (PreparedAtUse(snapshot, request)) {
         auto handle = PrepareShaderWithDiagnostics(request);
         snapshot.prepared->entries.push_back({codeOffset, handle});
         return handle;
@@ -241,11 +252,7 @@ ShaderRecompiler::PreparedShaderInvocation InvocationFor(const ShaderSnapshot& s
         invocationRequest.shader.code = ShaderRecompiler::GetPreparedCode(*entry.handle);
         if (auto invocation = ShaderRecompiler::PreparedShaderInvocation::TryCreate(invocationRequest, entry.handle, key)) return std::move(*invocation);
     }
-    if (snapshot.header.empty() || snapshot.prepared->deferred) {
-        if (snapshot.header.empty()) {
-            if (snapshot.type != 0 || request.shader.stage != ShaderRecompiler::ShaderStage::Compute) throw std::runtime_error("AGC driver: unregistered program is not a compute shader");
-            APS5_LOG_ERR("Compute shader 0x%llx was not registered; preparing its artifact at dispatch", static_cast<unsigned long long>(snapshot.codeAddress));
-        }
+    if (PreparedAtUse(snapshot, request)) {
         auto handle = PrepareShaderWithDiagnostics(request);
         invocationRequest = request;
         invocationRequest.shader.code = ShaderRecompiler::GetPreparedCode(*handle);
@@ -711,6 +718,26 @@ void ResolvePreparedGraphics(const ShaderSnapshot& front, const std::shared_ptr<
     transaction.Commit();
 }
 
+ShaderSnapshot PrepareNullPixelProgram(const VulkanDevice& device) {
+    ShaderSnapshot null{NullPixelProgramAddress(), reinterpret_cast<std::uintptr_t>(&NullPixelShader), NullPixelShader.type, {}, {}};
+    null.code.assign(std::begin(NullPixelCode), std::end(NullPixelCode));
+    null.header.resize(sizeof(Shader));
+    std::memcpy(null.header.data(), &NullPixelShader, sizeof(Shader));
+    auto nullRegisteredState = DecodeRegisteredState(null);
+    nullRegisteredState.shader.insert_or_assign(0x008u, static_cast<std::uint32_t>(null.codeAddress >> 8u));
+    nullRegisteredState.shader.insert_or_assign(0x009u, static_cast<std::uint32_t>(null.codeAddress >> 40u));
+    nullRegisteredState.shader.insert_or_assign(0x00bu, 0u);
+    nullRegisteredState.context.insert_or_assign(0x1b3u, 0x2u);
+    nullRegisteredState.context.insert_or_assign(0x1b4u, 0x2u);
+    null.registeredState = std::make_shared<const RegisteredShaderState>(std::move(nullRegisteredState));
+    QueueState nullState{};
+    nullState.shader = null.registeredState->shader;
+    nullState.context = null.registeredState->context;
+    nullState.userConfig = null.registeredState->userConfig;
+    null.prepared->entries = PrepareRegistered(null, device, nullState, true);
+    return null;
+}
+
 void Driver::RegisterShader(const Shader* shader) {
     PerformanceContext timingContext(FrameTiming::Preparation());
     PerformanceTimer timing("Shader.Register");
@@ -782,23 +809,7 @@ void Driver::RegisterShader(const Shader* shader) {
     rethrowFailure();
     PublishRegisteredShader(shaders, std::make_shared<const ShaderSnapshot>(std::move(snapshot)));
     if (shaders->find(NullPixelProgramAddress()) == shaders->end()) {
-        ShaderSnapshot null{NullPixelProgramAddress(), reinterpret_cast<std::uintptr_t>(&NullPixelShader), NullPixelShader.type, {}, {}};
-        null.code.assign(std::begin(NullPixelCode), std::end(NullPixelCode));
-        null.header.resize(sizeof(Shader));
-        std::memcpy(null.header.data(), &NullPixelShader, sizeof(Shader));
-        auto nullRegisteredState = DecodeRegisteredState(null);
-        nullRegisteredState.shader.emplace(0x008u, static_cast<std::uint32_t>(null.codeAddress >> 8u));
-        nullRegisteredState.shader.emplace(0x009u, static_cast<std::uint32_t>(null.codeAddress >> 40u));
-        nullRegisteredState.shader.emplace(0x00bu, 0u);
-        nullRegisteredState.context.emplace(0x1b3u, 0x2u);
-        nullRegisteredState.context.emplace(0x1b4u, 0x2u);
-        null.registeredState = std::make_shared<const RegisteredShaderState>(std::move(nullRegisteredState));
-        QueueState nullState{};
-        nullState.shader = null.registeredState->shader;
-        nullState.context = null.registeredState->context;
-        nullState.userConfig = null.registeredState->userConfig;
-        null.prepared->entries = PrepareRegistered(null, *localDevice, nullState, true);
-        PublishRegisteredShader(shaders, std::make_shared<const ShaderSnapshot>(std::move(null)));
+        PublishRegisteredShader(shaders, std::make_shared<const ShaderSnapshot>(PrepareNullPixelProgram(*localDevice)));
     }
     transaction.Commit();
 }

core/libs/prx/libSceAgcDriver/Graphics/src/ShaderInputState.cpp

@@ -158,7 +158,7 @@ ShaderRecompiler::ShaderPixelStageInfo DecodePixelStageInfo(const Registers& con
         .conservativeZExport = static_cast<ShaderRecompiler::ConservativeZExport>(conservativeZExport),
         .orderedPixelShader = ((shaderControl >> 16u) & 0x1u) != 0,
         .targetOutputMode = targetOutputMode,
-        .targetExportMapping = exportMappings
+        .targetExportMapping = nullProgram ? std::array<std::uint8_t, 8>{} : exportMappings
     };
 }
 

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

@@ -5,6 +5,7 @@
 #include "prx/libSceAgcDriver/Graphics/include/ShaderInputState.hpp"
 #include "Optimization/ResourceProgram.hpp"
 #include <spirv/unified1/spirv.hpp>
+#include <algorithm>
 #include <array>
 #include <vector>
 #include <cstring>
@@ -247,6 +248,44 @@ void Check(AgcDriver::VulkanDevice& device, AgcDriver::Graphics::ShaderPath path
     Reject([&] { DecodeGraphicsPrograms(invalid, queue, registry, true, true); }, "reserved graphics program address");
 }
 
+void NullPixelMatchesRegistration(AgcDriver::VulkanDevice& device) {
+    using namespace AgcDriver::DriverDetail;
+    Fixture front;
+    front.Initialize(2u);
+    const auto null = std::make_shared<ShaderSnapshot>(PrepareNullPixelProgram(device));
+    const auto registered = null->prepared->entries;
+    Require(!registered.empty(), "the null pixel program was not prepared at registration");
+    ShaderRegistry registry;
+    registry.emplace(front.snapshot->codeAddress, front.snapshot);
+    registry.emplace(null->codeAddress, null);
+    AgcDriver::QueueState queue{};
+    queue.context[0x8e] = 0xfu;
+    queue.context[0x8f] = 0xfu;
+    front.Bind(queue, 0xc8u, 0x8bu);
+    alignas(8) const std::array<std::uint32_t, 2> merged{};
+    const auto mergedAddress = reinterpret_cast<std::uintptr_t>(merged.data());
+    queue.shader[0x82u] = static_cast<std::uint32_t>(mergedAddress);
+    queue.shader[0x83u] = static_cast<std::uint32_t>(mergedAddress >> 32u);
+    queue.shader[0x008u] = 0u;
+    queue.shader[0x009u] = 0u;
+    queue.shader[0x00bu] = 4u << 1u;
+    queue.context[0x1b3u] = 0x30u;
+    queue.context[0x1b4u] = 0x30u;
+    DrawDecode draw{};
+    draw.state.stages.path = AgcDriver::Graphics::ShaderPath::Vertex;
+    draw.state.stages.vertexWaveSize = 32u;
+    draw.state.stages.fragmentWaveSize = 64u;
+    std::array<std::uint8_t, 8> mappings{};
+    mappings.fill(0xe4u);
+    draw.pixel = AgcDriver::Graphics::DecodePixelStageInfo(queue.context, mappings, true);
+    DecodeGraphicsPrograms(draw, queue, registry, false, true);
+    const auto& fragment = draw.programs.back();
+    Require(fragment.snapshot == null && fragment.userData.empty(), "a null pixel draw took the last pixel shader's user SGPRs");
+    const auto stages = PrepareGraphicsStages(draw, device.Target());
+    const auto& handle = stages.back().entry.handle;
+    Require(std::ranges::any_of(registered, [&](const auto& entry) { return entry.handle == handle; }), "a null pixel draw did not match the variant prepared at registration");
+}
+
 }
 
 int main(int argc, char** argv) {
@@ -255,6 +294,7 @@ int main(int argc, char** argv) {
         if (!device) return VulkanTestSkipped;
         Require(argc <= 2, "invalid test arguments");
         const auto dump = argc == 2 ? std::filesystem::path(argv[1]) : std::filesystem::path{};
+        NullPixelMatchesRegistration(*device);
         Check(*device, AgcDriver::Graphics::ShaderPath::Vertex, dump);
         Check(*device, AgcDriver::Graphics::ShaderPath::Geometry, dump);
         Check(*device, AgcDriver::Graphics::ShaderPath::Tessellation, dump);

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

@@ -2,6 +2,7 @@
 #include "Optimization/ResourceProgram.hpp"
 #include "ControlFlow/RequestSerializer.hpp"
 #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/Submit/include/Acb.hpp"
 #include <spirv/unified1/spirv.hpp>
@@ -63,7 +64,24 @@ void RunUnregistered(AgcDriver::VulkanDevice& device, ShaderRecompiler::Recompil
     }
 }
 
+void NullPixelAtDraw(AgcDriver::VulkanDevice& device) {
+    std::vector<std::uint32_t> code(64, 0);
+    code.front() = 0xbf810000u;
+    AgcDriver::DriverDetail::ShaderSnapshot snapshot{AgcDriver::DriverDetail::NullPixelProgramAddress(), 0, 1, code, {}};
+    snapshot.header.resize(sizeof(Shader));
+    std::array<std::uint32_t, 4> users{};
+    const auto pixel = AgcDriver::Graphics::DecodePixelStageInfo({}, {}, true);
+    const ShaderRecompiler::RecompileRequest registered{{ShaderRecompiler::ShaderStage::Fragment, snapshot.codeAddress, snapshot.code, 0, {}}, {64, 0, {}, {}, pixel, {}, {}}, device.Target(), {0, 0, 0, 128}};
+    snapshot.prepared->entries.push_back({0, ShaderRecompiler::PrepareShader(registered)});
+    ShaderRecompiler::RecompileRequest request{{ShaderRecompiler::ShaderStage::Fragment, snapshot.codeAddress, snapshot.code, 0, {}}, {32, 0, users, {}, pixel, {}, {}}, device.Target(), {0, 0, 20, 108}};
+    static_cast<void>(AgcDriver::DriverDetail::InvocationFor(snapshot, 0, request));
+    Require(snapshot.prepared->entries.size() == 2, "the null pixel program was not prepared at draw");
+    static_cast<void>(AgcDriver::DriverDetail::InvocationFor(snapshot, 0, request));
+    Require(AgcDriver::DriverDetail::SourceHandleFor(snapshot, 0, request) == snapshot.prepared->entries.back().handle && snapshot.prepared->entries.size() == 2, "the null pixel program was prepared again for the same draw");
+}
+
 void Run(AgcDriver::VulkanDevice& device) {
+    NullPixelAtDraw(device);
     alignas(256) std::array<std::uint32_t, 1> code{0xbf810000u};
     std::array<std::uint32_t, 4> users{};
     const ShaderRecompiler::ShaderComputeStageInfo compute{{1, 1, 1}, 0, {false, false, false}, false, 1, {}};