frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 21:14:14 +08:00
Code Issues Packages Projects Releases Wiki Activity

Enable Extension Support (#142028)

Browse Source 
Pull Request resolved: https://github.com/pytorch/pytorch/pull/142028
Approved by: https://github.com/ezyang, https://github.com/eqy
This commit is contained in:
drisspg
2024-12-03 22:29:29 -05:00
committed by PyTorch MergeBot
parent 38d10a1b17
commit 0582b32f6c
1 changed files with 7 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
torch/utils/cpp_extension.py
View File

@ -1991,10 +1991,12 @@ def _get_cuda_arch_flags(cflags: Optional[List[str]] = None) -> List[str]:
('Ampere', '8.0;8.6+PTX'),
('Ada', '8.9+PTX'),
('Hopper', '9.0+PTX'),
('Blackwell', '10.0+PTX'),
])
supported_arches = ['3.5', '3.7', '5.0', '5.2', '5.3', '6.0', '6.1', '6.2',
'7.0', '7.2', '7.5', '8.0', '8.6', '8.7', '8.9', '9.0', '9.0a']
'7.0', '7.2', '7.5', '8.0', '8.6', '8.7', '8.9', '9.0', '9.0a'
'10.0']
valid_arch_strings = supported_arches + [s + "+PTX" for s in supported_arches]
# The default is sm_30 for CUDA 9.x and 10.x
@ -2040,7 +2042,10 @@ def _get_cuda_arch_flags(cflags: Optional[List[str]] = None) -> List[str]:
if arch not in valid_arch_strings:
raise ValueError(f"Unknown CUDA arch ({arch}) or GPU not supported")
else:
num = arch[0] + arch[2:].split("+")[0]
# Handle both single and double-digit architecture versions
version = arch.split('+')[0] # Remove "+PTX" if present
major, minor = version.split('.')
num = f"{major}{minor}"
flags.append(f'-gencode=arch=compute_{num},code=sm_{num}')
if arch.endswith('+PTX'):
flags.append(f'-gencode=arch=compute_{num},code=compute_{num}')