(M, N, K, dtype, reps=100, warmup_reps=100)
| 730 | |
| 731 | |
| 732 | def bench(M, N, K, dtype, reps=100, warmup_reps=100): |
| 733 | a = torch.randn((M, K), device="cuda", dtype=torch.float16).to(dtype) |
| 734 | b = torch.randn((K, N), device="cuda", dtype=torch.float16).to(dtype) |
| 735 | bT = b |
| 736 | b = b.T.contiguous() |
| 737 | |
| 738 | if device_blas is not None: |
| 739 | blas_name = device_blas_name() |
| 740 | bench_fn(blas_name, reps, warmup_reps, device_blas_matmul, a, b) |
| 741 | if dtype == torch.float16: |
| 742 | bench_fn("torch", reps, warmup_reps, torch_matmul, a, b) |
| 743 | bench_fn("torch-nontransposed", reps, warmup_reps, torch_matmul_nontransposed, a, bT) |
| 744 | bench_fn("naive", reps, warmup_reps, matmul, a, b.T) |
| 745 | bench_fn("persistent", reps, warmup_reps, matmul_persistent, a, b.T) |
| 746 | warp_specialize = [False, True] if HAS_WARP_SPECIALIZE else [False] |
| 747 | for ws in warp_specialize: |
| 748 | ws_str = "_ws" if ws else "" |
| 749 | # disable on-host warpspec on Hopper |
| 750 | if HAS_HOST_TENSOR_DESC and not (is_hopper() and ws): |
| 751 | bench_fn( |
| 752 | f"tma_persistent{ws_str}", |
| 753 | reps, |
| 754 | warmup_reps, |
| 755 | lambda a, b: matmul_tma_persistent(a, b, ws), |
| 756 | a, |
| 757 | b, |
| 758 | ) |
| 759 | bench_fn(f"tma{ws_str}", reps, warmup_reps, lambda a, b: matmul_tma(a, b, ws), a, b) |
| 760 | if HAS_TENSOR_DESC: |
| 761 | bench_fn( |
| 762 | f"descriptor_persistent{ws_str}", |
| 763 | reps, |
| 764 | warmup_reps, |
| 765 | lambda a, b: matmul_descriptor_persistent(a, b, ws), |
| 766 | a, |
| 767 | b, |
| 768 | ) |
| 769 | |
| 770 | |
| 771 | def run_test(expect, fn, a, b, label, enabled=True): |
no test coverage detected