perf(vk): mirror the pools the GPU reads most, in place of colder ones

gpu/vk/vk_memory.cpp

@@ -70,6 +70,12 @@ struct Pool {
     uint64_t synced_epoch = 0;   // the submission that last brought it up to date
     uint64_t sync_seg = 0;       // (and the parallel replay's segment whose uploads did)
     uint64_t last_value = 0;     // (the newest submission that bound it: its mirror is in use until that one is done)
+    // (hookshot) its heat: the bytes of the ranges bound from it (every binding, the bind memo's hits too; relaxed, any thread),
+    // and those of the last 64-submission windows, halved each window (mirrors_swap_for_hot). made_epoch: its mirror's making
+    uint64_t bound_bytes = 0, heat_mark = 0, heat = 0, made_epoch = 0;
+    // (hookshot) its mirror may hold GPU writes no write-back copies to guest memory (memory_note_gpu_write: a bare entry's
+    // guessed range, the pointer table's, a wide binding; a written binding over RX_VK_WB_MAX_KB): never a swap's victim
+    bool gpu_only = false;
     uint64_t direct_write_epoch = 0;   // (the part that last wrote it in guest memory, before it had a mirror)
     uint64_t retired_value = 0;  // (retired: the submission that may still use it)
     bool was_mirrored = false;   // (its mirror was given back: its pages stay in the captures, for the next one)
@@ -2091,7 +2097,8 @@ Pool* pool_lookup(uint64_t guest, uint64_t size) {
 //   write generations) are copied in, in the recording's upload command buffer (submitted right before it).
 // - GPU -> guest: what the submission's shaders and copies may write into a mirror is copied back into the imported
 //   pool at the submission's end, so the CPU (and everything that reads guest memory) sees it.
-// RX_VK_MIRROR=0: no mirrors; RX_VK_MIRROR_MB: their budget (default 3072).
+// RX_VK_MIRROR=0: no mirrors; RX_VK_MIRROR_MB: their budget (default 3072); RX_VK_MIRROR_HEAT=0: no swaps by heat
+// (mirrors_swap_for_hot).
 bool mirror_on() {
     static const int on = [] {
         const char* e = getenv("RX_VK_MIRROR");
@@ -2459,6 +2466,11 @@ void memory_note_gpu_write(uint64_t gpu, uint64_t size) {
     }
     std::lock_guard lk(pools_mu);
     rb_hazard_locked(guest, guest + size);
+    // (hookshot: the mirrors it may write keep them, mirrors_swap_for_hot)
+    auto it = imported.upper_bound(guest);
+    if (it != imported.begin()) --it;
+    for (; it != imported.end() && it->second.lo < guest + size; ++it)
+        if (it->second.hi > guest && it->second.mstate == Pool::MIRROR_ON) it->second.gpu_only = true;
 }
 void memory_rb_recording_begun() {
     std::lock_guard lk(pools_mu);
@@ -2660,7 +2672,11 @@ bool bind_memo_hit(uint64_t guest, uint64_t size, uint64_t stamp, uint64_t cache
 struct BindMemo {
     uint64_t guest = 0, size = 0, stamp = 0;
     BufferRef ref{};
+    Pool* pool = nullptr;   // (hookshot: whose heat a hit adds to)
 };
+inline void heat_add(Pool* p, uint64_t size) {
+    if (p) std::atomic_ref<uint64_t>(p->bound_bytes).fetch_add(size, std::memory_order_relaxed);
+}
 thread_local BindMemo bind_memo[256];
 
 void note_direct_write(Pool& p, uint64_t off, uint64_t size) {   // (pools_mu held)
@@ -2681,15 +2697,18 @@ BufferRef gpu_buffer_bind(uint64_t gpu, uint64_t size, bool written) {
     }
     // (the same read, while that stamp still stands: the buffer and the offset are the ones found before. a miss
     // runs the lookup below, including a mirror sync, and remembers what it returned. a write is not remembered)
-    auto remember = [&](BufferRef ref) {
+    auto remember = [&](Pool* p, BufferRef ref) {
+        heat_add(p, size);
         if (!written && ref.buffer != VK_NULL_HANDLE && size)
-            bind_memo[bind_slot(guest, size)] = BindMemo{guest, size, gpu_bind_stamp(), ref};
+            bind_memo[bind_slot(guest, size)] = BindMemo{guest, size, gpu_bind_stamp(), ref, p};
         return ref;
     };
     if (!written) {
         const BindMemo& m = bind_memo[bind_slot(guest, size)];
-        if (bind_memo_hit(guest, size, gpu_bind_stamp(), m.guest, m.size, m.stamp, m.ref.buffer != VK_NULL_HANDLE))
+        if (bind_memo_hit(guest, size, gpu_bind_stamp(), m.guest, m.size, m.stamp, m.ref.buffer != VK_NULL_HANDLE)) {
+            heat_add(m.pool, size);
             return m.ref;
+        }
     }
     // (a read binding of a pool this thread found before and whose mirror is in use this submission or that has none,
     // without the lock: every thread's, the replay thread's too (each draw binds several uniform buffers: the pool's
@@ -2702,9 +2721,9 @@ BufferRef gpu_buffer_bind(uint64_t gpu, uint64_t size, bool written) {
         const uint64_t off = guest - p->lo;
         if (p->last_value < mem_rec_value) p->last_value = mem_rec_value;
         if (!mirror_on() || !mirror_upload || p->mstate == Pool::MIRROR_REFUSED || p->mstate == Pool::MIRROR_WANTED)
-            return remember({p->buffer, off, size, 0});
+            return remember(p, {p->buffer, off, size, 0});
         if (p->mstate == Pool::MIRROR_ON && p->synced_epoch == mirror_epoch && p->sync_seg <= texture_segment())
-            return remember({p->mirror, off, size, 0});
+            return remember(p, {p->mirror, off, size, 0});
     }
     ParLock plk;   // (before pools_mu, always: a mirror sync records into the shared upload command buffer)
     std::lock_guard lk(pools_mu);
@@ -2732,7 +2751,7 @@ BufferRef gpu_buffer_bind(uint64_t gpu, uint64_t size, bool written) {
         // (hookshot) written by the GPU in guest memory itself, past the mirror (bound outside a recording's uploads): the
         // write-back diff's shadow no longer knows those pages, and reads them again at its next write-back, after this write
         if (written) shadow_forget(*p, off, std::min(size, p->hi - p->lo - off));
-        return remember({p->buffer, off, size, 0});
+        return remember(p, {p->buffer, off, size, 0});
     }
     // (a pool's mirror is made at its first binding, in this part's uploads: bound in guest memory until the next
     // submission, the GPU read it there while the game's threads wrote it for their next frames, and a pass that drew
@@ -2759,7 +2778,7 @@ BufferRef gpu_buffer_bind(uint64_t gpu, uint64_t size, bool written) {
         e.first++, e.second = (int)p->mstate;
         direct_bytes_by_pool[p->lo] += size;
         if (written) note_direct_write(*p, off, size);
-        return remember({p->buffer, off, size, 0});
+        return remember(p, {p->buffer, off, size, 0});
     }
     if (p->synced_epoch != mirror_epoch || p->sync_seg > texture_segment()) mirror_sync(*p);
     if (written) {
@@ -2776,9 +2795,11 @@ BufferRef gpu_buffer_bind(uint64_t gpu, uint64_t size, bool written) {
         }();
         if (!wb_max || size <= wb_max)
             writebacks.push_back({p, off, std::min(off + size, p->hi - p->lo)});
+        else
+            p->gpu_only = true;   // (hookshot: not copied back, so the mirror is all there is of it)
         rb_hazard_locked(p->lo + off, p->lo + std::min(off + size, p->hi - p->lo));   // (memory_rb_refresh)
     }
-    return remember({p->mirror, off, size, 0});
+    return remember(p, {p->mirror, off, size, 0});
 }
 
 // ---- mirrors given back ----
@@ -2873,6 +2894,7 @@ void mirror_free(Pool& p) {
     st_mirror_bytes_freed += p.hi - p.lo;
     st_mirrors_freed++;
     p.mstate = Pool::MIRROR_NONE;
+    p.gpu_only = false;
     pools_gen++;   // (a parallel replay thread's pool cache: looked up again)
 }
 void mirrors_give_back(uint64_t want_bytes) {   // (pools_mu held)
@@ -2984,6 +3006,76 @@ void imports_give_back() {   // (pools_mu held)
              (long long)st_imports_evicted, (unsigned long long)(st_import_bytes_evicted >> 20));
 }
 
+// (hookshot) Mirrors by heat. A pool refused for the budget got a mirror only when one went unbound for kIdleEpochs submissions:
+// with the budget full of pools bound now and then but read little, the pool bound every submission stayed in guest memory for
+// good. On an RTX A1000 Laptop (4 GB; 495 MB of mirrors) in play, a 160 MB pool the GPU bound 513 times in 600 submissions,
+// 29 GB of ranges, read over PCIe, while the mirrors held 480 MB of colder pools: the GPU at ~70% and 15 of its 45 W, 17-25 fps.
+// Now a wanted pool that does not fit takes the mirrors of pools with a quarter of its heat or less (all of them together),
+// coldest first: only mirrors the GPU is done with (last_value, as mirrors_give_back), made kSwapHold submissions ago or more,
+// one swap per 16 submissions. And only mirrors guest memory has all of after their write-backs: freeing one the GPU wrote
+// past them (gpu_only: the foliage culling's records through a bare entry, bound every frame through 64 bytes and so cold)
+// sent its records back to stale guest memory, the missing foliage again. With RX_VK_FOLLOW_STORES (stores through the page
+// table, bound nowhere) any mirror may be one: no swaps. RX_VK_MIRROR_HEAT=0: as before
+// (heat = heat/2 + a window's bytes settles at twice a window's: the 64 MB floor is 32 MB of ranges bound a window)
+constexpr uint64_t kSwapHold = 600;
+int64_t st_swaps = 0;
+/// The swap's choice, as data (tests/test_mirror_heat.py): of the candidates' sizes and heats, coldest first, how many make
+/// room for want_bytes under the budget; 0 when they cannot, or when their heat together is more than a quarter of the
+/// wanted pool's (or its heat is under 64 MB, ~32 MB of ranges bound a window: not worth a mirror's copy)
+size_t mirror_swap_pick(const uint64_t* bytes, const uint64_t* heat, size_t count, uint64_t used, uint64_t budget,
+                        uint64_t want_bytes, uint64_t want_heat, uint64_t* freed_out, uint64_t* cold_heat_out) {
+    *freed_out = *cold_heat_out = 0;
+    if (want_bytes > budget || used + want_bytes <= budget || want_heat < (64ull << 20)) return 0;
+    uint64_t freed = 0, cold_heat = 0;
+    size_t k = 0;
+    while (k < count && used - freed + want_bytes > budget) freed += bytes[k], cold_heat += heat[k], k++;
+    if (used - freed + want_bytes > budget || cold_heat > want_heat / 4) return 0;
+    *freed_out = freed, *cold_heat_out = cold_heat;
+    return k;
+}
+void mirrors_swap_for_hot() {   // (pools_mu held)
+    static const bool on = (!getenv("RX_VK_MIRROR_HEAT") || !env_value_off(getenv("RX_VK_MIRROR_HEAT"))) &&
+                           !env_on("RX_VK_FOLLOW_STORES");
+    if (!on || mirror_wanted.empty()) return;
+    static uint64_t last_swap;
+    if (mirror_epoch - last_swap < 16) return;
+    // (the hottest wanted pool only: trying the others after it let small pools trade 16 MB mirrors back and forth, each swap
+    // making the mirrors young again, and the 160 MB pool never found kSwapHold-old ones to take. And only one read enough to
+    // matter: 64 MB of heat, ~32 MB of ranges a window)
+    auto hottest = std::max_element(mirror_wanted.begin(), mirror_wanted.end(),
+                                    [](const Pool* a, const Pool* b) { return a->heat < b->heat; });
+    {
+        Pool* w = *hottest;
+        const uint64_t n = w->hi - w->lo;
+        if (w->mstate != Pool::MIRROR_WANTED) return;
+        std::vector<Pool*> cold;
+        for (auto& [lo, p] : imported)
+            if (p.mstate == Pool::MIRROR_ON && &p != w && !p.inherit_pending && !p.gpu_only && p.last_value <= mem_completed &&
+                mirror_epoch - p.made_epoch >= kSwapHold && p.heat < w->heat / 4)
+                cold.push_back(&p);
+        std::sort(cold.begin(), cold.end(), [](const Pool* a, const Pool* b) { return a->heat < b->heat; });
+        std::vector<uint64_t> bytes(cold.size()), heat(cold.size());
+        for (size_t i = 0; i < cold.size(); i++) bytes[i] = cold[i]->hi - cold[i]->lo, heat[i] = cold[i]->heat;
+        uint64_t freed = 0, cold_heat = 0;
+        const size_t k = mirror_swap_pick(bytes.data(), heat.data(), cold.size(), mirror_bytes, mirror_budget(), n, w->heat,
+                                          &freed, &cold_heat);
+        if (!k) return;
+        for (size_t i = 0; i < k; i++) mirror_free(*cold[i]);
+        // (made here, first: the space freed is its, not a colder wanted pool's)
+        mirror_wanted.erase(hottest);
+        const bool made = mirror_make(*w);
+        last_swap = mirror_epoch;
+        if (!made)
+            logf("vk: pool %#llx (%llu MB): its mirror was not made after %zu colder ones (%llu MB) were given back for it",
+                 (unsigned long long)w->lo, (unsigned long long)(n >> 20), k, (unsigned long long)(freed >> 20));
+        else if (st_swaps++ < 32 || st_swaps % 64 == 0)
+            logf("vk: pool %#llx (%llu MB, ~%llu MB bound a window) takes the mirrors of %zu colder pools (%llu MB, ~%llu MB bound a "
+                 "window; %lld swaps so; RX_VK_MIRROR_HEAT=0: off)", (unsigned long long)w->lo, (unsigned long long)(n >> 20),
+                 (unsigned long long)(w->heat >> 21), k, (unsigned long long)(freed >> 20), (unsigned long long)(cold_heat >> 21),
+                 (long long)st_swaps);
+    }
+}
+
 void memory_begin_submission() {
     std::lock_guard lk(pools_mu);
     reap_mirrors();
@@ -3000,7 +3092,14 @@ void memory_begin_submission() {
         for (Pool* p : mirror_wanted) want += p->hi - p->lo;
         if (mirror_epoch % 64 == 0 || mirror_bytes + want > mirror_budget()) mirrors_give_back(want);
     }
+    if (mirror_epoch % 64 == 0)   // (hookshot: the heat's windows, mirrors_swap_for_hot)
+        for (auto& [lo, p] : imported) {
+            const uint64_t b = std::atomic_ref<uint64_t>(p.bound_bytes).load(std::memory_order_relaxed);
+            p.heat = p.heat / 2 + (b - p.heat_mark);
+            p.heat_mark = b;
+        }
     if (mirror_wanted.empty()) return;
+    mirrors_swap_for_hot();
     for (Pool* p : mirror_wanted) mirror_make(*p);
     mirror_wanted.clear();
     publish_cap_ranges();
@@ -3153,6 +3252,8 @@ bool mirror_make(Pool& pool) {
             }
         }
         p->mstate = Pool::MIRROR_ON;
+        p->made_epoch = mirror_epoch;
+        p->gpu_only = false;
         pools_gen.fetch_add(1, std::memory_order_relaxed);   // (the binding's buffer changes: guest memory, then this mirror)
         p->synced_gen = gen_now;
         p->synced_epoch = mirror_epoch;

tests/test_direct_buffer_writes.py

@@ -53,9 +53,10 @@ def main():
                              block(source, "struct TlPool {", True), block(source, "struct WriteBack {", True)])
     pages = source[source.index("std::mutex wb_pages_mu;"):source.index("void wb_pages_note(")]
     direct_note = block(source, "void note_direct_write(") if "void note_direct_write(" in source else ""
+    heat = block(source, "inline void heat_add(") if "inline void heat_add(" in source else ""   # (the mirrors' heat, #1377)
     actual = "\n\n".join([pages, block(source, "void wb_pages_note("), block(source, "uint64_t memory_writeback_value("),
         block(source, "uint64_t gpu_bind_stamp("), block(source, "uint32_t bind_slot("), block(source, "bool bind_memo_hit("),
-        block(source, "struct BindMemo {", True), "thread_local BindMemo bind_memo[256];", direct_note,
+        block(source, "struct BindMemo {", True), heat, "thread_local BindMemo bind_memo[256];", direct_note,
         block(source, "void shadow_forget(Pool& p, uint64_t at, uint64_t len) {"),
         block(source, "BufferRef gpu_buffer_bind(")])
     flags = subprocess.CREATE_NO_WINDOW | subprocess.BELOW_NORMAL_PRIORITY_CLASS if os.name == "nt" else 0

tests/test_mirror_heat.py

@@ -0,0 +1,110 @@
+"""Mirrors by heat: the pool the GPU reads most takes the mirrors of colder ones, and only when that is worth it.
+
+A pool refused for the mirrors' budget got a mirror only when one went unbound for 600 submissions, so with the budget
+full of pools bound now and then, the pool bound every submission was read over PCIe for good. mirror_swap_pick
+(gpu/vk/vk_memory.cpp) is extracted unmodified and run on the cases below; no GPU, game or Vulkan is involved.
+
+python -I -B tests/test_mirror_heat.py [--cxx g++]
+"""
+import argparse
+import subprocess
+import tempfile
+from pathlib import Path
+
+RUN_ALL_ARGS = ["--cxx", "g++"]
+
+ROOT = Path(__file__).resolve().parents[1]
+MB = 1 << 20
+
+FIXTURE = r"""
+#include <cstddef>
+#include <cstdint>
+#include <cstdio>
+
+/* ACTUAL_PICK */
+
+static int fails;
+static void check(const char* what, size_t got, size_t want) {
+    if (got != want) std::printf("FAIL %s: %zu, expected %zu\n", what, got, want), fails++;
+}
+
+int main() {
+    const uint64_t MB = 1ull << 20;
+    uint64_t freed, cold;
+    // (the A1000 case: 480 of 495 MB of mirrors, a 160 MB pool with ~3 GB of ranges a window; four 16 MB pools nearly unread
+    // and larger warm ones: the cold ones first, enough of them to make room)
+    {
+        const uint64_t bytes[] = {16 * MB, 16 * MB, 16 * MB, 16 * MB, 64 * MB, 96 * MB, 96 * MB};
+        const uint64_t heat[] = {0, 1 * MB, 5 * MB, 9 * MB, 20 * MB, 40 * MB, 60 * MB};
+        const size_t k = mirror_swap_pick(bytes, heat, 7, 480 * MB, 495 * MB, 160 * MB, 3000 * MB, &freed, &cold);
+        check("hot pool: the coldest that make room", k, 6);
+        check("hot pool: bytes freed", freed >> 20, 224);
+    }
+    // (it fits already: nothing taken)
+    {
+        const uint64_t bytes[] = {16 * MB};
+        const uint64_t heat[] = {0};
+        check("fits: no swap", mirror_swap_pick(bytes, heat, 1, 300 * MB, 495 * MB, 160 * MB, 3000 * MB, &freed, &cold), 0);
+    }
+    // (the candidates together read more than a quarter of what the wanted pool does: kept)
+    {
+        const uint64_t bytes[] = {80 * MB, 80 * MB};
+        const uint64_t heat[] = {150 * MB, 160 * MB};
+        check("warm victims: kept", mirror_swap_pick(bytes, heat, 2, 480 * MB, 495 * MB, 96 * MB, 1000 * MB, &freed, &cold), 0);
+    }
+    // (a pool that reads little: never worth a mirror's copy, however cold the others; v1 traded 16 MB mirrors back and forth so)
+    {
+        const uint64_t bytes[] = {16 * MB};
+        const uint64_t heat[] = {0};
+        check("cool wanted pool: no swap", mirror_swap_pick(bytes, heat, 1, 490 * MB, 495 * MB, 16 * MB, 30 * MB, &freed, &cold), 0);
+    }
+    // (not enough candidates to make room: none taken, rather than some for nothing)
+    {
+        const uint64_t bytes[] = {16 * MB, 16 * MB};
+        const uint64_t heat[] = {0, 0};
+        check("too few: none", mirror_swap_pick(bytes, heat, 2, 480 * MB, 495 * MB, 160 * MB, 3000 * MB, &freed, &cold), 0);
+        check("too few: nothing freed", freed, 0);
+    }
+    // (bigger than the whole budget: never)
+    {
+        const uint64_t bytes[] = {495 * MB};
+        const uint64_t heat[] = {0};
+        check("over the budget", mirror_swap_pick(bytes, heat, 1, 495 * MB, 495 * MB, 600 * MB, 9000 * MB, &freed, &cold), 0);
+    }
+    if (!fails) std::printf("mirror heat: all cases pass\n");
+    return fails != 0;
+}
+"""
+
+
+def block(text, marker):
+    start = text.index(marker)
+    end = text.index("{", start) + 1
+    depth = 1
+    while depth:
+        depth += (text[end] == "{") - (text[end] == "}")
+        end += 1
+    return text[start:end]
+
+
+def main():
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("--cxx", default="g++")
+    a = p.parse_args()
+    source = (ROOT / "gpu/vk/vk_memory.cpp").read_text(encoding="utf-8")
+    pick = block(source, "size_t mirror_swap_pick(")
+    swap = block(source, "void mirrors_swap_for_hot()")
+    # (the swap takes only mirrors the GPU is done with, old enough, that guest memory has all of after their write-backs,
+    # none with stores through the page table, and makes the wanted pool first)
+    for need in ("p.last_value <= mem_completed", "kSwapHold", "!p.gpu_only", "RX_VK_FOLLOW_STORES", "mirror_swap_pick(",
+                 "mirror_wanted.erase(hottest)", "mirror_make(*w)"):
+        assert need in swap, "mirrors_swap_for_hot: no " + need
+    with tempfile.TemporaryDirectory(prefix="hookshot-mirror-heat-") as temp:
+        src, exe = Path(temp) / "fixture.cpp", Path(temp) / "fixture.exe"
+        src.write_text(FIXTURE.replace("/* ACTUAL_PICK */", pick), encoding="utf-8")
+        subprocess.run([a.cxx, "-std=c++17", "-O0", "-Wall", "-Wextra", "-Werror", str(src), "-o", str(exe)], check=True)
+        subprocess.run([str(exe)], check=True)
+
+
+if __name__ == "__main__":
+    main()