[HOLD] perf(agc): pre-warm deferred shaders on SET_SH_REG packet execution

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

@@ -107,6 +107,7 @@ ShaderSnapshot PrepareNullPixelProgram(const VulkanDevice& device);
 std::optional<ShaderRecompiler::ShaderFloatMode> RegisteredFloatMode(const ShaderSnapshot& snapshot);
 std::shared_ptr<const ShaderSnapshot> RegisteredHeader(const ShaderRegistry& registry, const Shader* shader);
 std::shared_ptr<const ShaderSnapshot> RegisteredProgram(const ShaderRegistry& registry, std::uint64_t address, std::initializer_list<std::uint8_t> types);
+bool TriggerDeferredShaderPreparation(const ShaderRegistry& registry, const std::shared_ptr<VulkanDevice>& device, const QueueState& state, std::uint32_t programRegister, std::initializer_list<std::uint8_t> types);
 void PublishRegisteredShader(std::shared_ptr<ShaderRegistry>& registry, const std::shared_ptr<const ShaderSnapshot>& snapshot);
 void SubmitRegistrationPreparation(std::function<void()> task, const PreparedShaders* prepared = nullptr);
 void PrioritizeRegistrationPreparation(PreparedShaders& prepared);

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Packets/PacketExecution.cpp

@@ -8,13 +8,43 @@
 #include "prx/libSceAgcDriver/Execution/include/GuestMemory.hpp"
 #include "prx/libSceAgcDriver/Execution/include/PerformanceTimer.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Pm4.hpp"
+#include "prx/libSceAgcDriver/Execution/include/Driver/Shaders/ShaderRegistry.hpp"
 #include "prx/libSceAgcDriver/Eq/include/Event.hpp"
 #include <cstdlib>
 #include <functional>
 #include <shared_mutex>
 
 namespace AgcDriver::DriverDetail {
 
+void checkDeferredShadersOnShRegWrite(std::uint32_t opcode, std::span<const std::uint32_t> packet, std::size_t cursor, const Submission& submission, const std::shared_ptr<VulkanDevice>& device, const QueueState& queue) {
+    if (submission.shaders == nullptr || device == nullptr) return;
+    if (opcode == 0x76) {
+        const auto regOffset = packet[1] & ~0x70000000u;
+        const auto regCount = static_cast<std::uint32_t>(packet.size() - 2);
+        if (regOffset <= 0x20du && regOffset + regCount > 0x20cu) TriggerDeferredShaderPreparation(*submission.shaders, device, queue, 0x20cu, {0});
+        if (regOffset <= 0x009u && regOffset + regCount > 0x008u) TriggerDeferredShaderPreparation(*submission.shaders, device, queue, 0x008u, {1});
+        if (regOffset <= 0x0c9u && regOffset + regCount > 0x0c8u) TriggerDeferredShaderPreparation(*submission.shaders, device, queue, 0x0c8u, {2, 4});
+        if (regOffset <= 0x089u && regOffset + regCount > 0x088u) TriggerDeferredShaderPreparation(*submission.shaders, device, queue, 0x088u, {6});
+        if (regOffset <= 0x109u && regOffset + regCount > 0x108u) TriggerDeferredShaderPreparation(*submission.shaders, device, queue, 0x108u, {7});
+    } else if (opcode == 0x63) {
+        const auto& pairs = submission.registerLists.at(cursor);
+        bool checkCompute = false, checkFragment = false, checkVertex = false, checkMesh = false, checkTess = false;
+        for (std::size_t i = 0; i < pairs.size(); i += 2) {
+            const auto reg = pairs[i] & ~0x70000000u;
+            if (reg == 0x20cu || reg == 0x20du) checkCompute = true;
+            else if (reg == 0x008u || reg == 0x009u) checkFragment = true;
+            else if (reg == 0x0c8u || reg == 0x0c9u) checkVertex = true;
+            else if (reg == 0x088u || reg == 0x089u) checkMesh = true;
+            else if (reg == 0x108u || reg == 0x109u) checkTess = true;
+        }
+        if (checkCompute) TriggerDeferredShaderPreparation(*submission.shaders, device, queue, 0x20cu, {0});
+        if (checkFragment) TriggerDeferredShaderPreparation(*submission.shaders, device, queue, 0x008u, {1});
+        if (checkVertex) TriggerDeferredShaderPreparation(*submission.shaders, device, queue, 0x0c8u, {2, 4});
+        if (checkMesh) TriggerDeferredShaderPreparation(*submission.shaders, device, queue, 0x088u, {6});
+        if (checkTess) TriggerDeferredShaderPreparation(*submission.shaders, device, queue, 0x108u, {7});
+    }
+}
+
 template <typename TWork>
 void Driver::timed(double WorkerProfile::*bucket, TWork&& work) {
     static thread_local WorkerProfile profile;
@@ -305,6 +335,9 @@ void Driver::execute(const Submission& submission) {
                 if (opcode == 0x49 || opcode == 0x37) {
                     if (const auto label = Pm4::DecodeLabelWrite(packet)) noteLabelStore(label->address, label->Bytes(), ++eventSerial, submission.queue);
                 }
+                if (opcode == 0x76 || opcode == 0x63) {
+                    checkDeferredShadersOnShRegWrite(opcode, packet, cursor, submission, device.Load(), queue);
+                }
             }
             if (endOfPipeInterrupt && !interruptDeferred) AgcDriverDeliverEopInterrupt(submission.queue);
         }

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

@@ -915,6 +915,58 @@ void PrepareRegisteredInBackground(std::shared_ptr<const ShaderSnapshot> snapsho
     }, preparedKey);
 }
 
