Map a CUDA kernel name to an AutoKernel op type.
(kernel_name: str)
| 447 | # --------------------------------------------------------------------------- |
| 448 | |
| 449 | def classify_kernel(kernel_name: str) -> str: |
| 450 | """Map a CUDA kernel name to an AutoKernel op type.""" |
| 451 | name_lower = kernel_name.lower() |
| 452 | |
| 453 | for fragments, op_type in _KERNEL_CLASSIFICATION: |
| 454 | for frag in fragments: |
| 455 | if frag in name_lower: |
| 456 | return op_type |
| 457 | |
| 458 | # Check for standalone "mm" -- common in cuBLAS kernel names like |
| 459 | # "void cutlass::...sgemm..." or names containing "_mm_" or ending in "mm". |
| 460 | # Avoid false positives from words like "command", "summary", "commit". |
| 461 | if "mm" in name_lower: |
| 462 | if re.search(r"(?:^|[^a-z])mm(?:$|[^a-z])", name_lower): |
| 463 | return "matmul" |
| 464 | |
| 465 | return "other" |
| 466 | |
| 467 | |
| 468 | def is_autokernel_supported(op_type: str) -> bool: |