Change how ./configure and bazel finds CUDA paths.
Use a single python script (third_party/gpus/find_cuda_config.py) from configure.py and the different *_configure.bzl scripts to find the different CUDA library and header paths based on a set of environment variables.
PiperOrigin-RevId: 243669844
diff --git a/configure.py b/configure.py
index 2d919d8..c95b57f 100644
--- a/configure.py
+++ b/configure.py
@@ -862,154 +862,41 @@
return '.'.join(v[:sequence_count])
+def set_tf_cuda_paths(environ_cp):
+ """Set TF_CUDA_PATHS."""
+ ask_cuda_paths = (
+ 'Please specify the comma-separated list of base paths to look for CUDA '
+ 'libraries and headers. [Leave empty to use the default]: '
+ )
+ tf_cuda_paths = get_from_env_or_user_or_default(environ_cp, 'TF_CUDA_PATHS',
+ ask_cuda_paths,
+ '')
+ if tf_cuda_paths:
+ environ_cp['TF_CUDA_PATHS'] = tf_cuda_paths
+
+
def set_tf_cuda_version(environ_cp):
- """Set CUDA_TOOLKIT_PATH and TF_CUDA_VERSION."""
+ """Set TF_CUDA_VERSION."""
ask_cuda_version = (
'Please specify the CUDA SDK version you want to use. '
'[Leave empty to default to CUDA %s]: ') % _DEFAULT_CUDA_VERSION
-
- for _ in range(_DEFAULT_PROMPT_ASK_ATTEMPTS):
- # Configure the Cuda SDK version to use.
- tf_cuda_version = get_from_env_or_user_or_default(environ_cp,
- 'TF_CUDA_VERSION',
- ask_cuda_version,
- _DEFAULT_CUDA_VERSION)
- tf_cuda_version = reformat_version_sequence(str(tf_cuda_version), 2)
-
- # Find out where the CUDA toolkit is installed
- default_cuda_path = _DEFAULT_CUDA_PATH
- if is_windows() or is_cygwin():
- default_cuda_path = cygpath(
- environ_cp.get('CUDA_PATH', _DEFAULT_CUDA_PATH_WIN))
- elif is_linux():
- # If the default doesn't exist, try an alternative default.
- if (not os.path.exists(default_cuda_path)
- ) and os.path.exists(_DEFAULT_CUDA_PATH_LINUX):
- default_cuda_path = _DEFAULT_CUDA_PATH_LINUX
- ask_cuda_path = ('Please specify the location where CUDA %s toolkit is'
- ' installed. Refer to README.md for more details. '
- '[Default is %s]: ') % (tf_cuda_version, default_cuda_path)
- cuda_toolkit_path = get_from_env_or_user_or_default(environ_cp,
- 'CUDA_TOOLKIT_PATH',
- ask_cuda_path,
- default_cuda_path)
- if is_windows() or is_cygwin():
- cuda_toolkit_path = cygpath(cuda_toolkit_path)
-
- if is_windows():
- cuda_rt_lib_paths = ['lib/x64/cudart.lib']
- elif is_linux():
- cuda_rt_lib_paths = [
- '%s/libcudart.so.%s' % (x, tf_cuda_version) for x in [
- 'lib64',
- 'lib/powerpc64le-linux-gnu',
- 'lib/x86_64-linux-gnu',
- ]
- ]
- elif is_macos():
- cuda_rt_lib_paths = ['lib/libcudart.%s.dylib' % tf_cuda_version]
-
- cuda_toolkit_paths_full = [
- os.path.join(cuda_toolkit_path, x) for x in cuda_rt_lib_paths
- ]
- if any(os.path.exists(x) for x in cuda_toolkit_paths_full):
- break
-
- # Reset and retry
- print('Invalid path to CUDA %s toolkit. %s cannot be found' %
- (tf_cuda_version, cuda_toolkit_paths_full))
- environ_cp['TF_CUDA_VERSION'] = ''
- environ_cp['CUDA_TOOLKIT_PATH'] = ''
-
- else:
- raise UserInputError('Invalid TF_CUDA_SETTING setting was provided %d '
- 'times in a row. Assuming to be a scripting mistake.' %
- _DEFAULT_PROMPT_ASK_ATTEMPTS)
-
- # Set CUDA_TOOLKIT_PATH and TF_CUDA_VERSION
- environ_cp['CUDA_TOOLKIT_PATH'] = cuda_toolkit_path
- write_action_env_to_bazelrc('CUDA_TOOLKIT_PATH', cuda_toolkit_path)
+ tf_cuda_version = get_from_env_or_user_or_default(environ_cp,
+ 'TF_CUDA_VERSION',
+ ask_cuda_version,
+ _DEFAULT_CUDA_VERSION)
environ_cp['TF_CUDA_VERSION'] = tf_cuda_version
- write_action_env_to_bazelrc('TF_CUDA_VERSION', tf_cuda_version)
def set_tf_cudnn_version(environ_cp):
- """Set CUDNN_INSTALL_PATH and TF_CUDNN_VERSION."""
+ """Set TF_CUDNN_VERSION."""
ask_cudnn_version = (
'Please specify the cuDNN version you want to use. '
'[Leave empty to default to cuDNN %s]: ') % _DEFAULT_CUDNN_VERSION
-
- for _ in range(_DEFAULT_PROMPT_ASK_ATTEMPTS):
- tf_cudnn_version = get_from_env_or_user_or_default(environ_cp,
- 'TF_CUDNN_VERSION',
- ask_cudnn_version,
- _DEFAULT_CUDNN_VERSION)
- tf_cudnn_version = reformat_version_sequence(str(tf_cudnn_version), 1)
-
- default_cudnn_path = environ_cp.get('CUDA_TOOLKIT_PATH')
- ask_cudnn_path = (r'Please specify the location where cuDNN %s library is '
- 'installed. Refer to README.md for more details. [Default'
- ' is %s]: ') % (tf_cudnn_version, default_cudnn_path)
- cudnn_install_path = get_from_env_or_user_or_default(
- environ_cp, 'CUDNN_INSTALL_PATH', ask_cudnn_path, default_cudnn_path)
-
- # Result returned from "read" will be used unexpanded. That make "~"
- # unusable. Going through one more level of expansion to handle that.
- cudnn_install_path = os.path.realpath(
- os.path.expanduser(cudnn_install_path))
- if is_windows() or is_cygwin():
- cudnn_install_path = cygpath(cudnn_install_path)
-
- if is_windows():
- cuda_dnn_lib_path = 'lib/x64/cudnn.lib'
- cuda_dnn_lib_alt_path = 'lib/x64/cudnn.lib'
- elif is_linux():
- cuda_dnn_lib_path = 'lib64/libcudnn.so.%s' % tf_cudnn_version
- cuda_dnn_lib_alt_path = 'libcudnn.so.%s' % tf_cudnn_version
- elif is_macos():
- cuda_dnn_lib_path = 'lib/libcudnn.%s.dylib' % tf_cudnn_version
- cuda_dnn_lib_alt_path = 'libcudnn.%s.dylib' % tf_cudnn_version
-
- cuda_dnn_lib_path_full = os.path.join(cudnn_install_path, cuda_dnn_lib_path)
- cuda_dnn_lib_alt_path_full = os.path.join(cudnn_install_path,
- cuda_dnn_lib_alt_path)
- if os.path.exists(cuda_dnn_lib_path_full) or os.path.exists(
- cuda_dnn_lib_alt_path_full):
- break
-
- # Try another alternative for Linux
- if is_linux():
- ldconfig_bin = which('ldconfig') or '/sbin/ldconfig'
- cudnn_path_from_ldconfig = run_shell([ldconfig_bin, '-p'])
- cudnn_path_from_ldconfig = re.search('.*libcudnn.so .* => (.*)',
- cudnn_path_from_ldconfig)
- if cudnn_path_from_ldconfig:
- cudnn_path_from_ldconfig = cudnn_path_from_ldconfig.group(1)
- if os.path.exists('%s.%s' %
- (cudnn_path_from_ldconfig, tf_cudnn_version)):
- cudnn_install_path = os.path.dirname(cudnn_path_from_ldconfig)
- break
-
- # Reset and Retry
- print(
- 'Invalid path to cuDNN %s toolkit. None of the following files can be '
- 'found:' % tf_cudnn_version)
- print(cuda_dnn_lib_path_full)
- print(cuda_dnn_lib_alt_path_full)
- if is_linux():
- print('%s.%s' % (cudnn_path_from_ldconfig, tf_cudnn_version))
-
- environ_cp['TF_CUDNN_VERSION'] = ''
- else:
- raise UserInputError('Invalid TF_CUDNN setting was provided %d '
- 'times in a row. Assuming to be a scripting mistake.' %
- _DEFAULT_PROMPT_ASK_ATTEMPTS)
-
- # Set CUDNN_INSTALL_PATH and TF_CUDNN_VERSION
- environ_cp['CUDNN_INSTALL_PATH'] = cudnn_install_path
- write_action_env_to_bazelrc('CUDNN_INSTALL_PATH', cudnn_install_path)
+ tf_cudnn_version = get_from_env_or_user_or_default(environ_cp,
+ 'TF_CUDNN_VERSION',
+ ask_cudnn_version,
+ _DEFAULT_CUDNN_VERSION)
environ_cp['TF_CUDNN_VERSION'] = tf_cudnn_version
- write_action_env_to_bazelrc('TF_CUDNN_VERSION', tf_cudnn_version)
def is_cuda_compatible(lib, cuda_ver, cudnn_ver):
@@ -1041,253 +928,38 @@
return cudnn_ok and cuda_ok
-def set_tf_tensorrt_install_path(environ_cp):
- """Set TENSORRT_INSTALL_PATH and TF_TENSORRT_VERSION.
-
- Adapted from code contributed by Sami Kama (https://github.com/samikama).
-
- Args:
- environ_cp: copy of the os.environ.
-
- Raises:
- ValueError: if this method was called under non-Linux platform.
- UserInputError: if user has provided invalid input multiple times.
- """
+def set_tf_tensorrt_version(environ_cp):
+ """Set TF_TENSORRT_VERSION."""
if not is_linux():
raise ValueError('Currently TensorRT is only supported on Linux platform.')
- # Ask user whether to add TensorRT support.
+ # Backwards compatibility: early-out when TF_NEED_TENSORRT=0.
if str(int(get_var(environ_cp, 'TF_NEED_TENSORRT', 'TensorRT',
False))) != '1':
return
- for _ in range(_DEFAULT_PROMPT_ASK_ATTEMPTS):
- ask_tensorrt_path = (r'Please specify the location where TensorRT is '
- 'installed. [Default is %s]:') % (
- _DEFAULT_TENSORRT_PATH_LINUX)
- trt_install_path = get_from_env_or_user_or_default(
- environ_cp, 'TENSORRT_INSTALL_PATH', ask_tensorrt_path,
- _DEFAULT_TENSORRT_PATH_LINUX)
-
- # Result returned from "read" will be used unexpanded. That make "~"
- # unusable. Going through one more level of expansion to handle that.
- trt_install_path = os.path.realpath(os.path.expanduser(trt_install_path))
-
- def find_libs(search_path):
- """Search for libnvinfer.so in "search_path"."""
- fl = set()
- if os.path.exists(search_path) and os.path.isdir(search_path):
- fl.update([
- os.path.realpath(os.path.join(search_path, x))
- for x in os.listdir(search_path)
- if 'libnvinfer.so' in x
- ])
- return fl
-
- possible_files = find_libs(trt_install_path)
- possible_files.update(find_libs(os.path.join(trt_install_path, 'lib')))
- possible_files.update(find_libs(os.path.join(trt_install_path, 'lib64')))
- cuda_ver = convert_version_to_int(environ_cp['TF_CUDA_VERSION'])
- cudnn_ver = convert_version_to_int(environ_cp['TF_CUDNN_VERSION'])
- nvinfer_pattern = re.compile('.*libnvinfer.so.?(.*)$')
- highest_ver = [0, None, None]
-
- for lib_file in possible_files:
- if is_cuda_compatible(lib_file, cuda_ver, cudnn_ver):
- matches = nvinfer_pattern.search(lib_file)
- if not matches.groups():
- continue
- ver_str = matches.group(1)
- ver = convert_version_to_int(ver_str) if len(ver_str) else 0
- if ver > highest_ver[0]:
- highest_ver = [ver, ver_str, lib_file]
- if highest_ver[1] is not None:
- trt_install_path = os.path.dirname(highest_ver[2])
- tf_tensorrt_version = highest_ver[1]
- break
-
- # Try another alternative from ldconfig.
- ldconfig_bin = which('ldconfig') or '/sbin/ldconfig'
- ldconfig_output = run_shell([ldconfig_bin, '-p'])
- search_result = re.search('.*libnvinfer.so\\.?([0-9.]*).* => (.*)',
- ldconfig_output)
- if search_result:
- libnvinfer_path_from_ldconfig = search_result.group(2)
- if os.path.exists(libnvinfer_path_from_ldconfig):
- if is_cuda_compatible(libnvinfer_path_from_ldconfig, cuda_ver,
- cudnn_ver):
- trt_install_path = os.path.dirname(libnvinfer_path_from_ldconfig)
- tf_tensorrt_version = search_result.group(1)
- break
-
- # Reset and Retry
- if possible_files:
- print('TensorRT libraries found in one the following directories',
- 'are not compatible with selected cuda and cudnn installations')
- print(trt_install_path)
- print(os.path.join(trt_install_path, 'lib'))
- print(os.path.join(trt_install_path, 'lib64'))
- if search_result:
- print(libnvinfer_path_from_ldconfig)
- else:
- print(
- 'Invalid path to TensorRT. None of the following files can be found:')
- print(trt_install_path)
- print(os.path.join(trt_install_path, 'lib'))
- print(os.path.join(trt_install_path, 'lib64'))
- if search_result:
- print(libnvinfer_path_from_ldconfig)
-
- else:
- raise UserInputError('Invalid TF_TENSORRT setting was provided %d '
- 'times in a row. Assuming to be a scripting mistake.' %
- _DEFAULT_PROMPT_ASK_ATTEMPTS)
-
- # Set TENSORRT_INSTALL_PATH and TF_TENSORRT_VERSION
- environ_cp['TENSORRT_INSTALL_PATH'] = trt_install_path
- write_action_env_to_bazelrc('TENSORRT_INSTALL_PATH', trt_install_path)
+ ask_tensorrt_version = (
+ 'Please specify the TensorRT version you want to use. '
+ '[Leave empty to not use a specific version]: ')
+ tf_tensorrt_version = get_from_env_or_user_or_default(environ_cp,
+ 'TF_TENSORRT_VERSION',
+ ask_tensorrt_version,
+ '')
environ_cp['TF_TENSORRT_VERSION'] = tf_tensorrt_version
- write_action_env_to_bazelrc('TF_TENSORRT_VERSION', tf_tensorrt_version)
-def set_tf_nccl_install_path(environ_cp):
- """Set NCCL_INSTALL_PATH, NCCL_HDR_PATH and TF_NCCL_VERSION.
-
- Args:
- environ_cp: copy of the os.environ.
-
- Raises:
- ValueError: if this method was called under non-Linux platform.
- UserInputError: if user has provided invalid input multiple times.
- """
+def set_tf_nccl_version(environ_cp):
+ """Set TF_NCCL_VERSION."""
if not is_linux():
- raise ValueError('Currently NCCL is only supported on Linux platforms.')
+ raise ValueError('Currently NCCL is only supported on Linux platform.')
ask_nccl_version = (
'Please specify the locally installed NCCL version you want to use. '
- '[Default is to use https://github.com/nvidia/nccl]: ')
-
- for _ in range(_DEFAULT_PROMPT_ASK_ATTEMPTS):
- tf_nccl_version = get_from_env_or_user_or_default(environ_cp,
- 'TF_NCCL_VERSION',
- ask_nccl_version, '')
-
- if not tf_nccl_version:
- break # No need to get install path, building the open source code.
-
- tf_nccl_version = reformat_version_sequence(str(tf_nccl_version), 1)
-
- # Look with ldconfig first if we can find the library in paths
- # like /usr/lib/x86_64-linux-gnu and the header file in the corresponding
- # include directory. This is where the NCCL .deb packages install them.
-
- # First check to see if NCCL is in the ldconfig.
- # If its found, use that location.
- if is_linux():
- ldconfig_bin = which('ldconfig') or '/sbin/ldconfig'
- nccl2_path_from_ldconfig = run_shell([ldconfig_bin, '-p'])
- nccl2_path_from_ldconfig = re.search('.*libnccl.so .* => (.*)',
- nccl2_path_from_ldconfig)
- if nccl2_path_from_ldconfig:
- nccl2_path_from_ldconfig = nccl2_path_from_ldconfig.group(1)
- if os.path.exists('%s.%s' % (nccl2_path_from_ldconfig, tf_nccl_version)):
- nccl_install_path = os.path.dirname(nccl2_path_from_ldconfig)
- print('NCCL libraries found in ' + nccl2_path_from_ldconfig)
-
- # Check if this is the main system lib location
- if re.search('.*linux-gnu', nccl_install_path):
- trunc_nccl_install_path = '/usr'
- print('This looks like a system path.')
- else:
- trunc_nccl_install_path = nccl_install_path + '/..'
-
- # Look for header
- nccl_hdr_path = trunc_nccl_install_path + '/include'
- print('Assuming NCCL header path is ' + nccl_hdr_path)
- if os.path.exists(nccl_hdr_path + '/nccl.h'):
- # Set NCCL_INSTALL_PATH
- environ_cp['NCCL_INSTALL_PATH'] = nccl_install_path
- write_action_env_to_bazelrc('NCCL_INSTALL_PATH', nccl_install_path)
-
- # Set NCCL_HDR_PATH
- environ_cp['NCCL_HDR_PATH'] = nccl_hdr_path
- write_action_env_to_bazelrc('NCCL_HDR_PATH', nccl_hdr_path)
- break
- else:
- print(
- 'The header for NCCL2 cannot be found. Please install the libnccl-dev package.'
- )
- else:
- print('NCCL2 is listed by ldconfig but the library is not found. '
- 'Your ldconfig is out of date. Please run sudo ldconfig.')
- else:
- # NCCL is not found in ldconfig. Ask the user for the location.
- default_nccl_path = environ_cp.get('CUDA_TOOLKIT_PATH')
- ask_nccl_path = (
- r'Please specify the location where NCCL %s library is '
- 'installed. Refer to README.md for more details. [Default '
- 'is %s]:') % (tf_nccl_version, default_nccl_path)
- nccl_install_path = get_from_env_or_user_or_default(
- environ_cp, 'NCCL_INSTALL_PATH', ask_nccl_path, default_nccl_path)
-
- # Result returned from "read" will be used unexpanded. That make "~"
- # unusable. Going through one more level of expansion to handle that.
- nccl_install_path = os.path.realpath(
- os.path.expanduser(nccl_install_path))
- if is_windows() or is_cygwin():
- nccl_install_path = cygpath(nccl_install_path)
-
- nccl_lib_path = ''
- if is_windows():
- nccl_lib_path = 'lib/x64/nccl.lib'
- elif is_linux():
- nccl_lib_filename = 'libnccl.so.%s' % tf_nccl_version
- nccl_lpath = '%s/lib/%s' % (nccl_install_path, nccl_lib_filename)
- if not os.path.exists(nccl_lpath):
- for relative_path in NCCL_LIB_PATHS:
- path = '%s/%s%s' % (nccl_install_path, relative_path,
- nccl_lib_filename)
- if os.path.exists(path):
- print('NCCL found at ' + path)
- nccl_lib_path = path
- break
- else:
- nccl_lib_path = nccl_lpath
- elif is_macos():
- nccl_lib_path = 'lib/libnccl.%s.dylib' % tf_nccl_version
-
- nccl_lib_path = os.path.join(nccl_install_path, nccl_lib_path)
- nccl_hdr_path = os.path.join(
- os.path.dirname(nccl_lib_path), '../include/nccl.h')
- print('Assuming NCCL header path is ' + nccl_hdr_path)
- if os.path.exists(nccl_lib_path) and os.path.exists(nccl_hdr_path):
- # Set NCCL_INSTALL_PATH
- environ_cp['NCCL_INSTALL_PATH'] = os.path.dirname(nccl_lib_path)
- write_action_env_to_bazelrc('NCCL_INSTALL_PATH',
- os.path.dirname(nccl_lib_path))
-
- # Set NCCL_HDR_PATH
- environ_cp['NCCL_HDR_PATH'] = os.path.dirname(nccl_hdr_path)
- write_action_env_to_bazelrc('NCCL_HDR_PATH',
- os.path.dirname(nccl_hdr_path))
- break
-
- # Reset and Retry
- print(
- 'Invalid path to NCCL %s toolkit, %s or %s not found. Please use the '
- 'O/S agnostic package of NCCL 2' %
- (tf_nccl_version, nccl_lib_path, nccl_hdr_path))
-
- environ_cp['TF_NCCL_VERSION'] = ''
- else:
- raise UserInputError('Invalid TF_NCCL setting was provided %d '
- 'times in a row. Assuming to be a scripting mistake.' %
- _DEFAULT_PROMPT_ASK_ATTEMPTS)
-
- # Set TF_NCCL_VERSION
+ '[Leave empty to use https://github.com/nvidia/nccl]: ')
+ tf_nccl_version = get_from_env_or_user_or_default(environ_cp,
+ 'TF_NCCL_VERSION',
+ ask_nccl_version, '')
environ_cp['TF_NCCL_VERSION'] = tf_nccl_version
- write_action_env_to_bazelrc('TF_NCCL_VERSION', tf_nccl_version)
-
def get_native_cuda_compute_capabilities(environ_cp):
"""Get native cuda compute capabilities.
@@ -1720,11 +1392,69 @@
set_action_env_var(environ_cp, 'TF_NEED_CUDA', 'CUDA', False)
if (environ_cp.get('TF_NEED_CUDA') == '1' and
'TF_CUDA_CONFIG_REPO' not in environ_cp):
- set_tf_cuda_version(environ_cp)
- set_tf_cudnn_version(environ_cp)
- if is_linux():
- set_tf_tensorrt_install_path(environ_cp)
- set_tf_nccl_install_path(environ_cp)
+ for _ in range(_DEFAULT_PROMPT_ASK_ATTEMPTS):
+ set_tf_cuda_version(environ_cp)
+ set_tf_cudnn_version(environ_cp)
+ cuda_libraries = ['cuda', 'cudnn']
+ if is_linux():
+ set_tf_tensorrt_version(environ_cp)
+ if 'TF_TENSORRT_VERSION' in environ_cp: # if env variable exists
+ cuda_libraries.append('tensorrt')
+ set_tf_nccl_version(environ_cp)
+ if environ_cp['TF_NCCL_VERSION']: # if env variable not empty
+ cuda_libraries.append('nccl')
+
+ def maybe_encode_env(env):
+ """Encodes unicode in env to str on Windows python 2.x."""
+ if not is_windows() or sys.version_info[0] != 2:
+ return env
+ for k, v in env.items():
+ if isinstance(k, unicode):
+ k = k.encode('ascii')
+ if isinstance(v, unicode):
+ v = v.encode('ascii')
+ env[k] = v
+ return env
+
+ # Verify CUDA configuration by calling find_cuda_config.py.
+ proc = subprocess.Popen(
+ [
+ environ_cp['PYTHON_BIN_PATH'],
+ 'third_party/gpus/find_cuda_config.py'
+ ] + cuda_libraries,
+ stdout=subprocess.PIPE,
+ env=maybe_encode_env(environ_cp))
+
+ cuda_env_variable_names = [
+ 'TF_CUDA_VERSION', 'TF_CUBLAS_VERSION', 'TF_CUDNN_VERSION',
+ 'TF_TENSORRT_VERSION', 'TF_NCCL_VERSION', 'TF_CUDA_PATHS'
+ ]
+
+ if proc.wait():
+ print('\nCould not find all requested CUDA libraries or headers.\n')
+ for name in cuda_env_variable_names:
+ if name in environ_cp:
+ del environ_cp[name]
+ set_tf_cuda_paths(environ_cp)
+ continue
+
+ for line in proc.stdout:
+ match = re.match('cuda_toolkit_path: (.*)', line.decode('ascii'))
+ if match:
+ cuda_toolkit_path = match.group(1)
+
+ for name in cuda_env_variable_names:
+ if name in environ_cp:
+ write_action_env_to_bazelrc(name, environ_cp[name])
+
+ break
+ else:
+ raise UserInputError(
+ 'Invalid CUDA setting were provided %d '
+ 'times in a row. Assuming to be a scripting mistake.' %
+ _DEFAULT_PROMPT_ASK_ATTEMPTS)
+
+ environ_cp['CUDA_TOOLKIT_PATH'] = cuda_toolkit_path
set_tf_cuda_compute_capabilities(environ_cp)
if 'LD_LIBRARY_PATH' in environ_cp and environ_cp.get(
diff --git a/tensorflow/contrib/cmake/CMakeLists.txt b/tensorflow/contrib/cmake/CMakeLists.txt
index 2ad9ae4..44431d5 100644
--- a/tensorflow/contrib/cmake/CMakeLists.txt
+++ b/tensorflow/contrib/cmake/CMakeLists.txt
@@ -522,6 +522,7 @@
"#define CUDA_CUDA_CONFIG_H_\n"
"#define TF_CUDA_CAPABILITIES ${TF_CUDA_CAP}\n"
"#define TF_CUDA_VERSION \"64_${short_CUDA_VER}\"\n"
+ "#define TF_CUDA_LIB_VERSION \"64_${short_CUDA_VER}\"\n"
"#define TF_CUDNN_VERSION \"64_${CUDNN_VERSION}\"\n"
"#define TF_CUDA_TOOLKIT_PATH \"${CUDA_TOOLKIT_ROOT_DIR}\"\n"
"#endif // CUDA_CUDA_CONFIG_H_\n"
@@ -615,4 +616,4 @@
endif()
if(tensorflow_BUILD_CC_TESTS OR tensorflow_BUILD_PYTHON_TESTS)
include(tf_tests.cmake)
-endif()
\ No newline at end of file
+endif()
diff --git a/tensorflow/opensource_only.files b/tensorflow/opensource_only.files
index a2ceca5..b6b9d4c 100644
--- a/tensorflow/opensource_only.files
+++ b/tensorflow/opensource_only.files
@@ -102,6 +102,7 @@
tensorflow/third_party/gpus/rocm/BUILD.tpl
tensorflow/third_party/gpus/rocm/build_defs.bzl.tpl
tensorflow/third_party/gpus/cuda_configure.bzl
+tensorflow/third_party/gpus/find_cuda_config.py
tensorflow/third_party/gpus/rocm_configure.bzl
tensorflow/third_party/snappy.BUILD
tensorflow/third_party/cython.BUILD
diff --git a/tensorflow/stream_executor/platform/default/dso_loader.cc b/tensorflow/stream_executor/platform/default/dso_loader.cc
index 5a37873..ede9114 100644
--- a/tensorflow/stream_executor/platform/default/dso_loader.cc
+++ b/tensorflow/stream_executor/platform/default/dso_loader.cc
@@ -31,6 +31,7 @@
namespace {
string GetCudaVersion() { return TF_CUDA_VERSION; }
+string GetCudaLibVersion() { return TF_CUDA_LIB_VERSION; }
string GetCudnnVersion() { return TF_CUDNN_VERSION; }
port::StatusOr<void*> GetDsoHandle(const string& name, const string& version) {
@@ -75,15 +76,15 @@
}
port::StatusOr<void*> GetCublasDsoHandle() {
- return GetDsoHandle("cublas", GetCudaVersion());
+ return GetDsoHandle("cublas", GetCudaLibVersion());
}
port::StatusOr<void*> GetCufftDsoHandle() {
- return GetDsoHandle("cufft", GetCudaVersion());
+ return GetDsoHandle("cufft", GetCudaLibVersion());
}
port::StatusOr<void*> GetCurandDsoHandle() {
- return GetDsoHandle("curand", GetCudaVersion());
+ return GetDsoHandle("curand", GetCudaLibVersion());
}
port::StatusOr<void*> GetCuptiDsoHandle() {
diff --git a/third_party/gpus/cuda/cuda_config.h.tpl b/third_party/gpus/cuda/cuda_config.h.tpl
index 811b040..916315d 100644
--- a/third_party/gpus/cuda/cuda_config.h.tpl
+++ b/third_party/gpus/cuda/cuda_config.h.tpl
@@ -19,6 +19,7 @@
#define TF_CUDA_CAPABILITIES %{cuda_compute_capabilities}
#define TF_CUDA_VERSION "%{cuda_version}"
+#define TF_CUDA_LIB_VERSION "%{cuda_lib_version}"
#define TF_CUDNN_VERSION "%{cudnn_version}"
#define TF_CUDA_TOOLKIT_PATH "%{cuda_toolkit_path}"
diff --git a/third_party/gpus/cuda_configure.bzl b/third_party/gpus/cuda_configure.bzl
index 872400b..3e31143 100644
--- a/third_party/gpus/cuda_configure.bzl
+++ b/third_party/gpus/cuda_configure.bzl
@@ -11,12 +11,14 @@
* `TF_DOWNLOAD_CLANG`: Whether to download a recent release of clang
compiler and use it to build tensorflow. When this option is set
CLANG_CUDA_COMPILER_PATH is ignored.
- * `CUDA_TOOLKIT_PATH`: The path to the CUDA toolkit. Default is
+ * `TF_CUDA_PATHS`: The base paths to look for CUDA and cuDNN. Default is
+ `/usr/local/cuda,usr/`.
+ * `CUDA_TOOLKIT_PATH` (deprecated): The path to the CUDA toolkit. Default is
`/usr/local/cuda`.
* `TF_CUDA_VERSION`: The version of the CUDA toolkit. If this is blank, then
use the system default.
* `TF_CUDNN_VERSION`: The version of the cuDNN library.
- * `CUDNN_INSTALL_PATH`: The path to the cuDNN library. Default is
+ * `CUDNN_INSTALL_PATH` (deprecated): The path to the cuDNN library. Default is
`/usr/local/cuda`.
* `TF_CUDA_COMPUTE_CAPABILITIES`: The CUDA compute capabilities. Default is
`3.5,5.2`.
@@ -48,84 +50,8 @@
_TF_DOWNLOAD_CLANG = "TF_DOWNLOAD_CLANG"
_PYTHON_BIN_PATH = "PYTHON_BIN_PATH"
-_DEFAULT_CUDA_VERSION = ""
-_DEFAULT_CUDNN_VERSION = ""
-_DEFAULT_CUDA_TOOLKIT_PATH = "/usr/local/cuda"
-_DEFAULT_CUDNN_INSTALL_PATH = "/usr/local/cuda"
_DEFAULT_CUDA_COMPUTE_CAPABILITIES = ["3.5", "5.2"]
-# Lookup paths for CUDA / cuDNN libraries, relative to the install directories.
-#
-# Paths will be tried out in the order listed below. The first successful path
-# will be used. For example, when looking for the cudart libraries, the first
-# attempt will be lib64/cudart inside the CUDA toolkit.
-CUDA_LIB_PATHS = [
- "lib64/",
- "lib64/stubs/",
- "lib/powerpc64le-linux-gnu/",
- "lib/x86_64-linux-gnu/",
- "lib/x64/",
- "lib/",
- "",
-]
-
-# Lookup paths for cupti.h, relative to the CUDA toolkit directory.
-#
-# On most systems, the cupti library is not installed in the same directory as
-# the other CUDA libraries but rather in a special extras/CUPTI directory.
-CUPTI_HEADER_PATHS = [
- "extras/CUPTI/include/",
- "include/cuda/CUPTI/",
- "include/",
-]
-
-# Lookup paths for the cupti library, relative to the
-#
-# On most systems, the cupti library is not installed in the same directory as
-# the other CUDA libraries but rather in a special extras/CUPTI directory.
-CUPTI_LIB_PATHS = [
- "extras/CUPTI/lib64/",
- "lib/powerpc64le-linux-gnu/",
- "lib/x86_64-linux-gnu/",
- "lib64/",
- "extras/CUPTI/libx64/",
- "extras/CUPTI/lib/",
- "lib/",
-]
-
-# Lookup paths for CUDA headers (cuda.h) relative to the CUDA toolkit directory.
-CUDA_INCLUDE_PATHS = [
- "include/",
- "include/cuda/",
-]
-
-# Lookup paths for cudnn.h relative to the CUDNN install directory.
-CUDNN_INCLUDE_PATHS = [
- "",
- "include/",
- "include/cuda/",
-]
-
-# Lookup paths for NVVM libdevice relative to the CUDA directory toolkit.
-#
-# libdevice implements mathematical functions for GPU kernels, and is provided
-# in NVVM bitcode (a subset of LLVM bitcode).
-NVVM_LIBDEVICE_PATHS = [
- "nvvm/libdevice/",
- "share/cuda/",
- "lib/nvidia-cuda-toolkit/libdevice/",
-]
-
-# Files used to detect the NVVM libdevice path.
-NVVM_LIBDEVICE_FILES = [
- # CUDA 9.0 has a single file.
- "libdevice.10.bc",
-
- # CUDA 8.0 has separate files for compute versions 2.0, 3.0, 3.5 and 5.0.
- # Probing for one of them is sufficient.
- "libdevice.compute_20.10.bc",
-]
-
def _get_python_bin(repository_ctx):
"""Gets the python bin path."""
python_bin = repository_ctx.os.environ.get(_PYTHON_BIN_PATH)
@@ -414,31 +340,6 @@
return enable_cuda == "1"
return False
-def cuda_toolkit_path(repository_ctx):
- """Finds the cuda toolkit directory.
-
- Args:
- repository_ctx: The repository context.
-
- Returns:
- A speculative real path of the cuda toolkit install directory.
- """
- cuda_toolkit_path = _DEFAULT_CUDA_TOOLKIT_PATH
- if _CUDA_TOOLKIT_PATH in repository_ctx.os.environ:
- cuda_toolkit_path = repository_ctx.os.environ[_CUDA_TOOLKIT_PATH].strip()
- if not repository_ctx.path(cuda_toolkit_path).exists:
- auto_configure_fail("Cannot find cuda toolkit path.")
- return str(repository_ctx.path(cuda_toolkit_path).realpath)
-
-def _cudnn_install_basedir(repository_ctx):
- """Finds the cudnn install directory."""
- cudnn_install_path = _DEFAULT_CUDNN_INSTALL_PATH
- if _CUDNN_INSTALL_PATH in repository_ctx.os.environ:
- cudnn_install_path = repository_ctx.os.environ[_CUDNN_INSTALL_PATH].strip()
- if not repository_ctx.path(cudnn_install_path).exists:
- auto_configure_fail("Cannot find cudnn install path.")
- return cudnn_install_path
-
def matches_version(environ_version, detected_version):
"""Checks whether the user-specified version matches the detected version.
@@ -477,71 +378,7 @@
_NVCC_VERSION_PREFIX = "Cuda compilation tools, release "
-def _cuda_version(repository_ctx, cuda_toolkit_path, cpu_value):
- """Detects the version of CUDA installed on the system.
-
- Args:
- repository_ctx: The repository context.
- cuda_toolkit_path: The CUDA install directory.
-
- Returns:
- String containing the version of CUDA.
- """
-
- # Run nvcc --version and find the line containing the CUDA version.
- nvcc_path = repository_ctx.path("%s/bin/nvcc%s" % (
- cuda_toolkit_path,
- ".exe" if cpu_value == "Windows" else "",
- ))
- if not nvcc_path.exists:
- auto_configure_fail("Cannot find nvcc at %s" % str(nvcc_path))
- result = repository_ctx.execute([str(nvcc_path), "--version"])
- if result.stderr:
- auto_configure_fail("Error running nvcc --version: %s" % result.stderr)
- lines = result.stdout.splitlines()
- version_line = lines[len(lines) - 1]
- if version_line.find(_NVCC_VERSION_PREFIX) == -1:
- auto_configure_fail(
- "Could not parse CUDA version from nvcc --version. Got: %s" %
- result.stdout,
- )
-
- # Parse the CUDA version from the line containing the CUDA version.
- prefix_removed = version_line.replace(_NVCC_VERSION_PREFIX, "")
- parts = prefix_removed.split(",")
- if len(parts) != 2 or len(parts[0]) < 2:
- auto_configure_fail(
- "Could not parse CUDA version from nvcc --version. Got: %s" %
- result.stdout,
- )
- full_version = parts[1].strip()
- if full_version.startswith("V"):
- full_version = full_version[1:]
-
- # Check whether TF_CUDA_VERSION was set by the user and fail if it does not
- # match the detected version.
- environ_version = ""
- if _TF_CUDA_VERSION in repository_ctx.os.environ:
- environ_version = repository_ctx.os.environ[_TF_CUDA_VERSION].strip()
- if environ_version and not matches_version(environ_version, full_version):
- auto_configure_fail(
- ("CUDA version detected from nvcc (%s) does not match " +
- "TF_CUDA_VERSION (%s)") % (full_version, environ_version),
- )
-
- # We only use the version consisting of the major and minor version numbers.
- version_parts = full_version.split(".")
- if len(version_parts) < 2:
- auto_configure_fail("CUDA version detected from nvcc (%s) is incomplete.")
- if cpu_value == "Windows":
- version = "64_%s%s" % (version_parts[0], version_parts[1])
- else:
- version = "%s.%s" % (version_parts[0], version_parts[1])
- return version
-
_DEFINE_CUDNN_MAJOR = "#define CUDNN_MAJOR"
-_DEFINE_CUDNN_MINOR = "#define CUDNN_MINOR"
-_DEFINE_CUDNN_PATCHLEVEL = "#define CUDNN_PATCHLEVEL"
def find_cuda_define(repository_ctx, header_dir, header_file, define):
"""Returns the value of a #define in a header file.
@@ -611,60 +448,6 @@
version = version[:version_end].strip()
return version
-def _cudnn_version(repository_ctx, cudnn_install_basedir, cpu_value):
- """Detects the version of cuDNN installed on the system.
-
- Args:
- repository_ctx: The repository context.
- cpu_value: The name of the host operating system.
- cudnn_install_basedir: The cuDNN install directory.
-
- Returns:
- A string containing the version of cuDNN.
- """
- cudnn_header_dir = _find_cudnn_header_dir(
- repository_ctx,
- cudnn_install_basedir,
- )
- major_version = find_cuda_define(
- repository_ctx,
- cudnn_header_dir,
- "cudnn.h",
- _DEFINE_CUDNN_MAJOR,
- )
- minor_version = find_cuda_define(
- repository_ctx,
- cudnn_header_dir,
- "cudnn.h",
- _DEFINE_CUDNN_MINOR,
- )
- patch_version = find_cuda_define(
- repository_ctx,
- cudnn_header_dir,
- "cudnn.h",
- _DEFINE_CUDNN_PATCHLEVEL,
- )
- full_version = "%s.%s.%s" % (major_version, minor_version, patch_version)
-
- # Check whether TF_CUDNN_VERSION was set by the user and fail if it does not
- # match the detected version.
- environ_version = ""
- if _TF_CUDNN_VERSION in repository_ctx.os.environ:
- environ_version = repository_ctx.os.environ[_TF_CUDNN_VERSION].strip()
- if environ_version and not matches_version(environ_version, full_version):
- cudnn_h_path = repository_ctx.path(
- "%s/include/cudnn.h" % cudnn_install_basedir,
- )
- auto_configure_fail(("cuDNN version detected from %s (%s) does not match " +
- "TF_CUDNN_VERSION (%s)") %
- (str(cudnn_h_path), full_version, environ_version))
-
- # Only use the major version to match the SONAME of the library.
- version = major_version
- if cpu_value == "Windows":
- version = "64_" + version
- return version
-
def compute_capabilities(repository_ctx):
"""Returns a list of strings representing cuda compute capabilities."""
if _TF_CUDA_COMPUTE_CAPABILITIES not in repository_ctx.os.environ:
@@ -699,7 +482,7 @@
def _is_windows(repository_ctx):
"""Returns true if the host operating system is windows."""
- return get_cpu_value(repository_ctx) == "Windows"
+ return repository_ctx.os.name.lower().find("windows") >= 0
def lib_name(base_name, cpu_value, version = None, static = False):
"""Constructs the platform-specific name of a library.
@@ -777,57 +560,11 @@
Returns the path to the library.
"""
file_name = lib_name(lib, cpu_value, version, static)
- return find_lib(repository_ctx, [
- "%s/%s%s" % (basedir, path, file_name)
- for path in CUDA_LIB_PATHS
- ], check_soname = version and not static)
-
-def _find_cupti_header_dir(repository_ctx, cuda_config):
- """Returns the path to the directory containing cupti.h
-
- On most systems, the cupti library is not installed in the same directory as
- the other CUDA libraries but rather in a special extras/CUPTI directory.
-
- Args:
- repository_ctx: The repository context.
- cuda_config: The CUDA config as returned by _get_cuda_config
-
- Returns:
- The path of the directory containing the cupti header.
- """
- cuda_toolkit_path = cuda_config.cuda_toolkit_path
- for relative_path in CUPTI_HEADER_PATHS:
- if repository_ctx.path(
- "%s/%scupti.h" % (cuda_toolkit_path, relative_path),
- ).exists:
- return ("%s/%s" % (cuda_toolkit_path, relative_path))[:-1]
- auto_configure_fail("Cannot find cupti.h under %s" % ", ".join(
- [cuda_toolkit_path + "/" + s for s in CUPTI_HEADER_PATHS],
- ))
-
-def _find_cupti_lib(repository_ctx, cuda_config):
- """Finds the cupti library on the system.
-
- On most systems, the cupti library is not installed in the same directory as
- the other CUDA libraries but rather in a special extras/CUPTI directory.
-
- Args:
- repository_ctx: The repository context.
- cuda_config: The cuda configuration as returned by _get_cuda_config.
-
- Returns:
- Returns the path to the library.
- """
- file_name = lib_name(
- "cupti",
- cuda_config.cpu_value,
- cuda_config.cuda_version,
+ return find_lib(
+ repository_ctx,
+ ["%s/%s" % (basedir, file_name)],
+ check_soname = version and not static,
)
- basedir = cuda_config.cuda_toolkit_path
- return find_lib(repository_ctx, [
- "%s/%s%s" % (basedir, path, file_name)
- for path in CUPTI_LIB_PATHS
- ])
def _find_libs(repository_ctx, cuda_config):
"""Returns the CUDA and cuDNN libraries on the system.
@@ -840,26 +577,27 @@
Map of library names to structs of filename and path.
"""
cpu_value = cuda_config.cpu_value
+ stub_dir = "" if _is_windows(repository_ctx) else "/stubs"
return {
"cuda": _find_cuda_lib(
"cuda",
repository_ctx,
cpu_value,
- cuda_config.cuda_toolkit_path,
+ cuda_config.config["cuda_library_dir"] + stub_dir,
None,
),
"cudart": _find_cuda_lib(
"cudart",
repository_ctx,
cpu_value,
- cuda_config.cuda_toolkit_path,
+ cuda_config.config["cuda_library_dir"],
cuda_config.cuda_version,
),
"cudart_static": _find_cuda_lib(
"cudart_static",
repository_ctx,
cpu_value,
- cuda_config.cuda_toolkit_path,
+ cuda_config.config["cuda_library_dir"],
cuda_config.cuda_version,
static = True,
),
@@ -867,112 +605,75 @@
"cublas",
repository_ctx,
cpu_value,
- cuda_config.cuda_toolkit_path,
- cuda_config.cuda_version,
+ cuda_config.config["cublas_library_dir"],
+ cuda_config.cuda_lib_version,
),
"cusolver": _find_cuda_lib(
"cusolver",
repository_ctx,
cpu_value,
- cuda_config.cuda_toolkit_path,
- cuda_config.cuda_version,
+ cuda_config.config["cuda_library_dir"],
+ cuda_config.cuda_lib_version,
),
"curand": _find_cuda_lib(
"curand",
repository_ctx,
cpu_value,
- cuda_config.cuda_toolkit_path,
- cuda_config.cuda_version,
+ cuda_config.config["cuda_library_dir"],
+ cuda_config.cuda_lib_version,
),
"cufft": _find_cuda_lib(
"cufft",
repository_ctx,
cpu_value,
- cuda_config.cuda_toolkit_path,
- cuda_config.cuda_version,
+ cuda_config.config["cuda_library_dir"],
+ cuda_config.cuda_lib_version,
),
"cudnn": _find_cuda_lib(
"cudnn",
repository_ctx,
cpu_value,
- cuda_config.cudnn_install_basedir,
+ cuda_config.config["cudnn_library_dir"],
cuda_config.cudnn_version,
),
- "cupti": _find_cupti_lib(repository_ctx, cuda_config),
+ "cupti": _find_cuda_lib(
+ "cupti",
+ repository_ctx,
+ cpu_value,
+ cuda_config.config["cupti_library_dir"],
+ cuda_config.cuda_version,
+ ),
"cusparse": _find_cuda_lib(
"cusparse",
repository_ctx,
cpu_value,
- cuda_config.cuda_toolkit_path,
- cuda_config.cuda_version,
+ cuda_config.config["cuda_library_dir"],
+ cuda_config.cuda_lib_version,
),
}
-def _find_cuda_include_path(repository_ctx, cuda_config):
- """Returns the path to the directory containing cuda.h
-
- Args:
- repository_ctx: The repository context.
- cuda_config: The CUDA config as returned by _get_cuda_config
-
- Returns:
- The path of the directory containing the CUDA headers.
- """
- cuda_toolkit_path = cuda_config.cuda_toolkit_path
- for relative_path in CUDA_INCLUDE_PATHS:
- if repository_ctx.path(
- "%s/%scuda.h" % (cuda_toolkit_path, relative_path),
- ).exists:
- return ("%s/%s" % (cuda_toolkit_path, relative_path))[:-1]
- auto_configure_fail("Cannot find cuda.h under %s" % cuda_toolkit_path)
-
-def _find_cudnn_header_dir(repository_ctx, cudnn_install_basedir):
- """Returns the path to the directory containing cudnn.h
-
- Args:
- repository_ctx: The repository context.
- cudnn_install_basedir: The cudnn install directory as returned by
- _cudnn_install_basedir.
-
- Returns:
- The path of the directory containing the cudnn header.
- """
- for relative_path in CUDA_INCLUDE_PATHS:
- if repository_ctx.path(
- "%s/%scudnn.h" % (cudnn_install_basedir, relative_path),
- ).exists:
- return ("%s/%s" % (cudnn_install_basedir, relative_path))[:-1]
- if repository_ctx.path("/usr/include/cudnn.h").exists:
- return "/usr/include"
- auto_configure_fail("Cannot find cudnn.h under %s" % cudnn_install_basedir)
-
-def _find_nvvm_libdevice_dir(repository_ctx, cuda_config):
- """Returns the path to the directory containing libdevice in bitcode format.
-
- Args:
- repository_ctx: The repository context.
- cuda_config: The CUDA config as returned by _get_cuda_config
-
- Returns:
- The path of the directory containing the CUDA headers.
- """
- cuda_toolkit_path = cuda_config.cuda_toolkit_path
- for libdevice_file in NVVM_LIBDEVICE_FILES:
- for relative_path in NVVM_LIBDEVICE_PATHS:
- if repository_ctx.path("%s/%s%s" % (
- cuda_toolkit_path,
- relative_path,
- libdevice_file,
- )).exists:
- return ("%s/%s" % (cuda_toolkit_path, relative_path))[:-1]
- auto_configure_fail(
- "Cannot find libdevice*.bc files under %s" % cuda_toolkit_path,
- )
-
def _cudart_static_linkopt(cpu_value):
"""Returns additional platform-specific linkopts for cudart."""
return "" if cpu_value == "Darwin" else "\"-lrt\","
+# TODO(csigg): Only call once instead of from here, tensorrt_configure.bzl,
+# and nccl_configure.bzl.
+def find_cuda_config(repository_ctx, cuda_libraries):
+ """Returns CUDA config dictionary from running find_cuda_config.py"""
+ exec_result = repository_ctx.execute([
+ _get_python_bin(repository_ctx),
+ repository_ctx.path(Label("@org_tensorflow//third_party/gpus:find_cuda_config.py")),
+ ] + cuda_libraries)
+ if exec_result.return_code:
+ auto_configure_fail("Failed to run find_cuda_config.py: %s" % exec_result.stderr)
+ config = {}
+ for line in exec_result.stdout.splitlines():
+ elements = line.split(": ")
+ if len(elements) != 2:
+ auto_configure_fail("Unexpected output from find_cuda_config.py")
+ config[elements[0]] = elements[1]
+ return config
+
def _get_cuda_config(repository_ctx):
"""Detects and returns information about the CUDA installation on the system.
@@ -988,22 +689,30 @@
compute_capabilities: A list of the system's CUDA compute capabilities.
cpu_value: The name of the host operating system.
"""
+ config = find_cuda_config(repository_ctx, ["cuda", "cudnn"])
cpu_value = get_cpu_value(repository_ctx)
- toolkit_path = cuda_toolkit_path(repository_ctx)
- cuda_version = _cuda_version(repository_ctx, toolkit_path, cpu_value)
- cudnn_install_basedir = _cudnn_install_basedir(repository_ctx)
- cudnn_version = _cudnn_version(
- repository_ctx,
- cudnn_install_basedir,
- cpu_value,
- )
+ toolkit_path = config["cuda_toolkit_path"]
+
+ is_windows = _is_windows(repository_ctx)
+ cuda_version = config["cuda_version"].split(".")
+ cuda_version = ("64_%s%s" if is_windows else "%s.%s") % (cuda_version[0], cuda_version[1])
+ cudnn_version = ("64_%s" if is_windows else "%s") % config["cudnn_version"]
+
+ # cuda_lib_version is for libraries like cuBLAS, cuFFT, cuSOLVER, etc.
+ # It changed from 'x.y' to just 'x' in CUDA 10.1.
+ if (int(cuda_version[0]), int(cuda_version[1])) >= (10, 1):
+ cuda_lib_version = ("64_%s" if is_windows else "%s") % cuda_version[0]
+ else:
+ cuda_lib_version = cuda_version
+
return struct(
cuda_toolkit_path = toolkit_path,
- cudnn_install_basedir = cudnn_install_basedir,
cuda_version = cuda_version,
cudnn_version = cudnn_version,
+ cuda_lib_version = cuda_lib_version,
compute_capabilities = compute_capabilities(repository_ctx),
cpu_value = cpu_value,
+ config = config,
)
def _tpl(repository_ctx, tpl, substitutions = {}, out = None):
@@ -1108,13 +817,14 @@
repository_ctx,
"cuda:cuda_config.h",
{
- "%{cuda_version}": _DEFAULT_CUDA_VERSION,
- "%{cudnn_version}": _DEFAULT_CUDNN_VERSION,
+ "%{cuda_version}": "",
+ "%{cuda_lib_version}": "",
+ "%{cudnn_version}": "",
"%{cuda_compute_capabilities}": ",".join([
"CudaVersion(\"%s\")" % c
for c in _DEFAULT_CUDA_COMPUTE_CAPABILITIES
]),
- "%{cuda_toolkit_path}": _DEFAULT_CUDA_TOOLKIT_PATH,
+ "%{cuda_toolkit_path}": "",
},
"cuda/cuda/cuda_config.h",
)
@@ -1249,13 +959,11 @@
"""Creates the repository containing files set up to build with CUDA."""
cuda_config = _get_cuda_config(repository_ctx)
- cuda_include_path = _find_cuda_include_path(repository_ctx, cuda_config)
- cudnn_header_dir = _find_cudnn_header_dir(
- repository_ctx,
- cuda_config.cudnn_install_basedir,
- )
- cupti_header_dir = _find_cupti_header_dir(repository_ctx, cuda_config)
- nvvm_libdevice_dir = _find_nvvm_libdevice_dir(repository_ctx, cuda_config)
+ cuda_include_path = cuda_config.config["cuda_include_dir"]
+ cublas_include_path = cuda_config.config["cublas_include_dir"]
+ cudnn_header_dir = cuda_config.config["cudnn_include_dir"]
+ cupti_header_dir = cuda_config.config["cupti_include_dir"]
+ nvvm_libdevice_dir = cuda_config.config["nvvm_library_dir"]
# Create genrule to copy files from the installed CUDA toolkit into execroot.
copy_rules = [
@@ -1279,6 +987,24 @@
),
]
+ if cublas_include_path != cuda_include_path:
+ copy_rules.append(make_copy_files_rule(
+ repository_ctx,
+ name = "cublas-include",
+ srcs = [
+ cublas_include_path + "/cublas.h",
+ cublas_include_path + "/cublas_v2.h",
+ cublas_include_path + "/cublas_api.h",
+ ],
+ outs = [
+ "cuda/include/cublas.h",
+ "cuda/include/cublas_v2.h",
+ "cuda/include/cublas_api.h",
+ ],
+ ))
+ else:
+ copy_rules.append("filegroup(name = 'cublas-include')\n")
+
cuda_libs = _find_libs(repository_ctx, cuda_config)
cuda_lib_srcs = []
cuda_lib_outs = []
@@ -1340,7 +1066,9 @@
"%{cusparse_lib}": cuda_libs["cusparse"].basename,
"%{copy_rules}": "\n".join(copy_rules),
"%{cuda_headers}": (
- '":cuda-include",\n' + ' ":cudnn-include",'
+ '":cuda-include",\n' +
+ ' ":cublas-include",' +
+ ' ":cudnn-include",'
),
},
"cuda/BUILD",
@@ -1423,8 +1151,8 @@
"flag: \"-fno-canonical-system-headers\""
)
nvcc_path = str(
- repository_ctx.path("%s/bin/nvcc%s" % (
- cuda_config.cuda_toolkit_path,
+ repository_ctx.path("%s/nvcc%s" % (
+ cuda_config.config["cuda_binary_dir"],
".exe" if _is_windows(repository_ctx) else "",
)),
)
@@ -1472,6 +1200,7 @@
"cuda:cuda_config.h",
{
"%{cuda_version}": cuda_config.cuda_version,
+ "%{cuda_lib_version}": cuda_config.cuda_lib_version,
"%{cudnn_version}": cuda_config.cudnn_version,
"%{cuda_compute_capabilities}": ",".join([
"CudaVersion(\"%s\")" % c
@@ -1546,6 +1275,7 @@
_PYTHON_BIN_PATH,
"TMP",
"TMPDIR",
+ "TF_CUDA_PATHS",
],
)
diff --git a/third_party/gpus/find_cuda_config.py b/third_party/gpus/find_cuda_config.py
new file mode 100644
index 0000000..87be0b2
--- /dev/null
+++ b/third_party/gpus/find_cuda_config.py
@@ -0,0 +1,457 @@
+# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Prints CUDA library and header directories and versions found on the system.
+
+The script searches for CUDA library and header files on the system, inspects
+them to determine their version and prints the configuration to stdout.
+The paths to inspect and the required versions are specified through environment
+variables. If no valid configuration is found, the script prints to stderr and
+returns an error code.
+
+The list of libraries to find is specified as arguments. Supported libraries are
+CUDA (includes cuBLAS), cuDNN, NCCL, and TensorRT.
+
+The script takes a list of base directories specified by the TF_CUDA_PATHS
+environment variable as comma-separated glob list. The script looks for headers
+and library files in a hard-coded set of subdirectories from these base paths.
+If TF_CUDA_PATHS is not specified, a OS specific default is used:
+
+ Linux: /usr/local/cuda, /usr, and paths from 'ldconfig -p'.
+ Windows: CUDA_PATH environment variable, or
+ C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\*
+
+For backwards compatibility, some libraries also use alternative base
+directories from other environment variables if they are specified. List of
+library-specific environment variables:
+
+ Library Version env variable Additional base directories
+ ----------------------------------------------------------------
+ CUDA TF_CUDA_VERSION CUDA_TOOLKIT_PATH
+ cuBLAS TF_CUBLAS_VERSION CUDA_TOOLKIT_PATH
+ cuDNN TF_CUDNN_VERSION CUDNN_INSTALL_PATH
+ NCCL TF_NCCL_VERSION NCCL_INSTALL_PATH, NCCL_HDR_PATH
+ TensorRT TF_TENSORRT_VERSION TENSORRT_INSTALL_PATH
+
+Versions environment variables can be of the form 'x' or 'x.y' to request a
+specific version, empty or unspecified to accept any version.
+
+The output of a found library is of the form:
+tf_<library>_version: x.y.z
+tf_<library>_header_dir: ...
+tf_<library>_library_dir: ...
+"""
+
+import os
+import glob
+import platform
+import re
+import subprocess
+import sys
+
+# pylint: disable=g-import-not-at-top
+try:
+ from shutil import which
+except ImportError:
+ from distutils.spawn import find_executable as which
+# pylint: enable=g-import-not-at-top
+
+
+class ConfigError(Exception):
+ pass
+
+
+def _is_linux():
+ return platform.system() == "Linux"
+
+
+def _is_windows():
+ return platform.system() == "Windows"
+
+
+def _is_macos():
+ return platform.system() == "Darwin"
+
+
+def _matches_version(actual_version, required_version):
+ """Checks whether some version meets the requirements.
+
+ All elements of the required_version need to be present in the
+ actual_version.
+
+ required_version actual_version result
+ -----------------------------------------
+ 1 1.1 True
+ 1.2 1 False
+ 1.2 1.3 False
+ 1 True
+
+ Args:
+ required_version: The version specified by the user.
+ actual_version: The version detected from the CUDA installation.
+ Returns: Whether the actual version matches the required one.
+ """
+ if actual_version is None:
+ return False
+ return actual_version.startswith(required_version)
+
+
+def _at_least_version(actual_version, required_version):
+ actual = [int(v) for v in actual_version.split(".")]
+ required = [int(v) for v in required_version.split(".")]
+ return actual >= required
+
+
+def _get_header_version(path, name):
+ """Returns preprocessor defines in C header file."""
+ for line in open(path, "r").readlines():
+ match = re.match("#define %s (\d+)" % name, line)
+ if match:
+ return match.group(1)
+ return ""
+
+
+def _cartesian_product(first, second):
+ """Returns all path combinations of first and second."""
+ return [os.path.join(f, s) for f in first for s in second]
+
+
+def _get_ld_config_paths():
+ """Returns all directories from 'ldconfig -p'."""
+ if not _is_linux():
+ return []
+ ldconfig_path = which("ldconfig") or "/sbin/ldconfig"
+ output = subprocess.check_output([ldconfig_path, "-p"])
+ pattern = re.compile(".* => (.*)")
+ result = set()
+ for line in output.splitlines():
+ match = pattern.match(line.decode("ascii"))
+ if match:
+ result.add(os.path.dirname(match.group(1)))
+ return list(result)
+
+
+def _get_default_cuda_paths(cuda_version):
+ if not cuda_version:
+ cuda_version = "*"
+ elif not "." in cuda_version:
+ cuda_version = cuda_version + ".*"
+
+ if _is_windows():
+ return [
+ os.environ.get(
+ "CUDA_PATH",
+ "C:\\Program Files\\NVIDIA GPU Computing Toolkit\\CUDA\\v%s\\" %
+ cuda_version)
+ ]
+ return ["/usr/local/cuda-%s" % cuda_version, "/usr"] + _get_ld_config_paths()
+
+
+def _header_paths(base_paths):
+ return _cartesian_product(base_paths, [
+ "",
+ "include",
+ "include/cuda",
+ "include/*-linux-gnu",
+ "extras/CUPTI/include",
+ "include/cuda/CUPTI",
+ ])
+
+
+def _library_paths(base_paths):
+ return _cartesian_product(base_paths, [
+ "",
+ "lib64",
+ "lib",
+ "lib/*-linux-gnu",
+ "lib/x64",
+ "extras/CUPTI/*",
+ ])
+
+
+def _not_found_error(paths, filepattern):
+ return ConfigError(
+ "Could not find any %s in:%s" %
+ (filepattern, "".join(["\n %s" % path for path in sorted(paths)])))
+
+
+def _find_file(paths, filepattern):
+ for path in paths:
+ for file in glob.glob(os.path.join(path, filepattern)):
+ return file
+ raise _not_found_error(paths, filepattern)
+
+
+def _find_library(base_paths, library_name, required_version):
+ """Returns first valid path to the requested library."""
+ if _is_windows():
+ filepattern = library_name + ".lib"
+ elif _is_macos():
+ filepattern = "%s*.dylib" % (".".join(["lib" + library_name] +
+ required_version.split(".")[:1]))
+ else:
+ filepattern = ".".join(["lib" + library_name, "so"] +
+ required_version.split(".")[:1]) + "*"
+ return _find_file(_library_paths(base_paths), filepattern)
+
+
+def _find_versioned_file(paths, filepattern, required_version, get_version):
+ """Returns first valid path to a file that matches the requested version."""
+ for path in paths:
+ for file in glob.glob(os.path.join(path, filepattern)):
+ actual_version = get_version(file)
+ if _matches_version(actual_version, required_version):
+ return file, actual_version
+ raise _not_found_error(
+ paths, filepattern + " matching version '%s'" % required_version)
+
+
+def _find_header(base_paths, header_name, required_version, get_version):
+ """Returns first valid path to a header that matches the requested version."""
+ return _find_versioned_file(
+ _header_paths(base_paths), header_name, required_version, get_version)
+
+
+def _find_cuda_config(base_paths, required_version):
+
+ def get_header_version(path):
+ version = int(_get_header_version(path, "CUDA_VERSION"))
+ if not version:
+ return None
+ return "%d.%d" % (version // 1000, version % 1000 // 10)
+
+ cuda_header_path, header_version = _find_header(base_paths, "cuda.h",
+ required_version,
+ get_header_version)
+ cuda_version = header_version # x.y, see above.
+
+ cuda_library_path = _find_library(base_paths, "cudart", cuda_version)
+
+ def get_nvcc_version(path):
+ pattern = "Cuda compilation tools, release \d+\.\d+, V(\d+\.\d+\.\d+)"
+ for line in subprocess.check_output([path, "--version"]).splitlines():
+ match = re.match(pattern, line.decode("ascii"))
+ if match:
+ return match.group(1)
+ return None
+
+ nvcc_name = "nvcc.exe" if _is_windows() else "nvcc"
+ nvcc_path, nvcc_version = _find_versioned_file(
+ _cartesian_product(base_paths, [
+ "",
+ "bin",
+ ]), nvcc_name, cuda_version, get_nvcc_version)
+
+ nvvm_path = _find_file(
+ _cartesian_product(base_paths, [
+ "nvvm/libdevice",
+ "share/cuda",
+ "lib/nvidia-cuda-toolkit/libdevice",
+ ]), "libdevice*.10.bc")
+
+ cupti_header_path = _find_file(_header_paths(base_paths), "cupti.h")
+ cupti_library_path = _find_library(base_paths, "cupti", required_version)
+
+ cuda_binary_dir = os.path.dirname(nvcc_path)
+ nvvm_library_dir = os.path.dirname(nvvm_path)
+
+ # XLA requires the toolkit path to find ptxas and libdevice.
+ # TODO(csigg): pass in both directories instead.
+ cuda_toolkit_paths = (
+ os.path.normpath(os.path.join(cuda_binary_dir, "..")),
+ os.path.normpath(os.path.join(nvvm_library_dir, "../..")),
+ )
+ if cuda_toolkit_paths[0] != cuda_toolkit_paths[1]:
+ raise ConfigError("Inconsistent CUDA toolkit path: %s vs %s" %
+ cuda_toolkit_paths)
+
+ return {
+ "cuda_version": cuda_version,
+ "cuda_include_dir": os.path.dirname(cuda_header_path),
+ "cuda_library_dir": os.path.dirname(cuda_library_path),
+ "cuda_binary_dir": cuda_binary_dir,
+ "nvvm_library_dir": nvvm_library_dir,
+ "cupti_include_dir": os.path.dirname(cupti_header_path),
+ "cupti_library_dir": os.path.dirname(cupti_library_path),
+ "cuda_toolkit_path": cuda_toolkit_paths[0],
+ }
+
+
+def _find_cublas_config(base_paths, required_version, cuda_version):
+
+ if _at_least_version(cuda_version, "10.1"):
+
+ def get_header_version(path):
+ version = (
+ _get_header_version(path, name)
+ for name in ("CUBLAS_VER_MAJOR", "CUBLAS_VER_MINOR",
+ "CUBLAS_VER_PATCH"))
+ return ".".join(version)
+
+ header_path, header_version = _find_header(base_paths, "cublas_api.h",
+ required_version,
+ get_header_version)
+ # cuBLAS uses the major version only.
+ cublas_version = header_version.split(".")[0]
+
+ if not _matches_version(cuda_version, cublas_version):
+ raise ConfigError("cuBLAS version %s does not match CUDA version %s" %
+ (cublas_version, cuda_version))
+
+ else:
+ # There is no version info available before CUDA 10.1, just find the file.
+ header_path = _find_file(_header_paths(base_paths), "cublas_api.h")
+ # cuBLAS version is the same as CUDA version (x.y).
+ cublas_version = required_version
+
+ library_path = _find_library(base_paths, "cublas", cublas_version)
+
+ return {
+ "cublas_include_dir": os.path.dirname(header_path),
+ "cublas_library_dir": os.path.dirname(library_path),
+ }
+
+
+def _find_cudnn_config(base_paths, required_version):
+
+ def get_header_version(path):
+ version = (
+ _get_header_version(path, name)
+ for name in ("CUDNN_MAJOR", "CUDNN_MINOR", "CUDNN_PATCHLEVEL"))
+ return ".".join(version)
+
+ header_path, header_version = _find_header(base_paths, "cudnn.h",
+ required_version,
+ get_header_version)
+ cudnn_version = header_version.split(".")[0]
+
+ library_path = _find_library(base_paths, "cudnn", cudnn_version)
+
+ return {
+ "cudnn_version": cudnn_version,
+ "cudnn_include_dir": os.path.dirname(header_path),
+ "cudnn_library_dir": os.path.dirname(library_path),
+ }
+
+
+def _find_nccl_config(base_paths, required_version):
+
+ def get_header_version(path):
+ version = (
+ _get_header_version(path, name)
+ for name in ("NCCL_MAJOR", "NCCL_MINOR", "NCCL_PATCH"))
+ return ".".join(version)
+
+ header_path, header_version = _find_header(base_paths, "nccl.h",
+ required_version,
+ get_header_version)
+ nccl_version = header_version.split(".")[0]
+
+ library_path = _find_library(base_paths, "nccl", nccl_version)
+
+ return {
+ "nccl_version": nccl_version,
+ "nccl_include_dir": os.path.dirname(header_path),
+ "nccl_library_dir": os.path.dirname(library_path),
+ }
+
+
+def _find_tensorrt_config(base_paths, required_version):
+
+ def get_header_version(path):
+ version = (
+ _get_header_version(path, name)
+ for name in ("NV_TENSORRT_MAJOR", "NV_TENSORRT_MINOR",
+ "NV_TENSORRT_PATCH"))
+ return ".".join(version)
+
+ header_path, header_version = _find_header(base_paths, "NvInfer.h",
+ required_version,
+ get_header_version)
+ tensorrt_version = header_version.split(".")[0]
+
+ library_path = _find_library(base_paths, "nvinfer", tensorrt_version)
+
+ return {
+ "tensorrt_version": tensorrt_version,
+ "tensorrt_include_dir": os.path.dirname(header_path),
+ "tensorrt_library_dir": os.path.dirname(library_path),
+ }
+
+
+def _list_from_env(env_name, default=[]):
+ """Returns comma-separated list from environment variable."""
+ if env_name in os.environ:
+ return os.environ[env_name].split(",")
+ return default
+
+
+def _normalize_path(path):
+ """Returns normalized path, with forward slashes on Windows."""
+ path = os.path.normpath(path)
+ if _is_windows():
+ path = path.replace("\\", "/")
+ return path
+
+
+def find_cuda_config():
+ """Returns a dictionary of CUDA library and header file paths."""
+ libraries = [argv.lower() for argv in sys.argv[1:]]
+ cuda_version = os.environ.get("TF_CUDA_VERSION", "")
+ base_paths = _list_from_env("TF_CUDA_PATHS",
+ _get_default_cuda_paths(cuda_version))
+
+ result = {}
+ if "cuda" in libraries:
+ cuda_paths = _list_from_env("CUDA_TOOLKIT_PATH", base_paths)
+ result.update(_find_cuda_config(cuda_paths, cuda_version))
+
+ cuda_version = result["cuda_version"]
+ cublas_version = os.environ.get("TF_CUBLAS_VERSION", "")
+ result.update(_find_cublas_config(cuda_paths, cublas_version, cuda_version))
+
+ if "cudnn" in libraries:
+ cudnn_paths = _list_from_env("CUDNN_INSTALL_PATH", base_paths)
+ cudnn_version = os.environ.get("TF_CUDNN_VERSION", "")
+ result.update(_find_cudnn_config(cudnn_paths, cudnn_version))
+
+ if "nccl" in libraries:
+ nccl_paths = _list_from_env("NCCL_INSTALL_PATH",
+ base_paths) + _list_from_env("NCCL_HDR_PATH")
+ nccl_version = os.environ.get("TF_NCCL_VERSION", "")
+ result.update(_find_nccl_config(nccl_paths, nccl_version))
+
+ if "tensorrt" in libraries:
+ tensorrt_paths = _list_from_env("TENSORRT_INSTALL_PATH", base_paths)
+ tensorrt_version = os.environ.get("TF_TENSORRT_VERSION", "")
+ result.update(_find_tensorrt_config(tensorrt_paths, tensorrt_version))
+
+ for k, v in result.items():
+ if k.endswith("_dir") or k.endswith("_path"):
+ result[k] = _normalize_path(v)
+
+ return result
+
+
+def main():
+ try:
+ for key, value in sorted(find_cuda_config().items()):
+ print("%s: %s" % (key, value))
+ except ConfigError as e:
+ sys.stderr.write(str(e))
+ sys.exit(1)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/third_party/nccl/nccl_configure.bzl b/third_party/nccl/nccl_configure.bzl
index 07e4ad7..cfb2599 100644
--- a/third_party/nccl/nccl_configure.bzl
+++ b/third_party/nccl/nccl_configure.bzl
@@ -4,19 +4,20 @@
`nccl_configure` depends on the following environment variables:
* `TF_NCCL_VERSION`: Installed NCCL version or empty to build from source.
- * `NCCL_INSTALL_PATH`: The installation path of the NCCL library.
- * `NCCL_HDR_PATH`: The installation path of the NCCL header files.
+ * `NCCL_INSTALL_PATH` (deprecated): The installation path of the NCCL library.
+ * `NCCL_HDR_PATH` (deprecated): The installation path of the NCCL header
+ files.
+ * `TF_CUDA_PATHS`: The base paths to look for CUDA and cuDNN. Default is
+ `/usr/local/cuda,usr/`.
+
"""
load(
"//third_party/gpus:cuda_configure.bzl",
- "auto_configure_fail",
"compute_capabilities",
- "cuda_toolkit_path",
"enable_cuda",
- "find_cuda_define",
+ "find_cuda_config",
"get_cpu_value",
- "matches_version",
)
_CUDA_TOOLKIT_PATH = "CUDA_TOOLKIT_PATH"
@@ -59,63 +60,6 @@
def _label(file):
return Label("//third_party/nccl:{}".format(file))
-def _find_nccl_header(repository_ctx, nccl_install_path):
- """Finds the NCCL header on the system.
-
- Args:
- repository_ctx: The repository context.
- nccl_install_path: The NCCL library install directory.
-
- Returns:
- The path to the NCCL header.
- """
- header_path = repository_ctx.path("%s/include/nccl.h" % nccl_install_path)
- if not header_path.exists:
- auto_configure_fail("Cannot find %s" % str(header_path))
- return header_path
-
-def _check_nccl_version(repository_ctx, nccl_install_path, nccl_hdr_path, nccl_version):
- """Checks whether the header file matches the specified version of NCCL.
-
- Args:
- repository_ctx: The repository context.
- nccl_install_path: The NCCL library install directory.
- nccl_hdr_path: The NCCL header path.
- nccl_version: The expected NCCL version.
-
- Returns:
- A string containing the library version of NCCL.
- """
- header_path = repository_ctx.path("%s/nccl.h" % nccl_hdr_path)
- if not header_path.exists:
- header_path = _find_nccl_header(repository_ctx, nccl_install_path)
- header_dir = str(header_path.realpath.dirname)
- major_version = find_cuda_define(
- repository_ctx,
- header_dir,
- "nccl.h",
- _DEFINE_NCCL_MAJOR,
- )
- minor_version = find_cuda_define(
- repository_ctx,
- header_dir,
- "nccl.h",
- _DEFINE_NCCL_MINOR,
- )
- patch_version = find_cuda_define(
- repository_ctx,
- header_dir,
- "nccl.h",
- _DEFINE_NCCL_PATCH,
- )
- header_version = "%s.%s.%s" % (major_version, minor_version, patch_version)
- if not matches_version(nccl_version, header_version):
- auto_configure_fail(
- ("NCCL library version detected from %s/nccl.h (%s) does not " +
- "match TF_NCCL_VERSION (%s). To fix this rerun configure again.") %
- (header_dir, header_version, nccl_version),
- )
-
def _nccl_configure_impl(repository_ctx):
"""Implementation of the nccl_configure repository rule."""
if (not enable_cuda(repository_ctx) or
@@ -146,14 +90,8 @@
})
else:
# Create target for locally installed NCCL.
- nccl_install_path = repository_ctx.os.environ[_NCCL_INSTALL_PATH].strip()
- nccl_hdr_path = repository_ctx.os.environ[_NCCL_HDR_PATH].strip()
- _check_nccl_version(repository_ctx, nccl_install_path, nccl_hdr_path, nccl_version)
- repository_ctx.template("BUILD", _label("system.BUILD.tpl"), {
- "%{version}": nccl_version,
- "%{install_path}": nccl_install_path,
- "%{hdr_path}": nccl_hdr_path,
- })
+ config = find_cuda_config(repository_ctx, ["nccl"])
+ repository_ctx.template("BUILD", _label("system.BUILD.tpl"), config)
nccl_configure = repository_rule(
implementation = _nccl_configure_impl,
@@ -164,6 +102,7 @@
_TF_NCCL_VERSION,
_TF_CUDA_COMPUTE_CAPABILITIES,
_TF_NEED_CUDA,
+ "TF_CUDA_PATHS",
],
)
"""Detects and configures the NCCL configuration.
diff --git a/third_party/nccl/system.BUILD.tpl b/third_party/nccl/system.BUILD.tpl
index 970dddb..b970aee 100644
--- a/third_party/nccl/system.BUILD.tpl
+++ b/third_party/nccl/system.BUILD.tpl
@@ -5,7 +5,7 @@
cc_library(
name = "nccl",
- srcs = ["libnccl.so.%{version}"],
+ srcs = ["libnccl.so.%{nccl_version}"],
hdrs = ["nccl.h"],
include_prefix = "third_party/nccl",
visibility = ["//visibility:public"],
@@ -17,9 +17,12 @@
genrule(
name = "nccl-files",
outs = [
- "libnccl.so.%{version}",
+ "libnccl.so.%{nccl_version}",
"nccl.h",
],
- cmd = """cp "%{hdr_path}/nccl.h" "$(@D)/nccl.h" &&
- cp "%{install_path}/libnccl.so.%{version}" "$(@D)/libnccl.so.%{version}" """,
+ cmd = """
+cp "%{nccl_header_dir}/nccl.h" "$(@D)/nccl.h" &&
+cp "%{nccl_library_dir}/libnccl.so.%{nccl_version}" \
+ "$(@D)/libnccl.so.%{nccl_version}"
+""",
)
diff --git a/third_party/tensorrt/BUILD.tpl b/third_party/tensorrt/BUILD.tpl
index 4ff7da9..1db5903 100644
--- a/third_party/tensorrt/BUILD.tpl
+++ b/third_party/tensorrt/BUILD.tpl
@@ -11,15 +11,15 @@
cc_library(
name = "tensorrt_headers",
- hdrs = [%{tensorrt_headers}],
+ hdrs = [":tensorrt_include"],
visibility = ["//visibility:public"],
)
cc_library(
name = "tensorrt",
- srcs = %{tensorrt_libs},
+ srcs = [":tensorrt_lib"],
copts = cuda_default_copts(),
- data = %{tensorrt_libs},
+ data = [":tensorrt_lib"],
linkstatic = 1,
visibility = ["//visibility:public"],
deps = [
diff --git a/third_party/tensorrt/tensorrt_configure.bzl b/third_party/tensorrt/tensorrt_configure.bzl
index 9edcd60..004cc0e 100644
--- a/third_party/tensorrt/tensorrt_configure.bzl
+++ b/third_party/tensorrt/tensorrt_configure.bzl
@@ -9,14 +9,10 @@
load(
"//third_party/gpus:cuda_configure.bzl",
- "auto_configure_fail",
- "find_cuda_define",
- "find_lib",
+ "find_cuda_config",
"get_cpu_value",
"lib_name",
- "make_copy_dir_rule",
"make_copy_files_rule",
- "matches_version",
)
_TENSORRT_INSTALL_PATH = "TENSORRT_INSTALL_PATH"
@@ -30,108 +26,6 @@
_DEFINE_TENSORRT_SONAME_MINOR = "#define NV_TENSORRT_SONAME_MINOR"
_DEFINE_TENSORRT_SONAME_PATCH = "#define NV_TENSORRT_SONAME_PATCH"
-def _headers_exist(repository_ctx, path):
- """Returns whether all TensorRT header files could be found in 'path'.
-
- Args:
- repository_ctx: The repository context.
- path: The TensorRT include path to check.
-
- Returns:
- True if all TensorRT header files can be found in the path.
- """
- for h in _TF_TENSORRT_HEADERS:
- if not repository_ctx.path("%s/%s" % (path, h)).exists:
- return False
- return True
-
-def _find_trt_header_dir(repository_ctx, trt_install_path):
- """Returns the path to the directory containing headers of TensorRT.
-
- Args:
- repository_ctx: The repository context.
- trt_install_path: The TensorRT library install directory.
-
- Returns:
- The path of the directory containing the TensorRT header.
- """
- if trt_install_path == "/usr/lib/x86_64-linux-gnu":
- path = "/usr/include/x86_64-linux-gnu"
- if _headers_exist(repository_ctx, path):
- return path
- if trt_install_path == "/usr/lib/aarch64-linux-gnu":
- path = "/usr/include/aarch64-linux-gnu"
- if _headers_exist(repository_ctx, path):
- return path
- path = str(repository_ctx.path("%s/../include" % trt_install_path).realpath)
- if _headers_exist(repository_ctx, path):
- return path
- auto_configure_fail(
- "Cannot find NvInfer.h with TensorRT install path %s" % trt_install_path,
- )
-
-def _trt_lib_version(repository_ctx, trt_install_path):
- """Detects the library (e.g. libnvinfer) version of TensorRT.
-
- Args:
- repository_ctx: The repository context.
- trt_install_path: The TensorRT library install directory.
-
- Returns:
- A string containing the library version of TensorRT.
- """
- trt_header_dir = _find_trt_header_dir(repository_ctx, trt_install_path)
- major_version = find_cuda_define(
- repository_ctx,
- trt_header_dir,
- "NvInfer.h",
- _DEFINE_TENSORRT_SONAME_MAJOR,
- )
- minor_version = find_cuda_define(
- repository_ctx,
- trt_header_dir,
- "NvInfer.h",
- _DEFINE_TENSORRT_SONAME_MINOR,
- )
- patch_version = find_cuda_define(
- repository_ctx,
- trt_header_dir,
- "NvInfer.h",
- _DEFINE_TENSORRT_SONAME_PATCH,
- )
- full_version = "%s.%s.%s" % (major_version, minor_version, patch_version)
- environ_version = repository_ctx.os.environ[_TF_TENSORRT_VERSION].strip()
- if not matches_version(environ_version, full_version):
- auto_configure_fail(
- ("TensorRT library version detected from %s/%s (%s) does not match " +
- "TF_TENSORRT_VERSION (%s). To fix this rerun configure again.") %
- (trt_header_dir, "NvInfer.h", full_version, environ_version),
- )
-
- # Only use the major version to match the SONAME of the library.
- return major_version
-
-def _find_trt_libs(repository_ctx, cpu_value, trt_install_path, trt_lib_version):
- """Finds the given TensorRT library on the system.
-
- Adapted from code contributed by Sami Kama (https://github.com/samikama).
-
- Args:
- repository_ctx: The repository context.
- trt_install_path: The TensorRT library installation directory.
- trt_lib_version: The version of TensorRT library files as returned
- by _trt_lib_version.
-
- Returns:
- The path to the library.
- """
- result = {}
- for lib in _TF_TENSORRT_LIBS:
- file_name = lib_name(lib, cpu_value, trt_lib_version)
- path = find_lib(repository_ctx, ["%s/%s" % (trt_install_path, file_name)])
- result[file_name] = path
- return result
-
def _tpl(repository_ctx, tpl, substitutions):
repository_ctx.template(
tpl,
@@ -162,54 +56,33 @@
)
return
- if _TENSORRT_INSTALL_PATH not in repository_ctx.os.environ:
+ if _TF_TENSORRT_VERSION not in repository_ctx.os.environ:
_create_dummy_repository(repository_ctx)
return
+ config = find_cuda_config(repository_ctx, ["tensorrt"])
+ trt_version = config["tensorrt_version"]
cpu_value = get_cpu_value(repository_ctx)
- if (cpu_value != "Linux"):
- auto_configure_fail("TensorRT is supported only on Linux.")
- if _TF_TENSORRT_VERSION not in repository_ctx.os.environ:
- auto_configure_fail("TensorRT library (libnvinfer) version is not set.")
- trt_install_path = repository_ctx.os.environ[_TENSORRT_INSTALL_PATH].strip()
- if not repository_ctx.path(trt_install_path).exists:
- auto_configure_fail(
- "Cannot find TensorRT install path %s." % trt_install_path,
- )
- # Copy the library files.
- trt_lib_version = _trt_lib_version(repository_ctx, trt_install_path)
- trt_libs = _find_trt_libs(repository_ctx, cpu_value, trt_install_path, trt_lib_version)
- trt_lib_srcs = []
- trt_lib_outs = []
- for path in trt_libs.values():
- trt_lib_srcs.append(str(path))
- trt_lib_outs.append("tensorrt/lib/" + path.basename)
- copy_rules = [make_copy_files_rule(
- repository_ctx,
- name = "tensorrt_lib",
- srcs = trt_lib_srcs,
- outs = trt_lib_outs,
- )]
-
- # Copy the header files header files.
- trt_header_dir = _find_trt_header_dir(repository_ctx, trt_install_path)
- trt_header_srcs = [
- "%s/%s" % (trt_header_dir, header)
- for header in _TF_TENSORRT_HEADERS
- ]
- trt_header_outs = [
- "tensorrt/include/" + header
- for header in _TF_TENSORRT_HEADERS
- ]
- copy_rules.append(
+ # Copy the library and header files.
+ libraries = [lib_name(lib, cpu_value, trt_version) for lib in _TF_TENSORRT_LIBS]
+ library_dir = config["tensorrt_library_dir"] + "/"
+ headers = _TF_TENSORRT_HEADERS
+ include_dir = config["tensorrt_include_dir"] + "/"
+ copy_rules = [
+ make_copy_files_rule(
+ repository_ctx,
+ name = "tensorrt_lib",
+ srcs = [library_dir + library for library in libraries],
+ outs = ["tensorrt/lib/" + library for library in libraries],
+ ),
make_copy_files_rule(
repository_ctx,
name = "tensorrt_include",
- srcs = trt_header_srcs,
- outs = trt_header_outs,
+ srcs = [include_dir + header for header in headers],
+ outs = ["tensorrt/incude/" + header for header in headers],
),
- )
+ ]
# Set up config file.
_tpl(repository_ctx, "build_defs.bzl", {"%{if_tensorrt}": "if_true"})
@@ -217,8 +90,7 @@
# Set up BUILD file.
_tpl(repository_ctx, "BUILD", {
"%{copy_rules}": "\n".join(copy_rules),
- "%{tensorrt_headers}": '":tensorrt_include"',
- "%{tensorrt_libs}": str(trt_lib_outs),
+ "%{tensorrt_libs}": str(libraries),
})
tensorrt_configure = repository_rule(
@@ -226,6 +98,8 @@
environ = [
_TENSORRT_INSTALL_PATH,
_TF_TENSORRT_VERSION,
+ _TF_TENSORRT_CONFIG_REPO,
+ "TF_CUDA_PATHS",
],
)
"""Detects and configures the local CUDA toolchain.
diff --git a/third_party/toolchains/preconfig/centos6/cuda10.0-cudnn7/cuda/cuda/cuda_config.h b/third_party/toolchains/preconfig/centos6/cuda10.0-cudnn7/cuda/cuda/cuda_config.h
index 783d678..72a7cf7 100644
--- a/third_party/toolchains/preconfig/centos6/cuda10.0-cudnn7/cuda/cuda/cuda_config.h
+++ b/third_party/toolchains/preconfig/centos6/cuda10.0-cudnn7/cuda/cuda/cuda_config.h
@@ -19,6 +19,7 @@
#define TF_CUDA_CAPABILITIES CudaVersion("3.0"), CudaVersion("6.0")
#define TF_CUDA_VERSION "10.0"
+#define TF_CUDA_LIB_VERSION "10.0"
#define TF_CUDNN_VERSION "7"
#define TF_CUDA_TOOLKIT_PATH "/usr/local/cuda-10.0"
diff --git a/third_party/toolchains/preconfig/centos7/cuda10.0-cudnn7/cuda/cuda/cuda_config.h b/third_party/toolchains/preconfig/centos7/cuda10.0-cudnn7/cuda/cuda/cuda_config.h
index 783d678..72a7cf7 100644
--- a/third_party/toolchains/preconfig/centos7/cuda10.0-cudnn7/cuda/cuda/cuda_config.h
+++ b/third_party/toolchains/preconfig/centos7/cuda10.0-cudnn7/cuda/cuda/cuda_config.h
@@ -19,6 +19,7 @@
#define TF_CUDA_CAPABILITIES CudaVersion("3.0"), CudaVersion("6.0")
#define TF_CUDA_VERSION "10.0"
+#define TF_CUDA_LIB_VERSION "10.0"
#define TF_CUDNN_VERSION "7"
#define TF_CUDA_TOOLKIT_PATH "/usr/local/cuda-10.0"
diff --git a/third_party/toolchains/preconfig/ubuntu14.04/cuda10.0-cudnn7/cuda/cuda/cuda_config.h b/third_party/toolchains/preconfig/ubuntu14.04/cuda10.0-cudnn7/cuda/cuda/cuda_config.h
index 783d678..72a7cf7 100644
--- a/third_party/toolchains/preconfig/ubuntu14.04/cuda10.0-cudnn7/cuda/cuda/cuda_config.h
+++ b/third_party/toolchains/preconfig/ubuntu14.04/cuda10.0-cudnn7/cuda/cuda/cuda_config.h
@@ -19,6 +19,7 @@
#define TF_CUDA_CAPABILITIES CudaVersion("3.0"), CudaVersion("6.0")
#define TF_CUDA_VERSION "10.0"
+#define TF_CUDA_LIB_VERSION "10.0"
#define TF_CUDNN_VERSION "7"
#define TF_CUDA_TOOLKIT_PATH "/usr/local/cuda-10.0"
diff --git a/third_party/toolchains/preconfig/ubuntu14.04/cuda9.0-cudnn7/cuda/cuda/cuda_config.h b/third_party/toolchains/preconfig/ubuntu14.04/cuda9.0-cudnn7/cuda/cuda/cuda_config.h
index 5d0d301..09776bd 100755
--- a/third_party/toolchains/preconfig/ubuntu14.04/cuda9.0-cudnn7/cuda/cuda/cuda_config.h
+++ b/third_party/toolchains/preconfig/ubuntu14.04/cuda9.0-cudnn7/cuda/cuda/cuda_config.h
@@ -19,6 +19,7 @@
#define TF_CUDA_CAPABILITIES CudaVersion("3.0")
#define TF_CUDA_VERSION "9.0"
+#define TF_CUDA_LIB_VERSION "9.0"
#define TF_CUDNN_VERSION "7"
#define TF_CUDA_TOOLKIT_PATH "/usr/local/cuda-9.0"