shader_recompiler: linear-time breadth-first search; cache storage buffer tracking

Translating title 1's 28,102-instruction compute shader took 3.4 s on
the GPU thread (a freeze the first time it runs), three quarters of it
in GlobalMemoryToStorageBufferPass:
- BreadthFirstSearch checked "visited" with a linear scan per argument,
  quadratic once a search walks thousands of instructions; past 64 it
  now uses a hash set (same visit order, same result).
- The pass searched again for every global memory access; accesses
  through the same address now share the result (the IR does not
  change while storage buffers are collected).
That shader now takes 0.85 s; other new pipelines also cost less.

2 files changed, 40 insertions, 4 deletions

diff --git a/src/shader_recompiler/frontend/ir/breadth_first_search.h b/src/shader_recompiler/frontend/ir/breadth_first_search.h
index afe153f180..ad19917781 100644
--- a/src/shader_recompiler/frontend/ir/breadth_first_search.h
+++ b/src/shader_recompiler/frontend/ir/breadth_first_search.h
@@ -6,10 +6,12 @@
 
 #pragma once
 
+#include <algorithm>
 #include <functional>
 #include <optional>
 #include <type_traits>
 #include <queue>
+#include <unordered_set>
 
 #include <boost/container/small_vector.hpp>
 
@@ -28,6 +30,23 @@ auto BreadthFirstSearch(Value value, Pred&& pred, Resolve&& resolve = {})
     // Breadth-first search visiting the right most arguments first
     // Small vector has been determined from shaders in Super Smash Bros. Ultimate
     boost::container::small_vector<const Inst*, 2> visited;
+    // Once a search grows, a linear scan of `visited` per argument made it quadratic: a 28,000
+    // instruction Switch 2 compute shader took 3.4 s to translate, most of it here. Same visit
+    // order, so the same result.
+    std::unordered_set<const Inst*> visited_set;
+    const auto first_visit = [&](const Inst* inst) {
+        if (!visited_set.empty()) {
+            return visited_set.insert(inst).second;
+        }
+        if (std::ranges::find(visited, inst) != visited.end()) {
+            return false;
+        }
+        visited.push_back(inst);
+        if (visited.size() > 64) {
+            visited_set.insert(visited.begin(), visited.end());
+        }
+        return true;
+    };
     std::queue<const Inst*> queue;
     queue.push(value.InstRecursive());
 
@@ -47,8 +66,7 @@ auto BreadthFirstSearch(Value value, Pred&& pred, Resolve&& resolve = {})
             }
             // Queue instruction if it hasn't been visited
             const Inst* const arg_inst{arg_value.InstRecursive()};
-            if (std::ranges::find(visited, arg_inst) == visited.end()) {
-                visited.push_back(arg_inst);
+            if (first_visit(arg_inst)) {
                 queue.push(arg_inst);
             }
         }
diff --git a/src/shader_recompiler/ir_opt/global_memory_to_storage_buffer_pass.cpp b/src/shader_recompiler/ir_opt/global_memory_to_storage_buffer_pass.cpp
index f49d54e7ab..bdb4494a15 100644
--- a/src/shader_recompiler/ir_opt/global_memory_to_storage_buffer_pass.cpp
+++ b/src/shader_recompiler/ir_opt/global_memory_to_storage_buffer_pass.cpp
@@ -5,6 +5,7 @@
 // SPDX-License-Identifier: GPL-2.0-or-later
 
 #include <optional>
+#include <unordered_map>
 
 #include <boost/container/flat_set.hpp>
 #include <boost/container/small_vector.hpp>
@@ -54,6 +55,11 @@ struct StorageInfo {
     StorageBufferSet set;
     StorageInstVector to_replace;
     StorageWritesSet writes;
+    /// Track results by the instruction the search starts from, with and without the NVN bias:
+    /// accesses through the same address share them, and a search can walk thousands of
+    /// instructions (the IR does not change while storage buffers are collected).
+    std::unordered_map<const IR::Inst*, std::optional<StorageBufferAddr>> biased_tracks;
+    std::unordered_map<const IR::Inst*, std::optional<StorageBufferAddr>> unbiased_tracks;
 };
 
 /// Returns true when the instruction is a global memory instruction
@@ -387,10 +393,22 @@ void CollectStorageBuffers(IR::Block& block, IR::Inst& inst, StorageInfo& info)
     }
     // First try to find storage buffers in the NVN address
     const IR::U32 low_addr{low_addr_info->value};
-    std::optional<StorageBufferAddr> storage_buffer{Track(low_addr, &nvn_bias)};
+    const auto cached_track = [&](const Bias* bias, auto& cache) {
+        if (low_addr.IsImmediate()) {
+            return Track(low_addr, bias);
+        }
+        const IR::Inst* const start{low_addr.InstRecursive()};
+        if (const auto it = cache.find(start); it != cache.end()) {
+            return it->second;
+        }
+        const std::optional<StorageBufferAddr> result{Track(low_addr, bias)};
+        cache.emplace(start, result);
+        return result;
+    };
+    std::optional<StorageBufferAddr> storage_buffer{cached_track(&nvn_bias, info.biased_tracks)};
     if (!storage_buffer) {
         // If it fails, track without a bias
-        storage_buffer = Track(low_addr, nullptr);
+        storage_buffer = cached_track(nullptr, info.unbiased_tracks);
         if (!storage_buffer) {
             // If that also fails, use NVN fallbacks
             LOG_WARNING(Shader, "Storage buffer failed to track, using global memory fallbacks");