Reuse a guest pool's span without taking the pool lock

.github/workflows/source-regressions.yml

@@ -7,6 +7,8 @@ on:
       - nx2/nup.py
       - nx2/make_game.py
       - nx2/unpack.py
+      - gpu/vk/vk_memory.cpp
+      - tests/test_pool_span.py
       - nx2/test_nup_paths.py
       - tests/test_make_game_cache.py
       - tests/test_failure_reason.py
@@ -27,6 +29,8 @@ on:
       - nx2/nup.py
       - nx2/make_game.py
       - nx2/unpack.py
+      - gpu/vk/vk_memory.cpp
+      - tests/test_pool_span.py
       - nx2/test_nup_paths.py
       - tests/test_make_game_cache.py
       - tests/test_failure_reason.py
@@ -70,6 +74,9 @@ jobs:
           python-version: '3.12'
       - name: Check extraction boundaries
         run: python -I -B nx2/test_nup_paths.py
+      - name: Check a pool span is reused without the lock
+        if: runner.os == 'Linux'
+        run: python -I -B tests/test_pool_span.py
       - name: Check source-cache coherency
         run: python -I -B tests/test_make_game_cache.py
       - name: Check a failed step's error line

gpu/vk/vk_memory.cpp

@@ -59,6 +59,7 @@ struct GuestPool {
 };
 
 std::vector<GuestPool> guest_pools;   // by lo, from the host
+std::atomic<uint64_t> guest_pools_gen{1};   // (bumped when that table is replaced: a span cached without the lock goes stale)
 std::map<uint64_t, Pool> imported;    // by lo, made on first use
 Pool* last_pool = nullptr;            // (the last pool_of hit: draws in a row bind buffers of the same pools)
 extern Pool* pool_cache[4096];
