From 751e3356e1b2cb5421ad1ba7e4cfa4c4564915f4 Mon Sep 17 00:00:00 2001
From: Billy Laws <blaws05@gmail.com>
Date: Sun, 9 Oct 2022 13:52:40 +0100
Subject: [PATCH] Keep shader trap lock held for the duration of an execution

Avoids constant relocking on the GPFIFO thread (~0.5% of total time)
---
 .../maxwell_3d/pipeline_state.cpp             | 25 +++++++++++++++----
 .../interconnect/maxwell_3d/pipeline_state.h  |  8 ++++--
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/app/src/main/cpp/skyline/gpu/interconnect/maxwell_3d/pipeline_state.cpp b/app/src/main/cpp/skyline/gpu/interconnect/maxwell_3d/pipeline_state.cpp
index b3ef0875..491ab8d7 100644
--- a/app/src/main/cpp/skyline/gpu/interconnect/maxwell_3d/pipeline_state.cpp
+++ b/app/src/main/cpp/skyline/gpu/interconnect/maxwell_3d/pipeline_state.cpp
@@ -236,7 +236,7 @@ namespace skyline::gpu::interconnect::maxwell3d {
             return;
         }
 
-        auto [blockMapping, blockOffset]{ctx.channelCtx.asCtx->gmmu.LookupBlock(engine->programRegion + engine->pipeline.programOffset)};
+        auto[blockMapping, blockOffset]{ctx.channelCtx.asCtx->gmmu.LookupBlock(engine->programRegion + engine->pipeline.programOffset)};
 
         // Skip looking up the mirror if it is the same as the one used for the previous update
         if (!mirrorBlock.valid() || !mirrorBlock.contains(blockMapping)) {
@@ -246,8 +246,13 @@ namespace skyline::gpu::interconnect::maxwell3d {
                 auto newIt{mirrorMap.emplace(blockMapping.data(), std::make_unique<MirrorEntry>(ctx.memory.CreateMirror(blockMapping)))};
 
                 // We need to create the trap after allocating the entry so that we have an `invalid` pointer we can pass in
-                auto trapHandle{ctx.nce.CreateTrap(blockMapping, [](){}, [](){ return true; }, [dirty = &newIt.first->second->dirty, mutex = &trapMutex](){
-                    std::scoped_lock lock{*mutex}; // Don't use lock callback here since we need trapMutex to be always locked on accesses to prevent UAFs
+                auto trapHandle{ctx.nce.CreateTrap(blockMapping, [mutex = &trapMutex]() {
+                    std::scoped_lock lock{*mutex};
+                    return;
+                }, []() { return true; }, [dirty = &newIt.first->second->dirty, mutex = &trapMutex]() {
+                    std::unique_lock lock{*mutex, std::try_to_lock};
+                    if (!lock)
+                        return false;
                     *dirty = true;
                     return true;
                 })};
@@ -264,6 +269,9 @@ namespace skyline::gpu::interconnect::maxwell3d {
             mirrorBlock = blockMapping;
         }
 
+        if (!trapExecutionLock)
+            trapExecutionLock.emplace(trapMutex);
+
         // If the mirror entry has been written to, clear its shader binary cache and retrap to catch any future writes
         if (entry->dirty) {
             entry->cache.clear();
@@ -300,15 +308,20 @@ namespace skyline::gpu::interconnect::maxwell3d {
         entry->cache.insert({blockMapping.data() + blockOffset, CacheEntry{binary, hash}});
     }
 
-    // TODO: this needs to be checked every draw irresspective of pipeline dirtiness
     bool PipelineStageState::Refresh(InterconnectContext &ctx) {
-        std::scoped_lock lock{trapMutex};
+        if (!trapExecutionLock)
+            trapExecutionLock.emplace(trapMutex);
+
         if (entry && entry->dirty)
             return true;
 
         return false;
     }
 
+    void PipelineStageState::PurgeCaches() {
+        trapExecutionLock.reset();
+    }
+
     PipelineStageState::~PipelineStageState() {
         std::scoped_lock lock{trapMutex};
         //for (const auto &mirror : mirrorMap)
@@ -575,6 +588,8 @@ namespace skyline::gpu::interconnect::maxwell3d {
 
     void PipelineState::PurgeCaches() {
         pipeline = nullptr;
+        for (auto &stage : pipelineStages)
+            stage.MarkDirty(true);
     }
 
     std::shared_ptr<TextureView> PipelineState::GetColorRenderTargetForClear(InterconnectContext &ctx, size_t index) {
diff --git a/app/src/main/cpp/skyline/gpu/interconnect/maxwell_3d/pipeline_state.h b/app/src/main/cpp/skyline/gpu/interconnect/maxwell_3d/pipeline_state.h
index 82b0314c..7f33262b 100644
--- a/app/src/main/cpp/skyline/gpu/interconnect/maxwell_3d/pipeline_state.h
+++ b/app/src/main/cpp/skyline/gpu/interconnect/maxwell_3d/pipeline_state.h
@@ -55,7 +55,7 @@ namespace skyline::gpu::interconnect::maxwell3d {
         void Flush(InterconnectContext &ctx, PackedPipelineState &packedState);
     };
 
-    class PipelineStageState : dirty::RefreshableManualDirty {
+    class PipelineStageState : dirty::RefreshableManualDirty, dirty::CachedManualDirty {
       public:
         struct EngineRegisters {
             const engine::Pipeline &pipeline;
@@ -80,6 +80,7 @@ namespace skyline::gpu::interconnect::maxwell3d {
             tsl::robin_map<u8 *, CacheEntry> cache;
             std::optional<nce::NCE::TrapHandle> trap;
             bool dirty{};
+
             MirrorEntry(span<u8> alignedMirror) : mirror{alignedMirror} {}
         };
 
@@ -87,7 +88,8 @@ namespace skyline::gpu::interconnect::maxwell3d {
         engine::Pipeline::Shader::Type shaderType;
 
         tsl::robin_map<u8 *, std::unique_ptr<MirrorEntry>> mirrorMap;
-        SpinLock trapMutex; //!< Protects accesses from trap handlers to the mirror map
+        std::mutex trapMutex; //!< Protects accesses from trap handlers to the mirror map
+        std::optional<std::scoped_lock<std::mutex>> trapExecutionLock;
         MirrorEntry *entry{};
         span<u8> mirrorBlock{}; //!< Guest mapped memory block corresponding to `entry`
 
@@ -102,6 +104,8 @@ namespace skyline::gpu::interconnect::maxwell3d {
         void Flush(InterconnectContext &ctx);
 
         bool Refresh(InterconnectContext &ctx);
+
+        void PurgeCaches();
     };
 
     class VertexInputState : dirty::ManualDirty {