Merge pull request #96488 from danieldk/libtorch

libtorch-bin: init at 1.6.0

pkgs/development/libraries/science/math/libtorch/bin.nix

@@ -0,0 +1,111 @@
+{ callPackage
+, stdenv
+, fetchzip
+, lib
+
+, addOpenGLRunpath
+, patchelf
+, fixDarwinDylibNames
+
+, cudaSupport
+, nvidia_x11
+}:
+
+let
+  version = "1.6.0";
+  device = if cudaSupport then "cuda" else "cpu";
+  srcs = import ./binary-hashes.nix;
+  unavailable = throw "libtorch is not available for this platform";
+in stdenv.mkDerivation {
+  inherit version;
+  pname = "libtorch";
+
+  src = fetchzip srcs."${stdenv.targetPlatform.system}-${device}" or unavailable;
+
+  nativeBuildInputs =
+    if stdenv.isDarwin then [ fixDarwinDylibNames ]
+    else [ addOpenGLRunpath patchelf ]
+      ++ stdenv.lib.optionals cudaSupport [ addOpenGLRunpath ];
+
+  buildInputs = [
+    stdenv.cc.cc
+  ] ++ lib.optionals cudaSupport [ nvidia_x11 ];
+
+  dontBuild = true;
+  dontConfigure = true;
+  dontStrip = true;
+
+  installPhase = ''
+    # Copy headers and CMake files.
+    install -Dm755 -t $dev/lib lib/*.a
+    cp -r include $dev
+    cp -r share $dev
+
+    install -Dm755 -t $out/lib lib/*${stdenv.hostPlatform.extensions.sharedLibrary}*
+
+    # We do not care about Java support...
+    rm -f $out/lib/lib*jni* 2> /dev/null || true
+  '';
+
+  postFixup = let
+    libPaths = [ stdenv.cc.cc.lib ]
+      ++ stdenv.lib.optionals cudaSupport [ nvidia_x11 ];
+    rpath = stdenv.lib.makeLibraryPath libPaths;
+  in stdenv.lib.optionalString stdenv.isLinux ''
+    find $out/lib -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
+      echo "setting rpath for $lib..."
+      patchelf --set-rpath "${rpath}:$out/lib" "$lib"
+      ${lib.optionalString cudaSupport ''
+        addOpenGLRunpath "$lib"
+      ''}
+    done
+  '' + stdenv.lib.optionalString stdenv.isDarwin ''
+    install_name_tool -change @rpath/libshm.dylib $out/lib/libshm.dylib $out/lib/libtorch_python.dylib
+    install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libtorch_python.dylib
+    install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libtorch_python.dylib
+    install_name_tool -change @rpath/libtorch.dylib $out/lib/libtorch.dylib $out/lib/libtorch_python.dylib
+    install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libtorch_python.dylib
+
+    install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libtorch.dylib
+    install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libtorch.dylib
+    install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libtorch.dylib
+
+    install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libtorch_cpu.dylib
+    install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libtorch_cpu.dylib
+    install_name_tool -change @rpath/libtensorpipe.dylib $out/lib/libtensorpipe.dylib $out/lib/libtorch_cpu.dylib
+
+    install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libcaffe2_observers.dylib
+    install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libcaffe2_observers.dylib
+    install_name_tool -change @rpath/libtorch.dylib $out/lib/libtorch.dylib $out/lib/libcaffe2_observers.dylib
+    install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libcaffe2_observers.dylib
+
+    install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libcaffe2_module_test_dynamic.dylib
+    install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libcaffe2_module_test_dynamic.dylib
+    install_name_tool -change @rpath/libtorch.dylib $out/lib/libtorch.dylib $out/lib/libcaffe2_module_test_dynamic.dylib
+    install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libcaffe2_module_test_dynamic.dylib
+
+    install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libcaffe2_detectron_ops.dylib
+    install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libcaffe2_detectron_ops.dylib
+    install_name_tool -change @rpath/libtorch.dylib $out/lib/libtorch.dylib $out/lib/libcaffe2_detectron_ops.dylib
+    install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libcaffe2_detectron_ops.dylib
+
+    install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libshm.dylib
+    install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libshm.dylib
+    install_name_tool -change @rpath/libtorch.dylib $out/lib/libtorch.dylib $out/lib/libshm.dylib
+    install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libshm.dylib
+
+    install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libtorch_global_deps.dylib
+    install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libtorch_global_deps.dylib
+  '';
+
+  outputs = [ "out" "dev" ];
+
+  passthru.tests = callPackage ./test { };
+
+  meta = with stdenv.lib; {
+    description = "C++ API of the PyTorch machine learning framework";
+    homepage = "https://pytorch.org/";
+    license = licenses.unfree; # Includes CUDA and Intel MKL.
+    platforms = with platforms; linux ++ darwin;
+  };
+}

pkgs/development/libraries/science/math/libtorch/binary-hashes.nix

@@ -0,0 +1,14 @@
+{
+  x86_64-darwin-cpu = {
+    url = "https://download.pytorch.org/libtorch/cpu/libtorch-macos-1.6.0.zip";
+    sha256 = "0d4n7la31qzl4s9pwvm07la7q6lhcwiww0yjpfz3kw6nvx84p22r";
+  };
+  x86_64-linux-cpu = {
+    url = "https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-1.6.0%2Bcpu.zip";
+    sha256 = "1975b4zvyihzh89vnwspw0vf9qr05sxj8939vcrlmv3gzvdspcxz";
+  };
+  x86_64-linux-cuda = {
+    url = "https://download.pytorch.org/libtorch/cu102/libtorch-cxx11-abi-shared-with-deps-1.6.0.zip";
+    sha256 = "127qnfyi1faqbm40sbnsyqxjhrqj82bzwqyz7c1hs2bm0zgrrpya";
+  };
+}

pkgs/development/libraries/science/math/libtorch/test/CMakeLists.txt

@@ -0,0 +1,4 @@
+find_package(Torch REQUIRED)
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
+add_executable(test test.cpp)
+target_link_libraries(test "${TORCH_LIBRARIES}")

pkgs/development/libraries/science/math/libtorch/test/default.nix

@@ -0,0 +1,26 @@
+{ stdenv, cmake, libtorch-bin, symlinkJoin }:
+
+stdenv.mkDerivation {
+  pname = "libtorch-test";
+  version = libtorch-bin.version;
+
+  src = ./.;
+
+  postPatch = ''
+    cat CMakeLists.txt
+  '';
+
+  makeFlags = [ "VERBOSE=1" ];
+
+  nativeBuildInputs = [ cmake ];
+
+  buildInputs = [ libtorch-bin ];
+
+  installPhase = ''
+    touch $out
+  '';
+
+  checkPhase = ''
+    ./test
+  '';
+}

pkgs/development/libraries/science/math/libtorch/test/test.cpp

@@ -0,0 +1,7 @@
+#include <torch/torch.h>
+#include <iostream>
+
+int main() {
+  torch::Tensor tensor = torch::eye(3);
+  std::cout << tensor << std::endl;
+}

pkgs/top-level/all-packages.nix

@@ -1372,6 +1372,11 @@ in
     then python.pkgs.tensorflow.libtensorflow
     else libtensorflow-bin;
 
+  libtorch-bin = callPackage ../development/libraries/science/math/libtorch/bin.nix {
+    inherit (linuxPackages) nvidia_x11;
+    cudaSupport = pkgs.config.cudaSupport or false;
+  };
+
   behdad-fonts = callPackage ../data/fonts/behdad-fonts { };
 
   bless = callPackage ../applications/editors/bless { };