perf(agc-driver): prepare shader ABI variants before taking the preparation lock

A shader preparation transaction took the global writers lock when it
was constructed and kept it to its end. ResolveShaderAbi, reached from
sceAgcCreateInterpolantMapping and sceAgcLinkShaders, compiled its
variants inside such a transaction, so every game thread creating or
linking shaders waited behind one compile: Uncharted's stalled frames
spent about half a second each in shader registration.

The root transaction now takes the lock on its first Read, Edit or
Hold. ResolveShaderAbi prepares the variants before reading the
prepared state when its ABI was not committed yet. Registry
publication holds the lock as before: PublishRegisteredShader holds
it, and RegisterShader holds it before the driver mutex, so the lock
order is unchanged. Commit takes no lock: a transaction with changes
already holds it.

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

@@ -16,6 +16,7 @@ public:
     ShaderPreparationTransaction& operator=(const ShaderPreparationTransaction&) = delete;
     PreparedShaderState& Edit(const ShaderSnapshot& snapshot);
     const PreparedShaderState& Read(const ShaderSnapshot& snapshot) const;
+    void Hold();
     void Commit();
 
 private:

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

@@ -130,6 +130,12 @@ void AwaitRegisteredPreparation(PreparedShaders& prepared, std::unique_lock<std:
     if (prepared.failure != nullptr) std::rethrow_exception(prepared.failure);
 }
 
+bool CommittedAbi(const ShaderSnapshot& snapshot, std::span<const std::uint64_t> key) {
+    std::unique_lock lock(snapshot.prepared->mutex);
+    AwaitRegisteredPreparation(*snapshot.prepared, lock);
+    return std::ranges::any_of(snapshot.prepared->registeredAbis, [&](const auto& abi) { return std::ranges::equal(abi, key); });
+}
+
 }
 
 bool ProgramUsesWorkgroup(const ShaderSnapshot& snapshot, std::size_t codeOffset) {
@@ -157,9 +163,13 @@ struct ShaderPreparationTransaction::State {
     };
     static std::mutex writers;
     static thread_local State* active;
-    std::unique_lock<std::mutex> lock{writers};
+    std::unique_lock<std::mutex> lock{writers, std::defer_lock};
     std::map<PreparedShaders*, Change> changes;
     bool failed = false;
+
+    void Hold() {
+        if (!lock.owns_lock()) lock.lock();
+    }
 };
 
 std::mutex ShaderPreparationTransaction::State::writers;
