1 files changed, 143 insertions, 0 deletions
diff --git a/nixpkgs/pkgs/development/libraries/science/math/caffe2/default.nix b/nixpkgs/pkgs/development/libraries/science/math/caffe2/default.nix
new file mode 100644
index 000000000000..fe9a8b671459
--- /dev/null
+++ b/nixpkgs/pkgs/development/libraries/science/math/caffe2/default.nix
@@ -0,0 +1,143 @@
+{ stdenv, lib, config, fetchFromGitHub
+, cmake
+, glog, gflags, gtest
+, protobuf, snappy
+, python, future, six, python-protobuf, numpy, pydot
+, eigen
+, doxygen
+, useCuda ? (config.cudaSupport or false), cudatoolkit ? null
+, useCudnn ? (config.cudnnSupport or false), cudnn ? null
+, useOpenmp ? false, openmp ? null
+, useOpencv3 ? true, opencv3 ? null
+, useLeveldb ? false, leveldb ? null
+, useLmdb ? true, lmdb ? null
+, useRocksdb ? false, rocksdb ? null
+, useZeromq ? false, zeromq ? null
+, useMpi ? false, mpi ? null
+# TODO: distributed computations
+#, useGloo ? false
+#, useNccl ? false
+#, useNnpack ? false
+}:
+
+assert useCuda -> cudatoolkit != null;
+assert useCudnn -> (useCuda && cudnn != null);
+assert useOpencv3 -> opencv3 != null;
+assert useLeveldb -> leveldb != null;
+assert useLmdb -> lmdb != null;
+assert useRocksdb -> rocksdb != null;
+assert useZeromq -> zeromq != null;
+assert useMpi -> mpi != null;
+
+let
+  # Third party modules that caffe2 holds as git submodules.
+  # Download them and create symlinks from caffe2/third_party.
+  installExtraSrc = extra: ''
+    rmdir "third_party/${extra.dst}"
+    ln -s "${extra.src}" "third_party/${extra.dst}"
+  '';
+
+  cub = {
+    src = fetchFromGitHub {
+      owner  = "NVlabs";
+      repo   = "cub";
+      rev    = "v1.7.4";
+      sha256 = "0ksd5n1lxqhm5l5cd2lps4cszhjkf6gmzahaycs7nxb06qci8c66";
+    };
+    dst = "cub";
+  };
+
+  pybind11 = {
+    src = fetchFromGitHub {
+      owner  = "pybind";
+      repo   = "pybind11";
+      rev    = "86e2ad4f77442c3350f9a2476650da6bee253c52";
+      sha256 = "05gi58dirvc8fgm0avpydvidzsbh2zrzgfaq671ym09f6dz0bcgz";
+    };
+    dst = "pybind11";
+  };
+
+  ccVersion = lib.getVersion stdenv.cc;
+in
+
+stdenv.mkDerivation rec {
+  pname = "caffe2";
+  version = "0.8.1";
+  src = fetchFromGitHub {
+    owner = "caffe2";
+    repo = "caffe2";
+    rev = "v${version}";
+    sha256 = "18y7zjc69j6n5642l9caddl641b0djf3pjn4wacdsc1wk1jiyqk8";
+  };
+
+  nativeBuildInputs = [ cmake doxygen gtest ];
+  outputs = [ "bin" "out" ];
+  propagatedBuildOutputs = [ ]; # otherwise propagates out -> bin cycle
+
+  buildInputs = [ glog gflags protobuf snappy eigen ]
+    ++ lib.optional useCuda cudatoolkit
+    ++ lib.optional useCudnn cudnn
+    ++ lib.optional useOpenmp openmp
+    ++ lib.optional useOpencv3 opencv3
+    ++ lib.optional useLeveldb leveldb
+    ++ lib.optional useLmdb lmdb
+    ++ lib.optional useRocksdb rocksdb
+    ++ lib.optional useZeromq zeromq
+  ;
+  propagatedBuildInputs = [ numpy future six python-protobuf pydot ];
+
+  patches = lib.optional (stdenv.cc.isGNU && lib.versionAtLeast ccVersion "7.0.0") [
+    ./fix_compilation_on_gcc7.patch
+  ] ++ lib.optional stdenv.cc.isClang [ ./update_clang_cvtsh_bugfix.patch ];
+
+  cmakeFlags = [ "-DBUILD_TEST=OFF"
+                 "-DBUILD_PYTHON=ON"
+                 ''-DUSE_CUDA=${if useCuda then "ON"else "OFF"}''
+                 ''-DUSE_OPENMP=${if useOpenmp then "ON"else "OFF"}''
+                 ''-DUSE_OPENCV=${if useOpencv3 then "ON"else "OFF"}''
+                 ''-DUSE_MPI=${if useMpi then "ON"else "OFF"}''
+                 ''-DUSE_LEVELDB=${if useLeveldb then "ON"else "OFF"}''
+                 ''-DUSE_LMDB=${if useLmdb then "ON"else "OFF"}''
+                 ''-DUSE_ROCKSDB=${if useRocksdb then "ON"else "OFF"}''
+                 ''-DUSE_ZMQ=${if useZeromq  then "ON"else "OFF"}''
+                 "-DUSE_GLOO=OFF"
+                 "-DUSE_NNPACK=OFF"
+                 "-DUSE_NCCL=OFF"
+                 "-DUSE_REDIS=OFF"
+                 "-DUSE_FFMPEG=OFF"
+               ]
+               ++ lib.optional useCuda [
+                 "-DCUDA_TOOLKIT_ROOT_DIR=${cudatoolkit}"
+                 "-DCUDA_FAST_MATH=ON"
+                 "-DCUDA_HOST_COMPILER=${cudatoolkit.cc}/bin/gcc"
+               ];
+
+  preConfigure = ''
+    ${installExtraSrc cub}
+    ${installExtraSrc pybind11}
+    # XXX hack
+    export NIX_CFLAGS_COMPILE="-I ${eigen}/include/eigen3/ $NIX_CFLAGS_COMPILE"
+  '';
+
+  postInstall = ''
+    moveToOutput "bin" "$bin"
+    mkdir -p $out/lib/${python.libPrefix}
+    ln -s $out/ $out/${python.sitePackages}
+  '';
+
+  doCheck = false;
+
+  meta = {
+    homepage = "https://caffe2.ai/";
+    description = "A new lightweight, modular, and scalable deep learning framework";
+    longDescription = ''
+      Caffe2 aims to provide an easy and straightforward way for you to experiment
+      with deep learning and leverage community contributions of new models and
+      algorithms. You can bring your creations to scale using the power of GPUs in the
+      cloud or to the masses on mobile with Caffe2's cross-platform libraries.
+    '';
+    platforms = with lib.platforms; linux;
+    license = lib.licenses.asl20;
+    maintainers = with lib.maintainers; [ yuriaisaka ];
+  };
+}