wehub-resource-sync
132 lines
3.1 KiB
Python
132 lines
3.1 KiB
Python
#!/usr/bin/env python3
|
|
"""Step-by-step kernel benchmark. Run each step sequentially with gc."""
|
|
import sys, os, time, gc
|
|
import numpy as np
|
|
sys.path.insert(0, "~/work/cider/cider/lib")
|
|
import mlx.core as mx
|
|
import _cider_prim as ext
|
|
|
|
DEV = "~/work/cider/dev/step_kernels"
|
|
STEPS = ["step1", "step2", "step3", "step4", "step5"]
|
|
WARMUP = 5
|
|
REPEAT = 50
|
|
|
|
# Large tile shapes (M%128==0)
|
|
LARGE = [
|
|
(128, 4096, 4096),
|
|
(256, 4096, 4096),
|
|
(128, 3584, 18944),
|
|
(256, 3584, 18944),
|
|
]
|
|
|
|
# Small tile shapes (step3+ only)
|
|
SMALL = [
|
|
(1, 4096, 4096),
|
|
(16, 4096, 4096),
|
|
(32, 4096, 4096),
|
|
(64, 4096, 4096),
|
|
]
|
|
|
|
def bench_one(A_mx, B_mx, kernel_dir):
|
|
"""Warmup + benchmark, return median ms."""
|
|
for _ in range(WARMUP):
|
|
mx.eval(ext.int8_matmul_int32(A_mx, B_mx, kernel_dir))
|
|
times = []
|
|
for _ in range(REPEAT):
|
|
t0 = time.perf_counter()
|
|
mx.eval(ext.int8_matmul_int32(A_mx, B_mx, kernel_dir))
|
|
t1 = time.perf_counter()
|
|
times.append((t1 - t0) * 1000)
|
|
times.sort()
|
|
return times[len(times) // 2]
|
|
|
|
# === Large tile benchmark ===
|
|
print("=" * 90)
|
|
print("Part 1: Large Tile (M>=128)")
|
|
print("=" * 90)
|
|
|
|
large_results = {s: {} for s in STEPS}
|
|
|
|
for M, K, N in LARGE:
|
|
key = f"({M},{K},{N})"
|
|
np.random.seed(42)
|
|
A_np = np.random.randint(-127, 128, (M, K), dtype=np.int8)
|
|
B_np = np.random.randint(-127, 128, (K, N), dtype=np.int8)
|
|
A_mx = mx.array(A_np)
|
|
B_mx = mx.array(B_np)
|
|
mx.eval(A_mx, B_mx)
|
|
del A_np, B_np
|
|
gc.collect()
|
|
|
|
for step in STEPS:
|
|
kdir = os.path.join(DEV, step)
|
|
med = bench_one(A_mx, B_mx, kdir)
|
|
large_results[step][key] = med
|
|
print(f" {step} {key}: {med:.4f} ms")
|
|
|
|
del A_mx, B_mx
|
|
mx.clear_cache()
|
|
gc.collect()
|
|
|
|
# Print large table
|
|
print()
|
|
hdr = f"{'Shape':<22}"
|
|
for s in STEPS:
|
|
hdr += f" | {s:>10}"
|
|
print(hdr)
|
|
print("-" * 82)
|
|
for M, K, N in LARGE:
|
|
key = f"({M},{K},{N})"
|
|
row = f"{key:<22}"
|
|
for s in STEPS:
|
|
v = large_results[s].get(key, -1)
|
|
row += f" | {v:>10.3f}"
|
|
print(row)
|
|
|
|
# === Small tile benchmark (step3/4/5 only) ===
|
|
print()
|
|
print("=" * 60)
|
|
print("Part 2: Small Tile (M<128, step3/4/5)")
|
|
print("=" * 60)
|
|
|
|
small_steps = ["step3", "step4", "step5"]
|
|
small_results = {s: {} for s in small_steps}
|
|
|
|
for M, K, N in SMALL:
|
|
key = f"({M},{K},{N})"
|
|
np.random.seed(42)
|
|
A_np = np.random.randint(-127, 128, (M, K), dtype=np.int8)
|
|
B_np = np.random.randint(-127, 128, (K, N), dtype=np.int8)
|
|
A_mx = mx.array(A_np)
|
|
B_mx = mx.array(B_np)
|
|
mx.eval(A_mx, B_mx)
|
|
del A_np, B_np
|
|
gc.collect()
|
|
|
|
for step in small_steps:
|
|
kdir = os.path.join(DEV, step)
|
|
med = bench_one(A_mx, B_mx, kdir)
|
|
small_results[step][key] = med
|
|
print(f" {step} {key}: {med:.4f} ms")
|
|
|
|
del A_mx, B_mx
|
|
mx.clear_cache()
|
|
gc.collect()
|
|
|
|
# Print small table
|
|
print()
|
|
hdr = f"{'Shape':<22}"
|
|
for s in small_steps:
|
|
hdr += f" | {s:>10}"
|
|
print(hdr)
|
|
print("-" * 58)
|
|
for M, K, N in SMALL:
|
|
key = f"({M},{K},{N})"
|
|
row = f"{key:<22}"
|
|
for s in small_steps:
|
|
v = small_results[s].get(key, -1)
|
|
row += f" | {v:>10.3f}"
|
|
print(row)
|
|
|
|
print("\nDone!")
|