open-standards/jcp/tools/counter-resolution/gpu_rail_period.py

"""How long does the GPU Energy rail hold its value?

I declared 10 ms as its update period. PR #146 reports it "republishes in
steps and can hold 300-600 ms under load". Declaring too short a period is the
unsafe direction: it lets a span be graded that the counter never covered.
This measures the distribution of gaps between actual changes.
"""
import time, statistics, threading, sys
exec(open('rails.py').read().split('import time, subprocess, sys')[0])

def gpu_abs():
    """Absolute cumulative GPU nJ: delta against a fixed first sample."""
    return None

def gaps(duration_s, label, load=None):
    stop = threading.Event()
    t = None
    if load:
        t = threading.Thread(target=load, args=(stop,), daemon=True); t.start()
        time.sleep(0.3)                      # let the load spin up
    prev_s = sample(); last_change = time.perf_counter(); g = []; steps = []
    t_end = time.perf_counter() + duration_s
    while time.perf_counter() < t_end:
        time.sleep(0.002)                    # 2 ms sampling: finer than any claim
        cur_s = sample()
        d = IOR.IOReportCreateSamplesDelta(prev_s, cur_s, None)
        tot = sum(v for _g, _sg, n, u, v in rails(d) if n == "GPU Energy")
        if tot > 0:
            now = time.perf_counter()
            g.append((now - last_change) * 1e3); steps.append(tot / 1e6)
            last_change = now
            prev_s = cur_s                   # only advance the baseline on a change
    if load:
        stop.set(); t.join(timeout=2)
    if not g:
        print(f"{label:<16} no change in {duration_s}s"); return
    g.sort()
    print(f"{label:<16} n={len(g):<4} gap ms: p50={statistics.median(g):7.1f} "
          f"p90={g[int(.9*len(g))-1]:7.1f} max={g[-1]:7.1f}   median step={statistics.median(steps):8.2f} mJ")

def mlx_load(stop):
    import mlx.core as mx
    a = mx.random.normal((2048, 2048)); b = mx.random.normal((2048, 2048))
    while not stop.is_set():
        c = a @ b; mx.eval(c)

gaps(6.0, "idle")
gaps(6.0, "mlx matmul", mlx_load)