[BUG] [Metal] Wired memory is released shortly after GPU idle (residency never renewed) #4609

import subprocess
import time

import mlx.core as mx


def wired_gib():
    page = int(subprocess.check_output(["sysctl", "-n", "hw.pagesize"]))
    out = subprocess.check_output(["vm_stat"]).decode()
    line = next(l for l in out.splitlines() if "wired down" in l)
    return int(line.split(":")[1].strip(" .")) * page / 2**30


def timed_sum(a):
    start = time.perf_counter()
    mx.eval(a.sum())
    return (time.perf_counter() - start) * 1000


def idle(seconds, beat=None):
    # Wait, running a tiny GPU op every `beat` seconds if given.
    # Return the lowest wired figure seen, sampled every 0.5 s.
    one = mx.zeros((1,))
    lowest = wired_gib()
    start = last = time.monotonic()
    while time.monotonic() - start < seconds:
        if beat and time.monotonic() - last >= beat:
            mx.eval(one + 1)
            last = time.monotonic()
        lowest = min(lowest, wired_gib())
        time.sleep(0.5)
    return lowest


print("mlx", mx.__version__)
print(subprocess.check_output(["sysctl", "iogpu.disable_wired_collector"]).decode().strip())

a = mx.random.uniform(shape=(2**30,))  # 4 GiB of float32
mx.eval(a)
mx.set_wired_limit(mx.get_active_memory())
warm = timed_sum(a)
print(f"warm compute {warm:.1f} ms, wired {wired_gib():.1f} GiB")

for beat in (None, 1, 3):
    timed_sum(a)
    lowest = idle(12, beat)
    first = timed_sum(a)
    second = timed_sum(a)
    label = "idle 12 s" if beat is None else f"idle 12 s, op every {beat} s"
    print(f"{label:27} lowest wired {lowest:4.1f} GiB   "
          f"first compute {first:7.1f} ms   second {second:6.1f} ms")