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

Automatically set BUILD_SPLIT_CUDA for cpp exts (#52503)

Browse Source 
Summary:
Fixes https://github.com/pytorch/vision/pull/3418#issuecomment-781673110

Pull Request resolved: https://github.com/pytorch/pytorch/pull/52503

Reviewed By: malfet

Differential Revision: D26546857

Pulled By: janeyx99

fbshipit-source-id: a100b408e7cd28695145a1dda7f2fa081bb7f21f
This commit is contained in:
peterjc123
2021-02-19 12:20:42 -08:00
committed by Facebook GitHub Bot
parent b6ed05130e
commit 44ff79d849
1 changed files with 6 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
torch/utils/cpp_extension.py
View File

@ -22,16 +22,21 @@ from typing import List, Optional, Union
from setuptools.command.build_ext import build_ext
from pkg_resources import packaging # type: ignore
BUILD_SPLIT_CUDA = os.getenv('BUILD_SPLIT_CUDA')
IS_WINDOWS = sys.platform == 'win32'
LIB_EXT = '.pyd' if IS_WINDOWS else '.so'
EXEC_EXT = '.exe' if IS_WINDOWS else ''
CLIB_PREFIX = '' if IS_WINDOWS else 'lib'
CLIB_EXT = '.dll' if IS_WINDOWS else '.so'
SHARED_FLAG = '/DLL' if IS_WINDOWS else '-shared'
_HERE = os.path.abspath(__file__)
_TORCH_PATH = os.path.dirname(os.path.dirname(_HERE))
TORCH_LIB_PATH = os.path.join(_TORCH_PATH, 'lib')
BUILD_SPLIT_CUDA = os.getenv('BUILD_SPLIT_CUDA') or (os.path.exists(os.path.join(
TORCH_LIB_PATH, f'{CLIB_PREFIX}torch_cuda_cu{CLIB_EXT}')) and os.path.exists(os.path.join(TORCH_LIB_PATH, f'{CLIB_PREFIX}torch_cuda_cpp{CLIB_EXT}')))
# Taken directly from python stdlib < 3.9
# See https://github.com/pytorch/pytorch/issues/48617
def _nt_quote_args(args: Optional[List[str]]) -> List[str]: