colin/nixpkgs
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

caffe: add nccl support

Browse Source
This commit is contained in:
Nikolay Amiantov 2018-02-25 20:02:10 +03:00
parent 25079ad238
commit 6f96122d5b
1 changed files with 5 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
pkgs/applications/science/math/caffe/default.nix
View File

@ -14,10 +14,12 @@
, openblas , openblas
, cudaSupport ? true, cudatoolkit , cudaSupport ? true, cudatoolkit
, cudnnSupport ? false, cudnn ? null , cudnnSupport ? false, cudnn ? null
, ncclSupport ? false, nccl ? null
, pythonSupport ? false, python ? null, numpy ? null , pythonSupport ? false, python ? null, numpy ? null
}: }:
assert cudnnSupport -> cudaSupport; assert cudnnSupport -> cudaSupport;
assert ncclSupport -> cudaSupport;
assert pythonSupport -> (python != null && numpy != null); assert pythonSupport -> (python != null && numpy != null);
stdenv.mkDerivation rec { stdenv.mkDerivation rec {
@ -41,11 +43,13 @@ stdenv.mkDerivation rec {
] ++ (if cudaSupport then [ ] ++ (if cudaSupport then [
"-DCUDA_ARCH_NAME=All" "-DCUDA_ARCH_NAME=All"
"-DCUDA_HOST_COMPILER=${cudatoolkit.cc}/bin/cc" "-DCUDA_HOST_COMPILER=${cudatoolkit.cc}/bin/cc"
] else [ "-DCPU_ONLY=ON" ]); ] else [ "-DCPU_ONLY=ON" ])
++ lib.optional ncclSupport "-DUSE_NCCL=ON";
buildInputs = [ boost google-gflags glog protobuf hdf5-cpp lmdb leveldb snappy opencv3 openblas ] buildInputs = [ boost google-gflags glog protobuf hdf5-cpp lmdb leveldb snappy opencv3 openblas ]
++ lib.optional cudaSupport cudatoolkit ++ lib.optional cudaSupport cudatoolkit
++ lib.optional cudnnSupport cudnn ++ lib.optional cudnnSupport cudnn
++ lib.optional ncclSupport nccl
++ lib.optionals pythonSupport [ python numpy ]; ++ lib.optionals pythonSupport [ python numpy ];
propagatedBuildInputs = lib.optional pythonSupport python.pkgs.protobuf; propagatedBuildInputs = lib.optional pythonSupport python.pkgs.protobuf;