summaryrefslogtreecommitdiffstats
path: root/tools/guest-diag/guest_pin_stream.py
blob: 6c22e62b18562a3e7b15712c9b799e14435595f0 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
#!/usr/bin/env python3
# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 Open-IOV contributors
# guest_pin_stream.py - stream the CUDA client's push-buffer VA->PA pairs to
# the host over vsock, continuously and without the guest agent on the hot
# path. Runs in the guest as root (pagemap reads).
#
# Protocol (text lines to the host):
#   #pid <n>            the client pid just changed (host resets its state)
#   0x<va>,0x<pa>,<n>   n contiguous pages (va+i*0x1000 -> pa+i*0x1000)
#
# The host side is tools/guest-diag/pin_gr_stream.py.
#
# Usage: sudo python3 guest_pin_stream.py [host-cid] [port] [lo-hex] [hi-hex] [pa-max-hex]
import os
import re
import socket
import struct
import subprocess
import sys
import threading
import time

CID = int(sys.argv[1]) if len(sys.argv) > 1 else 2
PORT = int(sys.argv[2]) if len(sys.argv) > 2 else 45678
LO = int(sys.argv[3], 16) if len(sys.argv) > 3 else 0x200400000
HI = int(sys.argv[4], 16) if len(sys.argv) > 4 else 0x206000000
PAMAX = int(sys.argv[5], 16) if len(sys.argv) > 5 else 0x180000000
INTERVAL = float(os.environ.get("STREAM_INTERVAL", "0.1"))
BATCH = 4096
REGION = 0x200000
# The GR GPFIFO rings sit at 0x20020c000 + 0x3000*rank (region offset 0xc000
# ..0x45000) and the per-channel push allocations follow; 512 KB per 2 MB
# region covers both without scanning holes that hold nothing.
SPAN = int(os.environ.get("STREAM_SPAN", "0x80000"), 16)
# CE push-buffer window: one 1 MB allocation per CE channel at
# 0x120000000 + 0x100000*rank; SPAN2 covers the 1 MB boundary inside a
# 2 MB region. Extend HI2 to 0x121400000: the UVM semaphore pool (the
# tracker sems the client's context-init wait polls, pool word at
# 0x121000000 + 4*idx) lives above the push buffers and must reach the host
# through this fast stream - otherwise only the slow QMP chain resolution
# pins it, and the pin lands after the client's RM thread-state deadline.
LO2 = int(os.environ.get("STREAM_LO2", "0x120000000"), 16)
HI2 = int(os.environ.get("STREAM_HI2", "0x121400000"), 16)
SPAN2 = int(os.environ.get("STREAM_SPAN2", "0x180000"), 16)
# UVM semaphore pool window: the tracker pool word lives at
# 0x121000000 + 4*idx, above the CE window head span (LO2+SPAN2 =
# 0x120180000), so scan_window never reaches it. Scan every 4 KB page of
# [POOL_LO, POOL_HI) directly - cheap (16 pages) and host-visible as pins.
POOL_LO = int(os.environ.get("STREAM_POOL_LO", "0x121000000"), 16)
POOL_HI = int(os.environ.get("STREAM_POOL_HI", "0x121010000"), 16)
# UVM system-memory AVMAs: the client's sysmem buffers (the DtoH staging
# buffer among them) are /dev/zero mappings at high GPU VAs (e.g.
# 0x79c21c800000). They are not inside the two fixed windows, and only the
# GMMU page tables can translate them, so stream their present pages as
# pins too. Restricted to the high VA range and to nvidia/zero mappings;
# the PAMAX filter rejects device (BAR) PFNs.
VMA_LO = int(os.environ.get("STREAM_VMA_LO", "0x40000000000"), 16)
VMA_EVERY = int(os.environ.get("STREAM_VMA_EVERY", "1"))
VMA_MAX_PAGES = int(os.environ.get("STREAM_VMA_MAX_PAGES", "16384"))
VMA_PATHS = tuple(p for p in os.environ.get(
    "STREAM_VMA_PATHS", "zero,nvidia").split(",") if p)
TRACE = ("/sys/kernel/tracing/trace"
         if os.path.exists("/sys/kernel/tracing/trace")
         else "/sys/kernel/debug/tracing/trace")
