Entry point called by bench.py. Must match reference.matmul_ref signature.
(A: torch.Tensor, B: torch.Tensor)
| 244 | |
| 245 | |
| 246 | def kernel_fn(A: torch.Tensor, B: torch.Tensor) -> torch.Tensor: |
| 247 | """Entry point called by bench.py. Must match reference.matmul_ref signature.""" |
| 248 | assert A.is_cuda and B.is_cuda |
| 249 | |
| 250 | # Handle non-fp16 inputs by casting |
| 251 | orig_dtype = A.dtype |
| 252 | if A.dtype != torch.float16: |
| 253 | A = A.to(torch.float16) |
| 254 | if B.dtype != torch.float16: |
| 255 | B = B.to(torch.float16) |
| 256 | |
| 257 | mod = _get_module() |
| 258 | C = mod.matmul_cuda(A, B) |
| 259 | |
| 260 | # Cast back if needed |
| 261 | if orig_dtype != torch.float16: |
| 262 | C = C.to(orig_dtype) |
| 263 | |
| 264 | return C |
nothing calls this directly
no test coverage detected