+bool CanPlanDeferredShader(const ShaderSnapshot& snapshot, const QueueState& state) noexcept {
+    try {
+        std::uint32_t programRegister;
+        std::uint32_t resourceRegister;
+        switch (snapshot.type) {
+        case 0:
+            programRegister = 0x20c;
+            resourceRegister = 0x213;
+            if (!state.shader.contains(programRegister) || !state.shader.contains(programRegister + 1u)) return false;
+            if (!state.shader.contains(resourceRegister)) return false;
+            if (!state.shader.contains(0x207u) || !state.shader.contains(0x208u) || !state.shader.contains(0x209u)) return false;
+            if ((state.shader.find(0x207u)->second & 0xffffu) == 0 ||
+                (state.shader.find(0x208u)->second & 0xffffu) == 0 ||
+                (state.shader.find(0x209u)->second & 0xffffu) == 0) return false;
+            return true;
+        case 1:
+            programRegister = 0x008;
+            resourceRegister = 0x00b;
+            if (!state.shader.contains(programRegister) || !state.shader.contains(programRegister + 1u)) return false;
+            if (!state.shader.contains(resourceRegister)) return false;
+            if (!state.context.contains(0x1b6u) || !state.context.contains(0x1b3u) || !state.context.contains(0x1b4u)) return false;
+            return true;
+        case 2:
+        case 4:
+            programRegister = 0x0c8;
+            resourceRegister = 0x08b;
+            if (!state.shader.contains(programRegister) || !state.shader.contains(programRegister + 1u)) return false;
+            if (!state.shader.contains(resourceRegister)) return false;
+            if (!state.context.contains(0x2d5u)) return false;
+            return true;
+        case 6:
+            programRegister = 0x088;
+            resourceRegister = 0x08b;
+            if (!state.shader.contains(programRegister) || !state.shader.contains(programRegister + 1u)) return false;
+            if (!state.shader.contains(resourceRegister)) return false;
+            if (!state.context.contains(0x2d5u)) return false;
+            return true;
+        case 7:
+            programRegister = 0x108;
+            resourceRegister = 0x10b;
+            if (!state.shader.contains(programRegister) || !state.shader.contains(programRegister + 1u)) return false;
+            if (!state.shader.contains(resourceRegister)) return false;
+            if (!state.context.contains(0x2d5u)) return false;
+            return true;
+        default:
+            return false;
+        }
+    } catch (...) {
+        return false;
+    }
+}
+
 }
 
 void SubmitRegistrationPreparation(std::function<void()> task, const PreparedShaders* prepared) {
@@ -935,6 +987,48 @@ void CloseRegistrationPreparation() {
     if (auto* prebuild = RegistrationPreparePool::ExistingPrebuild()) prebuild->Close();
 }
 
+bool TriggerDeferredShaderPreparation(const ShaderRegistry& registry, const std::shared_ptr<VulkanDevice>& device, const QueueState& state, std::uint32_t programRegister, std::initializer_list<std::uint8_t> types) {
+    if (device == nullptr) return false;
+    const auto low = state.shader.find(programRegister);
+    const auto high = state.shader.find(programRegister + 1u);
+    if (low == state.shader.end() || high == state.shader.end()) return false;
+    if ((high->second & ~0xffu) != 0) return false;
+    const auto address = (static_cast<std::uint64_t>(low->second) << 8u) | (static_cast<std::uint64_t>(high->second) << 40u);
+    if (address == 0) return false;
+    std::shared_ptr<const ShaderSnapshot> snapshot;
+    try {
+        snapshot = RegisteredProgram(registry, address, types);
+    } catch (...) {
+        return false;
+    }
+    if (snapshot == nullptr) return false;
+    {
+        std::lock_guard lock(snapshot->prepared->mutex);
+        if (!snapshot->prepared->deferred || snapshot->prepared->pending || !snapshot->prepared->entries.empty()) return false;
+        if (!CanPlanDeferredShader(*snapshot, state)) return false;
+        snapshot->prepared->pending = true;
+    }
+    std::shared_ptr<RegisteredPreparation> plan;
+    try {
+        plan = PlanRegistered(*snapshot, *device, state, false);
+    } catch (...) {
+        std::lock_guard lock(snapshot->prepared->mutex);
+        snapshot->prepared->pending = false;
+        snapshot->prepared->settled.notify_all();
+        return false;
+    }
+    if (plan == nullptr) {
+        std::lock_guard lock(snapshot->prepared->mutex);
+        snapshot->prepared->pending = false;
+        snapshot->prepared->settled.notify_all();
+        return false;
+    }
+    plan->prebuildDevice = device;
+    if (snapshot->type == 0 || snapshot->type == 1) BuildRegisteredAbiKey(state, *device, plan->abiKey);
+    PrepareRegisteredInBackground(snapshot, std::move(plan));
+    return true;
+}
+
 std::vector<PreparedGraphicsStage> PrepareGraphicsStages(const DrawDecode& decoded, const ShaderRecompiler::SpirvTarget& target) {
     ShaderPreparationTransaction transaction;
     require(decoded.programs.size() == decoded.roles.size(), "graphics ABI program roles are incomplete");

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

@@ -581,6 +581,63 @@ void RegisteredVertexWorkgroup(const ShaderRecompiler::SpirvTarget& target, bool
     std::cout << "registered NGG wave64 preparation passed; subgroup=" << target.subgroupSize << '\n';
 }
 
+void DeferredShaderPrewarmedOnSetShReg(AgcDriver::VulkanDevice& device) {
+    const auto testDevice = std::shared_ptr<AgcDriver::VulkanDevice>(&device, [](auto*) {});
+    alignas(256) std::array<std::uint32_t, 1> code{0xbf810000u};
+    const auto address = reinterpret_cast<std::uintptr_t>(code.data());
+    struct Header {
+        Shader shader{};
+        ShaderSpecialRegs specials{};
+    } header;
+    const auto headerAddress = reinterpret_cast<std::uintptr_t>(&header);
+    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.type = 0;
+    header.shader.specials = &header.specials;
+    header.specials.dispatch_modifier = 0x8000;
+
+    auto snapshot = std::make_shared<AgcDriver::DriverDetail::ShaderSnapshot>(address, headerAddress, 0, std::vector<std::uint32_t>{code.begin(), code.end()}, std::vector<std::byte>{sizeof(header)});
+    std::memcpy(snapshot->header.data(), &header, sizeof(header));
+    snapshot->prepared->deferred = true;
+
+    auto registry = std::make_shared<AgcDriver::DriverDetail::ShaderRegistry>();
+    AgcDriver::DriverDetail::PublishRegisteredShader(registry, snapshot);
+
+    AgcDriver::QueueState queue{};
+    queue.shader[0x207] = 1;
+    queue.shader[0x208] = 1;
+    queue.shader[0x209] = 1;
+    queue.shader[0x212] = 0;
+    queue.shader[0x213] = 12;
+
+    Require(!AgcDriver::DriverDetail::TriggerDeferredShaderPreparation(*registry, testDevice, queue, 0x20cu, {0}), "deferred shader prepared without program address");
+    Require(snapshot->prepared->entries.empty(), "deferred shader has entries without program address");
+
+    queue.shader[0x20c] = static_cast<std::uint32_t>(address >> 8u);
+    queue.shader[0x20d] = static_cast<std::uint32_t>(address >> 40u);
+
+    queue.shader.erase(0x213);
+    Require(!AgcDriver::DriverDetail::TriggerDeferredShaderPreparation(*registry, testDevice, queue, 0x20cu, {0}), "deferred shader prepared without resource register");
+    Require(snapshot->prepared->entries.empty(), "deferred shader has entries without resource register");
+    queue.shader[0x213] = 12;
+
+    Require(AgcDriver::DriverDetail::TriggerDeferredShaderPreparation(*registry, testDevice, queue, 0x20cu, {0}), "deferred shader preparation was not triggered on SET_SH_REG");
+
+    {
+        std::unique_lock lock(snapshot->prepared->mutex);
+        const bool settled = snapshot->prepared->settled.wait_for(lock, std::chrono::seconds(3), [&] { return !snapshot->prepared->pending; });
+        Require(settled, "deferred shader preparation timed out");
+        if (snapshot->prepared->failure != nullptr) std::rethrow_exception(snapshot->prepared->failure);
+        Require(snapshot->prepared->entries.size() == 2, "deferred shader was not pre-warmed after SET_SH_REG trigger");
+        Require(snapshot->prepared->entries.front().handle != nullptr, "pre-warmed deferred shader has null handle");
+    }
+
+    Require(!AgcDriver::DriverDetail::TriggerDeferredShaderPreparation(*registry, testDevice, queue, 0x20cu, {0}), "already pre-warmed shader prepared again");
+}
+
 void DeferredRegistration() {
     alignas(256) std::array<std::uint32_t, 2> code{0xbe842104u, 0xbf810000u};
     struct Header {
@@ -650,7 +707,9 @@ int main(int argc, char** argv) {
             return 0;
         }
         if (argc == 2 && std::string_view(argv[1]) == "--deferred") {
-            if (!OpenVulkanTestDevice()) return VulkanTestSkipped;
+            const auto device = OpenVulkanTestDevice();
+            if (!device) return VulkanTestSkipped;
+            DeferredShaderPrewarmedOnSetShReg(*device);
             DeferredRegistration();
             std::cout << "deferred shader preparation tests passed\n";
             return 0;
@@ -665,6 +724,7 @@ int main(int argc, char** argv) {
         if (!device) return VulkanTestSkipped;
         Run(*device);
         RegisteredVertexWorkgroup(device->Target());
+        DeferredShaderPrewarmedOnSetShReg(*device);
         FailureCapture(*device);
         PendingRegistrationPreparation(*device);
         Require(argc != 2 || std::string_view(argv[1]) != "--fail-before-registration", "injected failure before registration");