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

   1 { callPackage
   2 , stdenv
   3 , fetchzip
   4 , lib
   5
   6 , addOpenGLRunpath
   7 , patchelf
   8 , fixDarwinDylibNames
   9
  10 , cudaSupport
  11 , cudatoolkit_11_1
  12 , cudnn_cudatoolkit_11_1
  13 }:
  14
  15 let
  16   # The binary libtorch distribution statically links the CUDA
  17   # toolkit. This means that we do not need to provide CUDA to
  18   # this derivation. However, we should ensure on version bumps
  19   # that the CUDA toolkit for `passthru.tests` is still
  20   # up-to-date.
  21   version = "1.8.0";
  22   device = if cudaSupport then "cuda" else "cpu";
  23   srcs = import ./binary-hashes.nix version;
  24   unavailable = throw "libtorch is not available for this platform";
  25 in stdenv.mkDerivation {
  26   inherit version;
  27   pname = "libtorch";
  28
  29   src = fetchzip srcs."${stdenv.targetPlatform.system}-${device}" or unavailable;
  30
  31   nativeBuildInputs =
  32     if stdenv.isDarwin then [ fixDarwinDylibNames ]
  33     else [ patchelf ] ++ lib.optionals cudaSupport [ addOpenGLRunpath ];
  34
  35   dontBuild = true;
  36   dontConfigure = true;
  37   dontStrip = true;
  38
  39   installPhase = ''
  40     # Copy headers and CMake files.
  41     mkdir -p $dev
  42     cp -r include $dev
  43     cp -r share $dev
  44
  45     install -Dm755 -t $out/lib lib/*${stdenv.hostPlatform.extensions.sharedLibrary}*
  46
  47     # We do not care about Java support...
  48     rm -f $out/lib/lib*jni* 2> /dev/null || true
  49
  50     # Fix up library paths for split outputs
  51     substituteInPlace $dev/share/cmake/Torch/TorchConfig.cmake \
  52       --replace \''${TORCH_INSTALL_PREFIX}/lib "$out/lib" \
  53
  54     substituteInPlace \
  55       $dev/share/cmake/Caffe2/Caffe2Targets-release.cmake \
  56       --replace \''${_IMPORT_PREFIX}/lib "$out/lib" \
  57   '';
  58
  59   postFixup = let
  60     rpath = lib.makeLibraryPath [ stdenv.cc.cc.lib ];
  61   in lib.optionalString stdenv.isLinux ''
  62     find $out/lib -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
  63       echo "setting rpath for $lib..."
  64       patchelf --set-rpath "${rpath}:$out/lib" "$lib"
  65       ${lib.optionalString cudaSupport ''
  66         addOpenGLRunpath "$lib"
  67       ''}
  68     done
  69   '' + lib.optionalString stdenv.isDarwin ''
  70     install_name_tool -change @rpath/libshm.dylib $out/lib/libshm.dylib $out/lib/libtorch_python.dylib
  71     install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libtorch_python.dylib
  72     install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libtorch_python.dylib
  73     install_name_tool -change @rpath/libtorch.dylib $out/lib/libtorch.dylib $out/lib/libtorch_python.dylib
  74     install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libtorch_python.dylib
  75
  76     install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libtorch.dylib
  77     install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libtorch.dylib
  78     install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libtorch.dylib
  79
  80     install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libtorch_cpu.dylib
  81     install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libtorch_cpu.dylib
  82     install_name_tool -change @rpath/libtensorpipe.dylib $out/lib/libtensorpipe.dylib $out/lib/libtorch_cpu.dylib
  83
  84     install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libcaffe2_observers.dylib
  85     install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libcaffe2_observers.dylib
  86     install_name_tool -change @rpath/libtorch.dylib $out/lib/libtorch.dylib $out/lib/libcaffe2_observers.dylib
  87     install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libcaffe2_observers.dylib
  88
  89     install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libcaffe2_module_test_dynamic.dylib
  90     install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libcaffe2_module_test_dynamic.dylib
  91     install_name_tool -change @rpath/libtorch.dylib $out/lib/libtorch.dylib $out/lib/libcaffe2_module_test_dynamic.dylib
  92     install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libcaffe2_module_test_dynamic.dylib
  93
  94     install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libcaffe2_detectron_ops.dylib
  95     install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libcaffe2_detectron_ops.dylib
  96     install_name_tool -change @rpath/libtorch.dylib $out/lib/libtorch.dylib $out/lib/libcaffe2_detectron_ops.dylib
  97     install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libcaffe2_detectron_ops.dylib
  98
  99     install_name_tool -change @rpath/libc10.dylib $out/lib/libc10.dylib $out/lib/libshm.dylib
 100     install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libshm.dylib
 101     install_name_tool -change @rpath/libtorch.dylib $out/lib/libtorch.dylib $out/lib/libshm.dylib
 102     install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libshm.dylib
 103
 104     install_name_tool -change @rpath/libiomp5.dylib $out/lib/libiomp5.dylib $out/lib/libtorch_global_deps.dylib
 105     install_name_tool -change @rpath/libtorch_cpu.dylib $out/lib/libtorch_cpu.dylib $out/lib/libtorch_global_deps.dylib
 106   '';
 107
 108   outputs = [ "out" "dev" ];
 109
 110   passthru.tests.cmake = callPackage ./test {
 111     inherit cudaSupport;
 112     cudatoolkit = cudatoolkit_11_1;
 113     cudnn = cudnn_cudatoolkit_11_1;
 114   };
 115
 116   meta = with lib; {
 117     description = "C++ API of the PyTorch machine learning framework";
 118     homepage = "https://pytorch.org/";
 119     license = licenses.unfree; # Includes CUDA and Intel MKL.
 120     maintainers = with maintainers; [ danieldk ];
 121     platforms = with platforms; linux ++ darwin;
 122   };
 123 }