guest-diag: semfix also resolves+pins the CE pushbuffer range (descending GPA layout)

1 files changed, 107 insertions, 33 deletions

diff --git a/tools/guest-diag/semfix_watch.py b/tools/guest-diag/semfix_watch.py
index fec0a87..0ded7c2 100644
--- a/tools/guest-diag/semfix_watch.py
+++ b/tools/guest-diag/semfix_watch.py
@@ -1,21 +1,22 @@
 #!/usr/bin/env python3
-# semfix_watch.py - keep UVM semaphore-pool VA->PA translations fresh from
-# guest truth (host-side dev tool, no guest changes).
+# semfix_watch.py - keep UVM VA->PA translations fresh from guest truth
+# (host-side dev tool, no guest changes).
 #
-# Why: the emulation's content-scanned VA->PA map can hold a *stale*
-# translation for a semaphore-pool page (same VA, previous client's page).
-# The GPU-tracking-semaphore waits in nvidia-uvm then never clear because the
-# release writes land in the old page. The guest kernel knows the real
-# backing: the tracking semaphore object's `page->memory->vas[]` array holds
-# both the kernel CPU VA and the GPU VA of the pool page, and QEMU can turn
-# the CPU VA into the GPA (`gva2gpa`). Pinning that pair with sema_map_add
-# makes the service's release path write the page the client actually polls.
+# Two jobs:
+#  1. Semaphore pool: the emulation's content-scanned map can hold a *stale*
+#     translation for a tracking-semaphore pool page (same VA, previous
+#     client's page), so GPU-tracking-semaphore waits never clear. The
+#     tracking semaphore object (kprobe `uvmt` gives its kernel pointer) has
+#     `semaphore.page->memory->vas[]` with the pool page's CPU and GPU VAs;
+#     QEMU turns the CPU VA into a GPA (`gva2gpa`) which we pin.
+#  2. CE pushbuffer: the pushes live in a UVM rm_mem with GPU VA
+#     `0x120000000` whose CPU VA is vmalloc-mapped guest RAM. Its pages map
+#     linearly *descending* in GPA (page i -> base_gpa - i*0x1000). Pinning
+#     the range lets the host pump read the pushes directly and execute the
+#     PTE memsets/copies instead of only acknowledging their releases.
 #
-# Sources:
-#   - guest: kprobe `uvmt` (uvm_gpu_tracking_semaphore_update_completed_value)
-#     trace lines carry the tracking-semaphore kernel pointer as `sem=0x...`.
-#   - host: QEMU QMP `human-monitor-command` for `x/..gx` memory reads and
-#     `gva2gpa`.
+# Sources: guest kprobe on uvm_gpu_tracking_semaphore_update_completed_value
+# (`sem=0x...`), QEMU QMP `x/..gx` reads and `gva2gpa`.
 #
 # Usage: tools/guest-diag/semfix_watch.py [interval-seconds]
 # Runs until killed; prints one line per newly resolved/pinned page.
@@ -35,6 +36,9 @@ QMP_SOCK = os.environ.get("QMP_SOCK", "/tmp/qemu-qmp.sock")
 
 GPU_VA_LO = 0x100000000
 GPU_VA_HI = 0x2000000000
+PUSHBUF_VA_LO = 0x120000000
+PUSHBUF_VA_HI = 0x120400000
+PUSHBUF_MAX_PAGES = 1024
 
 
 class Qmp:
@@ -46,7 +50,7 @@ class Qmp:
         self.cmd("qmp_capabilities")
 
     def _rd(self):
-        time.sleep(0.05)
+        time.sleep(0.03)
         try:
             return self.s.recv(1 << 20).decode()
         except Exception:
@@ -64,10 +68,10 @@ class Qmp:
 
     def read64(self, addr, n=1):
         out = self.hmp("x/%dgx 0x%x" % (n, addr))
-        if ":" not in out:
+        m = re.search(r":(.*)$", out, re.S)
+        if not m:
             return []
-        body = out.split(":", 1)[1]
-        return [int(m, 16) for m in re.findall(r"0x([0-9a-fA-F]+)", body)]
+        return [int(x, 16) for x in re.findall(r"0x([0-9a-fA-F]+)", m.group(1))][:n]
 
     def gva2gpa(self, va):
         out = self.hmp("gva2gpa 0x%x" % va)
@@ -75,13 +79,21 @@ class Qmp:
         return int(m.group(1), 16) if m else None
 
 
+def is_kern(v):
+    return v and (v >> 48) == 0xFFFF
+
+
+def is_gpu(v):
+    return v and GPU_VA_LO <= v < GPU_VA_HI
+
+
 def guest_sems():
     """Kernel pointers of tracking semaphores the client is spinning on."""
     cmd = ("sudo sh -c 'echo > /sys/kernel/tracing/trace; sleep 2; "
            "grep uvmt /sys/kernel/tracing/trace | tail -4'")
     try:
         out = subprocess.run([VMEXEC, "exec", cmd], capture_output=True,
-                             text=True, timeout=60).stdout
+                             text=True, timeout=90).stdout
     except Exception:
         return set()
     return {int(m, 16) for m in re.findall(r"sem=0x([0-9a-fA-F]+)", out)}
@@ -89,12 +101,16 @@ def guest_sems():
 
 def pin(va, pa):
     cmd = "printf '0x%x,0x%x' > %s/sema_map_add" % (va, pa, PARAM)
