| 63 | |
| 64 | |
| 65 | def benchmark(M=BENCH_SIZE, N=BENCH_SIZE, K=BENCH_SIZE, |
| 66 | warmup=10, iters=100): |
| 67 | lib = ctypes.CDLL(SO_PATH) |
| 68 | lib.benchmark_kernel.argtypes = ([ctypes.c_void_p] * 3 |
| 69 | + [ctypes.c_int] * 5) |
| 70 | lib.benchmark_kernel.restype = ctypes.c_float |
| 71 | |
| 72 | A = torch.randn(M, K, dtype=torch.bfloat16, device="cuda") |
| 73 | B = torch.randn(N, K, dtype=torch.bfloat16, device="cuda") |
| 74 | D = torch.empty(M, N, dtype=torch.bfloat16, device="cuda") |
| 75 | |
| 76 | ms = lib.benchmark_kernel(A.data_ptr(), B.data_ptr(), D.data_ptr(), |
| 77 | M, N, K, warmup, iters) |
| 78 | tflops = 2 * M * N * K / (ms * 1e-3) / 1e12 |
| 79 | |
| 80 | for _ in range(warmup): |
| 81 | torch.matmul(A, B.t()) |
| 82 | torch.cuda.synchronize() |
| 83 | start = torch.cuda.Event(enable_timing=True) |
| 84 | end = torch.cuda.Event(enable_timing=True) |
| 85 | start.record() |
| 86 | for _ in range(iters): |
| 87 | torch.matmul(A, B.t()) |
| 88 | end.record() |
| 89 | torch.cuda.synchronize() |
| 90 | torch_ms = start.elapsed_time(end) / iters |
| 91 | torch_tflops = 2 * M * N * K / (torch_ms * 1e-3) / 1e12 |
| 92 | |
| 93 | print(f"Benchmark ({M}x{N}x{K}):") |
| 94 | print(f" Ours: {ms:.3f} ms, {tflops:.1f} TFLOPS") |
| 95 | print(f" Torch: {torch_ms:.3f} ms, {torch_tflops:.1f} TFLOPS") |
| 96 | print(f" Ratio: {tflops / torch_tflops:.3f}x") |
| 97 | return ms, tflops, torch_ms, torch_tflops |
| 98 | |
| 99 | |
| 100 | if __name__ == "__main__": |