#!/usr/bin/env python3
# sem_counters.py - enumerate a guest process's per-channel GPU completion
# counters (the u32 slots the client polls to observe push completion).
#
# The counters live in one sysmem page mapped through /dev/nvidiactl (rw-s);
# each channel owns a 16-byte-strided u32 in it (observed layout: 20 counters
# at page offset 0xec0, one per GR channel). This translates the slots to
# guest PAs through /proc/<pid>/pagemap so the host side can register them
# with the sema_counter module parameter.
#
# Usage:
#   sudo python3 sem_counters.py <pid> <page-va> [count] [stride]
#   sudo python3 sem_counters.py <pid> <page-va> --pairs <h1,h2,...>
#
# --pairs zips the given channel handles with the counter indices in order
# (ascending handle = ascending counter index, the observed layout) and
# prints "<handle>,<va>,<pa>" lines: handle/pa for the sema_counter parameter,
# va/pa for the sema_map_add pin so pushes releasing the counter resolve.
#
# Needs root in the guest (pagemap reads). Uses only os/sys.
import os
import sys


def parse_args(argv):
    pid = int(argv[0])
    page_va = int(argv[1], 16)
    count = 20
    stride = 16
    offset = 0xec0
    handles = None
    rest = argv[2:]
    i = 0
    while i < len(rest):
        a = rest[i]
        if a == '--pairs':
            handles = [int(x, 16) for x in rest[i + 1].split(',') if x]
            i += 2
            continue
        if a.startswith('--count='):
            count = int(a.split('=', 1)[1], 0)
        elif a.startswith('--stride='):
            stride = int(a.split('=', 1)[1], 0)
        elif a.startswith('--offset='):
            offset = int(a.split('=', 1)[1], 0)
        elif i == 0:
            count = int(a, 0)
        elif i == 1:
            stride = int(a, 0)
        i += 1
    return pid, page_va, count, stride, offset, handles


def main():
    pid, page_va, count, stride, offset, handles = parse_args(sys.argv[1:])

    fd = os.open('/proc/%d/pagemap' % pid, os.O_RDONLY)
    try:
        pas = []
        for i in range(count):
            va = page_va + offset + i * stride
            os.lseek(fd, (va // 4096) * 8, 0)
            e = int.from_bytes(os.read(fd, 8), 'little')
            present = (e >> 63) & 1
            pfn = e & ((1 << 55) - 1)
            pa = (pfn << 12) | (va & 0xfff)
            pas.append((va, present, pa))
    finally:
        os.close(fd)

    if handles is not None:
        if len(handles) != len(pas):
            print('ERROR: %d handles but %d counters' %
                  (len(handles), len(pas)), file=sys.stderr)
            return 1
        for h, (va, present, pa) in zip(handles, pas):
            if not present:
                print('ERROR: VA=0x%x not present for handle 0x%x' %
                      (va, h), file=sys.stderr)
                return 1
            print('0x%x,0x%x,%d' % (h, va, pa))
    else:
        for i, (va, present, pa) in enumerate(pas):
            print('idx=%d VA=0x%x present=%d PA=0x%x' %
                  (i, va, present, pa))
    return 0


if __name__ == '__main__':
    sys.exit(main())