chore: import upstream snapshot with attribution

2026-07-13 12:34:46 +08:00
commit f4e68ed970
84 changed files with 14896 additions and 0 deletions
@@ -0,0 +1,131 @@
+#!/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!")