uboot: (firmwareOdroidC2/C4) don't invoke patch tool, use patches = [] instead
https://github.com/NixOS/nixpkgs/blob/master/pkgs/stdenv/generic/setup.sh#L948 this can do it nicely. Signed-off-by: Anton Arapov <anton@deadbeef.mx>
This commit is contained in:
commit
56de2bcd43
30691 changed files with 3076956 additions and 0 deletions
146
pkgs/development/python-modules/jaxlib/bin.nix
Normal file
146
pkgs/development/python-modules/jaxlib/bin.nix
Normal file
|
|
@ -0,0 +1,146 @@
|
|||
# For the moment we only support the CPU and GPU backends of jaxlib. The TPU
|
||||
# backend will require some additional work. Those wheels are located here:
|
||||
# https://storage.googleapis.com/jax-releases/libtpu_releases.html.
|
||||
|
||||
# For future reference, the easiest way to test the GPU backend is to run
|
||||
# NIX_PATH=.. nix-shell -p python3 python3Packages.jax "python3Packages.jaxlib.override { cudaSupport = true; }"
|
||||
# export XLA_FLAGS=--xla_gpu_force_compilation_parallelism=1
|
||||
# python -c "from jax.lib import xla_bridge; assert xla_bridge.get_backend().platform == 'gpu'"
|
||||
# python -c "from jax import random; random.PRNGKey(0)"
|
||||
# python -c "from jax import random; x = random.normal(random.PRNGKey(0), (100, 100)); x @ x"
|
||||
# There's no convenient way to test the GPU backend in the derivation since the
|
||||
# nix build environment blocks access to the GPU. See also:
|
||||
# * https://github.com/google/jax/issues/971#issuecomment-508216439
|
||||
# * https://github.com/google/jax/issues/5723#issuecomment-913038780
|
||||
|
||||
{ absl-py
|
||||
, addOpenGLRunpath
|
||||
, autoPatchelfHook
|
||||
, buildPythonPackage
|
||||
, config
|
||||
, cudnn
|
||||
, fetchurl
|
||||
, flatbuffers
|
||||
, isPy39
|
||||
, lib
|
||||
, python
|
||||
, scipy
|
||||
, stdenv
|
||||
# Options:
|
||||
, cudaSupport ? config.cudaSupport or false
|
||||
, cudaPackages ? {}
|
||||
}:
|
||||
|
||||
let
|
||||
inherit (cudaPackages) cudatoolkit cudnn;
|
||||
in
|
||||
|
||||
# There are no jaxlib wheels targeting cudnn <8.0.5, and although there are
|
||||
# wheels for cudatoolkit <11.1, we don't support them.
|
||||
assert cudaSupport -> lib.versionAtLeast cudatoolkit.version "11.1";
|
||||
assert cudaSupport -> lib.versionAtLeast cudnn.version "8.0.5";
|
||||
|
||||
let
|
||||
version = "0.3.0";
|
||||
|
||||
pythonVersion = python.pythonVersion;
|
||||
|
||||
# Find new releases at https://storage.googleapis.com/jax-releases. When
|
||||
# upgrading, you can get these hashes from prefetch.sh.
|
||||
cpuSrcs = {
|
||||
"3.9" = fetchurl {
|
||||
url = "https://storage.googleapis.com/jax-releases/nocuda/jaxlib-${version}-cp39-none-manylinux2010_x86_64.whl";
|
||||
hash = "sha256-AfBVqoqChEXlEC5PgbtQ5rQzcbwo558fjqCjSPEmN5Q=";
|
||||
};
|
||||
"3.10" = fetchurl {
|
||||
url = "https://storage.googleapis.com/jax-releases/nocuda/jaxlib-${version}-cp310-none-manylinux2010_x86_64.whl";
|
||||
hash = "sha256-9uBkFOO8LlRpO6AP+S8XK9/d2yRdyHxQGlbAjShqHRQ=";
|
||||
};
|
||||
};
|
||||
|
||||
gpuSrcs = {
|
||||
"3.9-805" = fetchurl {
|
||||
url = "https://storage.googleapis.com/jax-releases/cuda11/jaxlib-${version}+cuda11.cudnn805-cp39-none-manylinux2010_x86_64.whl";
|
||||
hash = "sha256-CArIhzM5FrQi3TkdqpUqCeDQYyDMVXlzKFgjNXjLJXw=";
|
||||
};
|
||||
"3.9-82" = fetchurl {
|
||||
url = "https://storage.googleapis.com/jax-releases/cuda11/jaxlib-${version}+cuda11.cudnn82-cp39-none-manylinux2010_x86_64.whl";
|
||||
hash = "sha256-Q0plVnA9pUNQ+gCHSXiLNs4i24xCg8gBGfgfYe3bot4=";
|
||||
};
|
||||
"3.10-805" = fetchurl {
|
||||
url = "https://storage.googleapis.com/jax-releases/cuda11/jaxlib-${version}+cuda11.cudnn805-cp310-none-manylinux2010_x86_64.whl";
|
||||
hash = "sha256-JopevCEAs0hgDngIId6NqbLam5YfcS8Lr9cEffBKp1U=";
|
||||
};
|
||||
"3.10-82" = fetchurl {
|
||||
url = "https://storage.googleapis.com/jax-releases/cuda11/jaxlib-${version}+cuda11.cudnn82-cp310-none-manylinux2010_x86_64.whl";
|
||||
hash = "sha256-2f5TwbdP7EfQNRM3ZcJXCAkS2VXBwNYH6gwT9pdu3Go=";
|
||||
};
|
||||
};
|
||||
in
|
||||
buildPythonPackage rec {
|
||||
pname = "jaxlib";
|
||||
inherit version;
|
||||
format = "wheel";
|
||||
|
||||
# At the time of writing (2022-03-03), there are releases for <=3.10.
|
||||
# Supporting all of them is a pain, so we focus on 3.9, the current nixpkgs
|
||||
# python3 version, and 3.10.
|
||||
disabled = !(pythonVersion == "3.9" || pythonVersion == "3.10");
|
||||
|
||||
src =
|
||||
if !cudaSupport then cpuSrcs."${pythonVersion}" else
|
||||
let
|
||||
# jaxlib wheels are currently provided for cudnn versions at least 8.0.5 and
|
||||
# 8.2. Try to use 8.2 whenever possible.
|
||||
cudnnVersion = if (lib.versionAtLeast cudnn.version "8.2") then "82" else "805";
|
||||
in
|
||||
gpuSrcs."${pythonVersion}-${cudnnVersion}";
|
||||
|
||||
# Prebuilt wheels are dynamically linked against things that nix can't find.
|
||||
# Run `autoPatchelfHook` to automagically fix them.
|
||||
nativeBuildInputs = [ autoPatchelfHook ] ++ lib.optional cudaSupport addOpenGLRunpath;
|
||||
# Dynamic link dependencies
|
||||
buildInputs = [ stdenv.cc.cc ];
|
||||
|
||||
# jaxlib contains shared libraries that open other shared libraries via dlopen
|
||||
# and these implicit dependencies are not recognized by ldd or
|
||||
# autoPatchelfHook. That means we need to sneak them into rpath. This step
|
||||
# must be done after autoPatchelfHook and the automatic stripping of
|
||||
# artifacts. autoPatchelfHook runs in postFixup and auto-stripping runs in the
|
||||
# patchPhase. Dependencies:
|
||||
# * libcudart.so.11.0 -> cudatoolkit_11.lib
|
||||
# * libcublas.so.11 -> cudatoolkit_11
|
||||
# * libcuda.so.1 -> opengl driver in /run/opengl-driver/lib
|
||||
preInstallCheck = lib.optional cudaSupport ''
|
||||
shopt -s globstar
|
||||
|
||||
addOpenGLRunpath $out/**/*.so
|
||||
|
||||
for file in $out/**/*.so; do
|
||||
rpath=$(patchelf --print-rpath $file)
|
||||
# For some reason `makeLibraryPath` on `cudatoolkit_11` maps to
|
||||
# <cudatoolkit_11.lib>/lib which is different from <cudatoolkit_11>/lib.
|
||||
patchelf --set-rpath "$rpath:${cudatoolkit}/lib:${lib.makeLibraryPath [ cudatoolkit.lib cudnn ]}" $file
|
||||
done
|
||||
'';
|
||||
|
||||
propagatedBuildInputs = [ absl-py flatbuffers scipy ];
|
||||
|
||||
# Note that cudatoolkit is snecessary since jaxlib looks for "ptxas" in $PATH.
|
||||
# See https://github.com/NixOS/nixpkgs/pull/164176#discussion_r828801621 for
|
||||
# more info.
|
||||
postInstall = lib.optional cudaSupport ''
|
||||
mkdir -p $out/bin
|
||||
ln -s ${cudatoolkit}/bin/ptxas $out/bin/ptxas
|
||||
'';
|
||||
|
||||
pythonImportsCheck = [ "jaxlib" ];
|
||||
|
||||
meta = with lib; {
|
||||
description = "XLA library for JAX";
|
||||
homepage = "https://github.com/google/jax";
|
||||
license = licenses.asl20;
|
||||
maintainers = with maintainers; [ samuela ];
|
||||
platforms = [ "x86_64-linux" ];
|
||||
};
|
||||
}
|
||||
296
pkgs/development/python-modules/jaxlib/default.nix
Normal file
296
pkgs/development/python-modules/jaxlib/default.nix
Normal file
|
|
@ -0,0 +1,296 @@
|
|||
{ lib
|
||||
, pkgs
|
||||
, stdenv
|
||||
|
||||
# Build-time dependencies:
|
||||
, addOpenGLRunpath
|
||||
, bazel_5
|
||||
, binutils
|
||||
, buildBazelPackage
|
||||
, buildPythonPackage
|
||||
, cython
|
||||
, fetchFromGitHub
|
||||
, git
|
||||
, jsoncpp
|
||||
, pybind11
|
||||
, setuptools
|
||||
, symlinkJoin
|
||||
, wheel
|
||||
, which
|
||||
|
||||
# Python dependencies:
|
||||
, absl-py
|
||||
, flatbuffers
|
||||
, numpy
|
||||
, scipy
|
||||
, six
|
||||
|
||||
# Runtime dependencies:
|
||||
, double-conversion
|
||||
, giflib
|
||||
, grpc
|
||||
, libjpeg_turbo
|
||||
, python
|
||||
, snappy
|
||||
, zlib
|
||||
|
||||
# CUDA flags:
|
||||
, cudaCapabilities ? [ "sm_35" "sm_50" "sm_60" "sm_70" "sm_75" "compute_80" ]
|
||||
, cudaSupport ? false
|
||||
, cudaPackages ? {}
|
||||
|
||||
# MKL:
|
||||
, mklSupport ? true
|
||||
}:
|
||||
|
||||
let
|
||||
|
||||
inherit (cudaPackages) cudatoolkit cudnn nccl;
|
||||
|
||||
pname = "jaxlib";
|
||||
version = "0.3.0";
|
||||
|
||||
meta = with lib; {
|
||||
description = "JAX is Autograd and XLA, brought together for high-performance machine learning research.";
|
||||
homepage = "https://github.com/google/jax";
|
||||
license = licenses.asl20;
|
||||
maintainers = with maintainers; [ ndl ];
|
||||
platforms = [ "x86_64-linux" "aarch64-darwin" "x86_64-darwin"];
|
||||
hydraPlatforms = ["x86_64-linux" ]; # Don't think anybody is checking the darwin builds
|
||||
};
|
||||
|
||||
cudatoolkit_joined = symlinkJoin {
|
||||
name = "${cudatoolkit.name}-merged";
|
||||
paths = [
|
||||
cudatoolkit.lib
|
||||
cudatoolkit.out
|
||||
] ++ lib.optionals (lib.versionOlder cudatoolkit.version "11") [
|
||||
# for some reason some of the required libs are in the targets/x86_64-linux
|
||||
# directory; not sure why but this works around it
|
||||
"${cudatoolkit}/targets/${stdenv.system}"
|
||||
];
|
||||
};
|
||||
|
||||
cudatoolkit_cc_joined = symlinkJoin {
|
||||
name = "${cudatoolkit.cc.name}-merged";
|
||||
paths = [
|
||||
cudatoolkit.cc
|
||||
binutils.bintools # for ar, dwp, nm, objcopy, objdump, strip
|
||||
];
|
||||
};
|
||||
|
||||
bazel-build = buildBazelPackage {
|
||||
name = "bazel-build-${pname}-${version}";
|
||||
|
||||
bazel = bazel_5;
|
||||
|
||||
src = fetchFromGitHub {
|
||||
owner = "google";
|
||||
repo = "jax";
|
||||
rev = "${pname}-v${version}";
|
||||
sha256 = "0ndpngx5k6lf6jqjck82bbp0gs943z0wh7vs9gwbyk2bw0da7w72";
|
||||
};
|
||||
|
||||
nativeBuildInputs = [
|
||||
cython
|
||||
pkgs.flatbuffers
|
||||
git
|
||||
setuptools
|
||||
wheel
|
||||
which
|
||||
];
|
||||
|
||||
buildInputs = [
|
||||
double-conversion
|
||||
giflib
|
||||
grpc
|
||||
jsoncpp
|
||||
libjpeg_turbo
|
||||
numpy
|
||||
pkgs.flatbuffers
|
||||
pkgs.protobuf
|
||||
pybind11
|
||||
scipy
|
||||
six
|
||||
snappy
|
||||
zlib
|
||||
] ++ lib.optionals cudaSupport [
|
||||
cudatoolkit
|
||||
cudnn
|
||||
];
|
||||
|
||||
postPatch = ''
|
||||
rm -f .bazelversion
|
||||
'';
|
||||
|
||||
bazelTarget = "//build:build_wheel";
|
||||
|
||||
removeRulesCC = false;
|
||||
|
||||
GCC_HOST_COMPILER_PREFIX = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin";
|
||||
GCC_HOST_COMPILER_PATH = lib.optionalString cudaSupport "${cudatoolkit_cc_joined}/bin/gcc";
|
||||
|
||||
preConfigure = ''
|
||||
# dummy ldconfig
|
||||
mkdir dummy-ldconfig
|
||||
echo "#!${stdenv.shell}" > dummy-ldconfig/ldconfig
|
||||
chmod +x dummy-ldconfig/ldconfig
|
||||
export PATH="$PWD/dummy-ldconfig:$PATH"
|
||||
cat <<CFG > ./.jax_configure.bazelrc
|
||||
build --strategy=Genrule=standalone
|
||||
build --repo_env PYTHON_BIN_PATH="${python}/bin/python"
|
||||
build --action_env=PYENV_ROOT
|
||||
build --python_path="${python}/bin/python"
|
||||
build --distinct_host_configuration=false
|
||||
'' + lib.optionalString cudaSupport ''
|
||||
build --action_env CUDA_TOOLKIT_PATH="${cudatoolkit_joined}"
|
||||
build --action_env CUDNN_INSTALL_PATH="${cudnn}"
|
||||
build --action_env TF_CUDA_PATHS="${cudatoolkit_joined},${cudnn},${nccl}"
|
||||
build --action_env TF_CUDA_VERSION="${lib.versions.majorMinor cudatoolkit.version}"
|
||||
build --action_env TF_CUDNN_VERSION="${lib.versions.major cudnn.version}"
|
||||
build:cuda --action_env TF_CUDA_COMPUTE_CAPABILITIES="${lib.concatStringsSep "," cudaCapabilities}"
|
||||
'' + ''
|
||||
CFG
|
||||
'';
|
||||
|
||||
# Copy-paste from TF derivation.
|
||||
# Most of these are not really used in jaxlib compilation but it's simpler to keep it
|
||||
# 'as is' so that it's more compatible with TF derivation.
|
||||
TF_SYSTEM_LIBS = lib.concatStringsSep "," [
|
||||
"absl_py"
|
||||
"astor_archive"
|
||||
"astunparse_archive"
|
||||
"boringssl"
|
||||
# Not packaged in nixpkgs
|
||||
# "com_github_googleapis_googleapis"
|
||||
# "com_github_googlecloudplatform_google_cloud_cpp"
|
||||
"com_github_grpc_grpc"
|
||||
"com_google_protobuf"
|
||||
# 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&)'
|
||||
# "com_googlesource_code_re2"
|
||||
"curl"
|
||||
"cython"
|
||||
"dill_archive"
|
||||
"double_conversion"
|
||||
"enum34_archive"
|
||||
"flatbuffers"
|
||||
"functools32_archive"
|
||||
"gast_archive"
|
||||
"gif"
|
||||
"hwloc"
|
||||
"icu"
|
||||
"jsoncpp_git"
|
||||
"libjpeg_turbo"
|
||||
"lmdb"
|
||||
"nasm"
|
||||
# "nsync" # not packaged in nixpkgs
|
||||
"opt_einsum_archive"
|
||||
"org_sqlite"
|
||||
"pasta"
|
||||
"pcre"
|
||||
"png"
|
||||
"pybind11"
|
||||
"six_archive"
|
||||
"snappy"
|
||||
"tblib_archive"
|
||||
"termcolor_archive"
|
||||
"typing_extensions_archive"
|
||||
"wrapt"
|
||||
"zlib"
|
||||
];
|
||||
|
||||
# Make sure Bazel knows about our configuration flags during fetching so that the
|
||||
# relevant dependencies can be downloaded.
|
||||
bazelFetchFlags = bazel-build.bazelBuildFlags;
|
||||
|
||||
bazelBuildFlags = [
|
||||
"-c opt"
|
||||
] ++ lib.optional (stdenv.targetPlatform.isx86_64 && stdenv.targetPlatform.isUnix) [
|
||||
"--config=avx_posix"
|
||||
] ++ lib.optional cudaSupport [
|
||||
"--config=cuda"
|
||||
] ++ lib.optional mklSupport [
|
||||
"--config=mkl_open_source_only"
|
||||
];
|
||||
|
||||
fetchAttrs = {
|
||||
sha256 =
|
||||
if cudaSupport then
|
||||
"0d2rqwk9n4a6c51m4g21rxymv85kw2sdksni30cdx3pdcdbqgic7"
|
||||
else
|
||||
"0q540mwmh7grig0qq48ynzqi0gynimxnrq7k97wribqpkx99k39d";
|
||||
};
|
||||
|
||||
buildAttrs = {
|
||||
outputs = [ "out" ];
|
||||
|
||||
# Note: we cannot do most of this patching at `patch` phase as the deps are not available yet.
|
||||
# 1) Fix pybind11 include paths.
|
||||
# 2) Force static protobuf linkage to prevent crashes on loading multiple extensions
|
||||
# in the same python program due to duplicate protobuf DBs.
|
||||
# 3) Patch python path in the compiler driver.
|
||||
# 4) Patch tensorflow sources to work with later versions of protobuf. See
|
||||
# https://github.com/google/jax/issues/9534. Note that this should be
|
||||
# removed on the next release after 0.3.0.
|
||||
preBuild = ''
|
||||
for src in ./jaxlib/*.{cc,h}; do
|
||||
sed -i 's@include/pybind11@pybind11@g' $src
|
||||
done
|
||||
sed -i 's@-lprotobuf@-l:libprotobuf.a@' ../output/external/org_tensorflow/third_party/systemlibs/protobuf.BUILD
|
||||
sed -i 's@-lprotoc@-l:libprotoc.a@' ../output/external/org_tensorflow/third_party/systemlibs/protobuf.BUILD
|
||||
substituteInPlace ../output/external/org_tensorflow/tensorflow/compiler/xla/python/pprof_profile_builder.cc \
|
||||
--replace "status.message()" "std::string{status.message()}"
|
||||
'' + lib.optionalString cudaSupport ''
|
||||
patchShebangs ../output/external/org_tensorflow/third_party/gpus/crosstool/clang/bin/crosstool_wrapper_driver_is_not_gcc.tpl
|
||||
'';
|
||||
|
||||
installPhase = ''
|
||||
./bazel-bin/build/build_wheel --output_path=$out --cpu=${stdenv.targetPlatform.linuxArch}
|
||||
'';
|
||||
};
|
||||
|
||||
inherit meta;
|
||||
};
|
||||
|
||||
in
|
||||
buildPythonPackage {
|
||||
inherit meta pname version;
|
||||
format = "wheel";
|
||||
|
||||
src = "${bazel-build}/jaxlib-${version}-cp${builtins.replaceStrings ["."] [""] python.pythonVersion}-none-manylinux2010_${stdenv.targetPlatform.linuxArch}.whl";
|
||||
|
||||
# Note that cudatoolkit is necessary since jaxlib looks for "ptxas" in $PATH.
|
||||
# See https://github.com/NixOS/nixpkgs/pull/164176#discussion_r828801621 for
|
||||
# more info.
|
||||
postInstall = lib.optionalString cudaSupport ''
|
||||
mkdir -p $out/bin
|
||||
ln -s ${cudatoolkit}/bin/ptxas $out/bin/ptxas
|
||||
|
||||
find $out -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
|
||||
addOpenGLRunpath "$lib"
|
||||
patchelf --set-rpath "${cudatoolkit}/lib:${cudatoolkit.lib}/lib:${cudnn}/lib:${nccl}/lib:$(patchelf --print-rpath "$lib")" "$lib"
|
||||
done
|
||||
'';
|
||||
|
||||
nativeBuildInputs = lib.optional cudaSupport addOpenGLRunpath;
|
||||
|
||||
propagatedBuildInputs = [
|
||||
absl-py
|
||||
double-conversion
|
||||
flatbuffers
|
||||
giflib
|
||||
grpc
|
||||
jsoncpp
|
||||
libjpeg_turbo
|
||||
numpy
|
||||
scipy
|
||||
six
|
||||
snappy
|
||||
];
|
||||
|
||||
pythonImportsCheck = [ "jaxlib" ];
|
||||
|
||||
# Without it there are complaints about libcudart.so.11.0 not being found
|
||||
# because RPATH path entries added above are stripped.
|
||||
dontPatchELF = cudaSupport;
|
||||
}
|
||||
7
pkgs/development/python-modules/jaxlib/prefetch.sh
Executable file
7
pkgs/development/python-modules/jaxlib/prefetch.sh
Executable file
|
|
@ -0,0 +1,7 @@
|
|||
version="$1"
|
||||
nix hash to-sri --type sha256 "$(nix-prefetch-url https://storage.googleapis.com/jax-releases/nocuda/jaxlib-${version}-cp39-none-manylinux2010_x86_64.whl)"
|
||||
nix hash to-sri --type sha256 "$(nix-prefetch-url https://storage.googleapis.com/jax-releases/nocuda/jaxlib-${version}-cp310-none-manylinux2010_x86_64.whl)"
|
||||
nix hash to-sri --type sha256 "$(nix-prefetch-url https://storage.googleapis.com/jax-releases/cuda11/jaxlib-${version}+cuda11.cudnn805-cp39-none-manylinux2010_x86_64.whl)"
|
||||
nix hash to-sri --type sha256 "$(nix-prefetch-url https://storage.googleapis.com/jax-releases/cuda11/jaxlib-${version}+cuda11.cudnn82-cp39-none-manylinux2010_x86_64.whl)"
|
||||
nix hash to-sri --type sha256 "$(nix-prefetch-url https://storage.googleapis.com/jax-releases/cuda11/jaxlib-${version}+cuda11.cudnn805-cp310-none-manylinux2010_x86_64.whl)"
|
||||
nix hash to-sri --type sha256 "$(nix-prefetch-url https://storage.googleapis.com/jax-releases/cuda11/jaxlib-${version}+cuda11.cudnn82-cp310-none-manylinux2010_x86_64.whl)"
|
||||
Loading…
Add table
Add a link
Reference in a new issue