python3Packages.pytorch: fix AArch64 build

aarch64-linux builds fail because of the use of opcodes in QNNPACK
that the GNU assembpler does not support. This change cherry-picks an
upstream patch that fixes this. See:

https://github.com/pytorch/pytorch/issues/33124
https://github.com/pytorch/pytorch/pull/40584
This commit is contained in:
Daniël de Kok 2020-06-30 07:56:38 +02:00
parent d144f2dfda
commit 5c80cdeb1a

View file

@ -121,6 +121,21 @@ in buildPythonPackage rec {
sha256 = "1xjbn4hi96m7xslv3p2jc6qcsng0fx3w1m6isqfah81piljf8wng";
};
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.6.0).
(fetchpatch {
name = "qnnpack-neon-fix.patch";
url = "https://github.com/pytorch/pytorch/commit/7676682584d0caf9243bce74ea0a88711ec4a807.diff";
sha256 = "13spncaqlpsp8qk2850yly7xqwmhhfwznhmzkk8jgpslkbx75vgq";
})
];
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++