import argparse
import time
import mlx.core as mx
import numpy as np
MLX_DTYPES = {
"float16": mx.float16,
"bfloat16": mx.bfloat16,
"float32": mx.float32,
}
def parse_cases(cases):
parsed = []
for spec in cases.split(","):
m, n, k, s = [int(x) for x in spec.split("x")]
parsed.append((m, n, k, s))
return parsed
def make_segments(k, num_segments, pattern, seed):
if pattern == "equal":
cuts = np.linspace(0, k, num_segments + 1, dtype=np.int64)
else:
rng = np.random.default_rng(seed)
cuts = rng.integers(0, k + 1, size=(num_segments - 1,), dtype=np.int64)
cuts = np.sort(cuts)
cuts = np.concatenate(([0], cuts, [k]))
return np.stack([cuts[:-1], cuts[1:]], axis=1).astype(np.uint32)
def numpy_segmented_mm_ref(a, b, segments):
out = []
for start, end in segments:
out.append(a[:, start:end] @ b[start:end, :])
return np.stack(out, axis=0)
def mlx_segmented_mm_loop(a, b, segments):
segments_list = segments.tolist()
out = []
for start, end in segments_list:
out.append(a[:, start:end] @ b[start:end, :])
return mx.stack(out, axis=0)
def bench_mlx(a, b, segments, warmup, iters):
for _ in range(warmup):
y = mx.segmented_mm(a, b, segments)
mx.eval(y)
mx.synchronize()
start = time.perf_counter()
for _ in range(iters):
y = mx.segmented_mm(a, b, segments)
mx.eval(y)
mx.synchronize()
end = time.perf_counter()
return (end - start) * 1e3 / iters
def bench_mlx_loop(a, b, segments, warmup, iters):
for _ in range(warmup):
y = mlx_segmented_mm_loop(a, b, segments)
mx.eval(y)
mx.synchronize()
start = time.perf_counter()
for _ in range(iters):
y = mlx_segmented_mm_loop(a, b, segments)
mx.eval(y)
mx.synchronize()
end = time.perf_counter()
return (end - start) * 1e3 / iters
def print_table(headers, rows):
widths = [len(h) for h in headers]
for row in rows:
for i, cell in enumerate(row):
widths[i] = max(widths[i], len(cell))
def fmt_row(row):
return (
"| "
+ " | ".join(f"{cell:<{widths[i]}}" for i, cell in enumerate(row))
+ " |"
)
sep = "|-" + "-|-".join("-" * w for w in widths) + "-|"
print(fmt_row(headers))
print(sep)
for row in rows:
print(fmt_row(row))
def main():
parser = argparse.ArgumentParser()
parser.add_argument(
"--cases",
default=(
"128x128x1024x16,"
"128x128x1024x32,"
"256x256x2048x16,"
"512x512x4096x32,"
"1024x1024x4096x32,"
"1024x1024x8192x64"
),
help="Comma-separated MxNxKxS list.",
)
parser.add_argument(
"--dtype",
default="float32",
choices=["float16", "bfloat16", "float32"],
)
parser.add_argument("--warmup", type=int, default=10)
parser.add_argument("--iters", type=int, default=50)
parser.add_argument(
"--segments",
choices=["equal", "random"],
default="random",
help="Segment generation pattern.",
)
parser.add_argument("--seed", type=int, default=0)
parser.add_argument("--no-check", action="store_true")
args = parser.parse_args()
mlx_dtype = MLX_DTYPES[args.dtype]
print(
f"dtype={args.dtype} warmup={args.warmup} iters={args.iters} segments={args.segments}"
)
headers = [
"Case",
"MLX ms",
"Loop ms",
"Speedup",
"MLX err",
"Loop err",
]
rows = []
cases = parse_cases(args.cases)
for idx, (m, n, k, s) in enumerate(cases):
rng = np.random.default_rng(args.seed + idx)
a_np = rng.standard_normal((m, k)).astype(np.float32)
b_np = rng.standard_normal((k, n)).astype(np.float32)
seg_np = make_segments(k, s, args.segments, args.seed + idx)
a_mx = mx.array(a_np, dtype=mlx_dtype)
b_mx = mx.array(b_np, dtype=mlx_dtype)
seg_mx = mx.array(seg_np, dtype=mx.uint32)
mx.eval(a_mx, b_mx, seg_mx)
mlx_err_str = ""
loop_err_str = ""
if not args.no_check:
y_mlx = mx.segmented_mm(a_mx, b_mx, seg_mx)
y_loop = mlx_segmented_mm_loop(a_mx, b_mx, seg_mx)
mx.eval(y_mlx, y_loop)
if args.dtype == "float32":
ref = numpy_segmented_mm_ref(
a_np.astype(np.float64),
b_np.astype(np.float64),
seg_np.tolist(),
)
mlx_err = np.max(np.abs(np.array(y_mlx, dtype=np.float64) - ref))
loop_err = np.max(np.abs(np.array(y_loop, dtype=np.float64) - ref))
else:
a_mx_f32 = mx.array(a_np, dtype=mx.float32)
b_mx_f32 = mx.array(b_np, dtype=mx.float32)
ref = mx.segmented_mm(a_mx_f32, b_mx_f32, seg_mx)
mx.eval(ref)
mlx_err = float(mx.max(mx.abs(ref - y_mlx.astype(mx.float32))).item())
loop_err = float(mx.max(mx.abs(ref - y_loop.astype(mx.float32))).item())
mlx_err_str = f"{mlx_err:.2e}"
loop_err_str = f"{loop_err:.2e}"
t_mlx = bench_mlx(a_mx, b_mx, seg_mx, args.warmup, args.iters)
t_loop = bench_mlx_loop(a_mx, b_mx, seg_mx, args.warmup, args.iters)
ratio = t_loop / t_mlx if t_mlx > 0 else float("inf")
rows.append(
[
f"{m}x{n}x{k}x{s}",
f"{t_mlx:.3f}",
f"{t_loop:.3f}",
f"{ratio:.2f}x",
mlx_err_str,
loop_err_str,
]
)
print_table(headers, rows)
if not args.no_check:
if args.dtype == "float32":
print("err: max|result - numpy_fp64_ref|")
else:
print("err: max|result - own_fp32_result|")
if __name__ == "__main__":
main()