fix(shader): bound the bytes the recompiler's result memos hold

Every shader source keeps up to 256 materialized results of
Recompile(request, capture) in its result memo, and as many patch
recipes in its masked memo, with no limit on their bytes. In Ghost of
Tsushima's opening the memos held about 500 MiB after 5 minutes and
kept growing, at a 13 % hit rate.

Each memo entry now counts the bytes its result's vectors keep (the
bindings' vectors included; the module is shared with the prepared
module cache), plus its patch recipe in the masked memo. A result held
by both tables counts once in each. Past a budget, 256 MiB unless
APS5_RESULT_MEMO_MB sets another (0 sets none), the least recently used
entries of both tables of all sources are dropped until three quarters
of the budget is left. The recency stamps come from one clock for all
sources and both tables so they compare across them. The profile line
also reports the trimmed and held entries and their MiB.

The trim takes its lock blocking and repeats, re-reading the bytes held,
until they are within the budget, so the insertion that pushed them
over has its trim run after any other; it stops when nothing is left to
drop.

SetResultMemoBudget and ResultMemoBytes let agc_shader_result_memo_budget
check that a memo over its budget keeps the most recently used results
and drops the oldest one, with the masked memo off. A one-shot hook
inside the trim (SetResultMemoTrimHook) lets
agc_shader_result_memo_budget_threads insert two results from other
threads after the trim has read the bytes held: the memo must still end
within its budget. agc_shader_result_memo_budget_masked, with the masked
memo on, checks that a first result counts its masked entry too and that
a trim empties the masked memo.

Co-Authored-By: Claude Opus 5.5 <[email protected]>

core/shader/recompiler/Recompiler.cpp