@@ -278,6 +279,7 @@ void memory_sync() {
     std::sort(next.begin(), next.end(), [](const GuestPool& a, const GuestPool& b) { return a.lo < b.lo; });
     std::lock_guard lk(pools_mu);
     guest_pools.swap(next);
+    guest_pools_gen.fetch_add(1, std::memory_order_relaxed);
     last_pool = nullptr;
     std::fill(std::begin(pool_cache), std::end(pool_cache), nullptr);
     static int logged;
@@ -1460,21 +1462,49 @@ static bool game_pool_of(uint64_t gpu, uint64_t& lo, uint64_t& hi) {   // (pools
     lo = it->lo, hi = it->hi;
     return true;
 }
+/// True when gpu sits in [lo, hi) and that span was cached from the guest-pool table this generation still is.
+bool guest_span_covers(uint64_t gpu, uint64_t lo, uint64_t hi, uint64_t cached_gen, uint64_t gen) {
+    return cached_gen == gen && lo < hi && gpu >= lo && gpu < hi;
+}
+struct GuestSpan { uint64_t lo = 1, hi = 0, gen = 0; };
+thread_local GuestSpan guest_span;
+bool guest_span_cached(uint64_t gpu, uint64_t* lo, uint64_t* hi) {
+    static const bool imported_only = getenv("RX_VK_POOL_END_IMPORTED") != nullptr;   // (A/B: the import's end, as before)
+    if (imported_only || gpu >= (1ull << 40)) return false;
+    const uint64_t gen = guest_pools_gen.load(std::memory_order_relaxed);
+    if (!guest_span_covers(gpu, guest_span.lo, guest_span.hi, guest_span.gen, gen)) return false;
+    *lo = guest_span.lo;
+    *hi = guest_span.hi;
+    return true;
+}
+void guest_span_store(uint64_t lo, uint64_t hi) {   // (pools_mu held: the generation matches the table just read)
+    guest_span = {lo, hi, guest_pools_gen.load(std::memory_order_relaxed)};
+}
 uint64_t gpu_pool_end(uint64_t gpu) {
+    uint64_t lo, hi;
+    // (storage-buffer tables ask this per entry, and the entries of one table sit in one pool: the locked lookup of a
+    // pool was ~8% of replay time before buffer binds cached it)
+    if (guest_span_cached(gpu, &lo, &hi)) return hi;
     ParLock plk;   // (before pools_mu, always)
     std::lock_guard lk(pools_mu);
-    uint64_t lo, hi;
     static const bool imported_only = getenv("RX_VK_POOL_END_IMPORTED") != nullptr;   // (A/B: the import's end, as before)
-    if (!imported_only && game_pool_of(gpu, lo, hi)) return hi;
+    if (!imported_only && game_pool_of(gpu, lo, hi)) {
+        guest_span_store(lo, hi);
+        return hi;
+    }
     const Pool* p = pool_of(gpu, 1);
     return p ? p->hi : 0;
 }
 uint64_t gpu_pool_start(uint64_t gpu) {   // (EvilOOT) where it begins
+    uint64_t lo, hi;
+    if (guest_span_cached(gpu, &lo, &hi)) return lo;
     ParLock plk;
     std::lock_guard lk(pools_mu);
-    uint64_t lo, hi;
     static const bool imported_only = getenv("RX_VK_POOL_END_IMPORTED") != nullptr;
-    if (!imported_only && game_pool_of(gpu, lo, hi)) return lo;
+    if (!imported_only && game_pool_of(gpu, lo, hi)) {
+        guest_span_store(lo, hi);
+        return lo;
+    }
     const Pool* p = pool_of(gpu, 1);
     return p ? p->lo : 0;
 }

tests/test_pool_span.py

@@ -0,0 +1,68 @@
+"""A repeated pool-end lookup in the same guest pool skips the lock.
+
+Storage-buffer entries ask where their pool ends. The locked lookup was about 8% of
+replay time before buffer binds grew a cache; the pool-end query still took the lock.
+python -I -B tests/test_pool_span.py [--cxx g++]
+"""
+import argparse
+import subprocess
+import tempfile
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+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 before_lock(fn, name):
+    lock = fn.index("ParLock")
+    head = fn[:lock]
+    if "guest_span_cached(" not in head:
+        raise SystemExit(f"{name} takes the pool lock before checking the span cache")
+
+
+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")
+    fn = block(source, "bool guest_span_covers(")
+    before_lock(block(source, "uint64_t gpu_pool_end("), "gpu_pool_end")
+    before_lock(block(source, "uint64_t gpu_pool_start("), "gpu_pool_start")
+    sync = block(source, "void memory_sync(")
+    swap = sync.index("guest_pools.swap(next);")
+    if "guest_pools_gen" not in sync[swap:swap + 200]:
+        raise SystemExit("refreshing the guest pools does not drop the span cache")
+    fixture = f'''#include <cstdint>
+#include <cstdio>
+#include <cstdlib>
+#define CHECK(x) do {{ if (!(x)) {{ std::fprintf(stderr, "FAIL:%d %s\\n", __LINE__, #x); std::exit(1); }} }} while (0)
+{fn}
+int main() {{
+    CHECK(guest_span_covers(150, 100, 200, 4, 4));
+    CHECK(guest_span_covers(100, 100, 200, 4, 4));
+    CHECK(!guest_span_covers(200, 100, 200, 4, 4));
+    CHECK(!guest_span_covers(150, 100, 200, 4, 5));
+    CHECK(!guest_span_covers(150, 200, 100, 4, 4));
+    CHECK(!guest_span_covers(150, 100, 100, 1, 1));
+    std::printf("PASS: a pool span is reused only inside that generation\\n");
+}}
+'''
+    with tempfile.TemporaryDirectory(prefix="pool-span-") as tmp:
+        cpp = Path(tmp) / "fixture.cpp"
+        exe = Path(tmp) / "fixture"
+        cpp.write_text(fixture, encoding="utf-8")
+        subprocess.run([a.cxx, "-std=c++20", "-Wall", "-Wextra", "-Werror", str(cpp), "-o", str(exe)], check=True)
+        subprocess.run([str(exe)], check=True)
+
+
+if __name__ == "__main__":
+    main()