Entry point called by bench.py. Must match reference.matmul_ref signature.
(A: torch.Tensor, B: torch.Tensor)
| 60 | |
| 61 | |
| 62 | def kernel_fn(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor: |
| 63 | """Entry point called by bench.py. Must match reference.matmul_ref signature.""" |
| 64 | assert A.is_cuda and B.is_cuda |
| 65 | M, K = A.shape |
| 66 | K2, N = B.shape |
| 67 | assert K == K2 |
| 68 | |
| 69 | C = torch.empty((M, N), device=A.device, dtype=A.dtype) |
| 70 | |
| 71 | BLOCK_SIZE_M = 64 |
| 72 | BLOCK_SIZE_N = 64 |
| 73 | BLOCK_SIZE_K = 32 |
| 74 | |
| 75 | grid = (triton.cdiv(M, BLOCK_SIZE_M), triton.cdiv(N, BLOCK_SIZE_N)) |
| 76 | |
| 77 | matmul_kernel[grid]( |
| 78 | A, B, C, |
| 79 | M, N, K, |
| 80 | A.stride(0), A.stride(1), |
| 81 | B.stride(0), B.stride(1), |
| 82 | C.stride(0), C.stride(1), |
| 83 | BLOCK_SIZE_M=BLOCK_SIZE_M, |
| 84 | BLOCK_SIZE_N=BLOCK_SIZE_N, |
| 85 | BLOCK_SIZE_K=BLOCK_SIZE_K, |
| 86 | ) |
| 87 | return C |
no outgoing calls
no test coverage detected