@@ -369,6 +369,7 @@ struct ResultMemoEntry {
     std::shared_ptr<const RecompileResult> result;
     std::uint64_t used = 0;
     std::shared_ptr<const MemoPatch> patch;
+    std::uint64_t bytes = 0;
 };
 
 struct EmissionFailure {
@@ -399,7 +400,6 @@ struct SourceEntry : std::enable_shared_from_this<SourceEntry> {
     std::vector<EmissionFailure> emissionFailures;
     std::unique_ptr<ResultMemoTable> memo;
     std::unique_ptr<ResultMemoTable> maskedMemo;
-    std::uint64_t memoClock = 0;
 };
 
 namespace {
@@ -1455,8 +1455,9 @@ ResultMemoEntry* memoSet(const std::unique_ptr<ResultMemoTable>& table, std::uin
 }
 
 struct ResultMemoCounters {
-    std::atomic<std::uint64_t> hits{0}, patched{0}, misses{0}, evictions{0}, populateNanoseconds{0}, patchNanoseconds{0}, mismatches{0};
+    std::atomic<std::uint64_t> hits{0}, patched{0}, misses{0}, evictions{0}, trimmed{0}, populateNanoseconds{0}, patchNanoseconds{0}, mismatches{0};
     std::atomic<std::int64_t> lastReport{0};
+    std::atomic<std::uint64_t> clock{0}, heldBytes{0}, heldEntries{0};
 };
 
 void applyMemoPatch(RecompileResult& result, const MemoPatch& patch, const ResourceSnapshot& snapshot) {
@@ -1503,6 +1504,123 @@ ResultMemoCounters& resultMemoCounters() {
     return counters;
 }
 
+std::uint64_t nextMemoUse() {
+    return resultMemoCounters().clock.fetch_add(1, std::memory_order_relaxed) + 1u;
+}
+
+template <typename T>
+std::uint64_t vectorBytes(const std::vector<T>& values) {
+    return static_cast<std::uint64_t>(values.capacity()) * sizeof(T);
+}
+
+std::uint64_t vectorBytes(const std::vector<bool>& values) {
+    return static_cast<std::uint64_t>(values.capacity() + 7u) / 8u;
+}
+
+std::uint64_t descriptorBindingBytes(const DescriptorBinding& binding) {
+    return vectorBytes(binding.guestDescriptor) + vectorBytes(binding.imageShapes) + vectorBytes(binding.imageSamplers) + vectorBytes(binding.bufferExtent) + vectorBytes(binding.samplerDepthCompare) + vectorBytes(binding.imageWritten) + vectorBytes(binding.imageDepthCompare) + vectorBytes(binding.imageAtomic) + vectorBytes(binding.imageAtomic64) + vectorBytes(binding.bufferAtomic) + vectorBytes(binding.bufferWritten) + vectorBytes(binding.samplerUnnormalized) + vectorBytes(binding.samplerIntegerBorder) + vectorBytes(binding.samplerCoordinatesIgnored) + vectorBytes(binding.imageUnnormalized) + vectorBytes(binding.imageUnnormalizedBaseLevel) + vectorBytes(binding.bufferRead);
+}
+
+std::uint64_t resultBytes(const RecompileResult& result) {
+    const auto bytes = [](const auto& values) { return vectorBytes(values); };
+    auto total = static_cast<std::uint64_t>(sizeof(RecompileResult)) + bytes(result.vertexInputPatches) + bytes(result.runtimeImageResources) + bytes(result.vertexInputs) + bytes(result.parameterExports) + bytes(result.fragmentParameters) + bytes(result.stencilReferences) + bytes(result.shaderDataUserWords) + bytes(result.specialization) + bytes(result.bindings) + bytes(result.pushConstants) + bytes(result.vertexAttributes) + bytes(result.imageTablePoison) + bytes(result.imageTableRanges) + bytes(result.imageTableReadRanges);
+    for (const auto& ranges : result.imageTableReadRanges) total += bytes(ranges);
+    for (const auto& binding : result.bindings) total += descriptorBindingBytes(binding);
+    return total;
+}
+
+std::uint64_t patchBytes(const MemoPatch& patch) {
+    auto total = static_cast<std::uint64_t>(sizeof(MemoPatch)) + vectorBytes(patch.buffers) + vectorBytes(patch.userWords);
+    for (const auto& buffer : patch.buffers) total += vectorBytes(buffer.second);
+    return total;
+}
+
+std::atomic<std::uint64_t>& resultMemoBudgetBytes() {
+    static std::atomic<std::uint64_t> budget{[] {
+        const char* value = std::getenv("APS5_RESULT_MEMO_MB");
+        const std::uint64_t megabytes = value == nullptr || *value == '\0' ? 256u : std::strtoull(value, nullptr, 10);
+        return megabytes << 20u;
+    }()};
+    return budget;
+}
+
+std::uint64_t resultMemoBudget() {
+    return resultMemoBudgetBytes().load(std::memory_order_relaxed);
+}
+
+struct ResultMemoRegistry {
+    std::mutex mutex;
+    std::vector<std::weak_ptr<SourceEntry>> sources;
+};
+
+ResultMemoRegistry& resultMemoRegistry() {
+    static ResultMemoRegistry registry;
+    return registry;
+}
+
+std::mutex& resultMemoTrimMutex() {
+    static std::mutex mutex;
+    return mutex;
+}
+
+std::function<void()>& resultMemoTrimHook() {
+    static std::function<void()> hook;
+    return hook;
+}
+
+void trimResultMemo(std::uint64_t budget) {
+    std::lock_guard trim(resultMemoTrimMutex());
+    auto& counters = resultMemoCounters();
+    for (auto held = counters.heldBytes.load(std::memory_order_relaxed); held > budget; held = counters.heldBytes.load(std::memory_order_relaxed)) {
+        if (auto hook = std::exchange(resultMemoTrimHook(), nullptr)) hook();
+        std::vector<std::shared_ptr<SourceEntry>> sources;
+        {
+            auto& registry = resultMemoRegistry();
+            std::lock_guard lock(registry.mutex);
+            sources.reserve(registry.sources.size());
+            for (const auto& weak : registry.sources) {
+                if (auto source = weak.lock()) sources.push_back(std::move(source));
+            }
+        }
+        std::vector<std::pair<std::uint64_t, std::uint64_t>> stamps;
+        for (const auto& source : sources) {
+            std::lock_guard lock(source->mutex);
+            for (const auto* table : {source->memo.get(), source->maskedMemo.get()}) {
+                if (table == nullptr) continue;
+                for (const auto& entry : *table) {
+                    if (entry.result != nullptr) stamps.emplace_back(entry.used, entry.bytes);
+                }
+            }
+        }
+        if (stamps.empty()) return;
+        std::sort(stamps.begin(), stamps.end());
+        const auto excess = held - budget / 4u * 3u;
+        std::uint64_t freed = 0;
+        std::uint64_t threshold = 0;
+        for (const auto& [used, bytes] : stamps) {
+            threshold = used;
+            freed += bytes;
+            if (freed >= excess) break;
+        }
+        std::uint64_t dropped = 0;
+        for (const auto& source : sources) {
+            std::lock_guard lock(source->mutex);
+            for (auto* table : {source->memo.get(), source->maskedMemo.get()}) {
+                if (table == nullptr) continue;
+                for (auto& entry : *table) {
+                    if (entry.result == nullptr || entry.used > threshold) continue;
+                    counters.heldBytes.fetch_sub(entry.bytes, std::memory_order_relaxed);
+                    counters.heldEntries.fetch_sub(1, std::memory_order_relaxed);
+                    counters.trimmed.fetch_add(1, std::memory_order_relaxed);
+                    entry = {};
+                    ++dropped;
+                }
+            }
+        }
+        if (dropped == 0) return;
+    }
+}
+
 void verifyPatchedResult(const RecompileResult& patched, const RecompileResult& full, std::uint64_t variantId) {
     const auto* difference = patchedResultDifference(patched, full);
     if (difference == nullptr) return;
@@ -1526,11 +1644,12 @@ void reportResultMemo() {
     const auto patched = counters.patched.exchange(0, std::memory_order_relaxed);
     const auto misses = counters.misses.exchange(0, std::memory_order_relaxed);
     const auto evictions = counters.evictions.exchange(0, std::memory_order_relaxed);
+    const auto trimmed = counters.trimmed.exchange(0, std::memory_order_relaxed);
     const auto populate = counters.populateNanoseconds.exchange(0, std::memory_order_relaxed);
     const auto patching = counters.patchNanoseconds.exchange(0, std::memory_order_relaxed);
     const auto mismatches = counters.mismatches.exchange(0, std::memory_order_relaxed);
     const auto total = static_cast<double>(hits + patched + misses);
-    std::fprintf(stderr, "[recompile] result memo (10 s): %llu hits, %llu patched, %llu misses (%.1f%% hits, %.1f%% patched), Populate %.1f us per miss / %.1f ms in total, patch %.1f us each / %.1f ms in total, %llu evictions, %llu patch mismatches\n", static_cast<unsigned long long>(hits), static_cast<unsigned long long>(patched), static_cast<unsigned long long>(misses), total != 0.0 ? 100.0 * static_cast<double>(hits) / total : 0.0, total != 0.0 ? 100.0 * static_cast<double>(patched) / total : 0.0, misses != 0 ? static_cast<double>(populate) / 1000.0 / static_cast<double>(misses) : 0.0, static_cast<double>(populate) / 1e6, patched != 0 ? static_cast<double>(patching) / 1000.0 / static_cast<double>(patched) : 0.0, static_cast<double>(patching) / 1e6, static_cast<unsigned long long>(evictions), static_cast<unsigned long long>(mismatches));
+    std::fprintf(stderr, "[recompile] result memo (10 s): %llu hits, %llu patched, %llu misses (%.1f%% hits, %.1f%% patched), Populate %.1f us per miss / %.1f ms in total, patch %.1f us each / %.1f ms in total, %llu evictions, %llu patch mismatches, %llu trimmed; %llu results held, %.1f MiB\n", static_cast<unsigned long long>(hits), static_cast<unsigned long long>(patched), static_cast<unsigned long long>(misses), total != 0.0 ? 100.0 * static_cast<double>(hits) / total : 0.0, total != 0.0 ? 100.0 * static_cast<double>(patched) / total : 0.0, misses != 0 ? static_cast<double>(populate) / 1000.0 / static_cast<double>(misses) : 0.0, static_cast<double>(populate) / 1e6, patched != 0 ? static_cast<double>(patching) / 1000.0 / static_cast<double>(patched) : 0.0, static_cast<double>(patching) / 1e6, static_cast<unsigned long long>(evictions), static_cast<unsigned long long>(mismatches), static_cast<unsigned long long>(trimmed), static_cast<unsigned long long>(counters.heldEntries.load(std::memory_order_relaxed)), static_cast<double>(counters.heldBytes.load(std::memory_order_relaxed)) / 1048576.0);
 }
 
 // Everything materializeResult reads besides the variant: the snapshot (the descriptor words, the
@@ -1644,7 +1763,7 @@ std::shared_ptr<const RecompileResult> materializePreparedMemoized(SourceEntry&
             for (std::size_t way = 0; way < ResultMemoWays; ++way) {
                 auto& entry = set[way];
                 if (entry.result == nullptr || entry.variantId != variantId || entry.hash != hash) continue;
-                entry.used = ++source.memoClock;
+                entry.used = nextMemoUse();
                 counters.hits.fetch_add(1, std::memory_order_relaxed);
                 if (memoHit != nullptr) *memoHit = true;
                 reportResultMemo();
@@ -1663,7 +1782,7 @@ std::shared_ptr<const RecompileResult> materializePreparedMemoized(SourceEntry&
             for (std::size_t way = 0; way < ResultMemoWays; ++way) {
                 auto& entry = set[way];
                 if (entry.result == nullptr || entry.variantId != variantId || entry.hash != maskedHash) continue;
-                entry.used = ++source.memoClock;
+                entry.used = nextMemoUse();
                 donor = entry.result;
                 donorPatch = entry.patch;
                 break;
@@ -1687,20 +1806,35 @@ std::shared_ptr<const RecompileResult> materializePreparedMemoized(SourceEntry&
     }
     result->cacheHit = cacheHit;
     std::shared_ptr<const RecompileResult> shared = std::move(result);
+    const auto bytes = resultBytes(*shared);
     const auto insert = [&](std::unique_ptr<ResultMemoTable>& table, std::uint64_t at, std::uint64_t key, std::shared_ptr<const MemoPatch> recipe, bool count) {
-        if (table == nullptr) table = std::make_unique<ResultMemoTable>();
+        if (table == nullptr) {
+            table = std::make_unique<ResultMemoTable>();
+            if (source.memo == nullptr || source.maskedMemo == nullptr) {
+                auto& registry = resultMemoRegistry();
+                std::lock_guard registryLock(registry.mutex);
+                registry.sources.push_back(source.weak_from_this());
+            }
+        }
         auto* set = memoSet(table, at);
         auto* victim = set;
         for (std::size_t way = 0; way < ResultMemoWays; ++way) {
             auto& entry = set[way];
             if (entry.result != nullptr && entry.variantId == variantId && entry.hash == key) {
-                entry.used = ++source.memoClock;
+                entry.used = nextMemoUse();
                 return entry.result;
             }
             if (entry.used < victim->used) victim = &entry;
         }
-        if (count && victim->result != nullptr) counters.evictions.fetch_add(1, std::memory_order_relaxed);
-        *victim = {variantId, key, shared, ++source.memoClock, std::move(recipe)};
+        if (victim->result != nullptr) {
+            if (count) counters.evictions.fetch_add(1, std::memory_order_relaxed);
+            counters.heldBytes.fetch_sub(victim->bytes, std::memory_order_relaxed);
+            counters.heldEntries.fetch_sub(1, std::memory_order_relaxed);
+        }
+        const auto held = bytes + (recipe != nullptr ? patchBytes(*recipe) : 0u);
+        *victim = {variantId, key, shared, nextMemoUse(), std::move(recipe), held};
+        counters.heldBytes.fetch_add(held, std::memory_order_relaxed);
+        counters.heldEntries.fetch_add(1, std::memory_order_relaxed);
         return shared;
     };
     {
@@ -1708,6 +1842,7 @@ std::shared_ptr<const RecompileResult> materializePreparedMemoized(SourceEntry&
         shared = insert(source.memo, index, hash, nullptr, true);
         if (patch != nullptr && patch->valid) static_cast<void>(insert(source.maskedMemo, maskedIndex, maskedHash, patch, false));
     }
+    if (const auto budget = resultMemoBudget(); budget != 0 && counters.heldBytes.load(std::memory_order_relaxed) > budget) trimResultMemo(budget);
     reportResultMemo();
     return shared;
 }
@@ -1940,4 +2075,18 @@ std::shared_ptr<const RecompileResult> Recompile(const RecompileRequest& request
     return recompileReporting(request, [&] { return RecompileImpl(request, capture, memoHit); });
 }
 
+void SetResultMemoBudget(std::uint64_t bytes) {
+    resultMemoBudgetBytes().store(bytes, std::memory_order_relaxed);
+    if (bytes != 0 && resultMemoCounters().heldBytes.load(std::memory_order_relaxed) > bytes) trimResultMemo(bytes);
+}
+
+void SetResultMemoTrimHook(std::function<void()> hook) {
+    std::lock_guard lock(resultMemoTrimMutex());
+    resultMemoTrimHook() = std::move(hook);
+}
+
+std::uint64_t ResultMemoBytes() {
+    return resultMemoCounters().heldBytes.load(std::memory_order_relaxed);
+}
+
 }