fix(agc): bound the SPIR-V that background compute pipeline builds compile at once

core/libs/prx/libSceAgcDriver/CMakeLists.txt

@@ -594,6 +594,17 @@ if(BUILD_TESTING)
                 "PATH=path_list_prepend:$<TARGET_FILE_DIR:libc>")
     endif()
 
+    add_test_executable(agc_compile_budget_tests tests/CompileBudget.cpp)
+    target_include_directories(agc_compile_budget_tests PRIVATE ${LIBS_INCLUDE_DIR})
+    target_link_libraries(agc_compile_budget_tests PRIVATE libc)
+    configure_windows_unwind(agc_compile_budget_tests)
+    add_test(NAME agc_compile_budget COMMAND agc_compile_budget_tests)
+    set_tests_properties(agc_compile_budget PROPERTIES TIMEOUT 30)
+    if(WIN32)
+        set_tests_properties(agc_compile_budget PROPERTIES ENVIRONMENT_MODIFICATION
+                "PATH=path_list_prepend:$<TARGET_FILE_DIR:libc>")
+    endif()
+
     add_test_executable(agc_capture_trace_tests tests/CaptureTrace.cpp)
     target_include_directories(agc_capture_trace_tests PRIVATE ${LIBS_INCLUDE_DIR})
     target_link_libraries(agc_capture_trace_tests PRIVATE libc)
@@ -3272,6 +3283,7 @@ if(BUILD_TESTING)
             "PATH=path_list_prepend:$<TARGET_FILE_DIR:libc>")
     endif()
     set_tests_properties(agc_driver_queued_index_read_order PROPERTIES TIMEOUT 60)
+    set_property(TEST agc_driver_compute_prebuild_share APPEND PROPERTY ENVIRONMENT "APS5_PREBUILD_THREADS=4")
 
     foreach(agcTest IN ITEMS agc_spirv_narrow_constant_store_tests agc_spirv13_image_tests agc_spirv_optimizer_tests agc_driver_constant_store_tests)
         if(NOT TARGET ${agcTest})

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

@@ -0,0 +1,77 @@
+#ifndef CORE_LIBS_PRX_LIBSCEAGCDRIVER_EXECUTION_INCLUDE_COMPILEBUDGET_HPP
+#define CORE_LIBS_PRX_LIBSCEAGCDRIVER_EXECUTION_INCLUDE_COMPILEBUDGET_HPP
+
+#include <condition_variable>
+#include <cstdint>
+#include <cstdlib>
+#include <mutex>
+
+namespace AgcDriver {
+
+class CompileBudget {
+public:
+    class Hold {
+    public:
+        Hold(CompileBudget& budget, std::uint64_t bytes) : budget(budget), bytes(bytes) {
+            std::unique_lock lock(budget.mutex);
+            if (bytes > budget.limit) {
+                ++budget.waitingAlone;
+                budget.released.wait(lock, [&] { return budget.held == 0; });
+                --budget.waitingAlone;
+            } else {
+                budget.released.wait(lock, [&] { return budget.waitingAlone == 0 && budget.held + bytes <= budget.limit; });
+            }
+            budget.held += bytes;
+            if (budget.held > budget.peak) budget.peak = budget.held;
+        }
+
+        ~Hold() {
+            {
+                std::lock_guard lock(budget.mutex);
+                budget.held -= bytes;
+            }
+            budget.released.notify_all();
+        }
+
+        Hold(const Hold&) = delete;
+        Hold& operator=(const Hold&) = delete;
+
+    private:
+        CompileBudget& budget;
+        std::uint64_t bytes;
+    };
+
+    explicit CompileBudget(std::uint64_t limit) : limit(limit) {}
+
+    std::uint64_t Peak() const {
+        std::lock_guard lock(mutex);
+        return peak;
+    }
+
+    std::uint64_t WaitingAlone() const {
+        std::lock_guard lock(mutex);
+        return waitingAlone;
+    }
+
+private:
+    mutable std::mutex mutex;
+    std::condition_variable released;
+    std::uint64_t limit;
+    std::uint64_t held = 0;
+    std::uint64_t peak = 0;
+    std::uint64_t waitingAlone = 0;
+};
+
+inline std::uint64_t BackgroundCompileLimitFor(const char* configured) {
+    constexpr std::uint64_t automaticKiB = 8192;
+    return (configured != nullptr ? std::strtoull(configured, nullptr, 10) : automaticKiB) * 1024u;
+}
+
+inline CompileBudget& BackgroundCompileBudget() {
+    static CompileBudget budget(BackgroundCompileLimitFor(std::getenv("APS5_BACKGROUND_COMPILE_KIB")));
+    return budget;
+}
+
+}
+
+#endif

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