TRACE_PIPE = ("/sys/kernel/tracing/trace_pipe"
              if os.path.exists("/sys/kernel/tracing/trace_pipe")
              else "/sys/kernel/debug/tracing/trace_pipe")
SEM_EVERY = int(os.environ.get("STREAM_SEM_EVERY", "4"))


class Kcore:
    """Kernel-truth reader: /proc/kcore for contents + manual page-table walk.

    The host's QMP-based struct walk intermittently reads garbage while the
    client spins in user mode (PTI hides the kernel half from the sampled
    vCPU), which left the CE push-buffer / semaphore-pool pages unpinned on
    ~half of all boots (the deterministic cudaMalloc stall). The guest can
    read its own kernel memory reliably: /proc/kcore has an ELF segment for
    every kernel mapping, and the direct-map segment lets us walk the page
    tables (init_top_pgt -> ... -> PTE) to turn the vmalloc CPU alias of a
    UVM rm_mem into the guest physical addresses the host needs.
    """

    def __init__(self):
        self.ok = False
        self.fd = None
        self.segs = []
        self.phys_base = 0
        self.page_offset_base = 0
        self.pgd_pa = 0
        self._pa_cache = {}

    def load(self):
        try:
            self.fd = open("/proc/kcore", "rb")
            hdr = self.fd.read(64)
            if hdr[:4] != b"\x7fELF":
                return False
            e_phoff = struct.unpack_from("<Q", hdr, 32)[0]
            e_phentsize, e_phnum = struct.unpack_from("<HH", hdr, 54)
            segs = []
            for i in range(e_phnum):
                self.fd.seek(e_phoff + i * e_phentsize)
                ph = self.fd.read(56)
                p_type, _fl, p_off, p_va, _pa, p_sz = struct.unpack_from(
                    "<IIQQQQ", ph, 0)
                if p_type == 1:
                    segs.append((p_va, p_va + p_sz, p_off))
            self.segs = segs
            syms = {}
            for line in open("/proc/kallsyms"):
                parts = line.split()
                if len(parts) == 3:
                    syms.setdefault(parts[2], int(parts[0], 16))
            need = ("phys_base", "page_offset_base", "init_top_pgt")
            if any(s not in syms or not syms[s] for s in need):
                return False
            self.phys_base = self.read64(syms["phys_base"])
            self.page_offset_base = self.read64(syms["page_offset_base"])
            if not self.phys_base or not self.page_offset_base:
                return False
            self.pgd_pa = (syms["init_top_pgt"] - 0xFFFFFFFF80000000 +
                           self.phys_base)
            self.ok = True
            return True
        except OSError:
            return False

    def read(self, va, n):
        for lo, hi, off in self.segs:
            if lo <= va and va + n <= hi:
                self.fd.seek(off + (va - lo))
                d = self.fd.read(n)
                return d if d and len(d) == n else None
        return None

    def read64(self, va):
        d = self.read(va, 8)
        return struct.unpack("<Q", d)[0] if d else None

    def phys64(self, pa):
        return self.read64(self.page_offset_base + (pa & 0xFFFFFFFFF))

    def va_to_pa(self, va):
        page = va & ~0xFFF
        cached = self._pa_cache.get(page)
        if cached is not None:
            return cached or None
        pa = self.pgd_pa
        for level, shift in enumerate((39, 30, 21, 12)):
            e = self.phys64(pa + ((va >> shift) & 0x1FF) * 8)
            if e is None or not (e & 1):
                self._pa_cache[page] = 0
                return None
            if level == 3:
                pa = e & 0x000FFFFFFFF000
                break
            if e & 0x80:
                if level == 1:
                    pa = (e & 0x000FFFFFC0000000) + (va & ((1 << 30) - 1))
                    break
                if level == 2:
                    pa = (e & 0x000FFFFFFFE00000) + (va & ((1 << 21) - 1))
                    break
                self._pa_cache[page] = 0
                return None
            pa = e & 0x000FFFFFFFF000
        self._pa_cache[page] = pa
        return pa


def is_kern_ptr(v):
    return v and (v >> 48) == 0xFFFF


