Function bench_fn

src/fused_mm_sampling/persistent_matmul.py:722–729

(label, reps, warmup_reps, fn, *args)

def bench_fn(label, reps, warmup_reps, fn, *args):
    print(f"Benchmarking {label}: ...", end="")
    for _ in range(warmup_reps):
        fn(*args)
    with proton_context():
        for _ in range(reps):
            fn(*args)
    print(f"\rBenchmarking {label}: done")


def bench(M, N, K, dtype, reps=100, warmup_reps=100):