@@ -224,6 +224,7 @@ void SubmitRegistrationPreparation(std::function<void()> task, const PreparedSha
 void PrioritizeRegistrationPreparation(PreparedShaders& prepared);
 void AwaitRegistrationPreparation();
 void SubmitLibraryPrebuild(std::function<void()> task);
+void SubmitComputePrebuild(std::function<void()> task);
 void AwaitLibraryPrebuilds();
 void CloseRegistrationPreparation();
 unsigned RegistrationPreparationThreads();

core/libs/prx/libSceAgcDriver/Execution/src/Driver/Dispatch/Dispatch.cpp

@@ -437,7 +437,7 @@ void Driver::dispatch(QueueState& queue, std::span<const std::uint32_t> packet,
         ++groupStats.resolved;
         if (cached) ++groupStats.resolvedHits;
         if (!cached && compiledResult != nullptr && !localDevice->HasComputePipeline(*compiledResult)) {
-            SubmitLibraryPrebuild([localDevice, compiledResult, address] {
+            SubmitComputePrebuild([localDevice, compiledResult, address] {
                 try { localDevice->PrebuildComputePipeline(*compiledResult, address); } catch (const std::exception&) {}
             });
         }

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

@@ -1544,13 +1544,13 @@ class RegistrationPreparePool {
 
 public:
 
-    void Submit(std::function<void()> task, const PreparedShaders* prepared = nullptr) {
+    void Submit(std::function<void()> task, const PreparedShaders* prepared = nullptr, bool budgeted = false) {
         {
             std::unique_lock lock(mutex);
             if (threads != 0 && !closed) {
                 const auto sequence = ++submitted;
-                if (prepared != nullptr && prepared->demanded.load()) tasks.push_front({std::move(task), prepared, sequence});
-                else tasks.push_back({std::move(task), prepared, sequence});
+                if (prepared != nullptr && prepared->demanded.load()) tasks.push_front({std::move(task), prepared, sequence, budgeted});
+                else tasks.push_back({std::move(task), prepared, sequence, budgeted});
                 unfinished.insert(sequence);
                 lock.unlock();
                 ready.notify_one();
@@ -1638,6 +1638,7 @@ class RegistrationPreparePool {
         }
         pthread_attr_destroy(&attributes);
 #endif
+        budgetedLimit = std::max(1u, threads / 2u);
     }
 
 #ifdef _WIN32
@@ -1654,26 +1655,34 @@ class RegistrationPreparePool {
             Task task;
             {
                 std::unique_lock lock(pool.mutex);
-                pool.ready.wait(lock, [&] { return !pool.tasks.empty() || pool.closed; });
-                if (pool.tasks.empty()) {
+                auto next = pool.tasks.end();
+                pool.ready.wait(lock, [&] {
+                    next = std::ranges::find_if(pool.tasks, [&](const Task& queued) { return !queued.budgeted || pool.budgetedRunning < pool.budgetedLimit; });
+                    return next != pool.tasks.end() || (pool.closed && pool.tasks.empty());
+                });
+                if (next == pool.tasks.end()) {
                     liveWorkers.fetch_sub(1, std::memory_order_acq_rel);
                     return nullptr;
                 }
-                task = std::move(pool.tasks.front());
-                pool.tasks.pop_front();
+                task = std::move(*next);
+                pool.tasks.erase(next);
                 ++pool.running;
+                if (task.budgeted) ++pool.budgetedRunning;
             }
             const auto sequence = task.sequence;
+            const bool budgeted = task.budgeted;
             task.function();
             task = {};
             bool drained = false;
             {
                 std::lock_guard lock(pool.mutex);
                 --pool.running;
+                if (budgeted) --pool.budgetedRunning;
                 pool.unfinished.erase(sequence);
                 drained = pool.tasks.empty() && pool.running == 0;
                 pool.idle.notify_all();
             }
+            if (budgeted) pool.ready.notify_all();
 #ifdef __GLIBC__
             if (drained) trimFreedMemory();
 #endif
@@ -1685,6 +1694,7 @@ class RegistrationPreparePool {
         std::function<void()> function;
         const PreparedShaders* prepared = nullptr;
         std::uint64_t sequence = 0;
+        bool budgeted = false;
     };
 
     std::mutex mutex;
@@ -1695,6 +1705,8 @@ class RegistrationPreparePool {
     unsigned threads = 0;
     inline static std::atomic<unsigned> liveWorkers{0};
     unsigned running = 0;
+    unsigned budgetedRunning = 0;
+    unsigned budgetedLimit = 1;
     std::uint64_t submitted = 0;
     std::set<std::uint64_t> unfinished;
     bool closed = false;
@@ -1750,6 +1762,10 @@ void SubmitLibraryPrebuild(std::function<void()> task) {
     RegistrationPreparePool::Prebuild().Submit(std::move(task));
 }
 
+void SubmitComputePrebuild(std::function<void()> task) {
+    RegistrationPreparePool::Prebuild().Submit(std::move(task), nullptr, true);
+}
+
 void AwaitLibraryPrebuilds() {
     if (auto* pool = RegistrationPreparePool::ExistingPrebuild()) pool->AwaitSubmitted();
 }

core/libs/prx/libSceAgcDriver/Execution/src/VulkanDevice.cpp

@@ -3,6 +3,7 @@
 #include "prx/libSceAgcDriver/Graphics/include/PipelineSpecialization.hpp"
 #include "prx/libSceAgcDriver/Execution/include/CaptureTrace.hpp"
 #include "prx/libSceAgcDriver/Execution/include/Mutex.hpp"
+#include "prx/libSceAgcDriver/Execution/include/CompileBudget.hpp"
 #include "BdaAbi.hpp"
 #if ANYPS5_ENABLE_SPIRV_TOOLS
 #include "SpirvBackend/SpirvOptimizer.hpp"
@@ -2329,6 +2330,11 @@ bool VulkanDevice::ReplayComputePipeline(const ShaderRecompiler::RecompileResult
         std::lock_guard pipelines(state->computePipelinesMutex);
         if (state->computePipelines.contains(key) || state->computePipelinesBuilding.contains(key)) return false;
     }
+    const CompileBudget::Hold hold(BackgroundCompileBudget(), shader.spirv.size() * sizeof(std::uint32_t));
+    {
+        std::lock_guard pipelines(state->computePipelinesMutex);
+        if (state->computePipelines.contains(key) || state->computePipelinesBuilding.contains(key)) return false;
+    }
     const auto context = graphicsContext();
     const std::array<Graphics::CompiledShader, 1> shaders{{{ShaderRecompiler::ShaderStage::Compute, &shader, 0}}};
     const auto layouts = Graphics::DescribeSetLayouts(context, shaders);
@@ -2353,6 +2359,8 @@ void VulkanDevice::PrebuildComputePipeline(const ShaderRecompiler::RecompileResu
     const auto pushStages = Graphics::PushConstantStages(shaders);
     if (pushStages != 0 && state->properties.limits.maxPushConstantsSize < Graphics::ComputePushConstantBytes) return;
     if (HasComputePipeline(shader)) return;
+    const CompileBudget::Hold hold(BackgroundCompileBudget(), shader.spirv.size() * sizeof(std::uint32_t));
+    if (HasComputePipeline(shader)) return;
     const auto layouts = Graphics::DescribeSetLayouts(context, shaders);
     static_cast<void>(computePipeline(context, shader, layouts.layouts.main, layouts.mainKey, pushStages, programAddress));
 }

core/libs/prx/libSceAgcDriver/tests/CompileBudget.cpp

@@ -0,0 +1,112 @@
+#include "prx/libSceAgcDriver/Execution/include/CompileBudget.hpp"
+#include <atomic>
+#include <chrono>
+#include <cstdint>
+#include <iostream>
+#include <stdexcept>
+#include <string>
+#include <thread>
+#include <vector>
+
+namespace {
+
+using namespace std::chrono_literals;
+using AgcDriver::BackgroundCompileLimitFor;
+using AgcDriver::CompileBudget;
+
+void Require(bool condition, const std::string& message) {
+    if (!condition) throw std::runtime_error(message);
+}
+
+void CheckLimitPolicy() {
+    Require(BackgroundCompileLimitFor(nullptr) == 8192u * 1024u, "an unset APS5_BACKGROUND_COMPILE_KIB is not 8 MiB");
+    Require(BackgroundCompileLimitFor("16") == 16u * 1024u, "APS5_BACKGROUND_COMPILE_KIB=16 is not 16 KiB");
+    Require(BackgroundCompileLimitFor("0") == 0u, "APS5_BACKGROUND_COMPILE_KIB=0 leaves room");
+}
+
+std::uint64_t MostHeldAtOnce(CompileBudget& budget, const std::vector<std::uint64_t>& sizes) {
+    std::atomic<std::uint64_t> held{0};
+    std::atomic<std::uint64_t> most{0};
+    std::vector<std::thread> threads;
+    for (const auto bytes : sizes) {
+        threads.emplace_back([&, bytes] {
+            const CompileBudget::Hold hold(budget, bytes);
+            const auto now = held.fetch_add(bytes) + bytes;
+            auto seen = most.load();
+            while (seen < now && !most.compare_exchange_weak(seen, now)) {}
+            std::this_thread::sleep_for(20ms);
+            held.fetch_sub(bytes);
+        });
+    }
+    for (auto& thread : threads) thread.join();
+    return most.load();
+}
+
+void CheckSharedLimit() {
+    CompileBudget budget(100);
+    const auto most = MostHeldAtOnce(budget, std::vector<std::uint64_t>(8, 40));
+    Require(most <= 80, "eight 40-byte holds under a 100-byte budget held " + std::to_string(most) + " bytes at once");
+    Require(budget.Peak() >= 40 && budget.Peak() <= 80, "eight 40-byte holds under a 100-byte budget peaked at " + std::to_string(budget.Peak()) + " bytes");
+}
+
+void CheckZeroLimitRunsOneAtATime() {
+    CompileBudget budget(0);
+    const auto most = MostHeldAtOnce(budget, std::vector<std::uint64_t>(6, 1));
+    Require(most == 1 && budget.Peak() == 1, "a zero budget let " + std::to_string(most) + " one-byte holds run at once");
+}
+
+void CheckLargerHoldRunsAlone() {
+    CompileBudget budget(100);
+    const auto most = MostHeldAtOnce(budget, {500, 10, 10, 10});
+    Require(most == 500 && budget.Peak() == 500, "a 500-byte hold under a 100-byte budget did not run alone: " + std::to_string(most) + " bytes were held at once");
+}
+
+void CheckWaitingLargerHoldKeepsNewHoldsOut() {
+    CompileBudget budget(100);
+    std::atomic<bool> firstHeld{false};
+    std::atomic<bool> releaseFirst{false};
+    std::atomic<int> entries{0};
+    std::atomic<int> largeEntry{-1};
+    std::atomic<int> smallEntry{-1};
+    std::thread first([&] {
+        const CompileBudget::Hold hold(budget, 40);
+        firstHeld = true;
+        while (!releaseFirst) std::this_thread::sleep_for(1ms);
+    });
+    while (!firstHeld) std::this_thread::sleep_for(1ms);
+    std::thread large([&] {
+        const CompileBudget::Hold hold(budget, 500);
+        largeEntry = entries++;
+        std::this_thread::sleep_for(10ms);
+    });
+    while (budget.WaitingAlone() == 0) std::this_thread::sleep_for(1ms);
+    std::thread small([&] {
+        const CompileBudget::Hold hold(budget, 40);
+        smallEntry = entries++;
+    });
+    std::this_thread::sleep_for(20ms);
+    const bool enteredEarly = smallEntry != -1;
+    releaseFirst = true;
+    first.join();
+    large.join();
+    small.join();
+    Require(!enteredEarly, "a 40-byte hold that fits entered while a 500-byte hold was waiting to run alone");
+    Require(largeEntry == 0 && smallEntry == 1, "the waiting 500-byte hold did not run before the 40-byte hold that came after it");
+}
+
+}
+
+int main() {
+    try {
+        CheckLimitPolicy();
+        CheckSharedLimit();
+        CheckZeroLimitRunsOneAtATime();
+        CheckLargerHoldRunsAlone();
+        CheckWaitingLargerHoldKeepsNewHoldsOut();
+        std::cout << "compile budget tests passed\n";
+        return 0;
+    } catch (const std::exception& error) {
+        std::cerr << error.what() << '\n';
+        return 1;
+    }
+}