tensorflow: add OpenGL path to find libcudart

3 files changed, 31 insertions, 3 deletions

diff --git a/pkgs/development/libraries/science/math/tensorflow/bin.nix b/pkgs/development/libraries/science/math/tensorflow/bin.nix
index 3db7a6f1d1e..f9f300883a1 100644
--- a/pkgs/development/libraries/science/math/tensorflow/bin.nix
+++ b/pkgs/development/libraries/science/math/tensorflow/bin.nix
@@ -1,6 +1,6 @@
 { stdenv
 , fetchurl
-, patchelf
+, addOpenGLRunpath
 , cudaSupport ? false, symlinkJoin, cudatoolkit, cudnn, nvidia_x11
 }:
 
@@ -35,6 +35,9 @@ let
     else ''
       patchelf --set-rpath "${rpath}:$out/lib" $out/lib/libtensorflow.so
       patchelf --set-rpath "${rpath}" $out/lib/libtensorflow_framework.so
+      ${optionalString cudaSupport ''
+        addOpenGLRunpath $out/lib/libtensorflow.so $out/lib/libtensorflow_framework.so
+      ''}
     '';
 
 in stdenv.mkDerivation rec {
@@ -43,6 +46,8 @@ in stdenv.mkDerivation rec {
 
   src = fetchurl url;
 
+  nativeBuildInputs = optional cudaSupport addOpenGLRunpath;
+
   # Patch library to use our libc, libstdc++ and others
   buildCommand = ''
     mkdir -pv $out
diff --git a/pkgs/development/python-modules/tensorflow/bin.nix b/pkgs/development/python-modules/tensorflow/bin.nix
index e46d0298953..db2308a0be8 100644
--- a/pkgs/development/python-modules/tensorflow/bin.nix
+++ b/pkgs/development/python-modules/tensorflow/bin.nix
@@ -26,6 +26,7 @@
 , symlinkJoin
 , keras-applications
 , keras-preprocessing
+, addOpenGLRunpath
 }:
 
 # We keep this binary build for two reasons:
@@ -76,6 +77,8 @@ in buildPythonPackage {
   ] ++ lib.optional (!isPy3k) mock
     ++ lib.optionals (pythonOlder "3.4") [ backports_weakref ];
 
+  nativeBuildInputs = lib.optional cudaSupport addOpenGLRunpath;
+
   # Upstream has a pip hack that results in bin/tensorboard being in both tensorflow
   # and the propageted input tensorflow-tensorboard which causes environment collisions.
   # another possibility would be to have tensorboard only in the buildInputs
@@ -94,7 +97,12 @@ in buildPythonPackage {
   lib.optionalString stdenv.isLinux ''
     rrPath="$out/${python.sitePackages}/tensorflow/:$out/${python.sitePackages}/tensorflow/contrib/tensor_forest/:${rpath}"
     internalLibPath="$out/${python.sitePackages}/tensorflow/python/_pywrap_tensorflow_internal.so"
-    find $out \( -name '*.so' -or -name '*.so.*' \) -exec patchelf --set-rpath "$rrPath" {} \;
+    find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
+      patchelf --set-rpath "$rrPath" "$lib"
+      ${lib.optionalString cudaSupport ''
+        addOpenGLRunpath "$lib"
+      ''}
+    done
   '';
 
 
diff --git a/pkgs/development/python-modules/tensorflow/default.nix b/pkgs/development/python-modules/tensorflow/default.nix
index 1895f2bfd0a..e180c9d70e9 100644
--- a/pkgs/development/python-modules/tensorflow/default.nix
+++ b/pkgs/development/python-modules/tensorflow/default.nix
@@ -1,4 +1,5 @@
 { stdenv, pkgs, buildBazelPackage, lib, fetchFromGitHub, fetchpatch, symlinkJoin
+, addOpenGLRunpath
 # Python deps
 , buildPythonPackage, isPy3k, pythonOlder, pythonAtLeast, python
 # Python libraries
@@ -112,7 +113,7 @@ let
 
     nativeBuildInputs = [
       swig which pythonEnv
-    ];
+    ] ++ lib.optional cudaSupport addOpenGLRunpath;
 
     buildInputs = [
       jemalloc
@@ -296,6 +297,12 @@ let
         bazel-bin/tensorflow/tools/pip_package/build_pip_package --src "$PWD/dist"
         cp -Lr "$PWD/dist" "$python"
       '';
+
+      postFixup = lib.optionalString cudaSupport ''
+        find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
+          addOpenGLRunpath "$lib"
+        done
+      '';
     };
 
     meta = with stdenv.lib; {
@@ -347,6 +354,14 @@ in buildPythonPackage {
     tensorflow-tensorboard
   ];
 
+  nativeBuildInputs = lib.optional cudaSupport addOpenGLRunpath;
+
+  postFixup = lib.optionalString cudaSupport ''
+    find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
+      addOpenGLRunpath "$lib"
+    done
+  '';
+
   # Actual tests are slow and impure.
   # TODO try to run them anyway
   # TODO better test (files in tensorflow/tools/ci_build/builds/*test)