def is_gpu_va(v):
    return v and 0x120000000 <= v < 0x200000000


def _direct_map_pa(kc, va):
    """Direct-map aliases translate linearly: pa = va - page_offset_base."""
    if kc.page_offset_base <= va < kc.page_offset_base + (1 << 40):
        return va - kc.page_offset_base
    return None


def _vmalloc_pages_for(kc, va):
    """Page count of the vmap/vmalloc area containing va (0 if unknown)."""
    now = time.time()
    if not hasattr(_vmalloc_pages_for, "areas") or \
            now - _vmalloc_pages_for.ts > 5.0:
        areas = []
        try:
            for line in open("/proc/vmallocinfo"):
                p = line.split()
                if len(p) < 3:
                    continue
                try:
                    lo, hi = p[0].split("-")
                    areas.append((int(lo, 16), int(hi, 16)))
                except ValueError:
                    continue
        except OSError:
            pass
        _vmalloc_pages_for.areas = areas
        _vmalloc_pages_for.ts = now
    base = va & ~0xFFF
    for lo, hi in _vmalloc_pages_for.areas:
        if lo <= base < hi:
            return (hi - base) >> 12
    return 0


def resolve_mem_pins(kc, mem, out, default_pages=16):
    """uvm_rm_mem kptr -> GPU-VA/GPA runs (direct map or page-table walk)."""
    if not is_kern_ptr(mem):
        return 0
    vas = [kc.read64(mem + 0x10 + i * 8) for i in range(16)]
    cpu = next((v for v in vas if is_kern_ptr(v)), None)
    gpus = [v for v in vas if is_gpu_va(v)]
    if not cpu or not gpus:
        return 0
    size = kc.read64(mem + 0x60) or 0
    dm = _direct_map_pa(kc, cpu)
    emitted = 0
    if dm is not None:
        npages = max(1, min(size >> 12, 4098)) if size else default_pages
        for gv in gpus:
            out.append("0x%x,0x%x,%d\n" % (gv, dm, npages))
            emitted += 1
        return emitted
    npages = (max(1, min(size >> 12, 4098)) if size
              else _vmalloc_pages_for(kc, cpu))
    if not npages:
        npages = 1
    runs = []
    cur = None
    for k in range(npages):
        pa = kc.va_to_pa(cpu + k * 0x1000)
        if pa is None:
            if cur:
                runs.append(cur)
                cur = None
            continue
        if cur and pa == cur[1] + cur[2] * 0x1000:
            cur[2] += 1
        else:
            if cur:
                runs.append(cur)
            cur = [k, pa, 1]
    if cur:
        runs.append(cur)
    for gv in gpus:
        for k, pa, n in runs:
            out.append("0x%x,0x%x,%d\n" % (gv + k * 0x1000, pa, n))
            emitted += 1
    return emitted


