109 lines
3.2 KiB
Python
109 lines
3.2 KiB
Python
"""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()
|