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

link specific versioned System NCCL, rather than generic file

Browse Source
This commit is contained in:
Soumith Chintala 2017-10-08 21:30:58 -07:00
parent 803afd58a0
commit b7e258f81e
2 changed files with 12 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
setup.py
View File

@ -15,7 +15,7 @@ import os
from tools.setup_helpers.env import check_env_flag from tools.setup_helpers.env import check_env_flag
from tools.setup_helpers.cuda import WITH_CUDA, CUDA_HOME from tools.setup_helpers.cuda import WITH_CUDA, CUDA_HOME
from tools.setup_helpers.cudnn import WITH_CUDNN, CUDNN_LIB_DIR, CUDNN_INCLUDE_DIR from tools.setup_helpers.cudnn import WITH_CUDNN, CUDNN_LIB_DIR, CUDNN_INCLUDE_DIR
from tools.setup_helpers.nccl import WITH_NCCL, WITH_SYSTEM_NCCL, NCCL_LIB_DIR, NCCL_INCLUDE_DIR, NCCL_ROOT_DIR from tools.setup_helpers.nccl import WITH_NCCL, WITH_SYSTEM_NCCL, NCCL_LIB_DIR, NCCL_INCLUDE_DIR, NCCL_ROOT_DIR, NCCL_SYSTEM_LIB
from tools.setup_helpers.nnpack import WITH_NNPACK, NNPACK_LIB_DIR, NNPACK_INCLUDE_DIRS from tools.setup_helpers.nnpack import WITH_NNPACK, NNPACK_LIB_DIR, NNPACK_INCLUDE_DIRS
from tools.setup_helpers.split_types import split_types from tools.setup_helpers.split_types import split_types
@ -215,7 +215,7 @@ class build_ext(setuptools.command.build_ext.build_ext):
print('-- Not using CUDA') print('-- Not using CUDA')
if WITH_NCCL and WITH_SYSTEM_NCCL: if WITH_NCCL and WITH_SYSTEM_NCCL:
print('-- Using system provided NCCL library at ' + print('-- Using system provided NCCL library at ' +
NCCL_LIB_DIR + ', ' + NCCL_INCLUDE_DIR) NCCL_SYSTEM_LIB + ', ' + NCCL_INCLUDE_DIR)
elif WITH_NCCL: elif WITH_NCCL:
print('-- Building NCCL library') print('-- Building NCCL library')
else: else:
@ -480,9 +480,8 @@ if WITH_CUDA:
if WITH_NCCL: if WITH_NCCL:
if WITH_SYSTEM_NCCL: if WITH_SYSTEM_NCCL:
main_libraries += ['nccl'] main_link_args += [NCCL_SYSTEM_LIB]
include_dirs.append(NCCL_INCLUDE_DIR) include_dirs.append(NCCL_INCLUDE_DIR)
library_dirs.append(NCCL_LIB_DIR)
else: else:
main_link_args += [NCCL_LIB] main_link_args += [NCCL_LIB]
extra_compile_args += ['-DWITH_NCCL'] extra_compile_args += ['-DWITH_NCCL']

9
tools/setup_helpers/nccl.py
View File

@ -18,6 +18,7 @@ conda_dir = os.path.join(os.path.dirname(sys.executable), '..')
WITH_NCCL = WITH_CUDA and platform.system() != 'Darwin' WITH_NCCL = WITH_CUDA and platform.system() != 'Darwin'
WITH_SYSTEM_NCCL = False WITH_SYSTEM_NCCL = False
NCCL_LIB_DIR = None NCCL_LIB_DIR = None
NCCL_SYSTEM_LIB = None
NCCL_INCLUDE_DIR = None NCCL_INCLUDE_DIR = None
NCCL_ROOT_DIR = None NCCL_ROOT_DIR = None
if WITH_CUDA and not check_env_flag('NO_SYSTEM_NCCL'): if WITH_CUDA and not check_env_flag('NO_SYSTEM_NCCL'):
@ -51,12 +52,20 @@ if WITH_CUDA and not check_env_flag('NO_SYSTEM_NCCL'):
if is_conda: if is_conda:
lib_paths.append(os.path.join(conda_dir, 'lib')) lib_paths.append(os.path.join(conda_dir, 'lib'))
for path in lib_paths: for path in lib_paths:
path = os.path.expanduser(path)
if path is None or not os.path.exists(path): if path is None or not os.path.exists(path):
continue continue
if glob.glob(os.path.join(path, 'libnccl*')): if glob.glob(os.path.join(path, 'libnccl*')):
NCCL_LIB_DIR = path NCCL_LIB_DIR = path
# try to find an exact versioned .so/.dylib, rather than libnccl.so
preferred_path = glob.glob(os.path.join(path, 'libnccl*[0-9]*'))
if len(preferred_path) == 0:
NCCL_SYSTEM_LIB = glob.glob(os.path.join(path, 'libnccl*'))[0]
else:
NCCL_SYSTEM_LIB = os.path.realpath(preferred_path[0])
break break
for path in include_paths: for path in include_paths:
path = os.path.expanduser(path)
if path is None or not os.path.exists(path): if path is None or not os.path.exists(path):
continue continue
if glob.glob(os.path.join(path, 'nccl.h')): if glob.glob(os.path.join(path, 'nccl.h')):