import os
import subprocess
import time
import matplotlib.pyplot as plt
import mlx.core as mx
import numpy as np
import torch
results_dir = "./results"
if not os.path.isdir(results_dir):
os.mkdir(results_dir)
device_name = subprocess.check_output(["sysctl", "-n", "machdep.cpu.brand_string"])
device_name = device_name.decode("utf-8").strip("\n")
N_warmup = 5
N_iter_bench = 50
N_iter_func = 20
out_vec_sizes = [128, 512, 2048, 4096]
in_vec_sizes = [128, 512, 2048, 4096]
benchmark_vector_lens = []
benchmark_vector_lens += [(i + 1) * 4096 for i in range(8)][::2]
benchmark_vector_lens += [(i + 1) * 4095 for i in range(8)][::2]
benchmark_vector_lens += [(i + 1) * 4097 for i in range(8)][::2]
benchmark_vector_lens += [64, 128, 512, 1024, 2048, 11008, 32000]
benchmark_vector_lens.sort()
def bench(f, m, v):
for i in range(N_warmup):
f(m, v)
torch.mps.synchronize()
s = time.perf_counter_ns()
for i in range(N_iter_bench):
f(m, v)
e = time.perf_counter_ns()
return (e - s) * 1e-9
def gemv_mlx(m, v):
ys = []
for i in range(N_iter_func):
y = m @ v
ys.append(y)
mx.eval(ys)
return ys
def gemv_t_mlx(m, v):
ys = []
for i in range(N_iter_func):
y = v @ m
ys.append(y)
mx.eval(ys)
return ys
@torch.no_grad()
def gemv_torch(m, v):
ys = []
for i in range(N_iter_func):
y = m @ v
ys.append(y)
torch.mps.synchronize()
return ys
@torch.no_grad()
def gemv_t_torch(m, v):
ys = []
for i in range(N_iter_func):
y = v @ m
ys.append(y)
torch.mps.synchronize()
return ys
def bench_lens(in_vec_len, out_vec_len, np_dtype, transpose=False):
shape_mat = (in_vec_len, out_vec_len) if transpose else (out_vec_len, in_vec_len)
shape_vec = (1, in_vec_len) if transpose else (in_vec_len, 1)
mat_npy = np.random.normal(0.0, 2.0 / in_vec_len, shape_mat).astype(np_dtype)
vec_npy = np.random.normal(0.0, 2.0 / in_vec_len, shape_vec).astype(np_dtype)
mat_mlx = mx.array(mat_npy)
vec_mlx = mx.array(vec_npy)
mat_trc = torch.from_numpy(mat_npy).to("mps")
vec_trc = torch.from_numpy(vec_npy).to("mps")
torch.mps.synchronize()
time_torch = (
bench(gemv_t_torch, mat_trc, vec_trc)
if transpose
else bench(gemv_torch, mat_trc, vec_trc)
)
time_mlx = (
bench(gemv_t_mlx, mat_mlx, vec_mlx)
if transpose
else bench(gemv_mlx, mat_mlx, vec_mlx)
)
c_mlx = (
np.asarray(vec_mlx @ mat_mlx) if transpose else np.asarray(mat_mlx @ vec_mlx)
)
c_npy = (vec_npy @ mat_npy) if transpose else (mat_npy @ vec_npy)
if not np.allclose(c_mlx, c_npy, atol=2e-5):
print(
f"Failed at {shape_mat} [transpose = {transpose}] with max(|a - b|) = {np.max(np.abs(c_npy - c_mlx))}"
)
return time_mlx, time_torch
def get_gflop_count(in_vec_len, out_vec_len):
return float(2.0 * N_iter_bench * N_iter_func * in_vec_len * out_vec_len) / float(
1024**3
)
def get_gbyte_size(in_vec_len, out_vec_len, np_dtype):
n_elem = in_vec_len * out_vec_len + in_vec_len + out_vec_len
item_size = 4 if np_dtype == np.float32 else 2
return float(N_iter_bench * N_iter_func * n_elem * item_size) / float(1024**3)
def bench_with_in_len(ax, in_vec_len, out_vector_lens, dtype, transpose):
np_dtype = getattr(np, dtype)
mlx_gb_s = []
mlx_gflops = []
pyt_gb_s = []
pyt_gflops = []
for out_vec_len in out_vector_lens:
gflop_count = get_gflop_count(in_vec_len, out_vec_len)
gbyte_size = get_gbyte_size(in_vec_len, out_vec_len, np_dtype)
time_mlx, time_torch = bench_lens(in_vec_len, out_vec_len, np_dtype, transpose)
mlx_gb_s.append(gbyte_size / time_mlx)
pyt_gb_s.append(gbyte_size / time_torch)
mlx_gflops.append(gflop_count / time_mlx)
pyt_gflops.append(gflop_count / time_torch)
if transpose:
title = f"gemv_t ([1, {in_vec_len}] [{in_vec_len}, out_vec_len]) | {dtype}"
else:
title = f"gemv ([out_vec_len, {in_vec_len}] X [{in_vec_len}, 1] ) | {dtype}"
ax.plot(out_vector_lens, mlx_gb_s, "tab:blue", label="MLX")
ax.plot(out_vector_lens, pyt_gb_s, "tab:red", label="Torch")
ax.set_title(title)
ax.set(xlabel="out_vector_len", ylabel="Performance (GB/s)")
ax.legend()
def bench_with_out_len(ax, out_vec_len, in_vector_lens, dtype, transpose):
np_dtype = getattr(np, dtype)
mlx_gb_s = []
mlx_gflops = []
pyt_gb_s = []
pyt_gflops = []
for in_vec_len in in_vector_lens:
gflop_count = get_gflop_count(in_vec_len, out_vec_len)
gbyte_size = get_gbyte_size(in_vec_len, out_vec_len, np_dtype)
time_mlx, time_torch = bench_lens(in_vec_len, out_vec_len, np_dtype, transpose)
mlx_gb_s.append(gbyte_size / time_mlx)
pyt_gb_s.append(gbyte_size / time_torch)
mlx_gflops.append(gflop_count / time_mlx)
pyt_gflops.append(gflop_count / time_torch)
if transpose:
title = f"([1, in_vec_len] [in_vec_len, {out_vec_len}])"
else:
title = f"([{out_vec_len}, in_vec_len] X [in_vec_len, 1] )"
ax.plot(in_vector_lens, mlx_gb_s, "tab:blue", label="MLX")
ax.plot(in_vector_lens, pyt_gb_s, "tab:red", label="Torch")
ax.set_title(title)
ax.set(xlabel="in_vector_len", ylabel="Performance (GB/s)")
ax.legend()
for transpose in (False, True):
for dtype in ("float32", "float16", "complex64"):
fig, axs = plt.subplots(
len(in_vec_sizes), 2, figsize=(8.5, 11), layout="constrained"
)
for i, in_vec_len in enumerate(in_vec_sizes):
bench_with_in_len(
axs[i][0], in_vec_len, benchmark_vector_lens, dtype, transpose
)
for i, out_vec_len in enumerate(out_vec_sizes):
bench_with_out_len(
axs[i][1], out_vec_len, benchmark_vector_lens, dtype, transpose
)
op_name = "gemv_t" if transpose else "gemv"
fig.suptitle(f"{device_name}: {dtype} {op_name}")
fig.savefig(
os.path.join(
results_dir, f"{device_name.replace(' ', '_')}_{dtype}_{op_name}.pdf"
)
)
plt.close(fig)