diff options
author | 2017-04-10 16:17:44 -0800 | |
---|---|---|
committer | 2017-04-10 17:28:17 -0700 | |
commit | 72c023d3967a3218cd3d830ce6e57f7c4d87a18c (patch) | |
tree | 6401e38c3e9cc64f472f31cc84c6b7b80712e718 /third_party/gpus/cuda_configure.bzl | |
parent | 2f5fde8dd977681110714a0564445a406506282c (diff) |
Add support for solvers based on the Nvidia cuSolver library.
Implement a GPU version of tf.cholesky as a first example.
Change: 152756373
Diffstat (limited to 'third_party/gpus/cuda_configure.bzl')
-rw-r--r-- | third_party/gpus/cuda_configure.bzl | 7 |
1 files changed, 7 insertions, 0 deletions
diff --git a/third_party/gpus/cuda_configure.bzl b/third_party/gpus/cuda_configure.bzl index 05ff584be0..e06092ab4e 100644 --- a/third_party/gpus/cuda_configure.bzl +++ b/third_party/gpus/cuda_configure.bzl @@ -541,6 +541,9 @@ def _find_libs(repository_ctx, cuda_config): "cublas": _find_cuda_lib( "cublas", repository_ctx, cpu_value, cuda_config.cuda_toolkit_path, cuda_config.cuda_version), + "cusolver": _find_cuda_lib( + "cusolver", 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), @@ -695,6 +698,7 @@ def _create_dummy_repository(repository_ctx): "%{cudart_static_linkopt}": _cudart_static_linkopt(cpu_value), "%{cudart_lib}": _lib_name("cudart", cpu_value), "%{cublas_lib}": _lib_name("cublas", cpu_value), + "%{cusolver_lib}": _lib_name("cusolver", cpu_value), "%{cudnn_lib}": _lib_name("cudnn", cpu_value), "%{cufft_lib}": _lib_name("cufft", cpu_value), "%{curand_lib}": _lib_name("curand", cpu_value), @@ -708,6 +712,7 @@ def _create_dummy_repository(repository_ctx): "%{cudart_static_linkopt}": _cudart_static_linkopt(cpu_value), "%{cudart_lib}": _lib_name("cudart", cpu_value), "%{cublas_lib}": _lib_name("cublas", cpu_value), + "%{cusolver_lib}": _lib_name("cusolver", cpu_value), "%{cudnn_lib}": _lib_name("cudnn", cpu_value), "%{cufft_lib}": _lib_name("cufft", cpu_value), "%{curand_lib}": _lib_name("curand", cpu_value), @@ -730,6 +735,7 @@ def _create_dummy_repository(repository_ctx): 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("cusolver", 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)) @@ -822,6 +828,7 @@ def _create_cuda_repository(repository_ctx): cuda_config.cpu_value), "%{cudart_lib}": cuda_libs["cudart"].file_name, "%{cublas_lib}": cuda_libs["cublas"].file_name, + "%{cusolver_lib}": cuda_libs["cusolver"].file_name, "%{cudnn_lib}": cuda_libs["cudnn"].file_name, "%{cufft_lib}": cuda_libs["cufft"].file_name, "%{curand_lib}": cuda_libs["curand"].file_name, |