Save inference parity implementation and evaluation harness

tools/mtplx-allocator-reference.py

@@ -0,0 +1,108 @@
+"""Model-free allocation/cache receipts from the installed pinned MTPLX runtime.
+
+No replacement allocator, timings, model loading or downloads. stdout is JSONL.
+Physical addresses are replaced with run-local identities before emission.
+"""
+import gc
+import json
+import os
+
+import mlx.core as mx
+import numpy as np
+
+
+def main():
+    assert mx.__version__ == "0.32.2"
+    arrays = {}
+    identities = {}
+    page = os.sysconf("SC_PAGE_SIZE")
+    mx.set_cache_limit(4 * 1024 * 1024)
+    mx.clear_cache()
+    assert mx.get_active_memory() == 0
+
+    def emit(op, **fields):
+        print(json.dumps(dict(op=op, active=mx.get_active_memory(),
+                              cached=mx.get_cache_memory(), **fields)), flush=True)
+
+    def alloc(name, size):
+        before = mx.get_active_memory()
+        arrays[name] = mx.array(np.zeros(size, dtype=np.uint8))
+        mx.eval(arrays[name])
+        mx.synchronize()
+        pointer = np.asarray(arrays[name]).__array_interface__["data"][0] if size else None
+        identity = identities.setdefault(pointer, len(identities)) if size else None
+        emit("alloc", name=name, size=size, buffer=identity,
+             storage=mx.get_active_memory() - before)
+
+    def free(name):
+        del arrays[name]
+        gc.collect()
+        mx.synchronize()
+        emit("free", name=name)
+
+    def clear():
+        assert not arrays
+        mx.clear_cache()
+        emit("clear")
+
+    def limit(size):
+        mx.set_cache_limit(size)
+        emit("limit", size=size)
+
+    emit("init", page=page, limit=4 * 1024 * 1024, runtime=mx.__version__)
+    for i, size in enumerate([0, 1, 255, 256, 257, page - 1, page, page + 1,
+                              2 * page - 1, 2 * page, 2 * page + 1]):
+        alloc(str(i), size)
+    for name in list(arrays):
+        free(name)
+    for i, size in enumerate([256, 1, page, page + 1, 2 * page + 1]):
+        alloc(str(i), size)
+    for name in list(arrays):
+        free(name)
+    clear()
+
+    # Equal-size multimap entries reuse the oldest insertion, not the newest.
+    for name in "abc":
+        alloc(name, 257)
+    for name in "bac":
+        free(name)
+    for name in "xyz":
+        alloc(name, 256)
+    for name in "xyz":
+        free(name)
+    clear()
+
+    # Strict upper bounds: min(2 * request, request + 2 * page).
+    for stored, requested in [(512, 256), (511, 256),
+                              (4 * page, 2 * page), (3 * page, 2 * page),
+                              (6 * page, 4 * page), (5 * page, 4 * page)]:
+        alloc("old", stored)
+        free("old")
+        alloc("new", requested)
+        free("new")
+        clear()
+
+    for name, size in zip("abcd", [64, 128, 256, 512]):
+        alloc(name, size)
+    for name in "cadb":
+        free(name)
+    limit(600)
+    alloc("e", 2048)  # trim from oldest, overshooting by complete buffers
+    free("e")       # free checks the pre-insertion size, allowing overshoot
+    alloc("f", 10000)
+    free("f")
+    limit(1)
+    alloc("g", 20000)  # >=90% requested release clears the whole pool
+    free("g")
+    clear()
+    limit(4096)
+    alloc("old", 64)
+    free("old")
+    limit(0)  # changing the limit does not clear; reuse is still attempted first
+    alloc("new", 64)
+    free("new")
+    clear()
+
+
+if __name__ == "__main__":
+    main()