Detect whether a kernel file uses the CUDA C++ or Triton backend. Returns 'cuda' or 'triton'.
(source: str)
| 50 | # --------------------------------------------------------------------------- |
| 51 | |
| 52 | def detect_backend(source: str) -> str: |
| 53 | """ |
| 54 | Detect whether a kernel file uses the CUDA C++ or Triton backend. |
| 55 | |
| 56 | Returns 'cuda' or 'triton'. |
| 57 | """ |
| 58 | # Explicit BACKEND declaration takes priority |
| 59 | match = re.search(r'^BACKEND\s*=\s*["\'](\w+)["\']', source, re.MULTILINE) |
| 60 | if match: |
| 61 | backend = match.group(1).lower() |
| 62 | if backend in ("cuda", "triton"): |
| 63 | return backend |
| 64 | |
| 65 | # Heuristic: look for CUDA indicators |
| 66 | has_cuda_src = "CUDA_SRC" in source |
| 67 | has_compile_cuda = "compile_cuda" in source |
| 68 | |
| 69 | if has_cuda_src or has_compile_cuda: |
| 70 | return "cuda" |
| 71 | |
| 72 | # Heuristic: look for Triton indicators |
| 73 | has_triton_import = "import triton" in source or "from triton" in source |
| 74 | has_triton_jit = "@triton.jit" in source or "@triton.autotune" in source |
| 75 | |
| 76 | if has_triton_import or has_triton_jit: |
| 77 | return "triton" |
| 78 | |
| 79 | # Default to triton if unclear |
| 80 | return "triton" |
| 81 | |
| 82 | |
| 83 | # --------------------------------------------------------------------------- |