diff options
author | David Z. Chen <dzc@google.com> | 2016-10-13 14:11:24 -0700 |
---|---|---|
committer | David Z. Chen <dzc@google.com> | 2016-11-30 01:18:14 -0800 |
commit | dfae1931365dae3354a8980829dfabcd59761232 (patch) | |
tree | 6ccd465eacc5b5821cd41de7b3d49ac8709bc5b3 | |
parent | 5657d0dee8d87f4594b3e5902ed3e3ca8d6dfc0a (diff) |
Autodetect cuda and cudnn versions and support more platforms.
This change implements cuda and cudnn version autodetection via running
nvcc --version and reading cudnn.h. As a result, TF_CUDA_VERSION and
TF_CUDNN_VERSION no longer have to be set. If TF_CUDA_VERSION and
TF_CUDNN_VERSION are set but do not match the detected versions,
cuda_configure will fail.
After the cuda and cudnn versions are detected, cuda_configure then
attempts to find the needed cuda and cudnn libraries in a set of known
installation directories.
-rw-r--r-- | tensorflow/core/platform/default/build_config/BUILD | 3 | ||||
-rw-r--r-- | third_party/gpus/cuda/BUILD.tpl | 52 | ||||
-rw-r--r-- | third_party/gpus/cuda/platform.bzl.tpl | 59 | ||||
-rw-r--r-- | third_party/gpus/cuda_configure.bzl | 416 |
4 files changed, 317 insertions, 213 deletions
diff --git a/tensorflow/core/platform/default/build_config/BUILD b/tensorflow/core/platform/default/build_config/BUILD index ca30603f51..bb52e75df3 100644 --- a/tensorflow/core/platform/default/build_config/BUILD +++ b/tensorflow/core/platform/default/build_config/BUILD @@ -9,7 +9,6 @@ exports_files(["LICENSE"]) load("//tensorflow:tensorflow.bzl", "tf_copts") load("//tensorflow:tensorflow.bzl", "tf_cuda_library") -load("@local_config_cuda//cuda:platform.bzl", "cuda_library_path") load("@local_config_sycl//sycl:platform.bzl", "sycl_library_path") cc_library( @@ -127,7 +126,7 @@ filegroup( cc_library( name = "cuda", data = [ - "@local_config_cuda//cuda:{}".format(cuda_library_path("cudart")), + "@local_config_cuda//cuda:cudart", ], linkopts = select({ "@local_config_cuda//cuda:darwin": [ diff --git a/third_party/gpus/cuda/BUILD.tpl b/third_party/gpus/cuda/BUILD.tpl index cb985c2251..82a96e3515 100644 --- a/third_party/gpus/cuda/BUILD.tpl +++ b/third_party/gpus/cuda/BUILD.tpl @@ -1,9 +1,5 @@ licenses(["restricted"]) # MPL2, portions GPL v3, LGPL v3, BSD-like -load("@local_config_cuda//cuda:platform.bzl", "cuda_library_path") -load("@local_config_cuda//cuda:platform.bzl", "cuda_static_library_path") -load("@local_config_cuda//cuda:platform.bzl", "cudnn_library_path") -load("@local_config_cuda//cuda:platform.bzl", "cupti_library_path") load("@local_config_cuda//cuda:platform.bzl", "readlink_command") package(default_visibility = ["//visibility:public"]) @@ -51,9 +47,7 @@ cc_library( cc_library( name = "cudart_static", - srcs = [ - cuda_static_library_path("cudart"), - ], + srcs = ["lib/%{cudart_static_lib}"], includes = ["include/"], linkopts = [ "-ldl", @@ -65,12 +59,8 @@ cc_library( cc_library( name = "cudart", - srcs = [ - cuda_library_path("cudart"), - ], - data = [ - cuda_library_path("cudart"), - ], + srcs = ["lib/%{cudart_lib}"], + data = ["lib/%{cudart_lib}"], includes = ["include/"], linkstatic = 1, visibility = ["//visibility:public"], @@ -78,12 +68,8 @@ cc_library( cc_library( name = "cublas", - srcs = [ - cuda_library_path("cublas"), - ], - data = [ - cuda_library_path("cublas"), - ], + srcs = ["lib/%{cublas_lib}"], + data = ["lib/%{cublas_lib}"], includes = ["include/"], linkstatic = 1, visibility = ["//visibility:public"], @@ -91,12 +77,8 @@ cc_library( cc_library( name = "cudnn", - srcs = [ - cudnn_library_path(), - ], - data = [ - cudnn_library_path(), - ], + srcs = ["lib/%{cudnn_lib}"], + data = ["lib/%{cudnn_lib}"], includes = ["include/"], linkstatic = 1, visibility = ["//visibility:public"], @@ -104,12 +86,8 @@ cc_library( cc_library( name = "cufft", - srcs = [ - cuda_library_path("cufft"), - ], - data = [ - cuda_library_path("cufft"), - ], + srcs = ["lib/%{cufft_lib}"], + data = ["lib/%{cufft_lib}"], includes = ["include/"], linkstatic = 1, visibility = ["//visibility:public"], @@ -117,12 +95,8 @@ cc_library( cc_library( name = "curand", - srcs = [ - cuda_library_path("curand"), - ], - data = [ - cuda_library_path("curand"), - ], + srcs = ["lib/%{curand_lib}"], + data = ["lib/%{curand_lib}"], includes = ["include/"], linkstatic = 1, visibility = ["//visibility:public"], @@ -155,8 +129,6 @@ cc_library( cc_library( name = "cupti_dsos", - data = [ - cupti_library_path(), - ], + data = ["lib/%{cupti_lib}"], visibility = ["//visibility:public"], ) diff --git a/third_party/gpus/cuda/platform.bzl.tpl b/third_party/gpus/cuda/platform.bzl.tpl index 53c02d6ea9..01ef24b94e 100644 --- a/third_party/gpus/cuda/platform.bzl.tpl +++ b/third_party/gpus/cuda/platform.bzl.tpl @@ -8,65 +8,6 @@ def cuda_sdk_version(): def cudnn_sdk_version(): return CUDNN_VERSION -def cuda_library_path(name, version = cuda_sdk_version()): - if PLATFORM == "Darwin": - if not version: - return "lib/lib{}.dylib".format(name) - else: - return "lib/lib{}.{}.dylib".format(name, version) - elif PLATFORM == "Windows": - if not version: - return "lib/x64/{}.lib".format(name) - else: - return "lib/x64/{}{}.lib".format(name, version) - else: - if not version: - return "lib64/lib{}.so".format(name) - else: - return "lib64/lib{}.so.{}".format(name, version) - -def cuda_static_library_path(name): - if PLATFORM == "Darwin": - return "lib/lib{}_static.a".format(name) - elif PLATFORM == "Windows": - return "lib/x64/{}_static.lib".format(name) - else: - return "lib64/lib{}_static.a".format(name) - -def cudnn_library_path(version = cudnn_sdk_version()): - if PLATFORM == "Darwin": - if not version: - return "lib/libcudnn.dylib" - else: - return "lib/libcudnn.{}.dylib".format(version) - elif PLATFORM == "Windows": - if not version: - return "lib/x64/cudnn.lib" - else: - return "lib/x64/cudnn{}.lib".format(version) - else: - if not version: - return "lib64/libcudnn.so" - else: - return "lib64/libcudnn.so.{}".format(version) - -def cupti_library_path(version = cuda_sdk_version()): - if PLATFORM == "Darwin": - if not version: - return "extras/CUPTI/lib/libcupti.dylib" - else: - return "extras/CUPTI/lib/libcupti.{}.dylib".format(version) - elif PLATFORM == "Windows": - if not version: - return "extras/CUPTI/libx64/cupti.lib" - else: - return "extras/CUPTI/libx64/cupti{}.lib".format(version) - else: - if not version: - return "extras/CUPTI/lib64/libcupti.so" - else: - return "extras/CUPTI/lib64/libcupti.so.{}".format(version) - def readlink_command(): if PLATFORM == "Darwin": return "greadlink" diff --git a/third_party/gpus/cuda_configure.bzl b/third_party/gpus/cuda_configure.bzl index eb24c0f8c7..955951a9b6 100644 --- a/third_party/gpus/cuda_configure.bzl +++ b/third_party/gpus/cuda_configure.bzl @@ -120,12 +120,11 @@ def _enable_cuda(repository_ctx): return False -def _cuda_toolkit_path(repository_ctx, cuda_version): +def _cuda_toolkit_path(repository_ctx): """Finds the cuda toolkit directory. Args: repository_ctx: The repository context. - cuda_version: The cuda toolkit version. Returns: A speculative real path of the cuda toolkit install directory. @@ -135,16 +134,6 @@ def _cuda_toolkit_path(repository_ctx, cuda_version): 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.") - - if cuda_version: - # Handle typical configuration where the real path is - # <basedir>/cuda-<version> and the provided path is <basedir>/cuda. - version_suffixed = "%s-%s" % (cuda_toolkit_path, cuda_version) - if repository_ctx.path(version_suffixed).exists: - cuda_toolkit_path = version_suffixed - # Returns the non-versioned path if cuda version is not provided or if the - # installation does not use a cuda- directory, such as on ArchLinux where - # CUDA installs directly to /opt/cuda. return str(repository_ctx.path(cuda_toolkit_path).realpath) @@ -158,20 +147,98 @@ def _cudnn_install_basedir(repository_ctx): return cudnn_install_path -def _cuda_version(repository_ctx): - """Detects the cuda version.""" +_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]) == 0: + auto_configure_fail( + "Could not parse CUDA version from nvcc --version. Got: %s" % + result.stdout) + version = parts[0].strip() + + # 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: - return repository_ctx.os.environ[_TF_CUDA_VERSION].strip() - else: - return "" + environ_version = repository_ctx.os.environ[_TF_CUDA_VERSION].strip() + if environ_version and version != environ_version: + auto_configure_fail( + "CUDA version detected from nvcc (%s) does not match " + + "TF_CUDA_VERSION (%s)" % (version, environ_version)) + return version + + +_DEFINE_CUDNN_MAJOR = "#define CUDNN_MAJOR" + +def _cudnn_version(repository_ctx, cudnn_install_basedir): + """Detects the version of cuDNN installed on the system. -def _cudnn_version(repository_ctx): - """Detects the cudnn version.""" + 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. + """ + # Find cudnn.h and grep for the line defining CUDNN_MAJOR. + cudnn_h_path = repository_ctx.path("%s/include/cudnn.h" % + cudnn_install_basedir) + if not cudnn_h_path.exists: + auto_configure_fail("Cannot find cudnn.h at %s" % str(cudnn_h_path)) + result = repository_ctx.execute([ + "grep", "-E", _DEFINE_CUDNN_MAJOR, str(cudnn_h_path)]) + if result.stderr: + auto_configure_fail("Error reading %s: %s" % + (result.stderr, str(cudnn_h_path))) + + # Parse the cuDNN major version from the line defining CUDNN_MAJOR + lines = result.stdout.splitlines() + if len(lines) == 0 or lines[0].find(_DEFINE_CUDNN_MAJOR) == -1: + auto_configure_fail("Cannot find line containing '%s' in %s" % + (_DEFINE_CUDNN_MAJOR, str(cudnn_h_path))) + version = lines[0].replace(_DEFINE_CUDNN_MAJOR, "").strip() + + # 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: - return repository_ctx.os.environ[_TF_CUDNN_VERSION].strip() - else: - return "" + environ_version = repository_ctx.os.environ[_TF_CUDNN_VERSION].strip() + if environ_version and version != environ_version: + auto_configure_fail( + "cuDNN version detected from %s (%s) does not match " + + "TF_CUDNN_VERSION (%s)" % (str(cudnn_h_path), version, environ_version)) + return version def _compute_capabilities(repository_ctx): @@ -191,6 +258,14 @@ def _compute_capabilities(repository_ctx): def _cpu_value(repository_ctx): + """Returns the name of the host operating system. + + Args: + repository_ctx: The repository context. + + Returns: + A string containing the name of the host operating system. + """ os_name = repository_ctx.os.name.lower() if os_name.startswith("mac os"): return "Darwin" @@ -200,75 +275,158 @@ def _cpu_value(repository_ctx): return result.stdout.strip() -def _cuda_symlink_files(cpu_value, cuda_version, cudnn_version): - """Returns a struct containing platform-specific paths. +def _lib_name(lib, cpu_value, version="", static=False): + """Constructs the platform-specific name of a library. Args: - cpu_value: The string representing the host OS. - cuda_version: The cuda version as returned by _cuda_version - cudnn_version: The cudnn version as returned by _cudnn_version + lib: The name of the library, such as "cudart" + cpu_value: The name of the host operating system. + version: The version of the library. + static: True the library is static or False if it is a shared object. + + Returns: + The platform-specific name of the library. """ - cuda_ext = ".%s" % cuda_version if cuda_version else "" - cudnn_ext = ".%s" % cudnn_version if cudnn_version else "" if cpu_value == "Linux": - return struct( - cuda_lib_path = "lib64", - cuda_rt_lib = "lib64/libcudart.so%s" % cuda_ext, - cuda_rt_lib_static = "lib64/libcudart_static.a", - cuda_blas_lib = "lib64/libcublas.so%s" % cuda_ext, - cuda_dnn_lib = "lib64/libcudnn.so%s" % cudnn_ext, - cuda_dnn_lib_alt = "libcudnn.so%s" % cudnn_ext, - cuda_rand_lib = "lib64/libcurand.so%s" % cuda_ext, - cuda_fft_lib = "lib64/libcufft.so%s" % cuda_ext, - cuda_cupti_lib = "extras/CUPTI/lib64/libcupti.so%s" % cuda_ext) - elif cpu_value == "Darwin": - return struct( - cuda_lib_path = "lib", - cuda_rt_lib = "lib/libcudart%s.dylib" % cuda_ext, - cuda_rt_lib_static = "lib/libcudart_static.a", - cuda_blas_lib = "lib/libcublas%s.dylib" % cuda_ext, - cuda_dnn_lib = "lib/libcudnn%s.dylib" % cudnn_ext, - cuda_dnn_lib_alt = "libcudnn%s.dylib" % cudnn_ext, - cuda_rand_lib = "lib/libcurand%s.dylib" % cuda_ext, - cuda_fft_lib = "lib/libcufft%s.dylib" % cuda_ext, - cuda_cupti_lib = "extras/CUPTI/lib/libcupti%s.dylib" % cuda_ext) + if static: + return "lib%s.a" % lib + else: + if version: + version = ".%s" % version + return "lib%s.so%s" % (lib, version) elif cpu_value == "Windows": - return struct( - cuda_lib_path = "lib", - cuda_rt_lib = "lib/x64/cudart%s.lib" % cuda_ext, - cuda_rt_lib_static = "lib/x64/cudart_static.lib", - cuda_blas_lib = "lib/x64/cublas%s.lib" % cuda_ext, - cuda_dnn_lib = "lib/x64/cudnn%s.lib" % cudnn_ext, - cuda_dnn_lib_alt = "cudnn%s.lib" % cudnn_ext, - cuda_rand_lib = "lib/x64/curand%s.lib" % cuda_ext, - cuda_fft_lib = "lib/x64/cufft%s.lib" % cuda_ext, - cuda_cupti_lib = "extras/CUPTI/libx64/cupti%s.lib" % cuda_ext) + return "%s.lib" % lib + elif cpu_value == "Darwin": + return "lib%s%s.dylib" % (lib, version) else: - auto_configure_fail("Not supported CPU value %s" % cpu_value) + auto_configure_fail("Invalid cpu_value: %s" % cpu_value) -def _check_lib(repository_ctx, cuda_toolkit_path, cuda_lib): - """Checks if cuda_lib exists under cuda_toolkit_path or fail if it doesn't. +def _find_cuda_lib(lib, repository_ctx, cpu_value, basedir, version="", + static=False): + """Finds the given CUDA or cuDNN library on the system. Args: + lib: The name of the library, such as "cudart" repository_ctx: The repository context. - cuda_toolkit_path: The cuda toolkit directory containing the cuda libraries. - cuda_lib: The library to look for under cuda_toolkit_path. + cpu_value: The name of the host operating system. + basedir: The install directory of CUDA or cuDNN. + version: The version of the library. + static: True if static library, False if shared object. + + Returns: + Returns a struct with the following fields: + file_name: The basename of the library found on the system. + path: The full path to the library. """ - lib_path = cuda_toolkit_path + "/" + cuda_lib - if not repository_ctx.path(lib_path).exists: - auto_configure_fail("Cannot find %s" % lib_path) + file_name = _lib_name(lib, cpu_value, version, static) + if cpu_value == "Linux": + path = repository_ctx.path("%s/lib64/%s" % (basedir, file_name)) + if path.exists: + return struct(file_name=file_name, path=str(path.realpath)) + path = repository_ctx.path( + "%s/lib/x86_64-linux-gnu/%s" % (basedir, file_name)) + if path.exists: + return struct(file_name=file_name, path=str(path.realpath)) + elif cpu_value == "Windows": + path = repository_ctx.path("%s/lib/x64/%s" % (basedir, file_name)) + if path.exists: + return struct(file_name=file_name, path=str(path.realpath)) -def _check_dir(repository_ctx, directory): - """Checks whether the directory exists and fail if it does not. + path = repository_ctx.path("%s/lib/%s" % (basedir, file_name)) + if path.exists: + return struct(file_name=file_name, path=str(path.realpath)) + path = repository_ctx.path("%s/%s" % (basedir, file_name)) + if path.exists: + return struct(file_name=file_name, path=str(path.realpath)) + + auto_configure_fail("Cannot find cuda library %s" % file_name) + + +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. - directory: The directory to check the existence of. + cuda_config: The cuda configuration as returned by _get_cuda_config. + + Returns: + Returns a struct with the following fields: + file_name: The basename of the library found on the system. + path: The full path to the library. """ - if not repository_ctx.path(directory).exists: - auto_configure_fail("Cannot find dir: %s" % directory) + file_name = _lib_name("cupti", cuda_config.cpu_value, + cuda_config.cuda_version) + if cuda_config.cpu_value == "Linux": + path = repository_ctx.path( + "%s/extras/CUPTI/lib64/%s" % (cuda_config.cuda_toolkit_path, file_name)) + if path.exists: + return struct(file_name=file_name, path=str(path.realpath)) + + path = repository_ctx.path( + "%s/lib/x86_64-linux-gnu/%s" % (cuda_config.cuda_toolkit_path, + file_name)) + if path.exists: + return struct(file_name=file_name, path=str(path.realpath)) + + elif cuda_config.cpu_value == "Windows": + path = repository_ctx.path( + "%s/extras/CUPTI/libx64/%s" % + (cuda_config.cuda_toolkit_path, file_name)) + if path.exists: + return struct(file_name=file_name, path=str(path.realpath)) + + path = repository_ctx.path( + "%s/extras/CUPTI/lib/%s" % (cuda_config.cuda_toolkit_path, file_name)) + if path.exists: + return struct(file_name=file_name, path=str(path.realpath)) + + path = repository_ctx.path( + "%s/lib/%s" % (cuda_config.cuda_toolkit_path, file_name)) + if path.exists: + return struct(file_name=file_name, path=str(path.realpath)) + + auto_configure_fail("Cannot find cupti library %s" % file_name) + +def _find_libs(repository_ctx, cuda_config): + """Returns the CUDA and cuDNN libraries on the system. + + Args: + repository_ctx: The repository context. + cuda_config: The CUDA config as returned by _get_cuda_config + + Returns: + Map of library names to structs of filename and path as returned by + _find_cuda_lib and _find_cupti_lib. + """ + cudnn_version = cuda_config.cudnn_version + cudnn_ext = ".%s" % cudnn_version if cudnn_version else "" + cpu_value = cuda_config.cpu_value + return { + "cudart": _find_cuda_lib( + "cudart", repository_ctx, cpu_value, cuda_config.cuda_toolkit_path, + cuda_config.cuda_version), + "cudart_static": _find_cuda_lib( + "cudart_static", repository_ctx, cpu_value, + cuda_config.cuda_toolkit_path, cuda_config.cuda_version, static=True), + "cublas": _find_cuda_lib( + "cublas", repository_ctx, cpu_value, cuda_config.cuda_toolkit_path, + cuda_config.cuda_version), + "curand": _find_cuda_lib( + "curand", repository_ctx, cpu_value, cuda_config.cuda_toolkit_path, + cuda_config.cuda_version), + "cufft": _find_cuda_lib( + "cufft", repository_ctx, cpu_value, cuda_config.cuda_toolkit_path, + cuda_config.cuda_version), + "cudnn": _find_cuda_lib( + "cudnn", repository_ctx, cpu_value, cuda_config.cudnn_install_basedir, + cuda_config.cudnn_version), + "cupti": _find_cupti_lib(repository_ctx, cuda_config), + } def _find_cudnn_header_dir(repository_ctx, cudnn_install_basedir): @@ -319,6 +477,34 @@ def _cudart_static_linkopt(cpu_value): """Returns additional platform-specific linkopts for cudart.""" return "" if cpu_value == "Darwin" else "\"-lrt\"," +def _get_cuda_config(repository_ctx): + """Detects and returns information about the CUDA installation on the system. + + Args: + repository_ctx: The repository context. + + Returns: + A struct containing the following fields: + cuda_toolkit_path: The CUDA toolkit installation directory. + cudnn_install_basedir: The cuDNN installation directory. + cuda_version: The version of CUDA on the system. + cudnn_version: The version of cuDNN on the system. + compute_capabilities: A list of the system's CUDA compute capabilities. + cpu_value: The name of the host operating system. + """ + cpu_value = _cpu_value(repository_ctx) + cuda_toolkit_path = _cuda_toolkit_path(repository_ctx) + cuda_version = _cuda_version(repository_ctx, cuda_toolkit_path, cpu_value) + cudnn_install_basedir = _cudnn_install_basedir(repository_ctx) + cudnn_version = _cudnn_version(repository_ctx, cudnn_install_basedir) + return struct( + cuda_toolkit_path = cuda_toolkit_path, + cudnn_install_basedir = cudnn_install_basedir, + cuda_version = cuda_version, + cudnn_version = cudnn_version, + compute_capabilities = _compute_capabilities(repository_ctx), + cpu_value = cpu_value) + def _tpl(repository_ctx, tpl, substitutions={}, out=None): if not out: @@ -365,8 +551,6 @@ error_gpu_disabled() def _create_dummy_repository(repository_ctx): cpu_value = _cpu_value(repository_ctx) - symlink_files = _cuda_symlink_files(cpu_value, _DEFAULT_CUDA_VERSION, - _DEFAULT_CUDNN_VERSION) # Set up BUILD file for cuda/. _file(repository_ctx, "cuda:build_defs.bzl") @@ -374,6 +558,18 @@ def _create_dummy_repository(repository_ctx): { "%{cudart_static_linkopt}": _cudart_static_linkopt(cpu_value), }) + _tpl(repository_ctx, "cuda:BUILD", + { + "%{cudart_static_lib}": _lib_name("cudart_static", cpu_value, + static=True), + "%{cudart_static_linkopt}": _cudart_static_linkopt(cpu_value), + "%{cudart_lib}": _lib_name("cudart", cpu_value), + "%{cublas_lib}": _lib_name("cublas", cpu_value), + "%{cudnn_lib}": _lib_name("cudnn", cpu_value), + "%{cufft_lib}": _lib_name("cufft", cpu_value), + "%{curand_lib}": _lib_name("curand", cpu_value), + "%{cupti_lib}": _lib_name("cupti", cpu_value), + }) _tpl(repository_ctx, "cuda:platform.bzl", { "%{cuda_version}": _DEFAULT_CUDA_VERSION, @@ -387,13 +583,13 @@ def _create_dummy_repository(repository_ctx): repository_ctx.file("cuda/include/cublas.h", "") repository_ctx.file("cuda/include/cudnn.h", "") repository_ctx.file("cuda/extras/CUPTI/include/cupti.h", "") - repository_ctx.file("cuda/%s" % symlink_files.cuda_rt_lib, "") - repository_ctx.file("cuda/%s" % symlink_files.cuda_rt_lib_static, "") - repository_ctx.file("cuda/%s" % symlink_files.cuda_blas_lib, "") - repository_ctx.file("cuda/%s" % symlink_files.cuda_dnn_lib, "") - repository_ctx.file("cuda/%s" % symlink_files.cuda_rand_lib, "") - repository_ctx.file("cuda/%s" % symlink_files.cuda_fft_lib, "") - repository_ctx.file("cuda/%s" % symlink_files.cuda_cupti_lib, "") + repository_ctx.file("cuda/lib/%s" % _lib_name("cudart", cpu_value)) + repository_ctx.file("cuda/lib/%s" % _lib_name("cudart_static", cpu_value)) + repository_ctx.file("cuda/lib/%s" % _lib_name("cublas", cpu_value)) + repository_ctx.file("cuda/lib/%s" % _lib_name("cudnn", cpu_value)) + repository_ctx.file("cuda/lib/%s" % _lib_name("curand", cpu_value)) + repository_ctx.file("cuda/lib/%s" % _lib_name("cufft", cpu_value)) + repository_ctx.file("cuda/lib/%s" % _lib_name("cupti", cpu_value)) # Set up cuda_config.h, which is used by # tensorflow/stream_executor/dso_loader.cc. @@ -428,56 +624,51 @@ def _symlink_dir(repository_ctx, src_dir, dest_dir): def _create_cuda_repository(repository_ctx): """Creates the repository containing files set up to build with CUDA.""" - cuda_version = _cuda_version(repository_ctx) - cuda_toolkit_path = _cuda_toolkit_path(repository_ctx, cuda_version) - cudnn_install_basedir = _cudnn_install_basedir(repository_ctx) - cudnn_version = _cudnn_version(repository_ctx) - compute_capabilities = _compute_capabilities(repository_ctx) - - cpu_value = _cpu_value(repository_ctx) - symlink_files = _cuda_symlink_files(cpu_value, cuda_version, cudnn_version) - _check_lib(repository_ctx, cuda_toolkit_path, symlink_files.cuda_rt_lib) - _check_lib(repository_ctx, cuda_toolkit_path, symlink_files.cuda_cupti_lib) - _check_dir(repository_ctx, cudnn_install_basedir) + cuda_config = _get_cuda_config(repository_ctx) cudnn_header_dir = _find_cudnn_header_dir(repository_ctx, - cudnn_install_basedir) - cudnn_lib_path = _find_cudnn_lib_path(repository_ctx, cudnn_install_basedir, - symlink_files) + cuda_config.cudnn_install_basedir) # Set up symbolic links for the cuda toolkit. We link at the individual file # level not at the directory level. This is because the external library may # have a different file layout from our desired structure. + cuda_toolkit_path = cuda_config.cuda_toolkit_path _symlink_dir(repository_ctx, cuda_toolkit_path + "/include", "cuda/include") - _symlink_dir(repository_ctx, - cuda_toolkit_path + "/" + symlink_files.cuda_lib_path, - "cuda/" + symlink_files.cuda_lib_path) _symlink_dir(repository_ctx, cuda_toolkit_path + "/bin", "cuda/bin") _symlink_dir(repository_ctx, cuda_toolkit_path + "/nvvm", "cuda/nvvm") _symlink_dir(repository_ctx, cuda_toolkit_path + "/extras/CUPTI/include", "cuda/extras/CUPTI/include") - repository_ctx.symlink(cuda_toolkit_path + "/" + symlink_files.cuda_cupti_lib, - "cuda/" + symlink_files.cuda_cupti_lib) + + cuda_libs = _find_libs(repository_ctx, cuda_config) + for lib in cuda_libs.values(): + repository_ctx.symlink(lib.path, "cuda/lib/" + lib.file_name) # Set up the symbolic links for cudnn if cudnn was was not installed to # CUDA_TOOLKIT_PATH. if not repository_ctx.path("cuda/include/cudnn.h").exists: repository_ctx.symlink(cudnn_header_dir + "/cudnn.h", "cuda/include/cudnn.h") - if not repository_ctx.path("cuda/" + symlink_files.cuda_dnn_lib).exists: - repository_ctx.symlink(cudnn_lib_path, "cuda/" + symlink_files.cuda_dnn_lib) # Set up BUILD file for cuda/ _file(repository_ctx, "cuda:build_defs.bzl") _tpl(repository_ctx, "cuda:BUILD", { - "%{cudart_static_linkopt}": _cudart_static_linkopt(cpu_value), + "%{cudart_static_lib}": cuda_libs["cudart_static"].file_name, + "%{cudart_static_linkopt}": _cudart_static_linkopt( + cuda_config.cpu_value), + "%{cudart_lib}": cuda_libs["cudart"].file_name, + "%{cublas_lib}": cuda_libs["cublas"].file_name, + "%{cudnn_lib}": cuda_libs["cudnn"].file_name, + "%{cufft_lib}": cuda_libs["cufft"].file_name, + "%{curand_lib}": cuda_libs["curand"].file_name, + "%{cupti_lib}": cuda_libs["cupti"].file_name, }) + _tpl(repository_ctx, "cuda:platform.bzl", { - "%{cuda_version}": cuda_version, - "%{cudnn_version}": cudnn_version, - "%{platform}": cpu_value, + "%{cuda_version}": cuda_config.cuda_version, + "%{cudnn_version}": cuda_config.cudnn_version, + "%{platform}": cuda_config.cpu_value, }) # Set up crosstool/ @@ -486,7 +677,7 @@ def _create_cuda_repository(repository_ctx): gcc_host_compiler_includes = _gcc_host_compiler_includes(repository_ctx, cc) _tpl(repository_ctx, "crosstool:CROSSTOOL", { - "%{cuda_include_path}": cuda_toolkit_path + '/include', + "%{cuda_include_path}": cuda_config.cuda_toolkit_path + '/include', "%{gcc_host_compiler_includes}": gcc_host_compiler_includes, }) _tpl(repository_ctx, @@ -495,17 +686,18 @@ def _create_cuda_repository(repository_ctx): "%{cpu_compiler}": str(cc), "%{gcc_host_compiler_path}": str(cc), "%{cuda_compute_capabilities}": ", ".join( - ["\"%s\"" % c for c in compute_capabilities]), + ["\"%s\"" % c for c in cuda_config.compute_capabilities]), }) # Set up cuda_config.h, which is used by # tensorflow/stream_executor/dso_loader.cc. _tpl(repository_ctx, "cuda:cuda_config.h", { - "%{cuda_version}": cuda_version, - "%{cudnn_version}": cudnn_version, + "%{cuda_version}": cuda_config.cuda_version, + "%{cudnn_version}": cuda_config.cudnn_version, "%{cuda_compute_capabilities}": ",".join( - ["CudaVersion(\"%s\")" % c for c in compute_capabilities]), + ["CudaVersion(\"%s\")" % c + for c in cuda_config.compute_capabilities]), }) |