guest-diag: resolve and stream the PMA scrubber buffer words (third-client scrub freeze)

guest_pin_stream: find the scrubber channel buffer among the nvidia vmap
areas (size in the CeUtils channel set, 906F ring shape at +0x64000 when
non-empty), walk its pages with kcore and emit '#scrub <fin_pa>,<hsem_pa>'
(the finish payload at +0x6c004 and the host semaphore at +0x6c000).

pin_gr_stream: handle '#scrub' lines and register the PAs through the
sema_scrubbuf_add module parameter, so the kernel's svc_scrub_pass can
author the guest's scrub progress.

2 files changed, 101 insertions, 0 deletions

diff --git a/tools/guest-diag/guest_pin_stream.py b/tools/guest-diag/guest_pin_stream.py
index 6607ead..10a9aba 100644
--- a/tools/guest-diag/guest_pin_stream.py
+++ b/tools/guest-diag/guest_pin_stream.py
@@ -1099,6 +1099,75 @@ def new_sems(seen):
     return found
 
 
+# --------------------------------------------------------------- scrubber
+# PMA scrubber channel discovery (third-client scrub freeze). The guest's
+# PMA scrubber is an internal CeUtils CE channel (4096 GPFIFO entries,
+# pbGpuVA 0x120000000, ring at pbGpuVA+0x64000, finish payload at
+# pbGpuVA+0x6c004). It lives in a kernel vmap area, not in any user pagemap,
+# and its GPU VAs collide numerically with the client CE push-buffer window
+# in the host's global translation map - so the host cannot resolve it and
+# must be told the finish/host-semaphore word PAs directly. Resolve them
+# here (vmallocinfo + manual kernel page-table walk) and emit one
+#   #scrub <fin_pa>,<hsem_pa>
+# line; pin_gr_stream.py registers them through sema_scrubbuf_add.
+SCRUB_AREA_SIZES = (0x6e000, 0x6d000, 0x6c000, 0x6b000, 0x6f000,
+                    0x70000, 0x71000, 0x72000)
+_scrub_sent = {}
+
+
+def _scrub_ring_ok(kc, ring_pa):
+    """A 906F GPFIFO ring page: all-zero (fresh) or entries whose GET names a
+    push inside the channel's own 0x120000000 window (encoded as
+    push_va - 0x100000000, 4-byte aligned)."""
+    buf = kc.read(kc.page_offset_base + ring_pa, 0x1000)
+    if not buf:
+        return False
+    for off in range(0, 0x100, 8):
+        e0, _e1 = struct.unpack_from('<II', buf, off)
+        if e0 == 0:
+            continue
+        if not (0x20000000 <= e0 < 0x20080000) or (e0 & 3):
+            return False
+    return True
+
+
+def scrub_scan(kc, lines):
+    """Emit the scrubber buffer words' PAs once (re-emitted if they move)."""
+    try:
+        f = open('/proc/vmallocinfo')
+    except OSError:
+        return
+    with f:
+        for line in f:
+            if 'nvidia' not in line or 'ioremap' in line or 'vmap' not in line:
+                continue
+            parts = line.split()
+            try:
+                a, b = parts[0].split('-')
+                lo, hi = int(a, 16), int(b, 16)
+            except (ValueError, IndexError):
+                continue
+            if (hi - lo) not in SCRUB_AREA_SIZES:
+                continue
+            fin_pa = kc.va_to_pa(lo + 0x6c000)
+            ring_pa = kc.va_to_pa(lo + 0x64000)
+            if fin_pa is None or ring_pa is None:
+                continue
+            if not _scrub_ring_ok(kc, ring_pa):
+                dbg('scrub candidate 0x%x rejected (ring shape)' % lo)
+                continue
+            fin = fin_pa + 4
+            hsem = fin_pa
+            key = (fin, hsem)
+            if _scrub_sent.get('key') == key:
+                return
+            _scrub_sent['key'] = key
+            lines.append('#scrub 0x%x,0x%x\n' % (fin, hsem))
+            dbg('scrub buffer emitted lo=0x%x fin=0x%x hsem=0x%x'
+                % (lo, fin, hsem))
+            return
+
+
 def nvidia_fds(pid):
     """CUDA device fds this pid holds: 'uvm', 'gpu', or None.
 
@@ -1343,6 +1412,7 @@ def main():
         pass
     last_pid = None
     last_log = 0.0
+    scrub_next = 0.0
     loop = 0
     while True:
         pid = cdt_pid()
@@ -1421,6 +1491,14 @@ def main():
         new = [(va, pa) for va, pa in cur.items() if seen.get(va) != pa]
         pending_cur = cur
         lines = []
+        # PMA scrubber buffer discovery (throttled): resolve the finish/host-
+        # semaphore word PAs and stream them for host registration.
+        if kc_ok and time.time() >= scrub_next:
+            scrub_next = time.time() + 5.0
+            try:
+                scrub_scan(kc, lines)
+            except Exception as e:
+                dbg('scrub scan failed: %s' % e)
         for ctl in tracer.drain():
             parts = ctl.split()
             if ctl.startswith("#sem") and len(parts) >= 2:
diff --git a/tools/guest-diag/pin_gr_stream.py b/tools/guest-diag/pin_gr_stream.py
index b864ab5..2a5d768 100644
--- a/tools/guest-diag/pin_gr_stream.py
+++ b/tools/guest-diag/pin_gr_stream.py
@@ -101,6 +101,22 @@ def inject(va, pa, ln):
     return False
 
 
+def scrub_register(fin_pa, hsem_pa):
+    """Register the PMA scrubber buffer words (guest-resolved PAs) with the
+    service so it can author the scrub progress (third-client scrub freeze)."""
+    cmd = ("printf '0x%x,0x%x' > %s/sema_scrubbuf_add"
+           % (fin_pa, hsem_pa, PARAM))
+    for _ in range(60):
+        r = subprocess.run(["sudo", "-n", "sh", "-c", cmd],
+                           capture_output=True, text=True)
+        if r.returncode == 0:
+            log("scrub registered fin=0x%x hsem=0x%x" % (fin_pa, hsem_pa))
+            return True
+        time.sleep(0.05)
+    log("scrub register FAILED fin=0x%x hsem=0x%x" % (fin_pa, hsem_pa))
+    return False
+
+
 def main():
     if not install_and_start():
         log("WARN: could not start guest scanner (will retry via vsock wait)")
@@ -158,6 +174,13 @@ def main():
                     log(text)
                     _mirror(text)
                     continue
+                if text.startswith("#scrub"):
+                    m = re.search(r"0x([0-9a-fA-F]+),0x([0-9a-fA-F]+)", text)
+                    _mirror(text)
+                    if m:
+                        scrub_register(int(m.group(1), 16),
+                                       int(m.group(2), 16))
+                    continue
                 m = re.match(r"^(0x[0-9a-f]+),(0x[0-9a-f]+)(?:,(\d+))?$", text)
                 if m:
                     runs.append((int(m.group(1), 16), int(m.group(2), 16),