def _emit_ring_pins(kc, gpu_va, entries_cpu, entries, out):
    """Pin the GPFIFO ring pages (and the pages after it, where the GPPUT
    buffer usually sits) from the ring's CPU mapping."""
    emitted = 0
    if not is_kern_ptr(entries_cpu) or not (32 <= entries <= 0x100000):
        return 0
    npages = min(((entries * 8) + 0xFFF) // 0x1000 + 2, 8)
    runs = []
    cur = None
    for k in range(npages):
        pa = kc.va_to_pa(entries_cpu + k * 0x1000)
        if pa is None:
            if cur:
                runs.append(cur)
                cur = None
            continue
        if cur and pa == cur[1] + cur[2] * 0x1000:
            cur[2] += 1
        else:
            if cur:
                runs.append(cur)
            cur = [k, pa, 1]
    if cur:
        runs.append(cur)
    for k, pa, n in runs:
        out.append("0x%x,0x%x,%d\n" % (gpu_va + k * 0x1000, pa, n))
        emitted += 1
    return emitted


def resolve_channel_from_sem(kc, sem, out):
    """Tracking sem (embedded in uvm_channel_t) -> ring/GPPUT pins.

    uvch events do not fire on this driver build, but the channel's tracking
    semaphore is a uvm_channel_t member updated on every push. Locate
    channel_info by scanning around the sem for its gpFifoGpuVa value (a
    page-aligned GPU VA in the GR ring window), then read the ring/GPPUT
    CPU pointers (UvmGpuChannelInfo offsets: gpPut +0x08, gpFifoEntries
    +0x10, numGpFifoEntries +0x18, gpFifoGpuVa +0x58, gpPutGpuVa +0x60).
    """
    if not is_kern_ptr(sem):
        return 0
    lo = sem - 0x2000
    data = kc.read(lo, 0x4000)
    if not data:
        return 0
    for off in range(0, len(data) - 8, 8):
        v = struct.unpack_from("<Q", data, off)[0]
        if not (GR_RING_LO <= v < GR_RING_HI) or (v & 0xFFF):
            continue
        pos = lo + off
        entries_cpu = kc.read64(pos - 0x48)
        entries = kc.read64(pos - 0x40)
        if entries:
            entries &= 0xFFFFFFFF
        emitted = _emit_ring_pins(kc, v, entries_cpu, entries, out)
        gp_put = kc.read64(pos - 0x50)
        gp_put_va = kc.read64(pos + 0x08)
        if is_kern_ptr(gp_put) and gp_put_va and \
                GR_RING_LO <= gp_put_va < 0x201000000:
            pa = kc.va_to_pa(gp_put)
            if pa is not None:
                out.append("0x%x,0x%x,1\n" % (gp_put_va & ~0xFFF,
                                               pa & ~0xFFF))
                emitted += 1
        return emitted
    return 0


def resolve_sem_pins(kc, sem, out):
    vals = [kc.read64(sem + i * 8) for i in range(12)]
    if not vals[0] or not is_kern_ptr(vals[0]):
        return 0
    mem = kc.read64(vals[0])
    if not mem or not is_kern_ptr(mem):
        return 0
    n = resolve_mem_pins(kc, mem, out, default_pages=16)
    try:
        n += resolve_channel_from_sem(kc, sem, out)
    except Exception:
        pass
    return n


def resolve_pb_pins(kc, pb, out):
    mem = kc.read64(pb + 8)
    if not mem:
        return 0
    return resolve_mem_pins(kc, mem, out, default_pages=64)


# UvmGpuChannelInfo field offsets (nv_uvm_types.h, x86_64):
#   +0x00 gpGet, +0x08 gpPut, +0x10 gpFifoEntries, +0x18 numGpFifoEntries,
#   +0x58 gpFifoGpuVa, +0x60 gpPutGpuVa, +0x68 gpGetGpuVa.
# The offset of channel_info inside uvm_channel_t is config-dependent, so the
# struct is located by searching the channel for its gpFifoGpuVa value (a GPU
# VA in the GR ring window, page aligned).
GR_RING_LO = 0x200200000
GR_RING_HI = 0x200400000
CHAN_INFO_GPFIFO_VA = 0x58


def resolve_channel_pins(kc, ch, out):
    """uvm_channel_t kptr -> pin the GPFIFO ring + GPPUT buffers.

    The ring and GP_PUT buffers are CPU-mapped in the guest kernel (UVM
    writes pushes and the put counter directly), so their guest-physical
    backing is reachable with the page-table walk. Without these pins the
    host cannot read GR channel entries (the stall class: RM event loop
    waits for a completion whose ring entry was never processed).
    """
    if not is_kern_ptr(ch):
        return 0
    base = None
    gpu_va = None
    data = kc.read(ch, 0x800)
    if not data:
        return 0
    for off in range(0, len(data) - 8, 8):
        v = struct.unpack_from("<Q", data, off)[0]
        if GR_RING_LO <= v < GR_RING_HI and (v & 0xFFF) == 0:
            base = ch + off - CHAN_INFO_GPFIFO_VA
            gpu_va = v
            break
    if base is None:
        return 0
    entries = kc.read64(base + 0x18)
    if entries:
        entries &= 0xFFFFFFFF
    if not entries or entries > 0x100000:
        return 0
    emitted = 0
    entries_cpu = kc.read64(base + 0x10)
    if is_kern_ptr(entries_cpu):
        nbytes = entries * 8
        npages = min((nbytes + 0xFFF) // 0x1000 + 2, 8)
        runs = []
        cur = None
        for k in range(npages):
            pa = kc.va_to_pa(entries_cpu + k * 0x1000)
            if pa is None:
                if cur:
                    runs.append(cur)
                    cur = None
                continue
            if cur and pa == cur[1] + cur[2] * 0x1000:
                cur[2] += 1
            else:
                if cur:
                    runs.append(cur)
                cur = [k, pa, 1]
        if cur:
            runs.append(cur)
        for k, pa, n in runs:
            out.append("0x%x,0x%x,%d\n" % (gpu_va + k * 0x1000, pa, n))
            emitted += 1
    gp_put = kc.read64(base + 0x08)
    gp_put_va = kc.read64(base + 0x60)
    if is_kern_ptr(gp_put) and gp_put_va and \
            GR_RING_LO <= gp_put_va < 0x201000000:
        pa = kc.va_to_pa(gp_put)
        if pa is not None:
            out.append("0x%x,0x%x,1\n" % (gp_put_va & ~0xFFF, pa & ~0xFFF))
            emitted += 1
    gp_get = kc.read64(base + 0x00)
    gp_get_va = kc.read64(base + 0x68)
    if is_kern_ptr(gp_get) and gp_get_va and \
            GR_RING_LO <= gp_get_va < 0x201000000:
        pa = kc.va_to_pa(gp_get)
        if pa is not None:
            out.append("0x%x,0x%x,1\n" % (gp_get_va & ~0xFFF, pa & ~0xFFF))
            emitted += 1
    return emitted



def clear_trace():
    """Drop stale ring events so the fallback tail reader cannot re-learn a
    previous client's freed sems after a client switch. The new client's
    events are delivered fresh through the consuming trace_pipe reader."""
    try:
        with open(TRACE, "w") as f:
            f.write("")
    except OSError:
        pass


class TraceReader:
    """Stream uvmt sems / uvbp pushbuffers from the guest trace via trace_pipe.

    trace_pipe is a consuming stream: every event is delivered exactly once,
    with no ring-window race. The old tail-of-trace reader polled the last
    256 KB every 0.3 s, so a one-shot sem= event could age out of the window
    when the client's spin flooded the ring before the next poll - that is
    the second-client stall class (the previous client's freed sem was
    re-learned from the stale ring while the new client's sem was missed).
    """

    def __init__(self):
        self.lines = []
        self.seen = set()
        self.lock = threading.Lock()
        self.stop = False
        self.fd = None
        self.t = threading.Thread(target=self._run, daemon=True)
        self.t.start()

    def _open(self):
        try:
            self.fd = os.open(TRACE_PIPE, os.O_RDONLY)
        except OSError:
            self.fd = None

    def _run(self):
        while not self.stop:
            if self.fd is None:
                self._open()
                if self.fd is None:
                    time.sleep(1)
                    continue
            try:
                data = os.read(self.fd, 65536)
            except OSError:
                try:
                    os.close(self.fd)
                except OSError:
                    pass
                self.fd = None
                time.sleep(0.5)
                continue
            if not data:
                continue
            with self.lock:
                for m in re.finditer(rb"sem=0x([0-9a-fA-F]+)", data):
                    v = int(m.group(1), 16)
                    if ("s", v) not in self.seen:
                        self.seen.add(("s", v))
                        self.lines.append("#sem 0x%x\n" % v)
                for m in re.finditer(rb"pb=0x([0-9a-fA-F]+)", data):
                    v = int(m.group(1), 16)
                    if ("p", v) not in self.seen:
                        self.seen.add(("p", v))
                        self.lines.append("#pb 0x%x\n" % v)
                for m in re.finditer(rb"ch=0x([0-9a-fA-F]+)", data):
                    v = int(m.group(1), 16)
                    if ("c", v) not in self.seen:
                        self.seen.add(("c", v))
                        self.lines.append("#ch 0x%x\n" % v)

    def drain(self):
        with self.lock:
            out, self.lines = self.lines, []
        return out


def new_sems(seen):
    """New tracking-sem kptrs (uvmt lines) in the trace tail.

    Read in the guest so the host does not need a guest-agent round trip for
    every semaphore discovery; the trace is only tailed while a client runs.
    """
    try:
        out = subprocess.run(["tail", "-c", "262144", TRACE],
                             capture_output=True, timeout=10).stdout
    except Exception:
        return []
    found = []
    for m in re.finditer(rb"sem=0x([0-9a-fA-F]+)", out):
        v = int(m.group(1), 16)
        if v not in seen:
            seen.add(v)
            found.append(v)
    return found


def cdt_pid():
    prefixes = tuple(p for p in os.environ.get("CLIENT_PREFIXES",
                                               "cdt,va,victory").split(",") if p)
    best = None
    for d in os.listdir('/proc'):
        if not d.isdigit():
            continue
        try:
            with open('/proc/%s/comm' % d) as f:
                name = f.read().strip()
            if not any(name.startswith(p) for p in prefixes):
                continue
            # Skip zombies: a defunct previous client keeps its comm and
            # would otherwise shadow the live client (pagemap reads then fail
            # and no pins are ever streamed).
            with open('/proc/%s/stat' % d) as f:
                st = f.read()
            rp = st.rfind(')')
            if rp >= 0 and st[rp + 2:rp + 3] == 'Z':
                continue
            best = int(d)
        except OSError:
            pass
    return best


def scan_window(fd, lo, hi, span, out):
    base = lo & ~(REGION - 1)
    while base < hi:
        va = max(base, lo) & ~0xFFF
        end = min(base + span, hi)
        while va < end:
            n = min(BATCH, (end - va) // 0x1000)
            if n <= 0:
                break
            os.lseek(fd, (va // 4096) * 8, 0)
            data = os.read(fd, n * 8)
            got = len(data) // 8
            for i in range(got):
                e = int.from_bytes(data[i * 8:i * 8 + 8], 'little')
                if not ((e >> 63) & 1):
                    continue
                pfn = e & ((1 << 55) - 1)
                pa = pfn << 12
                if not pa or (PAMAX and pa >= PAMAX):
                    continue
                out[va + i * 0x1000] = pa
            va += got * 0x1000
        base += REGION


def scan_range(fd, lo, hi, out):
    """Present pages of one explicit VA range, every 4 KB page.

    Same pagemap present/PFN/PAMAX filtering as scan_window, but with no
    REGION alignment or head-span limit: walks [lo, hi) directly. Used for
    the UVM semaphore pool window, which sits past the CE head span.
    """
    va = lo & ~0xFFF
    while va < hi:
        n = min(BATCH, (hi - va) // 0x1000)
        if n <= 0:
            break
        os.lseek(fd, (va // 4096) * 8, 0)
        data = os.read(fd, n * 8)
        got = len(data) // 8
        if got <= 0:
            break
        for i in range(got):
            e = int.from_bytes(data[i * 8:i * 8 + 8], 'little')
            if not ((e >> 63) & 1):
                continue
            pfn = e & ((1 << 55) - 1)
            pa = pfn << 12
            if not pa or (PAMAX and pa >= PAMAX):
                continue
            out[va + i * 0x1000] = pa
        va += got * 0x1000


def scan_vmas(fd, pid, out):
    """Present pages of the client's high-VA nvidia//dev/zero mappings.

    These are UVM system-memory mappings (sysmem buffers, DtoH staging);
    the fast windows do not cover them and the host cannot resolve them
    from guest RAM alone. `fd` is the pagemap fd of the client.
    """
    n = 0
    try:
        maps = open('/proc/%d/maps' % pid)
    except OSError:
        return
    try:
        for line in maps:
            parts = line.split()
            if len(parts) < 6:
                continue
            path = parts[5]
            if not any(p in path for p in VMA_PATHS):
                continue
            try:
                a, b = parts[0].split('-')
                lo, hi = int(a, 16), int(b, 16)
            except ValueError:
                continue
            if lo < VMA_LO:
                continue
            va = (lo + 0xFFF) & ~0xFFF
            while va < hi and n < VMA_MAX_PAGES:
                cnt = min(256, (hi - va) // 0x1000)
                if cnt <= 0:
                    break
                os.lseek(fd, (va // 4096) * 8, 0)
                data = os.read(fd, cnt * 8)
                got = len(data) // 8
                for i in range(got):
                    e = int.from_bytes(data[i * 8:i * 8 + 8], 'little')
                    if not ((e >> 63) & 1):
                        continue
                    pfn = e & ((1 << 55) - 1)
                    pa = pfn << 12
                    if not pa or (PAMAX and pa >= PAMAX):
                        continue
                    out[va + i * 0x1000] = pa
                va += got * 0x1000
                n += got
    finally:
        maps.close()


def scan(pid):
    """Present pages in the first SPAN bytes of every REGION, as {va: pa}.

    The pushes sit at the start of each per-channel allocation (2/4 MB
    aligned); scanning only the allocation heads keeps the per-loop page
    table work small, which is what makes the stream fast while the client
    is mapping and the guest is loaded.

    Two windows: the GR window (rings/pushbuffers/counter page) and the CE
    push-buffer window (0x120000000+, one 1 MB allocation per CE channel).
    The CE window uses a wider span because the per-channel allocations are
    1 MB apart inside a 2 MB region.
    """
    fd = os.open('/proc/%d/pagemap' % pid, os.O_RDONLY)
    out = {}
    scan.loop += 1
    try:
        scan_window(fd, LO, HI, SPAN, out)
        if HI2 > LO2:
            scan_window(fd, LO2, HI2, SPAN2, out)
        if POOL_HI > POOL_LO:
            scan_range(fd, POOL_LO, POOL_HI, out)
        if VMA_EVERY and scan.loop % VMA_EVERY == 0:
            scan_vmas(fd, pid, out)
    finally:
        os.close(fd)
    return out


scan.loop = 0


def main():
    with open('/tmp/guest_pin_stream.pid', 'w') as f:
        f.write('%d\n' % os.getpid())
    sock = None
    seen = {}
    seen_sems = set()
    seen_pbs = set()
    seen_chs = set()
    resolved_sems = set()
    resolved_pbs = set()
    resolved_chs = set()
    retry_at = {}
    tracer = TraceReader()
    kc = Kcore()
    kc_ok = kc.load()
    try:
        with open('/tmp/guest_pin_stream.log', 'a') as f:
            f.write('%s resolver: kcore=%s (sem->pool/pushbuffer pins)\n'
                    % (time.strftime('%T'), 'ON' if kc_ok else 'OFF'))
    except OSError:
        pass
    last_pid = None
    last_log = 0.0
    loop = 0
    while True:
        pid = cdt_pid()
        loop += 1
        if pid is None:
            last_pid = None
            seen = {}
            seen_sems = set()
            seen_pbs = set()
            seen_chs = set()
            time.sleep(INTERVAL)
            continue
        if pid != last_pid:
            last_pid = pid
            seen = {}
            seen_sems = set()
            seen_pbs = set()
            seen_chs = set()
            resolved_sems = set()
            resolved_pbs = set()
            resolved_chs = set()
            retry_at = {}
            clear_trace()
            if sock is not None:
                try:
                    sock.sendall(("#pid %d\n" % pid).encode())
                except OSError:
                    sock.close()
                    sock = None
        t0 = time.time()
        try:
            cur = scan(pid)
        except OSError:
            time.sleep(INTERVAL)
            continue
        dt = time.time() - t0
        # Diagnostic: the GR completion-counter page is the one page the
        # client's context-init wait polls; log its pagemap entry every loop
        # so a scan blind spot is visible in the guest log.
        cpm = None
        try:
            fdp = os.open('/proc/%d/pagemap' % pid, os.O_RDONLY)
            os.lseek(fdp, (0x20400fec0 // 4096) * 8, 0)
            cpm = int.from_bytes(os.read(fdp, 8), 'little')
            os.close(fdp)
        except OSError:
            cpm = None
        diag = ('cpm=%s' % ('present' if (cpm or 0) >> 63 & 1 else 'absent')
                if cpm is not None else 'cpm=ERR')
        # Guaranteed fast path for the GR completion-counter page: if the
        # window scan misses it (observed once), add it explicitly so the host
        # can register the slots before the client's first wait.
        if cpm is not None and (cpm >> 63) & 1:
            cpa = (cpm & ((1 << 55) - 1)) << 12
            if cpa and (not PAMAX or cpa < PAMAX):
                cur[0x20400f000] = cpa
        new = [(va, pa) for va, pa in cur.items() if seen.get(va) != pa]
        pending_cur = cur
        lines = []
        for ctl in tracer.drain():
            parts = ctl.split()
            if ctl.startswith("#sem") and len(parts) >= 2:
                seen_sems.add(int(parts[1], 16))
            elif ctl.startswith("#pb") and len(parts) >= 2:
                seen_pbs.add(int(parts[1], 16))
            elif ctl.startswith("#ch") and len(parts) >= 2:
                seen_chs.add(int(parts[1], 16))
            lines.append(ctl)
        if kc_ok:
            now_r = time.time()
            added = 0
            for sem in list(seen_sems):
                if sem in resolved_sems or now_r < retry_at.get(sem, 0):
                    continue
                got = []
                try:
                    resolve_sem_pins(kc, sem, got)
                except Exception:
                    got = []
                if got:
                    lines.extend(got)
                    resolved_sems.add(sem)
                    added += len(got)
                else:
                    retry_at[sem] = now_r + 1.0
            for pb in list(seen_pbs):
                if pb in resolved_pbs or now_r < retry_at.get(pb, 0):
                    continue
                got = []
                try:
                    resolve_pb_pins(kc, pb, got)
                except Exception:
                    got = []
                if got:
                    lines.extend(got)
                    resolved_pbs.add(pb)
                    added += len(got)
                else:
                    retry_at[pb] = now_r + 1.0
            for ch in list(seen_chs):
                if ch in resolved_chs or now_r < retry_at.get(ch, 0):
                    continue
                got = []
                try:
                    resolve_channel_pins(kc, ch, got)
                except Exception:
                    got = []
                if got:
                    lines.extend(got)
                    resolved_chs.add(ch)
                    added += len(got)
                else:
                    retry_at[ch] = now_r + 1.0
            if added:
                try:
                    with open('/tmp/guest_pin_stream.log', 'a') as f:
                        f.write('%s resolver pins: +%d run(s) '
                                '(sems=%d/%d pbs=%d/%d chs=%d/%d)\n'
                                % (time.strftime('%T'), added,
                                   len(resolved_sems), len(seen_sems),
                                   len(resolved_pbs), len(seen_pbs),
                                   len(resolved_chs), len(seen_chs)))
                except OSError:
                    pass
        now = time.time()
        if now - last_log > 1.0:
            last_log = now
            try:
                with open('/tmp/guest_pin_stream.log', 'a') as f:
                    f.write('%s pid=%s present=%d new=%d sem=%d scan=%.3fs %s\n'
                            % (time.strftime('%T'), pid, len(cur), len(new),
                               len(seen_sems), dt, diag))
            except OSError:
                pass
        if not new and not lines:
            seen = pending_cur
            time.sleep(INTERVAL)
            continue
        if new:
            new.sort()
            va0, pa0, n = new[0][0], new[0][1], 1
            for va, pa in new[1:]:
                if va == va0 + n * 0x1000 and pa == pa0 + n * 0x1000:
                    n += 1
                    continue
                lines.append('0x%x,0x%x,%d\n' % (va0, pa0, n))
                va0, pa0, n = va, pa, 1
            lines.append('0x%x,0x%x,%d\n' % (va0, pa0, n))
        payload = ''.join(lines).encode()
        if sock is None:
            try:
                sock = socket.socket(socket.AF_VSOCK, socket.SOCK_STREAM)
                sock.settimeout(5)
                sock.connect((CID, PORT))
                sock.settimeout(None)
                if last_pid is not None:
                    sock.sendall(("#pid %d\n" % last_pid).encode())
            except OSError:
                if sock:
                    sock.close()
                sock = None
                time.sleep(INTERVAL)
                continue
        try:
            sock.sendall(payload)
            seen = pending_cur
        except OSError:
            sock.close()
            sock = None
        time.sleep(INTERVAL)


if __name__ == '__main__':
    main()