From 6f96122d5b1f2a56d945ce537e87dd432b507eb4 Mon Sep 17 00:00:00 2001
From: Nikolay Amiantov <ab@fmap.me>
Date: Sun, 25 Feb 2018 20:02:10 +0300
Subject: [PATCH] caffe: add nccl support

---
 pkgs/applications/science/math/caffe/default.nix | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/pkgs/applications/science/math/caffe/default.nix b/pkgs/applications/science/math/caffe/default.nix
index dc9c5f7359ec..d7357c5048d5 100644
--- a/pkgs/applications/science/math/caffe/default.nix
+++ b/pkgs/applications/science/math/caffe/default.nix
@@ -14,10 +14,12 @@
 , openblas
 , cudaSupport ? true, cudatoolkit
 , cudnnSupport ? false, cudnn ? null
+, ncclSupport ? false, nccl ? null
 , pythonSupport ? false, python ? null, numpy ? null
 }:
 
 assert cudnnSupport -> cudaSupport;
+assert ncclSupport -> cudaSupport;
 assert pythonSupport -> (python != null && numpy != null);
 
 stdenv.mkDerivation rec {
@@ -41,11 +43,13 @@ stdenv.mkDerivation rec {
     ] ++ (if cudaSupport then [
            "-DCUDA_ARCH_NAME=All"
            "-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 ]
                 ++ lib.optional cudaSupport cudatoolkit
                 ++ lib.optional cudnnSupport cudnn
+                ++ lib.optional ncclSupport nccl
                 ++ lib.optionals pythonSupport [ python numpy ];
 
   propagatedBuildInputs = lib.optional pythonSupport python.pkgs.protobuf;