Generate -gencode flags for the current GPU architecture.
()
| 42 | # --------------------------------------------------------------------------- |
| 43 | |
| 44 | def _get_arch_flags() -> list: |
| 45 | """Generate -gencode flags for the current GPU architecture.""" |
| 46 | if not torch.cuda.is_available(): |
| 47 | return [] |
| 48 | |
| 49 | cap = torch.cuda.get_device_capability() |
| 50 | major, minor = cap |
| 51 | sm = f"{major}{minor}" |
| 52 | |
| 53 | flags = [ |
| 54 | # Compile for the exact GPU |
| 55 | f"-gencode=arch=compute_{sm},code=sm_{sm}", |
| 56 | # Also embed PTX for forward compatibility |
| 57 | f"-gencode=arch=compute_{sm},code=compute_{sm}", |
| 58 | ] |
| 59 | return flags |
| 60 | |
| 61 | |
| 62 | # --------------------------------------------------------------------------- |