Angelos Katharopoulos
110 lines
3.5 KiB
Python
110 lines
3.5 KiB
Python
# Copyright © 2023-2024 Apple Inc.
|
|
|
|
import argparse
|
|
|
|
import mlx.core as mx
|
|
import torch
|
|
from time_utils import measure_runtime
|
|
|
|
|
|
def benchmark_slice_update_mlx(dst_shape, slice_shape, slice_range, dtype, iters=10):
|
|
def slice_update(arguments):
|
|
for i in range(iters):
|
|
arguments["dst"] = (
|
|
arguments["dst"].at[slice_range].add(arguments["updates"])
|
|
)
|
|
mx.eval(arguments)
|
|
|
|
dtype = getattr(mx, dtype)
|
|
arguments = {
|
|
"dst": mx.random.normal(dst_shape).astype(dtype),
|
|
"updates": mx.random.normal(slice_shape).astype(dtype),
|
|
}
|
|
|
|
runtime = measure_runtime(slice_update, arguments=arguments)
|
|
bytes_processed = (
|
|
arguments["dst"][slice_range].nbytes * 2 + arguments["updates"].nbytes
|
|
) * iters
|
|
bandwidth_gb_s = bytes_processed / runtime / 1e6
|
|
return runtime, bandwidth_gb_s
|
|
|
|
|
|
def benchmark_slice_update_torch(
|
|
dst_shape, slice_shape, slice_range, device, dtype, iters=10
|
|
):
|
|
def slice_update(dst, updates, slice_range):
|
|
for i in range(iters):
|
|
dst[slice_range] = dst[slice_range] + updates
|
|
if device == torch.device("mps"):
|
|
torch.mps.synchronize()
|
|
|
|
dtype = getattr(torch, dtype)
|
|
updates = torch.randn(slice_shape, dtype=dtype).to(device)
|
|
dst = torch.randn(dst_shape, dtype=dtype).to(device)
|
|
|
|
runtime = measure_runtime(
|
|
slice_update, dst=dst, updates=updates, slice_range=slice_range
|
|
)
|
|
bytes_processed = (dst[slice_range].nbytes * 2 + updates.nbytes) * iters
|
|
bandwidth_gb_s = bytes_processed / runtime / 1e6
|
|
return runtime, bandwidth_gb_s
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser("Slice update benchmarks.")
|
|
parser.add_argument("--cpu", action="store_true", help="Use the CPU.")
|
|
args = parser.parse_args()
|
|
|
|
if args.cpu:
|
|
mx.set_default_device(mx.cpu)
|
|
device = torch.device("cpu")
|
|
elif torch.mps.is_available():
|
|
device = torch.device("mps")
|
|
elif torch.cuda.is_available():
|
|
device = torch.device("cuda")
|
|
else:
|
|
raise ValueError()
|
|
|
|
dtypes = ["float32", "bfloat16"]
|
|
|
|
test_cases = [
|
|
((10_000_000,), slice(0, 1_000_000), (1_000_000,)),
|
|
((100_000,), slice(10_000, 20_000), (10_000,)),
|
|
((1000, 64), slice(100, 200), (100, 64)),
|
|
((100, 100, 64), slice(20, 40), (20, 100, 64)),
|
|
(
|
|
(2048, 2048, 128),
|
|
(slice(500, 1500), slice(200, 1200), slice(32, 96)),
|
|
(1000, 1000, 64),
|
|
),
|
|
(
|
|
(2048, 2048, 128),
|
|
(slice(1800, 1850), slice(100, 200), slice(64, 128)),
|
|
(50, 100, 64),
|
|
),
|
|
(
|
|
(2048, 2048, 128),
|
|
(slice(1000, 1010), slice(1000, 1010), slice(64, 128)),
|
|
(10, 10, 64),
|
|
),
|
|
]
|
|
|
|
print(
|
|
f"{'Dtype':<12} {'Dst Shape':<25} {'Update Shape':<20} "
|
|
f"{'MLX (ms)':<12} {'MLX GB/s':<12} {'Torch (ms)':<12} {'Torch GB/s':<12}"
|
|
)
|
|
print("-" * 110)
|
|
|
|
for dtype in dtypes:
|
|
for dst_shape, slice_range, update_shape in test_cases:
|
|
mlx_time, mlx_bw = benchmark_slice_update_mlx(
|
|
dst_shape, update_shape, slice_range, dtype
|
|
)
|
|
torch_time, torch_bw = benchmark_slice_update_torch(
|
|
dst_shape, update_shape, slice_range, device, dtype
|
|
)
|
|
print(
|
|
f"{dtype:<12} {str(dst_shape):<25} {str(update_shape):<20} "
|
|
f"{mlx_time:<12.3f} {mlx_bw:<12.2f} {torch_time:<12.3f} {torch_bw:<12.2f}"
|
|
)
|