1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
|
{ stdenv, buildBazelPackage, lib, fetchFromGitHub, fetchpatch, symlinkJoin
, buildPythonPackage, isPy3k, pythonOlder, pythonAtLeast
, which, swig, binutils, glibcLocales
, python, jemalloc, openmpi
, numpy, six, protobuf, tensorflow-tensorboard, backports_weakref, mock, enum34, absl-py
, cudaSupport ? false, nvidia_x11 ? null, cudatoolkit ? null, cudnn ? null
# XLA without CUDA is broken
, xlaSupport ? cudaSupport
# Default from ./configure script
, cudaCapabilities ? [ "3.5" "5.2" ]
, sse42Support ? builtins.elem (stdenv.hostPlatform.platform.gcc.arch or "default") ["westmere" "sandybridge" "ivybridge" "haswell" "broadwell" "skylake" "skylake-avx512"]
, avx2Support ? builtins.elem (stdenv.hostPlatform.platform.gcc.arch or "default") [ "haswell" "broadwell" "skylake" "skylake-avx512"]
, fmaSupport ? builtins.elem (stdenv.hostPlatform.platform.gcc.arch or "default") [ "haswell" "broadwell" "skylake" "skylake-avx512"]
}:
assert cudaSupport -> nvidia_x11 != null
&& cudatoolkit != null
&& cudnn != null;
# unsupported combination
assert ! (stdenv.isDarwin && cudaSupport);
let
withTensorboard = pythonOlder "3.6";
cudatoolkit_joined = symlinkJoin {
name = "${cudatoolkit.name}-unsplit";
paths = [ cudatoolkit.out cudatoolkit.lib ];
};
tfFeature = x: if x then "1" else "0";
version = "1.5.0";
pkg = buildBazelPackage rec {
name = "tensorflow-build-${version}";
src = fetchFromGitHub {
owner = "tensorflow";
repo = "tensorflow";
rev = "v${version}";
sha256 = "1c4djsaip901nasm7a6dsimr02bsv70a7b1g0kysb4n39qpdh22q";
};
patches = [
# Fix build with Bazel >= 0.10
(fetchpatch {
url = "https://github.com/tensorflow/tensorflow/commit/6fcfab770c2672e2250e0f5686b9545d99eb7b2b.patch";
sha256 = "0p61za1mx3a7gj1s5lsps16fcw18iwnvq2b46v1kyqfgq77a12vb";
})
(fetchpatch {
url = "https://github.com/tensorflow/tensorflow/commit/3f57956725b553d196974c9ad31badeb3eabf8bb.patch";
sha256 = "11dja5gqy0qw27sc9b6yw9r0lfk8dznb32vrqqfcnypk2qmv26va";
})
];
nativeBuildInputs = [ swig which ];
buildInputs = [ python jemalloc openmpi glibcLocales numpy ]
++ lib.optionals cudaSupport [ cudatoolkit cudnn nvidia_x11 ];
preConfigure = ''
patchShebangs configure
export PYTHON_BIN_PATH="${python.interpreter}"
export PYTHON_LIB_PATH="$NIX_BUILD_TOP/site-packages"
export TF_NEED_GCP=1
export TF_NEED_HDFS=1
export TF_ENABLE_XLA=${tfFeature xlaSupport}
export CC_OPT_FLAGS=" "
# https://github.com/tensorflow/tensorflow/issues/14454
export TF_NEED_MPI=${tfFeature cudaSupport}
export TF_NEED_CUDA=${tfFeature cudaSupport}
${lib.optionalString cudaSupport ''
export CUDA_TOOLKIT_PATH=${cudatoolkit_joined}
export TF_CUDA_VERSION=${cudatoolkit.majorVersion}
export CUDNN_INSTALL_PATH=${cudnn}
export TF_CUDNN_VERSION=${cudnn.majorVersion}
export GCC_HOST_COMPILER_PATH=${cudatoolkit.cc}/bin/gcc
export TF_CUDA_COMPUTE_CAPABILITIES=${lib.concatStringsSep "," cudaCapabilities}
''}
mkdir -p "$PYTHON_LIB_PATH"
'';
NIX_LDFLAGS = lib.optionals cudaSupport [ "-lcublas" "-lcudnn" "-lcuda" "-lcudart" ];
hardeningDisable = [ "all" ];
bazelFlags = [ "--config=opt" ]
++ lib.optional sse42Support "--copt=-msse4.2"
++ lib.optional avx2Support "--copt=-mavx2"
++ lib.optional fmaSupport "--copt=-mfma"
++ lib.optional cudaSupport "--config=cuda";
bazelTarget = "//tensorflow/tools/pip_package:build_pip_package";
fetchAttrs = {
preInstall = ''
rm -rf $bazelOut/external/{bazel_tools,\@bazel_tools.marker,local_*,\@local_*}
'';
sha256 = "1nc98aqrp14q7llypcwaa0kdn9xi7r0p1mnd3vmmn1m299py33ca";
};
buildAttrs = {
preBuild = ''
patchShebangs .
find -type f -name CROSSTOOL\* -exec sed -i \
-e 's,/usr/bin/ar,${binutils.bintools}/bin/ar,g' \
{} \;
'';
installPhase = ''
sed -i 's,.*bdist_wheel.*,cp -rL . "$out"; exit 0,' bazel-bin/tensorflow/tools/pip_package/build_pip_package
bazel-bin/tensorflow/tools/pip_package/build_pip_package $PWD/dist
'';
};
dontFixup = true;
};
in buildPythonPackage rec {
pname = "tensorflow";
inherit version;
src = pkg;
installFlags = lib.optional (!withTensorboard) "--no-dependencies";
postPatch = lib.optionalString (pythonAtLeast "3.4") ''
sed -i '/enum34/d' setup.py
'';
propagatedBuildInputs = [ numpy six protobuf absl-py ]
++ lib.optional (!isPy3k) mock
++ lib.optionals (pythonOlder "3.4") [ backports_weakref enum34 ]
++ lib.optional withTensorboard tensorflow-tensorboard;
# Actual tests are slow and impure.
checkPhase = ''
${python.interpreter} -c "import tensorflow"
'';
meta = with stdenv.lib; {
description = "Computation using data flow graphs for scalable machine learning";
homepage = http://tensorflow.org;
license = licenses.asl20;
maintainers = with maintainers; [ jyp abbradar ];
platforms = platforms.linux;
broken = !(xlaSupport -> cudaSupport);
};
}
|