#!/usr/bin/env python3
# reg_sampler.py - read-only BAR0 register sampler for GR/FECS/SKED ground truth.
#
# Purpose
#   Sample the physical GPU's PGRAPH/FECS/SKED/PFIFO registers at high rate
#   while a GR launch is in flight, so the engine-side state transitions
#   (context load, CWD table occupancy, SKED activity, GRFIFO progress) are
#   observable without the driver's help. Used to capture the ground-truth
#   behaviour of a working CUDA channel (distro driver) and to inspect the
#   ladder's forwarded launches.
#
# Safety
#   Read-only MMIO. Safe to run while any driver owns the device. Needs root
#   (resource0 mmap). Do not run concurrently with a driver rebind/FLR.
#
# Usage
#   sudo python3 reg_sampler.py <seconds> <Hz> <outfile> [label]
#   python3 analyze_regs.py <outfile>        # value distribution + transitions
#
# Notes
#   - The sampler writes one TSV line per sample: t and one u32 per register.
#   - `badfxxxx` values are the GPU returning its "bad fence/clock-gated"
#     pattern during driver handoff windows; filter them in analysis.
#   - Offset source: kernel/gsp_phys_booter.c reg-probe table (TU104).
import mmap
import os
import struct
import sys
import time

DEV = "/sys/bus/pci/devices/0000:01:00.0/resource0"
REGS = [
    (0x000200, "PMC_ENABLE"),
    (0x400100, "PGRAPH_INTR"),
    (0x400108, "PGRAPH_EXCEPTION"),
    (0x400118, "PGRAPH_EXCEPTION1"),
    (0x40011C, "PGRAPH_EXCEPTION2"),
    (0x40013C, "PGRAPH_INTR_EN"),
    (0x502C90, "GPC0_EXCEPTION"),
    (0x502C84, "GPC0_ACTIVITY1"),
    (0x504730, "G0T0S0_WARP_ESR"),
    (0x504734, "G0T0S0_ESR_MASK"),
    (0x504738, "G0T0S0_PC"),
    (0x50473C, "G0T0S0_PC_HI"),
    (0x50476C, "G0T0S0_ESR_INFO"),
    (0x5047B0, "G0T0S1_WARP_ESR"),
    (0x504830, "G0T0S2_WARP_ESR"),
    (0x5048B0, "G0T0S3_WARP_ESR"),
    (0x500324, "GPC0_GPCMMU_ESR"),
    (0x50AC90, "GPC1_EXCEPTION"),
    (0x50C730, "G1T0S0_WARP_ESR"),
    (0x50C734, "G1T0S0_ESR_MASK"),
    (0x50C738, "G1T0S0_PC"),
    (0x50C73C, "G1T0S0_PC_HI"),
    (0x50C76C, "G1T0S0_ESR_INFO"),
    (0x508324, "GPC1_GPCMMU_ESR"),
    (0x512C90, "GPC2_EXCEPTION"),
    (0x51AC90, "GPC3_EXCEPTION"),
    (0x522C90, "GPC4_EXCEPTION"),
    (0x52AC90, "GPC5_EXCEPTION"),
    (0x400700, "PGRAPH_STATUS"),
    (0x400504, "GRFIFO_STATUS"),
    (0x400704, "TRAPPED_ADDR"),
    (0x400708, "TRAPPED_DATA_LO"),
    (0x404000, "FE_HWW_ESR"),
    (0x405B00, "CWD_FS"),
    (0x405B2C, "CWD_PARTITION_CTL"),
    (0x405B30, "CWD_PARTITION_DATA"),
    (0x407020, "SKED_HWW_ESR"),
    (0x407024, "SKED_HWW_ESR_EN"),
    (0x407028, "SKED_HWW_INFO_1"),
    (0x40702C, "SKED_HWW_INFO_2"),
    (0x407030, "SKED_HWW_INFO_3"),
    (0x407040, "SKED_STATUS_0"),
    (0x407054, "SKED_ACTIVITY"),
    (0x407060, "CWD_TABLE_0"),
    (0x407064, "CWD_TABLE_1"),
    (0x407068, "CWD_TABLE_2"),
    (0x40706C, "CWD_TABLE_3"),
    (0x407070, "SKED_CACHE_0"),
    (0x407074, "SKED_CACHE_1"),
    (0x407090, "SKED_L15_CACHE_0"),
    (0x407080, "TASKTABLE_0"),
    (0x407084, "TASKTABLE_1"),
    (0x407088, "TASKTABLE_2"),
    (0x409008, "FECS_IRQSTAT"),
    (0x409100, "FECS_CPUCTL"),
    (0x409400, "FECS_CTXSW_STATUS_1"),
    (0x409b00, "FECS_CURRENT_CTX"),
    (0x409b04, "FECS_NEW_CTX"),
    (0x409c00, "FECS_CTXSW_STATUS_FE_0"),
    (0x409c18, "FECS_HOST_INT_STATUS"),
    (0x409040, "FECS_MAILBOX0"),
    (0x409044, "FECS_MAILBOX1"),
    (0x40a300, "FE_QMD_VIRT_BASEA_0"),
    (0x40a304, "FE_QMD_VIRT_BASEA_1"),
    (0x40a400, "FE_QMD_VIRT_BASEB_0"),
    (0x40a500, "FE_QMD_VIRT_CONTROL_0"),
    (0x40a504, "FE_QMD_VIRT_CONTROL_1"),
    (0x40a508, "FE_QMD_VIRT_CONTROL_2"),
    (0x40a50c, "FE_QMD_VIRT_CONTROL_3"),
    (0x41a040, "GPCCS_MAILBOX0"),
    (0x41a100, "GPCCS_CPUCTL"),
]


def main():
    secs = float(sys.argv[1])
    hz = float(sys.argv[2])
    out = sys.argv[3]
    label = sys.argv[4] if len(sys.argv) > 4 else "sample"
    fd = os.open(DEV, os.O_RDONLY | os.O_SYNC)
    m = mmap.mmap(fd, 0x1000000, mmap.MAP_SHARED, mmap.PROT_READ)
    period = 1.0 / hz
    t0 = time.monotonic()
    n = 0
    with open(out, "w") as f:
        f.write("# %s t=%.3f\n" % (label, t0))
        f.write("# ts " + " ".join(nm for _, nm in REGS) + "\n")
        while time.monotonic() - t0 < secs:
            ts = time.monotonic() - t0
            vals = [struct.unpack_from("<I", m, off)[0] for off, _ in REGS]
            f.write("%.6f " % ts + " ".join("%08x" % v for v in vals) + "\n")
            n += 1
            dt = period - (time.monotonic() - t0 - ts)
            if dt > 0:
                time.sleep(dt)
    print("samples=%d" % n)


if __name__ == "__main__":
    main()