aboutsummaryrefslogtreecommitdiffhomepage
path: root/third_party/gpus/cuda_configure.bzl
diff options
context:
space:
mode:
authorGravatar A. Unique TensorFlower <gardener@tensorflow.org>2017-04-10 16:17:44 -0800
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2017-04-10 17:28:17 -0700
commit72c023d3967a3218cd3d830ce6e57f7c4d87a18c (patch)
tree6401e38c3e9cc64f472f31cc84c6b7b80712e718 /third_party/gpus/cuda_configure.bzl
parent2f5fde8dd977681110714a0564445a406506282c (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.bzl7
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,