|
|
|
|
|
|
|
|
import argparse |
|
|
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) |
|
|
|
