diff options
Diffstat (limited to 'nixpkgs/pkgs/development/python-modules/pytorch/default.nix')
| -rw-r--r-- | nixpkgs/pkgs/development/python-modules/pytorch/default.nix | 83 |
1 files changed, 53 insertions, 30 deletions
diff --git a/nixpkgs/pkgs/development/python-modules/pytorch/default.nix b/nixpkgs/pkgs/development/python-modules/pytorch/default.nix index 80c3c6d72931..4635f813c3a6 100644 --- a/nixpkgs/pkgs/development/python-modules/pytorch/default.nix +++ b/nixpkgs/pkgs/development/python-modules/pytorch/default.nix @@ -4,7 +4,7 @@ openMPISupport ? false, openmpi ? null, buildDocs ? false, cudaArchList ? null, - numpy, pyyaml, cffi, click, typing, cmake, oneDNN, hypothesis, numactl, psutil, + numpy, pyyaml, cffi, click, typing, cmake, hypothesis, numactl, psutil, linkFarm, symlinkJoin, # virtual pkg that consistently instantiates blas across nixpkgs @@ -25,15 +25,11 @@ assert !openMPISupport || openmpi != null; assert !cudaSupport || cudatoolkit != null; assert cudnn == null || cudatoolkit != null; assert !cudaSupport || (let majorIs = lib.versions.major cudatoolkit.version; - in majorIs == "9" || majorIs == "10"); + in majorIs == "9" || majorIs == "10" || majorIs == "11"); -let - hasDependency = dep: pkg: lib.lists.any (inp: inp == dep) pkg.buildInputs; - matchesCudatoolkit = hasDependency cudatoolkit; -in # confirm that cudatoolkits are sync'd across dependencies -assert !(openMPISupport && cudaSupport) || matchesCudatoolkit openmpi; -assert !cudaSupport || matchesCudatoolkit magma; +assert !(openMPISupport && cudaSupport) || openmpi.cudatoolkit == cudatoolkit; +assert !cudaSupport || magma.cudatoolkit == cudatoolkit; let cudatoolkit_joined = symlinkJoin { @@ -104,13 +100,16 @@ let "LD_LIBRARY_PATH=${cudaStub}\${LD_LIBRARY_PATH:+:}$LD_LIBRARY_PATH "; in buildPythonPackage rec { - version = "1.5.0"; pname = "pytorch"; + # Don't forget to update pytorch-bin to the same version. + version = "1.6.0"; + disabled = !isPy3k; outputs = [ "out" # output standard python package - "dev" # output libtorch only + "dev" # output libtorch headers + "lib" # output libtorch libraries ]; src = fetchFromGitHub { @@ -118,9 +117,31 @@ in buildPythonPackage rec { repo = "pytorch"; rev = "v${version}"; fetchSubmodules = true; - sha256 = "19qyrjd72mc0llcfn50av8ym05f2iwa38gv068wykji4ph7qjlv2"; + sha256 = "14hhjsi6fnpaw9m1a3bhvdinsks6fhss6bbcrfk6jgns64abqdaz"; }; + patches = lib.optionals stdenv.isAarch64 [ + # GNU aarch64 assembler does not support 4s on neon mov: + # https://github.com/pytorch/pytorch/issues/33124 + # + # Fix from: + # https://github.com/pytorch/pytorch/pull/40584 + # + # This patch can be removed with the next major version (1.7.0). + (fetchpatch { + name = "qnnpack-neon-fix.patch"; + url = "https://github.com/pytorch/pytorch/commit/7676682584d0caf9243bce74ea0a88711ec4a807.diff"; + sha256 = "13spncaqlpsp8qk2850yly7xqwmhhfwznhmzkk8jgpslkbx75vgq"; + }) + ] ++ lib.optionals stdenv.isDarwin [ + # pthreadpool added support for Grand Central Dispatch in April + # 2020. However, this relies on functionality (DISPATCH_APPLY_AUTO) + # that is available starting with macOS 10.13. However, our current + # base is 10.12. Until we upgrade, we can fall back on the older + # pthread support. + ./pthreadpool-disable-gcd.diff + ]; + preConfigure = lib.optionalString cudaSupport '' export TORCH_CUDA_ARCH_LIST="${lib.strings.concatStringsSep ";" final_cudaArchList}" export CC=${cudatoolkit.cc}/bin/gcc CXX=${cudatoolkit.cc}/bin/g++ @@ -136,9 +157,9 @@ in buildPythonPackage rec { USE_MKL = blas.implementation == "mkl"; - # Unlike MKL, MKLDNN is FOSS, so we enable support for it by default. Note - # that this was renamed to dnnl and then renamed again to oneDNN upstream, but - # pytorch still calls it by the old name mkldnn. + # Unlike MKL, oneDNN (née MKLDNN) is FOSS, so we enable support for + # it by default. PyTorch currently uses its own vendored version + # of oneDNN through Intel iDeep. USE_MKLDNN = mklDnnSupport; USE_MKLDNN_CBLAS = mklDnnSupport; @@ -187,7 +208,7 @@ in buildPythonPackage rec { ninja ] ++ lib.optionals cudaSupport [ cudatoolkit_joined ]; - buildInputs = [ blas blas.provider oneDNN ] + buildInputs = [ blas blas.provider ] ++ lib.optionals cudaSupport [ cudnn magma nccl ] ++ lib.optionals stdenv.isLinux [ numactl ]; @@ -224,33 +245,35 @@ in buildPythonPackage rec { ]; postInstall = '' mkdir $dev - cp -r $out/${python.sitePackages}/torch/lib $dev/lib cp -r $out/${python.sitePackages}/torch/include $dev/include cp -r $out/${python.sitePackages}/torch/share $dev/share + + mkdir $lib + cp -r $out/${python.sitePackages}/torch/lib $lib/lib ''; postFixup = stdenv.lib.optionalString stdenv.isDarwin '' - for f in $(ls $dev/lib/*.dylib); do - install_name_tool -id $dev/lib/$(basename $f) $f || true + for f in $(ls $lib/lib/*.dylib); do + install_name_tool -id $lib/lib/$(basename $f) $f || true done - install_name_tool -change @rpath/libshm.dylib $dev/lib/libshm.dylib $dev/lib/libtorch_python.dylib - install_name_tool -change @rpath/libtorch.dylib $dev/lib/libtorch.dylib $dev/lib/libtorch_python.dylib - install_name_tool -change @rpath/libc10.dylib $dev/lib/libc10.dylib $dev/lib/libtorch_python.dylib + install_name_tool -change @rpath/libshm.dylib $lib/lib/libshm.dylib $lib/lib/libtorch_python.dylib + install_name_tool -change @rpath/libtorch.dylib $lib/lib/libtorch.dylib $lib/lib/libtorch_python.dylib + install_name_tool -change @rpath/libc10.dylib $lib/lib/libc10.dylib $lib/lib/libtorch_python.dylib - install_name_tool -change @rpath/libc10.dylib $dev/lib/libc10.dylib $dev/lib/libtorch.dylib + install_name_tool -change @rpath/libc10.dylib $lib/lib/libc10.dylib $lib/lib/libtorch.dylib - install_name_tool -change @rpath/libtorch.dylib $dev/lib/libtorch.dylib $dev/lib/libcaffe2_observers.dylib - install_name_tool -change @rpath/libc10.dylib $dev/lib/libc10.dylib $dev/lib/libcaffe2_observers.dylib + install_name_tool -change @rpath/libtorch.dylib $lib/lib/libtorch.dylib $lib/lib/libcaffe2_observers.dylib + install_name_tool -change @rpath/libc10.dylib $lib/lib/libc10.dylib $lib/lib/libcaffe2_observers.dylib - install_name_tool -change @rpath/libtorch.dylib $dev/lib/libtorch.dylib $dev/lib/libcaffe2_module_test_dynamic.dylib - install_name_tool -change @rpath/libc10.dylib $dev/lib/libc10.dylib $dev/lib/libcaffe2_module_test_dynamic.dylib + install_name_tool -change @rpath/libtorch.dylib $lib/lib/libtorch.dylib $lib/lib/libcaffe2_module_test_dynamic.dylib + install_name_tool -change @rpath/libc10.dylib $lib/lib/libc10.dylib $lib/lib/libcaffe2_module_test_dynamic.dylib - install_name_tool -change @rpath/libtorch.dylib $dev/lib/libtorch.dylib $dev/lib/libcaffe2_detectron_ops.dylib - install_name_tool -change @rpath/libc10.dylib $dev/lib/libc10.dylib $dev/lib/libcaffe2_detectron_ops.dylib + install_name_tool -change @rpath/libtorch.dylib $lib/lib/libtorch.dylib $lib/lib/libcaffe2_detectron_ops.dylib + install_name_tool -change @rpath/libc10.dylib $lib/lib/libc10.dylib $lib/lib/libcaffe2_detectron_ops.dylib - install_name_tool -change @rpath/libtorch.dylib $dev/lib/libtorch.dylib $dev/lib/libshm.dylib - install_name_tool -change @rpath/libc10.dylib $dev/lib/libc10.dylib $dev/lib/libshm.dylib + install_name_tool -change @rpath/libtorch.dylib $lib/lib/libtorch.dylib $lib/lib/libshm.dylib + install_name_tool -change @rpath/libc10.dylib $lib/lib/libc10.dylib $lib/lib/libshm.dylib ''; |