OSSForks/pytorch
1
0
Fork 0
mirror of https://github.com/zebrajr/pytorch.git synced 2025-12-06 12:20:52 +01:00
Code Issues Actions 38 Packages Projects Releases Wiki Activity

Suppress warning when using native arch for jit loading cuda extensions. (#156923)

Browse Source 
Previeusly, if users want to let pytorch determine the cuda arch when jit loading cuda extensions, they should left environment variable `TORCH_CUDA_ARCH_LIST` empty, but which will raise an warning. This commit add an option to set `TORCH_CUDA_ARCH_LIST=native`, to tell pytorch users want to use native cuda arch intentionally.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/156923
Approved by: https://github.com/ezyang
This commit is contained in:
Hao Zhang(张浩) 2025-07-09 02:51:20 +00:00 committed by PyTorch MergeBot
parent bc6e0661a6
commit ab8874bd26
1 changed files with 6 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
torch/utils/cpp_extension.py
View File

@ -2419,11 +2419,12 @@ def _get_cuda_arch_flags(cflags: Optional[list[str]] = None) -> list[str]:
# See cmake/Modules_CUDA_fix/upstream/FindCUDA/select_compute_arch.cmake
_arch_list = os.environ.get('TORCH_CUDA_ARCH_LIST', None)
# If not given, determine what's best for the GPU / CUDA version that can be found
if not _arch_list:
logger.warning(
"TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation. \n"
"If this is not desired, please set os.environ['TORCH_CUDA_ARCH_LIST'] to specific architectures.")
# If not given or set as native, determine what's best for the GPU / CUDA version that can be found
if not _arch_list or _arch_list == "native":
if not _arch_list:
logger.warning(
"TORCH_CUDA_ARCH_LIST is not set, all archs for visible cards are included for compilation. \n"
"If this is not desired, please set os.environ['TORCH_CUDA_ARCH_LIST'] to specific architectures.")
arch_list = []
# the assumption is that the extension should run on any of the currently visible cards,
# which could be of different types - therefore all archs for visible cards should be included