(M, N, K, dtype)
| 780 | |
| 781 | |
| 782 | def validate(M, N, K, dtype): |
| 783 | print(f"{M=}, {N=}, {K=}, verification naive vs: ") |
| 784 | a = torch.randn((M, K), device="cuda", dtype=torch.float16).to(dtype) |
| 785 | b = torch.randn((K, N), device="cuda", dtype=torch.float16).to(dtype) |
| 786 | bT = b |
| 787 | b = b.T.contiguous() |
| 788 | |
| 789 | naive_result = matmul(a, b.T).to(torch.float16) |
| 790 | run_test(naive_result, torch_matmul, a, b, "Torch", enabled=dtype == torch.float16) |
| 791 | run_test( |
| 792 | naive_result, |
| 793 | torch_matmul_nontransposed, |
| 794 | a, |
| 795 | bT, |
| 796 | "Torch (Non-Transposed)", |
| 797 | enabled=dtype == torch.float16, |
| 798 | ) |
| 799 | run_test( |
| 800 | naive_result, device_blas_matmul, a, b, device_blas_name(), enabled=device_blas is not None |
| 801 | ) |
| 802 | run_test(naive_result, matmul_persistent, a, b.T, "Persistent") |
| 803 | |
| 804 | kernels = [ |
| 805 | (matmul_tma, "TMA", HAS_HOST_TENSOR_DESC), |
| 806 | (matmul_tma_persistent, "TMA Persistent", HAS_HOST_TENSOR_DESC), |
| 807 | (matmul_descriptor_persistent, "Tensor Descriptor Persistent", HAS_TENSOR_DESC), |
| 808 | ] |
| 809 | warp_specialize = [False, True] if HAS_WARP_SPECIALIZE else [False] |
| 810 | |
| 811 | for (kernel, label, enabled), warp_specialize in itertools.product(kernels, warp_specialize): |
| 812 | label = f"{label} (warp_specialize={warp_specialize})" |
| 813 | # skip if hopper and warp_specialize and not on-device |
| 814 | skipped = is_hopper() and warp_specialize and kernel != matmul_descriptor_persistent |
| 815 | enabled = enabled and (not warp_specialize or HAS_TENSOR_DESC) and (not skipped) |
| 816 | run_test(naive_result, lambda a, b: kernel(a, b, warp_specialize), a, b, label, enabled) |
| 817 | print() |
| 818 | |
| 819 | |
| 820 | def show_profile(precision, profile_name): |
no test coverage detected