tjh.dev
1{
2 lib,
3 pkgs,
4 stdenv,
5
6 # Build-time dependencies:
7 addDriverRunpath,
8 autoAddDriverRunpath,
9 bazel_6,
10 binutils,
11 buildBazelPackage,
12 buildPythonPackage,
13 cctools,
14 curl,
15 cython,
16 fetchFromGitHub,
17 git,
18 jsoncpp,
19 nsync,
20 openssl,
21 pybind11,
22 setuptools,
23 symlinkJoin,
24 wheel,
25 build,
26 which,
27
28 # Python dependencies:
29 absl-py,
30 flatbuffers,
31 ml-dtypes,
32 numpy,
33 scipy,
34 six,
35
36 # Runtime dependencies:
37 double-conversion,
38 giflib,
39 libjpeg_turbo,
40 python,
41 snappy-cpp,
42 zlib,
43
44 config,
45 # CUDA flags:
46 cudaSupport ? config.cudaSupport,
47 cudaPackages,
48
49 # MKL:
50 mklSupport ? true,
51}@inputs:
52
53let
54 inherit (cudaPackages)
55 cudaMajorMinorVersion
56 flags
57 nccl
58 ;
59
60 pname = "jaxlib";
61 version = "0.4.28";
62
63 # It's necessary to consistently use backendStdenv when building with CUDA
64 # support, otherwise we get libstdc++ errors downstream
65 stdenv = throw "Use effectiveStdenv instead";
66 effectiveStdenv = if cudaSupport then cudaPackages.backendStdenv else inputs.stdenv;
67
68 meta = with lib; {
69 description = "Source-built JAX backend. JAX is Autograd and XLA, brought together for high-performance machine learning research";
70 homepage = "https://github.com/google/jax";
71 license = licenses.asl20;
72 maintainers = with maintainers; [ ndl ];
73
74 # Make this platforms.unix once Darwin is supported.
75 # The top-level jaxlib now falls back to jaxlib-bin on unsupported platforms.
76 # aarch64-darwin is broken because of https://github.com/bazelbuild/rules_cc/pull/136
77 # however even with that fix applied, it doesn't work for everyone:
78 # https://github.com/NixOS/nixpkgs/pull/184395#issuecomment-1207287129
79 platforms = platforms.linux;
80 };
81
82 # Bazel wants a merged cudnn at configuration time
83 cudnnMerged = symlinkJoin {
84 name = "cudnn-merged";
85 paths = with cudaPackages; [
86 (lib.getDev cudnn)
87 (lib.getLib cudnn)
88 ];
89 };
90
91 # These are necessary at build time and run time.
92 cuda_libs_joined = symlinkJoin {
93 name = "cuda-joined";
94 paths = with cudaPackages; [
95 (lib.getLib cuda_cudart) # libcudart.so
96 (lib.getLib cuda_cupti) # libcupti.so
97 (lib.getLib libcublas) # libcublas.so
98 (lib.getLib libcufft) # libcufft.so
99 (lib.getLib libcurand) # libcurand.so
100 (lib.getLib libcusolver) # libcusolver.so
101 (lib.getLib libcusparse) # libcusparse.so
102 ];
103 };
104 # These are only necessary at build time.
105 cuda_build_deps_joined = symlinkJoin {
106 name = "cuda-build-deps-joined";
107 paths = with cudaPackages; [
108 cuda_libs_joined
109
110 # Binaries
111 (lib.getBin cuda_nvcc) # nvcc
112
113 # Archives
114 (lib.getOutput "static" cuda_cudart) # libcudart_static.a
115
116 # Headers
117 (lib.getDev cuda_cccl) # block_load.cuh
118 (lib.getDev cuda_cudart) # cuda.h
119 (lib.getDev cuda_cupti) # cupti.h
120 (lib.getDev cuda_nvcc) # See https://github.com/google/jax/issues/19811
121 (lib.getDev cuda_nvml_dev) # nvml.h
122 (lib.getDev cuda_nvtx) # nvToolsExt.h
123 (lib.getDev libcublas) # cublas_api.h
124 (lib.getDev libcufft) # cufft.h
125 (lib.getDev libcurand) # curand.h
126 (lib.getDev libcusolver) # cusolver_common.h
127 (lib.getDev libcusparse) # cusparse.h
128 ];
129 };
130
131 backend_cc_joined = symlinkJoin {
132 name = "cuda-cc-joined";
133 paths = [
134 effectiveStdenv.cc
135 binutils.bintools # for ar, dwp, nm, objcopy, objdump, strip
136 ];
137 };
138
139 # Copy-paste from TF derivation.
140 # Most of these are not really used in jaxlib compilation but it's simpler to keep it
141 # 'as is' so that it's more compatible with TF derivation.
142 tf_system_libs = [
143 "absl_py"
144 "astor_archive"
145 "astunparse_archive"
146 # Not packaged in nixpkgs
147 # "com_github_googleapis_googleapis"
148 # "com_github_googlecloudplatform_google_cloud_cpp"
149 # Issue with transitive dependencies after https://github.com/grpc/grpc/commit/f1d14f7f0b661bd200b7f269ef55dec870e7c108
150 # "com_github_grpc_grpc"
151 # ERROR: /build/output/external/bazel_tools/tools/proto/BUILD:25:6: no such target '@com_google_protobuf//:cc_toolchain':
152 # target 'cc_toolchain' not declared in package '' defined by /build/output/external/com_google_protobuf/BUILD.bazel
153 # "com_google_protobuf"
154 # Fails with the error: external/org_tensorflow/tensorflow/core/profiler/utils/tf_op_utils.cc:46:49: error: no matching function for call to 're2::RE2::FullMatch(absl::lts_2020_02_25::string_view&, re2::RE2&)'
155 # "com_googlesource_code_re2"
156 "curl"
157 "cython"
158 "dill_archive"
159 "double_conversion"
160 "flatbuffers"
161 "functools32_archive"
162 "gast_archive"
163 "gif"
164 "hwloc"
165 "icu"
166 "jsoncpp_git"
167 "libjpeg_turbo"
168 "lmdb"
169 "nasm"
170 "opt_einsum_archive"
171 "org_sqlite"
172 "pasta"
173 "png"
174 # ERROR: /build/output/external/pybind11/BUILD.bazel: no such target '@pybind11//:osx':
175 # target 'osx' not declared in package '' defined by /build/output/external/pybind11/BUILD.bazel
176 # "pybind11"
177 "six_archive"
178 "snappy"
179 "tblib_archive"
180 "termcolor_archive"
181 "typing_extensions_archive"
182 "wrapt"
183 "zlib"
184 ];
185
186 arch =
187 # KeyError: ('Linux', 'arm64')
188 if effectiveStdenv.hostPlatform.isLinux && effectiveStdenv.hostPlatform.linuxArch == "arm64" then
189 "aarch64"
190 else
191 effectiveStdenv.hostPlatform.linuxArch;
192
193 xla = effectiveStdenv.mkDerivation {
194 pname = "xla-src";
195 version = "unstable";
196
197 src = fetchFromGitHub {
198 owner = "openxla";
199 repo = "xla";
200 # Update this according to https://github.com/google/jax/blob/jaxlib-v${version}/third_party/xla/workspace.bzl.
201 rev = "e8247c3ea1d4d7f31cf27def4c7ac6f2ce64ecd4";
202 hash = "sha256-ZhgMIVs3Z4dTrkRWDqaPC/i7yJz2dsYXrZbjzqvPX3E=";
203 };
204
205 dontBuild = true;
206
207 # This is necessary for patchShebangs to know the right path to use.
208 nativeBuildInputs = [ python ];
209
210 # Main culprits we're targeting are third_party/tsl/third_party/gpus/crosstool/clang/bin/*.tpl
211 postPatch = ''
212 patchShebangs .
213 '';
214
215 installPhase = ''
216 cp -r . $out
217 '';
218 };
219
220 bazel-build = buildBazelPackage rec {
221 name = "bazel-build-${pname}-${version}";
222
223 # See https://github.com/google/jax/blob/main/.bazelversion for the latest.
224 bazel = bazel_6;
225
226 src = fetchFromGitHub {
227 owner = "google";
228 repo = "jax";
229 # google/jax contains tags for jax and jaxlib. Only use jaxlib tags!
230 rev = "refs/tags/${pname}-v${version}";
231 hash = "sha256-qSHPwi3is6Ts7pz5s4KzQHBMbcjGp+vAOsejW3o36Ek=";
232 };
233
234 nativeBuildInputs = [
235 cython
236 pkgs.flatbuffers
237 git
238 setuptools
239 wheel
240 build
241 which
242 ]
243 ++ lib.optionals effectiveStdenv.hostPlatform.isDarwin [ cctools ];
244
245 buildInputs = [
246 curl
247 double-conversion
248 giflib
249 jsoncpp
250 libjpeg_turbo
251 numpy
252 openssl
253 pkgs.flatbuffers
254 pkgs.protobuf
255 pybind11
256 scipy
257 six
258 snappy-cpp
259 zlib
260 ]
261 ++ lib.optionals (!effectiveStdenv.hostPlatform.isDarwin) [ nsync ];
262
263 # We don't want to be quite so picky regarding bazel version
264 postPatch = ''
265 rm -f .bazelversion
266 '';
267
268 bazelRunTarget = "//jaxlib/tools:build_wheel";
269 runTargetFlags = [
270 "--output_path=$out"
271 "--cpu=${arch}"
272 # This has no impact whatsoever...
273 "--jaxlib_git_hash='12345678'"
274 ];
275
276 removeRulesCC = false;
277
278 GCC_HOST_COMPILER_PREFIX = lib.optionalString cudaSupport "${backend_cc_joined}/bin";
279 GCC_HOST_COMPILER_PATH = lib.optionalString cudaSupport "${backend_cc_joined}/bin/gcc";
280
281 # The version is automatically set to ".dev" if this variable is not set.
282 # https://github.com/google/jax/commit/e01f2617b85c5bdffc5ffb60b3d8d8ca9519a1f3
283 JAXLIB_RELEASE = "1";
284
285 preConfigure =
286 # Dummy ldconfig to work around "Can't open cache file /nix/store/<hash>-glibc-2.38-44/etc/ld.so.cache" error
287 ''
288 mkdir dummy-ldconfig
289 echo "#!${effectiveStdenv.shell}" > dummy-ldconfig/ldconfig
290 chmod +x dummy-ldconfig/ldconfig
291 export PATH="$PWD/dummy-ldconfig:$PATH"
292 ''
293 +
294
295 # Construct .jax_configure.bazelrc. See https://github.com/google/jax/blob/b9824d7de3cb30f1df738cc42e486db3e9d915ff/build/build.py#L259-L345
296 # for more info. We assume
297 # * `cpu = None`
298 # * `enable_nccl = True`
299 # * `target_cpu_features = "release"`
300 # * `rocm_amdgpu_targets = None`
301 # * `enable_rocm = False`
302 # * `build_gpu_plugin = False`
303 # * `use_clang = False` (Should we use `effectiveStdenv.cc.isClang` instead?)
304 #
305 # Note: We should try just running https://github.com/google/jax/blob/ceb198582b62b9e6f6bdf20ab74839b0cf1db16e/build/build.py#L259-L266
306 # instead of duplicating the logic here. Perhaps we can leverage the
307 # `--configure_only` flag (https://github.com/google/jax/blob/ceb198582b62b9e6f6bdf20ab74839b0cf1db16e/build/build.py#L544-L548)?
308 ''
309 cat <<CFG > ./.jax_configure.bazelrc
310 build --strategy=Genrule=standalone
311 build --repo_env PYTHON_BIN_PATH="${python}/bin/python"
312 build --action_env=PYENV_ROOT
313 build --python_path="${python}/bin/python"
314 build --distinct_host_configuration=false
315 build --define PROTOBUF_INCLUDE_PATH="${pkgs.protobuf}/include"
316 ''
317 + lib.optionalString cudaSupport ''
318 build --config=cuda
319 build --action_env CUDA_TOOLKIT_PATH="${cuda_build_deps_joined}"
320 build --action_env CUDNN_INSTALL_PATH="${cudnnMerged}"
321 build --action_env TF_CUDA_PATHS="${cuda_build_deps_joined},${cudnnMerged},${lib.getDev nccl}"
322 build --action_env TF_CUDA_VERSION="${cudaMajorMinorVersion}"
323 build --action_env TF_CUDNN_VERSION="${lib.versions.major cudaPackages.cudnn.version}"
324 build:cuda --action_env TF_CUDA_COMPUTE_CAPABILITIES="${builtins.concatStringsSep "," flags.realArches}"
325 ''
326 +
327 # Note that upstream conditions this on `wheel_cpu == "x86_64"`. We just
328 # rely on `effectiveStdenv.hostPlatform.avxSupport` instead. So far so
329 # good. See https://github.com/google/jax/blob/b9824d7de3cb30f1df738cc42e486db3e9d915ff/build/build.py#L322
330 # for upstream's version.
331 lib.optionalString (effectiveStdenv.hostPlatform.avxSupport && effectiveStdenv.hostPlatform.isUnix)
332 ''
333 build --config=avx_posix
334 ''
335 + lib.optionalString mklSupport ''
336 build --config=mkl_open_source_only
337 ''
338 + ''
339 CFG
340 '';
341
342 # Make sure Bazel knows about our configuration flags during fetching so that the
343 # relevant dependencies can be downloaded.
344 bazelFlags = [
345 "-c opt"
346 # See https://bazel.build/external/advanced#overriding-repositories for
347 # information on --override_repository flag.
348 "--override_repository=xla=${xla}"
349 ]
350 ++ lib.optionals effectiveStdenv.cc.isClang [
351 # bazel depends on the compiler frontend automatically selecting these flags based on file
352 # extension but our clang doesn't.
353 # https://github.com/NixOS/nixpkgs/issues/150655
354 "--cxxopt=-x"
355 "--cxxopt=c++"
356 "--host_cxxopt=-x"
357 "--host_cxxopt=c++"
358 ];
359
360 # We intentionally overfetch so we can share the fetch derivation across all the different configurations
361 fetchAttrs = {
362 TF_SYSTEM_LIBS = lib.concatStringsSep "," tf_system_libs;
363 # we have to force @mkl_dnn_v1 since it's not needed on darwin
364 bazelTargets = [
365 bazelRunTarget
366 "@mkl_dnn_v1//:mkl_dnn"
367 ];
368 bazelFlags =
369 bazelFlags
370 ++ [
371 "--config=avx_posix"
372 "--config=mkl_open_source_only"
373 ]
374 ++ lib.optionals cudaSupport [
375 # ideally we'd add this unconditionally too, but it doesn't work on darwin
376 # we make this conditional on `cudaSupport` instead of the system, so that the hash for both
377 # the cuda and the non-cuda deps can be computed on linux, since a lot of contributors don't
378 # have access to darwin machines
379 "--config=cuda"
380 ];
381
382 sha256 =
383 (
384 if cudaSupport then
385 { x86_64-linux = "sha256-Uf0VMRE0jgaWEYiuphWkWloZ5jMeqaWBl3lSvk2y1HI="; }
386 else
387 {
388 x86_64-linux = "sha256-NzJJg6NlrPGMiR8Fn8u4+fu0m+AulfmN5Xqk63Um6sw=";
389 aarch64-linux = "sha256-Ro3qzrUxSR+3TH6ROoJTq+dLSufrDN/9oEo2MRkx7wM=";
390 }
391 ).${effectiveStdenv.system} or (throw "jaxlib: unsupported system: ${effectiveStdenv.system}");
392
393 # Non-reproducible fetch https://github.com/NixOS/nixpkgs/issues/321920#issuecomment-2184940546
394 preInstall = ''
395 cat << \EOF > "$bazelOut/external/go_sdk/versions.json"
396 []
397 EOF
398 '';
399 };
400
401 buildAttrs = {
402 outputs = [ "out" ];
403
404 TF_SYSTEM_LIBS = lib.concatStringsSep "," (
405 tf_system_libs
406 ++ lib.optionals (!effectiveStdenv.hostPlatform.isDarwin) [
407 "nsync" # fails to build on darwin
408 ]
409 );
410
411 # Note: we cannot do most of this patching at `patch` phase as the deps
412 # are not available yet.
413 preBuild = lib.optionalString effectiveStdenv.hostPlatform.isDarwin ''
414 substituteInPlace ../output/external/rules_cc/cc/private/toolchain/osx_cc_wrapper.sh.tpl \
415 --replace "/usr/bin/install_name_tool" "${cctools}/bin/install_name_tool"
416 substituteInPlace ../output/external/rules_cc/cc/private/toolchain/unix_cc_configure.bzl \
417 --replace "/usr/bin/libtool" "${cctools}/bin/libtool"
418 '';
419 };
420
421 inherit meta;
422 };
423 platformTag =
424 if effectiveStdenv.hostPlatform.isLinux then
425 "manylinux2014_${arch}"
426 else if effectiveStdenv.system == "x86_64-darwin" then
427 "macosx_10_9_${arch}"
428 else if effectiveStdenv.system == "aarch64-darwin" then
429 "macosx_11_0_${arch}"
430 else
431 throw "Unsupported target platform: ${effectiveStdenv.hostPlatform}";
432in
433buildPythonPackage {
434 inherit pname version;
435 format = "wheel";
436
437 src =
438 let
439 cp = "cp${builtins.replaceStrings [ "." ] [ "" ] python.pythonVersion}";
440 in
441 "${bazel-build}/jaxlib-${version}-${cp}-${cp}-${platformTag}.whl";
442
443 # Note that jaxlib looks for "ptxas" in $PATH. See https://github.com/NixOS/nixpkgs/pull/164176#discussion_r828801621
444 # for more info.
445 postInstall = lib.optionalString cudaSupport ''
446 mkdir -p $out/bin
447 ln -s ${lib.getExe' cudaPackages.cuda_nvcc "ptxas"} $out/bin/ptxas
448
449 find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
450 patchelf --add-rpath "${
451 lib.makeLibraryPath [
452 cuda_libs_joined
453 (lib.getLib cudaPackages.cudnn)
454 nccl
455 ]
456 }" "$lib"
457 done
458 '';
459
460 nativeBuildInputs = lib.optionals cudaSupport [ autoAddDriverRunpath ];
461
462 dependencies = [
463 absl-py
464 curl
465 double-conversion
466 flatbuffers
467 giflib
468 jsoncpp
469 libjpeg_turbo
470 ml-dtypes
471 numpy
472 scipy
473 six
474 ];
475
476 buildInputs = [
477 snappy-cpp
478 ];
479
480 pythonImportsCheck = [
481 "jaxlib"
482 # `import jaxlib` loads surprisingly little. These imports are actually bugs that appeared in the 0.4.11 upgrade.
483 "jaxlib.cpu_feature_guard"
484 "jaxlib.xla_client"
485 ];
486
487 # Without it there are complaints about libcudart.so.11.0 not being found
488 # because RPATH path entries added above are stripped.
489 dontPatchELF = cudaSupport;
490
491 passthru = {
492 # Note "bazel.*.tar.gz" can be accessed as `jaxlib.bazel-build.deps`
493 inherit bazel-build;
494 };
495
496 inherit meta;
497}