Merge pull request #4404 from eljackjoestar-cmd/codex/demanded-shader-preparation

fix(agc): prioritize demanded shader preparation

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

@@ -67,6 +67,7 @@ struct PreparedShaders : PreparedShaderState {
     std::mutex mutex;
     std::condition_variable settled;
     std::exception_ptr failure;
+    std::atomic<bool> demanded{false};
     bool pending = false;
 };
 
@@ -107,7 +108,8 @@ std::optional<ShaderRecompiler::ShaderFloatMode> RegisteredFloatMode(const Shade
 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);
 void PublishRegisteredShader(std::shared_ptr<ShaderRegistry>& registry, const std::shared_ptr<const ShaderSnapshot>& snapshot);
-void SubmitRegistrationPreparation(std::function<void()> task);
+void SubmitRegistrationPreparation(std::function<void()> task, const PreparedShaders* prepared = nullptr);
+void PrioritizeRegistrationPreparation(PreparedShaders& prepared);
 void SubmitLibraryPrebuild(std::function<void()> task);
 void CloseRegistrationPreparation();
 

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

@@ -141,6 +141,7 @@ bool RawGraphicsDrawn() {
 namespace {
 
 void AwaitRegisteredPreparation(PreparedShaders& prepared, std::unique_lock<std::mutex>& lock) {
+    if (prepared.pending) PrioritizeRegistrationPreparation(prepared);
     prepared.settled.wait(lock, [&] { return !prepared.pending; });
     if (prepared.failure != nullptr) std::rethrow_exception(prepared.failure);
 }
@@ -764,10 +765,18 @@ std::vector<PreparedShaders::Entry> PrepareRegistered(const ShaderSnapshot& snap
 class RegistrationPreparePool {
 public:
     static RegistrationPreparePool& Instance() {
-        static auto* pool = new RegistrationPreparePool(std::max(2u, std::thread::hardware_concurrency() / 2u));
+        static auto* pool = [] {
+            auto* created = new RegistrationPreparePool(std::max(2u, std::thread::hardware_concurrency() / 2u));
+            PreparationSlot().store(created);
+            return created;
+        }();
         return *pool;
     }
 
+    static RegistrationPreparePool* Existing() {
+        return PreparationSlot().load();
+    }
+
     static RegistrationPreparePool& Prebuild() {
         static auto* pool = [] {
             const char* text = std::getenv("APS5_PREBUILD_THREADS");
@@ -783,6 +792,11 @@ public:
     }
 
 private:
+    static std::atomic<RegistrationPreparePool*>& PreparationSlot() {
+        static std::atomic<RegistrationPreparePool*> slot{nullptr};
+        return slot;
+    }
+
     static std::atomic<RegistrationPreparePool*>& PrebuildSlot() {
         static std::atomic<RegistrationPreparePool*> slot{nullptr};
         return slot;
@@ -790,11 +804,12 @@ private:
 
 public:
 
-    void Submit(std::function<void()> task) {
+    void Submit(std::function<void()> task, const PreparedShaders* prepared = nullptr) {
         {
             std::unique_lock lock(mutex);
             if (threads != 0 && !closed) {
-                tasks.push_back(std::move(task));
+                if (prepared != nullptr && prepared->demanded.load()) tasks.push_front({std::move(task), prepared});
+                else tasks.push_back({std::move(task), prepared});
                 lock.unlock();
                 ready.notify_one();
                 return;
@@ -803,6 +818,14 @@ public:
         task();
     }
 
+    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));
+    }
+
     void Close() {
         std::unique_lock lock(mutex);
         closed = true;
@@ -828,7 +851,7 @@ private:
         auto& pool = *static_cast<RegistrationPreparePool*>(argument);
         if (PerformanceTrace::Configured()) PerformanceTrace::BindThread("shader prepare", TraceThread::DefaultCapacity / 8);
         for (;;) {
-            std::function<void()> task;
+            Task task;
             {
                 std::unique_lock lock(pool.mutex);
                 pool.ready.wait(lock, [&] { return !pool.tasks.empty(); });
@@ -836,8 +859,8 @@ private:
                 pool.tasks.pop_front();
                 ++pool.running;
             }
-            task();
-            task = nullptr;
+            task.function();
+            task = {};
             {
                 std::lock_guard lock(pool.mutex);
                 --pool.running;
@@ -847,10 +870,15 @@ private:
         return nullptr;
     }
 
+    struct Task {
+        std::function<void()> function;
+        const PreparedShaders* prepared = nullptr;
+    };
+
     std::mutex mutex;
     std::condition_variable ready;
     std::condition_variable idle;
-    std::deque<std::function<void()>> tasks;
+    std::deque<Task> tasks;
     unsigned threads = 0;
     unsigned running = 0;
     bool closed = false;
@@ -862,6 +890,7 @@ bool AsyncRegistrationPrepare() {
 }
 
 void PrepareRegisteredInBackground(std::shared_ptr<const ShaderSnapshot> snapshot, std::shared_ptr<RegisteredPreparation> plan) {
+    const auto* preparedKey = snapshot->prepared.get();
     SubmitRegistrationPreparation([snapshot = std::move(snapshot), plan = std::move(plan)] {
         PerformanceContext timingContext(FrameTiming::Preparation());
         std::vector<PreparedShaders::Entry> entries;
@@ -883,13 +912,18 @@ void PrepareRegisteredInBackground(std::shared_ptr<const ShaderSnapshot> snapsho
             prepared.pending = false;
         }
         prepared.settled.notify_all();
-    });
+    }, preparedKey);
 }
 
 }
 
-void SubmitRegistrationPreparation(std::function<void()> task) {
-    RegistrationPreparePool::Instance().Submit(std::move(task));
+void SubmitRegistrationPreparation(std::function<void()> task, const PreparedShaders* prepared) {
+    RegistrationPreparePool::Instance().Submit(std::move(task), prepared);
+}
+
+void PrioritizeRegistrationPreparation(PreparedShaders& prepared) {
+    prepared.demanded.store(true);
+    if (auto* pool = RegistrationPreparePool::Existing()) pool->Prioritize(prepared);
 }
 
 void SubmitLibraryPrebuild(std::function<void()> task) {

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

@@ -4,13 +4,16 @@
 #include "prx/libSceAgcDriver/Execution/include/Driver/Shaders/ShaderRegistry.hpp"
 #include "Optimization/ResourceProgram.hpp"
 #include <array>
+#include <algorithm>
 #include <atomic>
 #include <barrier>
 #include <chrono>
 #include <future>
 #include <cstring>
 #include <iostream>
+#include <latch>
 #include <optional>
+#include <semaphore>
 #include <thread>
 
 namespace {
@@ -300,6 +303,41 @@ void TransactionLock() {
     Require(excluded, "a transaction ran beside one that had read prepared state");
 }
 
+void PreparationPriority(bool beforeSubmit) {
+    const auto workers = std::max(2u, std::thread::hardware_concurrency() / 2u);
+    std::barrier started(static_cast<std::ptrdiff_t>(workers + 1u));
+    std::counting_semaphore<> release(0);
+    std::latch finished(static_cast<std::ptrdiff_t>(workers));
+    for (unsigned index = 0; index < workers; ++index) {
+        SubmitRegistrationPreparation([&] {
+            started.arrive_and_wait();
+            release.acquire();
+            finished.count_down();
+        });
+    }
+    started.arrive_and_wait();
+    PreparedShaders prepared;
+    std::atomic<unsigned> order{0};
+    std::promise<unsigned> background;
+    std::promise<unsigned> demanded;
+    auto backgroundDone = background.get_future();
+    auto demandedDone = demanded.get_future();
+    SubmitRegistrationPreparation([&] { background.set_value(order.fetch_add(1u)); });
+    if (beforeSubmit) PrioritizeRegistrationPreparation(prepared);
+    SubmitRegistrationPreparation([&] { demanded.set_value(order.fetch_add(1u)); }, &prepared);
+    if (!beforeSubmit) {
+        PrioritizeRegistrationPreparation(prepared);
+        PrioritizeRegistrationPreparation(prepared);
+    }
+    release.release();
+    const bool settled = demandedDone.wait_for(std::chrono::seconds(5)) == std::future_status::ready;
+    const bool backgroundSettled = backgroundDone.wait_for(std::chrono::seconds(5)) == std::future_status::ready;
+    release.release(static_cast<std::ptrdiff_t>(workers - 1u));
+    finished.wait();
+    Require(settled && backgroundSettled, "preparation queue did not complete its work");
+    Require(demandedDone.get() == 0u && backgroundDone.get() == 1u && order.load() == 2u, "demanded preparation waited behind background work or ran twice");
+}
+
 }
 
 int main() {
@@ -308,6 +346,8 @@ int main() {
         HeaderSelection();
         FusedHullHeaderSelection();
         TransactionLock();
+        PreparationPriority(false);
+        PreparationPriority(true);
         auto device = OpenVulkanTestDevice();
         if (!device) return VulkanTestSkipped;
         Registration(*device);