python3Packages.pytorch: convert boolean environment variables

We relied on the stdenv serialization for booleans. However, the
PyTorch configure script expects 0/1. So convert the options to 0/1
values.
This commit is contained in:
Daniël de Kok 2021-04-10 08:08:08 +02:00
parent 7b02acac3e
commit 4f876309fe

View file

@ -38,6 +38,7 @@ assert !(MPISupport && cudaSupport) || mpi.cudatoolkit == cudatoolkit;
assert !cudaSupport || magma.cudatoolkit == cudatoolkit;
let
setBool = v: if v then "1" else "0";
cudatoolkit_joined = symlinkJoin {
name = "${cudatoolkit.name}-unsplit";
# nccl is here purely for semantic grouping it could be moved to nativeBuildInputs
@ -160,16 +161,16 @@ in buildPythonPackage rec {
# Use pytorch's custom configurations
dontUseCmakeConfigure = true;
BUILD_NAMEDTENSOR = true;
BUILD_DOCS = buildDocs;
BUILD_NAMEDTENSOR = setBool true;
BUILD_DOCS = setBool buildDocs;
USE_MKL = blas.implementation == "mkl";
USE_MKL = setBool (blas.implementation == "mkl");
# 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;
USE_MKLDNN = setBool mklDnnSupport;
USE_MKLDNN_CBLAS = setBool mklDnnSupport;
preBuild = ''
export MAX_JOBS=$NIX_BUILD_CORES
@ -198,7 +199,7 @@ in buildPythonPackage rec {
PYTORCH_BUILD_VERSION = version;
PYTORCH_BUILD_NUMBER = 0;
USE_SYSTEM_NCCL=useSystemNccl; # don't build pytorch's third_party NCCL
USE_SYSTEM_NCCL=setBool useSystemNccl; # don't build pytorch's third_party NCCL
# Suppress a weird warning in mkl-dnn, part of ideep in pytorch
# (upstream seems to have fixed this in the wrong place?)