NixOS · Oct 23, 2019 · Oct 23, 2019 · Oct 23, 2019 · Oct 23, 2019
Showing with 184 additions and 30 deletions.

+5 −0 maintainers/maintainer-list.nix

+20 −2 pkgs/development/libraries/openmpi/default.nix

+4 −0 pkgs/development/libraries/science/math/magma/default.nix

+155 −28 pkgs/development/python-modules/pytorch/default.nix
diff --git a/maintainers/maintainer-list.nix b/maintainers/maintainer-list.nix
@@ -6708,6 +6708,11 @@
     githubId = 483735;
     name = "Dmitry Geurkov";
   };
+  tscholak = {
+    email = "torsten.scholak@googlemail.com";
+    github = "tscholak";
+    name = "Torsten Scholak";
+  };
   tstrobel = {
     email = "4ZKTUB6TEP74PYJOPWIR013S2AV29YUBW5F9ZH2F4D5UMJUJ6S@hash.domains";
     name = "Thomas Strobel";

diff --git a/pkgs/development/libraries/openmpi/default.nix b/pkgs/development/libraries/openmpi/default.nix
@@ -1,5 +1,8 @@
 { stdenv, fetchurl, fetchpatch, gfortran, perl, libnl
-, rdma-core, zlib, numactl, libevent, hwloc, pkgsTargetTarget
+, rdma-core, zlib, numactl, libevent, hwloc, pkgsTargetTarget, symlinkJoin
+
+# Enable CUDA support
+, cudaSupport ? false, cudatoolkit ? null
 
 # Enable the Sun Grid Engine bindings
 , enableSGE ? false
@@ -8,9 +11,15 @@
 , enablePrefix ? false
 }:
 
+assert !cudaSupport || cudatoolkit != null;
+
 let
   version = "4.0.2";
 
+  cudatoolkit_joined = symlinkJoin {
+    name = "${cudatoolkit.name}-unsplit";
+    paths = [ cudatoolkit.out cudatoolkit.lib ];
+  };
 in stdenv.mkDerivation rec {
   pname = "openmpi";
   inherit version;
@@ -33,15 +42,20 @@ in stdenv.mkDerivation rec {
 
   buildInputs = with stdenv; [ gfortran zlib ]
     ++ lib.optionals isLinux [ libnl numactl ]
+    ++ lib.optionals cudaSupport [ cudatoolkit ]
     ++ [ libevent hwloc ]
     ++ lib.optional (isLinux || isFreeBSD) rdma-core;
 
   nativeBuildInputs = [ perl ];
 
-  configureFlags = with stdenv; [ "--disable-mca-dso" ]
+  configureFlags = with stdenv; lib.optional (!cudaSupport) "--disable-mca-dso"
     ++ lib.optional isLinux  "--with-libnl=${libnl.dev}"
     ++ lib.optional enableSGE "--with-sge"
     ++ lib.optional enablePrefix "--enable-mpirun-prefix-by-default"
+    # TODO: add UCX support, which is recommended to use with cuda for the most robust OpenMPI build
+    # https://github.com/openucx/ucx
+    # https://www.open-mpi.org/faq/?category=buildcuda
+    ++ lib.optionals cudaSupport [ "--with-cuda=${cudatoolkit_joined}" "--enable-dlopen" ]
     ;
 
   enableParallelBuilding = true;
@@ -69,6 +83,10 @@ in stdenv.mkDerivation rec {
 
   doCheck = true;
 
+  passthru = {
+    inherit cudaSupport cudatoolkit;
+  };
+
   meta = with stdenv.lib; {
     homepage = https://www.open-mpi.org/;
     description = "Open source MPI-3 implementation";

diff --git a/pkgs/development/libraries/science/math/magma/default.nix b/pkgs/development/libraries/science/math/magma/default.nix
@@ -24,6 +24,10 @@ in stdenv.mkDerivation {
 
   MKLROOT = optionalString mklSupport mkl;
 
+  preConfigure = ''
+    export CC=${cudatoolkit.cc}/bin/gcc CXX=${cudatoolkit.cc}/bin/g++
+  '';
+
   enableParallelBuilding=true;
   buildFlags = [ "magma" "magma_sparse" ];
 

diff --git a/pkgs/development/python-modules/pytorch/default.nix b/pkgs/development/python-modules/pytorch/default.nix
@@ -1,18 +1,101 @@
-{ stdenv, fetchurl, buildPythonPackage, pythonOlder,
-  cudaSupport ? false, cudatoolkit ? null, cudnn ? null,
-  fetchFromGitHub, lib, numpy, pyyaml, cffi, typing, cmake, hypothesis, numactl,
+{ stdenv, fetchurl, fetchgit, buildPythonPackage, python, pythonOlder,
+  cudaSupport ? false, cudatoolkit ? null, cudnn ? null, nccl ? null, magma ? null,
+  mklSupport ? false, mkl ? null,
+  openMPISupport ? false, openmpi ? null,
+  buildNamedTensor ? false,
+  buildBinaries ? false,
+  cudaArchList ? null,
+  fetchFromGitHub, lib, numpy, pyyaml, cffi, click, typing, cmake, hypothesis, numactl,
   linkFarm, symlinkJoin,
+
+  # ninja (https://ninja-build.org) must be available to run C++ extensions tests,
+  ninja,
+
+  # dependencies for torch.utils.tensorboard
+  tensorboardSupport ? true, pillow, six, future, tensorflow-tensorboard,
+
   utillinux, which, isPy3k }:
 
-assert cudnn == null || cudatoolkit != null;
+assert !openMPISupport || openmpi != null;
+assert !tensorboardSupport || tensorflow-tensorboard != null;
+
+# assert that everything needed for cuda is present and that the correct cuda versions are used
 assert !cudaSupport || cudatoolkit != null;
+assert cudnn == null || cudatoolkit != null;
+assert !cudaSupport || (let majorIs = lib.versions.major cudatoolkit.version;
+                        in majorIs == "9" || majorIs == "10");
+
+let
+  hasDependency = dep: pkg: lib.lists.any (inp: inp == dep) pkg.buildInputs;
+  matchesCudatoolkit = hasDependency cudatoolkit;
+  matchesMkl = hasDependency mkl;
+in
+# confirm that cudatoolkits are sync'd across dependencies
+assert !(openMPISupport && cudaSupport) || matchesCudatoolkit openmpi;
+assert !cudaSupport || matchesCudatoolkit magma;
+
+# confirm that mkl is sync'd across dependencies
+assert !mklSupport || mkl != null;
+assert !(mklSupport && cudaSupport) || matchesMkl magma;
+assert !mklSupport || (numpy.blasImplementation == "mkl" && numpy.blas == mkl);
 
 let
   cudatoolkit_joined = symlinkJoin {
     name = "${cudatoolkit.name}-unsplit";
-    paths = [ cudatoolkit.out cudatoolkit.lib ];
+    # nccl is here purely for semantic grouping it could be moved to nativeBuildInputs
+    paths = [ cudatoolkit.out cudatoolkit.lib nccl.dev nccl.out ];
   };
 
+  # Give an explicit list of supported architectures for the build, See:
+  # - pytorch bug report: https://github.com/pytorch/pytorch/issues/23573
+  # - pytorch-1.2.0 build on nixpks: https://github.com/NixOS/nixpkgs/pull/65041
+  #
+  # This list was selected by omitting the TORCH_CUDA_ARCH_LIST parameter,
+  # observing the fallback option (which selected all architectures known
+  # from cudatoolkit_10_0, pytorch-1.2, and python-3.6), and doing a binary
+  # searching to find offending architectures.
+  #
+  # NOTE: Because of sandboxing, this derivation can't auto-detect the hardware's
+  # cuda architecture, so there is also now a problem around new architectures
+  # not being supported until explicitly added to this derivation.
+  #
+  # FIXME: CMake is throwing the following warning on python-1.2:
+  #
+  # ```
+  # CMake Warning at cmake/public/utils.cmake:172 (message):
+  #   In the future we will require one to explicitly pass TORCH_CUDA_ARCH_LIST
+  #   to cmake instead of implicitly setting it as an env variable.  This will
+  #   become a FATAL_ERROR in future version of pytorch.
+  # ```
+  # If this is causing problems for your build, this derivation may have to strip
+  # away the standard `buildPythonPackage` and use the
+  # [*Adjust Build Options*](https://github.com/pytorch/pytorch/tree/v1.2.0#adjust-build-options-optional)
+  # instructions. This will also add more flexibility around configurations
+  # (allowing FBGEMM to be built in pytorch-1.1), and may future proof this
+  # derivation.
+  brokenArchs = [ "3.0" ]; # this variable is only used as documentation.
+  cuda9ArchList = [
+    "3.5"
+    "5.0"
+    "5.2"
+    "6.0"
+    "6.1"
+    "7.0"
+    "7.0+PTX"  # I am getting a "undefined architecture compute_75" on cuda 9
+               # which leads me to believe this is the final cuda-9-compatible architecture.
+  ];
+  cuda10ArchList = cuda9ArchList ++ [
+    "7.5"
+    "7.5+PTX"  # < most recent architecture as of cudatoolkit_10_0 and pytorch-1.2.0
+  ];
+  final_cudaArchList =
+    if !cudaSupport || cudaArchList != null
+    then cudaArchList
+    else
+      if lib.versions.major cudatoolkit.version == "9"
+      then cuda9ArchList
+      else cuda10ArchList; # the assert above removes any ambiguity here.
+
   # Normally libcuda.so.1 is provided at runtime by nvidia-x11 via
   # LD_LIBRARY_PATH=/run/opengl-driver/lib.  We only use the stub
   # libcuda.so from cudatoolkit for running tests, so that we don’t have
@@ -25,28 +108,27 @@ let
     "LD_LIBRARY_PATH=${cudaStub}\${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH} ";
 
 in buildPythonPackage rec {
-  version = "1.0.0";
+  version = "1.2.0";
   pname = "pytorch";
   disabled = !isPy3k;
 
+  outputs = [
+    "out"   # output standard python package
+    "dev"   # output libtorch only
+  ];
+
   src = fetchFromGitHub {
     owner  = "pytorch";
     repo   = "pytorch";
     rev    = "v${version}";
     fetchSubmodules = true;
-    sha256 = "076cpbig4sywn9vv674c0xdg832sdrd5pk1d0725pjkm436kpvlm";
+    sha256 = "1biyq2p48chakf2xw7hazzqmr5ps1nx475ql8vkmxjg5zaa071cz";
   };
 
-  patches =
-    [ # Skips two tests that are only meant to run on multi GPUs
-      (fetchurl {
-        url = "https://github.com/pytorch/pytorch/commit/bfa666eb0deebac21b03486e26642fd70d66e478.patch";
-        sha256 = "1fgblcj02gjc0y62svwc5gnml879q3x2z7m69c9gax79dpr37s9i";
-      })
-    ];
-
   dontUseCmakeConfigure = true;
+
   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++
   '' + lib.optionalString (cudaSupport && cudnn != null) ''
     export CUDNN_INCLUDE_DIR=${cudnn}/include
@@ -61,7 +143,6 @@ in buildPythonPackage rec {
       RP_NEW=$(join_by : ''${RP[@]:2})
       patchelf --set-rpath \$ORIGIN:''${RP_NEW} "$1"
     }
-
     for f in $(find ''${out} -name 'libcaffe2*.so')
     do
       strip2 $f
@@ -74,38 +155,84 @@ in buildPythonPackage rec {
   PYTORCH_BUILD_VERSION = version;
   PYTORCH_BUILD_NUMBER = 0;
 
+  BUILD_NAMEDTENSOR = buildNamedTensor;  # experimental feature
+  USE_SYSTEM_NCCL=true;                  # 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?)
   # https://github.com/intel/mkl-dnn/commit/8134d346cdb7fe1695a2aa55771071d455fae0bc
-  NIX_CFLAGS_COMPILE = lib.optionals (numpy.blasImplementation == "mkl") [ "-Wno-error=array-bounds" ];
+  # https://github.com/pytorch/pytorch/issues/22346
+  #
+  # Also of interest: pytorch ignores CXXFLAGS uses CFLAGS for both C and C++:
+  # https://github.com/pytorch/pytorch/blob/v1.2.0/setup.py#L17
+  NIX_CFLAGS_COMPILE = lib.optionals (numpy.blas == mkl) [ "-Wno-error=array-bounds" ];
 
   nativeBuildInputs = [
-     cmake
-     utillinux
-     which
+    cmake
+    utillinux
+    which
+    ninja
   ] ++ lib.optionals cudaSupport [ cudatoolkit_joined ];
 
   buildInputs = [
-     numpy.blas
-  ] ++ lib.optionals cudaSupport [ cudnn ]
+    numpy.blas
+  ] ++ lib.optionals cudaSupport [ cudnn magma nccl ]
     ++ lib.optionals stdenv.isLinux [ numactl ];
 
   propagatedBuildInputs = [
     cffi
+    click
     numpy
     pyyaml
-  ] ++ lib.optional (pythonOlder "3.5") typing;
+  ] ++ lib.optionals openMPISupport [ openmpi ]
+    ++ lib.optional (pythonOlder "3.5") typing
+    ++ lib.optionals tensorboardSupport [pillow six future tensorflow-tensorboard];
+
+  checkInputs = [ hypothesis ninja ];
+
+  doCheck = false; # tests take a long time for channel release, so doCheck should be overridden only when developing
+  checkPhase = "${cudaStubEnv}python test/run_test.py"
+    + " --exclude utils" # utils requires git, which is not allowed in the check phase
+
+    # Other tests which have been disabled in previous nix derivations of pytorch.
+    # --exclude dataloader sparse torch utils thd_distributed distributed cpp_extensions
+    ;
+  postInstall = ''
+    mkdir $dev
+    cp -r $out/${python.sitePackages}/torch/lib     $dev/lib
+    cp -r $out/${python.sitePackages}/torch/include $dev/include
+  '';
+
+  postFixup = stdenv.lib.optionalString stdenv.isDarwin ''
+    for f in $(ls $dev/lib/*.dylib); do
+        install_name_tool -id $dev/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
 
-  checkInputs = [ hypothesis ];
-  checkPhase = ''
-    ${cudaStubEnv}python test/run_test.py --exclude dataloader sparse torch utils thd_distributed distributed cpp_extensions
+    install_name_tool -change @rpath/libc10.dylib $dev/lib/libc10.dylib $dev/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 $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 $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 $dev/lib/libtorch.dylib $dev/lib/libshm.dylib
+    install_name_tool -change @rpath/libc10.dylib $dev/lib/libc10.dylib $dev/lib/libshm.dylib
   '';
 
+
   meta = {
     description = "Open source, prototype-to-production deep learning platform";
     homepage    = https://pytorch.org/;
     license     = lib.licenses.bsd3;
-    platforms   = lib.platforms.linux;
-    maintainers = with lib.maintainers; [ teh thoughtpolice ];
+    platforms   = with lib.platforms; linux ++ lib.optionals (!cudaSupport) darwin;
+    maintainers = with lib.maintainers; [ teh thoughtpolice stites tscholak ]; # tscholak esp. for darwin-related builds
   };
 }