-    return subprocess.run(["sudo", "-n", "sh", "-c", cmd],
-                          capture_output=True, text=True).returncode == 0
+    for _ in range(40):
+        r = subprocess.run(["sudo", "-n", "sh", "-c", cmd],
+                           capture_output=True, text=True)
+        if r.returncode == 0 and "timed out" not in (r.stderr or ""):
+            return True
+        time.sleep(0.05)  # async queue full: drain and retry
+    return False
 
 
-def resolve(q, sem):
-    """(gpu_va, gpa) pairs for the pool page behind this tracking semaphore."""
+def resolve_sem_pool(q, sem):
     page = q.read64(sem)
     if not page or not page[0]:
         return []
@@ -102,39 +118,97 @@ def resolve(q, sem):
     if not mem or not mem[0]:
         return []
     vas = q.read64(mem[0] + 0x10, 16)
-    if not vas:
-        return []
-    cpu_kva = next((v for v in vas if v and (v >> 56) == 0xFF), None)
-    gpu_vas = [v for v in vas if v and GPU_VA_LO <= v < GPU_VA_HI]
-    if not cpu_kva or not gpu_vas:
+    cpu = next((v for v in vas if is_kern(v)), None)
+    gpu_vas = [v for v in vas if is_gpu(v)]
+    if not cpu or not gpu_vas:
         return []
-    gpa = q.gva2gpa(cpu_kva)
+    gpa = q.gva2gpa(cpu)
     if not gpa:
         return []
     return [(gv, gpa) for gv in gpu_vas]
 
 
+def resolve_pushbuffer(q, sem):
+    """Find the pushbuffer rm_mem near the channel and return (gpu_vas, cpu, gpa)."""
+    base = sem & ~0xFFF
+    words = []
+    for i in range(9):
+        chunk = q.read64(base - 0x8000 + i * 0x1000, 512)
+        if not chunk:
+            break
+        words += chunk
+    for M in dict.fromkeys(w for w in words if is_kern(w)):
+        pb = q.read64(M + 8)
+        if not pb or not is_kern(pb[0]):
+            continue
+        mem = q.read64(pb[0] + 8)
+        if not mem or not is_kern(mem[0]):
+            continue
+        vas = q.read64(mem[0] + 0x10, 16)
+        gpu_vas = [v for v in vas if PUSHBUF_VA_LO <= v < PUSHBUF_VA_HI]
+        if not gpu_vas:
+            continue
+        cpu = next((v for v in vas if is_kern(v)), None)
+        if not cpu:
+            continue
+        gpa = q.gva2gpa(cpu)
+        if not gpa:
+            continue
+        return gpu_vas, cpu, gpa
+    return [], None, None
+
+
+def pin_pushbuffer(q, gpu_vas, cpu, gpa):
+    pinned = 0
+    for gv in gpu_vas:
+        for i in range(PUSHBUF_MAX_PAGES):
+            if i in (0, 1, 2, 3, 8, 32, 128, 512):
+                g = q.gva2gpa(cpu + i * 0x1000)
+                if g != gpa - i * 0x1000:
+                    print("semfix: pushbuf va=0x%x: page %d not linear (gpa=%s), stop at %d"
+                          % (gv, i, hex(g) if g else None, pinned), flush=True)
+                    return pinned
+            if not pin(gv + i * 0x1000, gpa - i * 0x1000):
+                print("semfix: pushbuf va=0x%x: pin failed at page %d" % (gv, i), flush=True)
+                return pinned
+            pinned += 1
+        print("semfix: pushbuf va=0x%x cpu=0x%x gpa=0x%x: pinned %d page(s)"
+              % (gv, cpu, gpa, pinned), flush=True)
+    return pinned
+
+
 def main():
     interval = float(sys.argv[1]) if len(sys.argv) > 1 else 3.0
     q = Qmp(QMP_SOCK)
     seen = set()
+    pb_done = False
     print("semfix: watching (interval %.1fs)" % interval, flush=True)
     while True:
         for sem in guest_sems() - seen:
-            seen.add(sem)
             try:
-                pairs = resolve(q, sem)
+                pairs = resolve_sem_pool(q, sem)
             except Exception as e:
                 print("semfix: sem=0x%x resolve failed: %s" % (sem, e), flush=True)
                 continue
             if not pairs:
                 print("semfix: sem=0x%x unresolved (retry)" % sem, flush=True)
-                seen.discard(sem)
                 continue
+            seen.add(sem)
             for gv, gpa in pairs:
                 ok = pin(gv, gpa)
                 print("semfix: sem=0x%x va=0x%x -> gpa=0x%x pin=%s"
                       % (sem, gv, gpa, "ok" if ok else "FAILED"), flush=True)
+            if not pb_done:
+                try:
+                    gpu_vas, cpu, gpa = resolve_pushbuffer(q, sem)
+                except Exception as e:
+                    print("semfix: pushbuf resolve failed: %s" % e, flush=True)
+                    gpu_vas = []
+                if gpu_vas:
+                    pin_pushbuffer(q, gpu_vas, cpu, gpa)
+                    pb_done = True
+                else:
+                    print("semfix: pushbuf not found near sem=0x%x (retry)" % sem, flush=True)
         time.sleep(interval)