perf(agc): wait at submit only for the preparations of the programs a submission binds

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

@@ -12,6 +12,7 @@
 #include <map>
 #include <memory>
 #include <set>
+#include <span>
 #include <thread>
 #include <unordered_map>
 #include <unordered_set>
@@ -58,6 +59,13 @@ struct QueueWorker {
     std::thread thread;
 };
 
+struct ReferencedProgramScan {
+    std::vector<std::uint64_t> addresses;
+    bool complete = true;
+};
+
+ReferencedProgramScan ReferencedPrograms(std::span<const std::uint32_t> commands, const std::map<std::size_t, std::vector<std::uint32_t>>& registerLists);
+
 }
 
 #endif

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

@@ -332,6 +332,7 @@ void PublishRegisteredShader(std::shared_ptr<ShaderRegistry>& registry, const st
 void SubmitRegistrationPreparation(std::function<void()> task, const PreparedShaders* prepared = nullptr);
 void PrioritizeRegistrationPreparation(PreparedShaders& prepared);
 void AwaitRegistrationPreparation();
+void AwaitReferencedPreparations(const ShaderRegistry& registry, std::span<const std::uint64_t> addresses);
 void SubmitLibraryPrebuild(std::function<void()> task, std::uint64_t program = 0);
 bool LibraryPrebuildsWaiting();
 void PrioritizeLibraryPrebuilds(std::span<const std::uint64_t> programs);

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Queues/Submission.cpp

@@ -14,6 +14,7 @@
 #include "ThreadOwned.hpp"
 #include "prx/libc/include/HostThreadLocal.hpp"
 #include <algorithm>
+#include <array>
 #include <atomic>
 #include <bit>
 #include <cmath>
@@ -225,6 +226,60 @@ bool Driver::copySegment(Submission& submission, const std::uint32_t* guest, std
     return false;
 }
 
+ReferencedProgramScan ReferencedPrograms(std::span<const std::uint32_t> commands, const std::map<std::size_t, std::vector<std::uint32_t>>& registerLists) {
+    constexpr std::array<std::uint32_t, 7> programRegisters{0x008u, 0x048u, 0x088u, 0x0c8u, 0x108u, 0x148u, 0x20cu};
+    std::array<std::uint32_t, programRegisters.size() * 2> values{};
+    std::array<bool, programRegisters.size() * 2> written{};
+    std::array<bool, programRegisters.size()> bound{};
+    ReferencedProgramScan scan;
+    const auto write = [&](std::uint32_t offset, std::uint32_t value, bool conditional) {
+        const auto base = std::ranges::find_if(programRegisters, [&](std::uint32_t program) { return offset - program < 2u; });
+        if (base == programRegisters.end()) return;
+        if (conditional) {
+            scan.complete = false;
+            return;
+        }
+        const auto slot = static_cast<std::size_t>(base - programRegisters.begin()) * 2u + (offset - *base);
+        values[slot] = value;
+        written[slot] = true;
+        const auto low = slot & ~std::size_t{1};
+        if (written[low] && written[low + 1]) {
+            bound[low / 2] = true;
+            scan.addresses.push_back((static_cast<std::uint64_t>(values[low]) << 8u) | (static_cast<std::uint64_t>(values[low + 1] & 0xffu) << 40u));
+        }
+    };
+    std::size_t guardedEnd = 0;
+    for (std::size_t cursor = 0; cursor < commands.size(); cursor += Pm4::PacketWords(commands[cursor])) {
+        const auto header = commands[cursor];
+        if ((header >> 30u) != 3u) continue;
+        const auto opcode = (header >> 8u) & 0xffu;
+        const auto words = std::min<std::size_t>(Pm4::PacketWords(header), commands.size() - cursor);
+        const bool conditional = (header & 1u) != 0 || cursor < guardedEnd;
+        if (opcode == 0x22u && words == 5) guardedEnd = std::max(guardedEnd, cursor + words + Pm4::ConditionalWords(commands.subspan(cursor, words)));
+        if (Pm4::DrawOpcode(opcode) && !std::any_of(bound.begin(), bound.end() - 1, [](bool pair) { return pair; })) scan.complete = false;
+        if ((opcode == 0x15u || opcode == 0x16u) && !bound.back()) scan.complete = false;
+        if (opcode == 0x76u && words >= 3) {
+            const auto packet = commands.subspan(cursor, words);
+            const auto offset = packet[1] & ~0x70000000u;
+            for (std::size_t i = 2; i < words; ++i) write(offset + static_cast<std::uint32_t>(i - 2), packet[i], conditional);
+        } else if (opcode == 0x63u) {
+            const auto list = registerLists.find(cursor);
+            if (list == registerLists.end()) continue;
+            if (!Pm4::IndirectRegistersValid(list->second)) {
+                scan.complete = false;
+                continue;
+            }
+            for (std::size_t i = 0; i + 1 < list->second.size(); i += 2) write(list->second[i] & ~0x70000000u, list->second[i + 1], conditional);
+        }
+    }
+    for (std::size_t slot = 0; slot < written.size(); slot += 2) {
+        if (written[slot] != written[slot + 1]) scan.complete = false;
+    }
+    std::ranges::sort(scan.addresses);
+    scan.addresses.erase(std::unique(scan.addresses.begin(), scan.addresses.end()), scan.addresses.end());
+    return scan;
+}
+
 void Driver::readRegisterLists(Submission& submission) {
     submission.registerLists.clear();
     for (std::size_t cursor = 0; cursor < submission.commands.size(); cursor += Pm4::PacketWords(submission.commands[cursor])) {
@@ -511,7 +566,17 @@ void Driver::Submit(const Packet* packet, std::uint32_t queue, bool cpuFlip) {
     if (APS5_ENABLE_TIMING_LOG) submission.validatedAt = std::chrono::steady_clock::now();
     waitForFlipRoom(submission);
     holdFlipBehindWorker(submission, cpuFlip);
-    AwaitRegistrationPreparation();
+    const auto scan = ReferencedPrograms(submission.commands, submission.registerLists);
+    if (!scan.complete || submission.rewindTail != nullptr) {
+        AwaitRegistrationPreparation();
+    } else {
+        std::shared_ptr<const ShaderRegistry> registry;
+        {
+            std::lock_guard lock(mutex);
+            registry = shaders;
+        }
+        if (registry != nullptr) AwaitReferencedPreparations(*registry, scan.addresses);
+    }
     if (LibraryPrebuildsWaiting()) PrioritizeLibraryPrebuilds(BoundGraphicsPrograms(submission));
     if (APS5_ENABLE_TIMING_LOG) submission.roomReadyAt = std::chrono::steady_clock::now();
     static const bool trace = std::getenv("APS5_TRACE_GPU") != nullptr;

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

@@ -52,6 +52,7 @@
 #include <sstream>
 #include <stdexcept>
 #include <thread>
+#include <unordered_map>
 #include "Optimization/ShaderStageInputInfo.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Driver/Shaders/RegistrationPreparationBudget.hpp"
 #include "prx/libSceAgcDriver/Execution/include/PerformanceOverlay.hpp"
@@ -1640,6 +1641,7 @@ class RegistrationPreparePool {
                 else tasks.push_back({std::move(task), prepared, sequence, budgeted, program});
                 if (program != 0) waitingPrograms.fetch_add(1, std::memory_order_relaxed);
                 unfinished.insert(sequence);
+                if (prepared != nullptr) owned.emplace(prepared, sequence);
                 lock.unlock();
                 ready.notify_one();
                 return;
@@ -1650,10 +1652,17 @@ class RegistrationPreparePool {
 
     void Prioritize(const PreparedShaders& prepared) {
         std::lock_guard lock(mutex);
-        const auto found = std::ranges::find_if(tasks, [&](const auto& task) { return task.prepared == &prepared; });
-        if (found == tasks.end()) return;
-        if (std::getenv("APS5_TRACE_PREPARATION_WAIT") != nullptr) std::fprintf(stderr, "[shader-prepare] demanded queued task at position %zu of %zu\n", static_cast<std::size_t>(std::distance(tasks.begin(), found)), tasks.size());
-        std::rotate(tasks.begin(), found, std::next(found));
+        moveToFront(prepared);
+    }
+
+    void AwaitOwned(const PreparedShaders& prepared) {
+        std::unique_lock lock(mutex);
+        const auto [first, last] = owned.equal_range(&prepared);
+        if (first == last) return;
+        std::vector<std::uint64_t> sequences;
+        for (auto entry = first; entry != last; ++entry) sequences.push_back(entry->second);
+        moveToFront(prepared);
+        idle.wait(lock, [&] { return std::ranges::none_of(sequences, [&](const auto sequence) { return unfinished.contains(sequence); }); });
     }
 
     void Prioritize(std::span<const std::uint64_t> programs) {
@@ -1681,7 +1690,10 @@ class RegistrationPreparePool {
             if (dropQueued) {
                 dropped.swap(tasks);
                 waitingPrograms.store(0, std::memory_order_relaxed);
-                for (const auto& task : dropped) unfinished.erase(task.sequence);
+                for (const auto& task : dropped) {
+                    unfinished.erase(task.sequence);
+                    disown(task.prepared, task.sequence);
+                }
                 idle.notify_all();
             }
             idle.wait(lock, [&] { return tasks.empty() && running == 0; });
@@ -1694,6 +1706,20 @@ class RegistrationPreparePool {
 private:
     static constexpr std::size_t StackBytes = 64u << 20u;
 
+    void moveToFront(const PreparedShaders& prepared) {
+        const auto found = std::ranges::find_if(tasks, [&](const auto& task) { return task.prepared == &prepared; });
+        if (found == tasks.end()) return;
+        if (std::getenv("APS5_TRACE_PREPARATION_WAIT") != nullptr) std::fprintf(stderr, "[shader-prepare] demanded queued task at position %zu of %zu\n", static_cast<std::size_t>(std::distance(tasks.begin(), found)), tasks.size());
+        std::rotate(tasks.begin(), found, std::next(found));
+    }
+
+    void disown(const PreparedShaders* prepared, std::uint64_t sequence) {
+        if (prepared == nullptr) return;
+        const auto [first, last] = owned.equal_range(prepared);
+        const auto found = std::find_if(first, last, [&](const auto& entry) { return entry.second == sequence; });
+        if (found != last) owned.erase(found);
+    }
+
 #ifdef _WIN32
     struct Worker {
         HANDLE handle;
@@ -1771,6 +1797,7 @@ class RegistrationPreparePool {
             }
             const auto sequence = task.sequence;
             const bool budgeted = task.budgeted;
+            const auto* const prepared = task.prepared;
             task.function();
             task = {};
             bool drained = false;
@@ -1779,6 +1806,7 @@ class RegistrationPreparePool {
                 --pool.running;
                 if (budgeted) --pool.budgetedRunning;
                 pool.unfinished.erase(sequence);
+                pool.disown(prepared, sequence);
                 drained = pool.tasks.empty() && pool.running == 0;
                 pool.idle.notify_all();
             }
@@ -1811,6 +1839,7 @@ class RegistrationPreparePool {
     unsigned budgetedLimit = 1;
     std::uint64_t submitted = 0;
     std::set<std::uint64_t> unfinished;
+    std::unordered_multimap<const PreparedShaders*, std::uint64_t> owned;
     bool closed = false;
 };
 
@@ -1887,6 +1916,24 @@ void AwaitRegistrationPreparation() {
     if (auto* pool = RegistrationPreparePool::Existing()) pool->AwaitSubmitted();
 }
 
+void AwaitReferencedPreparations(const ShaderRegistry& registry, std::span<const std::uint64_t> addresses) {
+    for (const auto address : addresses) {
+        auto found = registry.upper_bound(address);
+        if (found == registry.begin()) continue;
+        --found;
+        for (const auto& snapshot : found->second) {
+            if (address - snapshot->codeAddress >= snapshot->code.size() * sizeof(std::uint32_t)) continue;
+            auto& prepared = *snapshot->prepared;
+            {
+                std::lock_guard lock(prepared.mutex);
+                if (!prepared.pending) continue;
+            }
+            prepared.demanded.store(true);
+            if (auto* pool = RegistrationPreparePool::Existing()) pool->AwaitOwned(prepared);
+        }
+    }
+}
+
 void SubmitLibraryPrebuild(std::function<void()> task, std::uint64_t program) {
     RegistrationPreparePool::Prebuild().Submit(std::move(task), nullptr, false, program);
 }