libtensorflow: support building on OS X and add optional cuda support

This commit is contained in:
Bas van Dijk 2018-07-23 11:07:09 +02:00
parent 6275a31f33
commit 7098bcac27
2 changed files with 54 additions and 12 deletions

View file

@ -1,11 +1,50 @@
{ stdenv, fetchurl, patchelf }:
stdenv.mkDerivation rec {
{ stdenv
, fetchurl
, patchelf
, cudaSupport ? false, symlinkJoin, cudatoolkit, cudnn, nvidia_x11
}:
with stdenv.lib;
let
tfType = if cudaSupport then "gpu" else "cpu";
system =
if stdenv.isx86_64
then if stdenv.isLinux then "linux-x86_64"
else if stdenv.isDarwin then "darwin-x86_64" else unavailable
else unavailable;
unavailable = throw "libtensorflow is not available for this platform!";
cudatoolkit_joined = symlinkJoin {
name = "unsplit_cudatoolkit";
paths = [ cudatoolkit.out
cudatoolkit.lib ];};
rpath = makeLibraryPath ([stdenv.cc.libc stdenv.cc.cc.lib] ++
optionals cudaSupport [ cudatoolkit_joined cudnn nvidia_x11 ]);
patchLibs =
if stdenv.isDarwin
then ''
install_name_tool -id $out/lib/libtensorflow.so $out/lib/libtensorflow.so
install_name_tool -id $out/lib/libtensorflow_framework.so $out/lib/libtensorflow_framework.so
''
else ''
${patchelf}/bin/patchelf --set-rpath "${rpath}:$out/lib" $out/lib/libtensorflow.so
${patchelf}/bin/patchelf --set-rpath "${rpath}" $out/lib/libtensorflow_framework.so
'';
in stdenv.mkDerivation rec {
pname = "libtensorflow";
version = "1.8.0";
name = "${pname}-${version}";
src = fetchurl {
url = "https://storage.googleapis.com/tensorflow/${pname}/${pname}-cpu-linux-x86_64-${version}.tar.gz";
sha256 = "0qzy15rc3x961cyi3bqnygrcnw4x69r28xkwhpwrv1r0gi6k73ha";
url = "https://storage.googleapis.com/tensorflow/${pname}/${pname}-${tfType}-${system}-${version}.tar.gz";
sha256 =
if system == "linux-x86_64" then
if cudaSupport
then "0m1g4sqr9as0jgfx7wlyay2nkad6wgvsyk2gvhfkqkq5sm1vbx85"
else "0qzy15rc3x961cyi3bqnygrcnw4x69r28xkwhpwrv1r0gi6k73ha"
else if system == "darwin-x86_64" then
if cudaSupport
then unavailable
else "0q8lmyj8l50hl6l48c640ixanvhqf2836bicyl9p2x8sj97b7y8l"
else unavailable;
};
# Patch library to use our libc, libstdc++ and others
@ -15,18 +54,16 @@ stdenv.mkDerivation rec {
tar -C $out -xzf $src
chmod +w $out/lib/libtensorflow.so
chmod +w $out/lib/libtensorflow_framework.so
${patchelf}/bin/patchelf --set-rpath "${stdenv.cc.libc}/lib:${stdenv.cc.cc.lib}/lib:$out/lib" $out/lib/libtensorflow.so
${patchelf}/bin/patchelf --set-rpath "${stdenv.cc.libc}/lib:${stdenv.cc.cc.lib}/lib" $out/lib/libtensorflow_framework.so
${patchLibs}
chmod -w $out/lib/libtensorflow.so
chmod -w $out/lib/libtensorflow_framework.so
'';
meta = with stdenv.lib; {
inherit version;
meta = {
description = "C API for TensorFlow";
license = licenses.asl20;
maintainers = [maintainers.basvandijk];
platforms = platforms.linux;
homepage = https://www.tensorflow.org/versions/master/install/install_c;
license = licenses.asl20;
platforms = with platforms; linux ++ darwin;
maintainers = [maintainers.basvandijk];
};
}

View file

@ -858,7 +858,12 @@ with pkgs;
tensor = libsForQt5.callPackage ../applications/networking/instant-messengers/tensor { };
libtensorflow = callPackage ../development/libraries/libtensorflow { };
libtensorflow = callPackage ../development/libraries/libtensorflow {
cudaSupport = config.cudaSupport or false;
inherit (linuxPackages) nvidia_x11;
cudatoolkit = cudatoolkit_9_0;
cudnn = cudnn_cudatoolkit_9_0;
};
blink1-tool = callPackage ../tools/misc/blink1-tool { };