guest-diag: demand-pin content resolver searches +/-16 pages and learns the winning per-area offset (multi-MB CE staging pushes never resolved at +/-1)

1 files changed, 32 insertions, 16 deletions

diff --git a/tools/guest-diag/guest_pin_stream.py b/tools/guest-diag/guest_pin_stream.py
index 3f4183c..6012dd1 100644
--- a/tools/guest-diag/guest_pin_stream.py
+++ b/tools/guest-diag/guest_pin_stream.py
@@ -300,32 +300,47 @@ def _nv_vmap_areas(kc):
     return areas
 
 
+# Learned page-offset per nv_vm_map_pages area for the demand content match
+# (2026-10-16b): the first successful resolution records which page index the
+# VA maps to, and later pushes in the same allocation retry that offset first
+# instead of walking the whole +/-16 neighbourhood.
+_area_off = {}
+
+
 def resolve_va_by_content(kc, va, out):
     """Demand pin: find the page carrying GPU VA 'va' by content match.
 
     The host reports the CE push-buffer VAs whose GPFIFO entries it cannot
     read (dmesg 'stalled VA #N: 0x...'). The guest walks every
     nv_vm_map_pages vmap area and tests the candidate leaf page at the VA's
-    in-window offset (plus one neighbour on either side, for the per-mem
-    header-page offset ambiguity) against the host's method-stream prefix
-    check. On a match it emits the exact pair, so the host can translate
-    and consume the entry - no kprobe event, no object-chain discovery, no
-    resolution lottery. Only the requested VA is pinned, so a wrong
-    candidate cannot poison other VAs.
+    in-window offset (plus a neighbourhood, for the per-mem header-page
+    offset ambiguity) against the host's method-stream prefix check. On a
+    match it emits the exact pair, so the host can translate and consume the
+    entry - no kprobe event, no object-chain discovery, no resolution
+    lottery. Only the requested VA is pinned, so a wrong candidate cannot
+    poison other VAs.
+
+    2026-10-16b: the offset search window was +/-1 page. Pushes of a client
+    whose pushbuffer vmap carries a larger per-allocation header (observed
+    with multi-MB CE staging: pushes at 0x120000f58+ never resolved and the
+    demand loop re-sent the same VA for minutes while the client crawled)
+    fall outside it. Widen to +/-16 pages and LEARN the winning offset per
+    vmap area, because consecutive pushes advance sequentially inside the
+    same allocation and share the offset.
     """
     if not (PB_CE_LO <= va < PB_CE_HI):
         return False
-    # 2026-10-16 fix: the CE push-buffer window is 16 MB (PB_CE_LO..PB_CE_HI)
-    # and each rank is a 1 MB slice of ONE nv_vm_map_pages vmap. The old
-    # `& 0xFFFFF` folded every 1 MB boundary to offset 0, so ranks 0x120700000,
-    # 0x120800000 and 0x120a00000 all resolved to page 0 of the first area -
-    # the same wrong PA - and the host read garbage for every rank but the
-    # first. Keep the full 16 MB offset: the page index is the true offset
-    # into the mem, and the per-area pages bound rejects areas too small.
     off = (va - PB_CE_LO) & 0xFFFFFF
     k = off >> 12
-    for (lo, _hi, pages) in _nv_vmap_areas(kc):
-        for adj in (0, 1, -1):
+    areas = _nv_vmap_areas(kc)
+    for (lo, _hi, pages) in areas:
+        learned = _area_off.get(lo)
+        if learned is not None:
+            adjs = [learned]
+        else:
+            adjs = [0, 1, -1]
+            adjs += [d for m in (2, 4, 8, 16) for d in (m, -m)]
+        for adj in adjs:
             idx = k + adj
             if idx < 0 or idx >= pages:
                 continue
@@ -338,12 +353,13 @@ def resolve_va_by_content(kc, va, out):
             pa = kc.va_to_pa(cpu)
             if pa is None:
                 continue
+            _area_off[lo] = adj
             out.append("0x%x,0x%x,1\n" % (va & ~0xFFF, pa))
             dbg("demand-va: 0x%x -> pa=0x%x (area off=%d adj=%+d)"
                 % (va & ~0xFFF, pa, idx, adj))
             return True
     dbg("demand-va: 0x%x not found by content (%d areas)"
-        % (va & ~0xFFF, len(_nv_vmap_areas(kc))))
+        % (va & ~0xFFF, len(areas)))
     return False