|
|
|
|
|
|
|
|
import argparse |
|
|
|
|
|
import mlx.core as mx |
|
|
from time_utils import time_fn |
|
|
|
|
|
B = 8 |
|
|
T = 1024 |
|
|
D = 512 |
|
|
|
|
|
|
|
|
def time_batch_matmul(): |
|
|
mx.random.seed(3) |
|
|
a = mx.random.uniform(shape=(B, T, D)) |
|
|
b = mx.random.uniform(shape=(D, D)) |
|
|
c = mx.random.uniform(shape=(B, T, D)) |
|
|
mx.eval(a, b, c) |
|
|
|
|
|
time_fn(mx.matmul, a, b) |
|
|
|
|
|
def batch_vjp_first(): |
|
|
return mx.vjp(mx.matmul, [a, b], [c])[1][0] |
|
|
|
|
|
time_fn(batch_vjp_first) |
|
|
|
|
|
def batch_vjp_second(): |
|
|
return mx.vjp(mx.matmul, [a, b], [c])[1][1] |
|
|
|
|
|
time_fn(batch_vjp_second) |
|
|
|
|
|
|
|
|
def time_unbatch_matmul(): |
|
|
mx.random.seed(3) |
|
|
a = mx.random.uniform(shape=(B * T, D)) |
|
|
b = mx.random.uniform(shape=(D, D)) |
|
|
c = mx.random.uniform(shape=(B * T, D)) |
|
|
mx.eval(a, b, c) |
|
|
time_fn(mx.matmul, a, b) |
|
|
|
|
|
def unbatch_vjp_first(): |
|
|
return mx.matmul(c, mx.transpose(b)) |
|
|
|
|
|
time_fn(unbatch_vjp_first) |
|
|
|
|
|
def unbatch_vjp_second(): |
|
|
return mx.matmul(mx.transpose(a), c) |
|
|
|
|
|
time_fn(unbatch_vjp_second) |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
parser = argparse.ArgumentParser("MLX benchmarks.") |
|
|
parser.add_argument("--gpu", action="store_true", help="Use the Metal back-end.") |
|
|
args = parser.parse_args() |
|
|
if args.gpu: |
|
|
mx.set_default_device(mx.gpu) |
|
|
else: |
|
|
mx.set_default_device(mx.cpu) |
|
|
|
|
|
time_batch_matmul() |
|
|
time_unbatch_matmul() |
|
|
|