@@ -180,6 +190,7 @@ ShaderPreparationTransaction::~ShaderPreparationTransaction() {
 
 PreparedShaderState& ShaderPreparationTransaction::Edit(const ShaderSnapshot& snapshot) {
     require(!committed, "shader preparation transaction is already committed");
+    root->Hold();
     auto& changes = root->changes;
     const auto found = changes.find(snapshot.prepared.get());
     if (found != changes.end()) {
@@ -193,6 +204,7 @@ PreparedShaderState& ShaderPreparationTransaction::Edit(const ShaderSnapshot& sn
 
 const PreparedShaderState& ShaderPreparationTransaction::Read(const ShaderSnapshot& snapshot) const {
     require(!committed, "shader preparation transaction is already committed");
+    root->Hold();
     auto& changes = root->changes;
     const auto found = changes.find(snapshot.prepared.get());
     if (found != changes.end()) return found->second.prepared;
@@ -201,6 +213,11 @@ const PreparedShaderState& ShaderPreparationTransaction::Read(const ShaderSnapsh
     return changes.emplace(snapshot.prepared.get(), State::Change{snapshot.prepared, *snapshot.prepared, false, snapshot.prepared->entries.size(), snapshot.prepared->registeredAbis.size()}).first->second.prepared;
 }
 
+void ShaderPreparationTransaction::Hold() {
+    require(!committed, "shader preparation transaction is already committed");
+    root->Hold();
+}
+
 void ShaderPreparationTransaction::Commit() {
     require(!committed && !root->failed, "shader preparation transaction was aborted");
     if (state != nullptr) {
@@ -294,6 +311,7 @@ std::shared_ptr<const ShaderSnapshot> RegisteredProgram(const ShaderRegistry& re
 
 void PublishRegisteredShader(std::shared_ptr<ShaderRegistry>& registry, const std::shared_ptr<const ShaderSnapshot>& snapshot) {
     ShaderPreparationTransaction transaction;
+    transaction.Hold();
     auto published = snapshot;
     if (registry != nullptr) {
         const auto found = registry->find(snapshot->codeAddress);
@@ -935,12 +953,14 @@ void Driver::ResolveShaderAbi(const Shader* shader, std::span<const ShaderRegist
     struct RegisteredAbiKeyStorage {};
     auto& key = HostThreadLocal<std::vector<std::uint64_t>, RegisteredAbiKeyStorage>();
     BuildRegisteredAbiKey(state, *localDevice, key);
+    const bool committedAbi = CommittedAbi(*snapshot, key);
+    auto entries = committedAbi ? std::vector<PreparedShaders::Entry>{} : PrepareRegistered(*snapshot, *localDevice, state, false);
     const auto& registeredAbis = transaction.Read(*snapshot).registeredAbis;
     if (std::ranges::find(registeredAbis, key) != registeredAbis.end()) {
         transaction.Commit();
         return;
     }
-    auto entries = PrepareRegistered(*snapshot, *localDevice, state, false);
+    if (committedAbi) entries = PrepareRegistered(*snapshot, *localDevice, state, false);
     auto& prepared = transaction.Edit(*snapshot);
     for (auto& entry : entries) {
         const auto duplicate = std::ranges::any_of(prepared.entries, [&](const auto& existing) { return existing.handle->artifact == entry.handle->artifact; });
@@ -1094,6 +1114,7 @@ void Driver::RegisterShader(const Shader* shader) {
     snapshot.header.resize(fields.header_size);
     std::memcpy(snapshot.header.data(), static_cast<const void*>(shader), fields.header_size);
 
+    transaction.Hold();
     {
         std::lock_guard lock(mutex);
         rethrowFailure();

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

@@ -10,6 +10,7 @@
 #include <future>
 #include <cstring>
 #include <iostream>
+#include <optional>
 #include <thread>
 
 namespace {
@@ -277,6 +278,28 @@ void Registration(AgcDriver::VulkanDevice& device) {
     Require(registry->at(0x10000).back() == metadata, "changed shader metadata was treated as identical");
 }
 
+void TransactionLock() {
+    ShaderSnapshot shader{};
+    const auto publish = [&] {
+        ShaderPreparationTransaction other;
+        other.Edit(shader).rectangleRequested = true;
+        other.Commit();
+    };
+    std::optional<ShaderPreparationTransaction> open(std::in_place);
+    auto unread = std::async(std::launch::async, publish);
+    const bool unblocked = unread.wait_for(std::chrono::seconds(10)) == std::future_status::ready;
+    if (!unblocked) open.reset();
+    unread.get();
+    Require(unblocked, "an open transaction that read nothing blocked another one");
+    static_cast<void>(open->Read(shader));
+    auto waiting = std::async(std::launch::async, publish);
+    const bool excluded = waiting.wait_for(std::chrono::milliseconds(200)) == std::future_status::timeout;
+    open->Commit();
+    open.reset();
+    waiting.get();
+    Require(excluded, "a transaction ran beside one that had read prepared state");
+}
+
 }
 
 int main() {
@@ -284,6 +307,7 @@ int main() {
         Transactions();
         HeaderSelection();
         FusedHullHeaderSelection();
+        TransactionLock();
         auto device = OpenVulkanTestDevice();
         if (!device) return VulkanTestSkipped;
         Registration(*device);