Victor Engmark
c4f5dfad0d
To avoid confusion and spurious newlines in the output. Found by searching for a backslash followed by a newline, some indentation, and then the end of multi-line string marker, using the following extended regex: ```regex \\\n +''; ```
135 lines
3.7 KiB
Nix
135 lines
3.7 KiB
Nix
{
|
|
callPackage,
|
|
stdenv,
|
|
fetchzip,
|
|
lib,
|
|
llvmPackages,
|
|
config,
|
|
|
|
autoAddDriverRunpath,
|
|
autoPatchelfHook,
|
|
patchelf,
|
|
fixDarwinDylibNames,
|
|
|
|
cudaSupport ? config.cudaSupport,
|
|
cudaPackages_13,
|
|
libz,
|
|
}:
|
|
|
|
let
|
|
# The binary libtorch distribution statically links the CUDA
|
|
# toolkit. This means that we do not need to provide CUDA to
|
|
# this derivation. However, we should ensure on version bumps
|
|
# that the CUDA toolkit for `passthru.tests` is still
|
|
# up-to-date.
|
|
version = "2.9.0";
|
|
device = if cudaSupport then "cuda" else "cpu";
|
|
srcs = import ./binary-hashes.nix version;
|
|
unavailable = throw "libtorch is not available for this platform";
|
|
in
|
|
stdenv.mkDerivation {
|
|
inherit version;
|
|
pname = "libtorch";
|
|
|
|
src = fetchzip srcs."${stdenv.hostPlatform.system}-${device}" or unavailable;
|
|
|
|
nativeBuildInputs =
|
|
if stdenv.hostPlatform.isDarwin then
|
|
[ fixDarwinDylibNames ]
|
|
else
|
|
[
|
|
patchelf
|
|
autoPatchelfHook
|
|
]
|
|
++ lib.optionals cudaSupport [ autoAddDriverRunpath ];
|
|
|
|
dontBuild = true;
|
|
dontConfigure = true;
|
|
dontStrip = true;
|
|
|
|
installPhase = ''
|
|
# Copy headers and CMake files.
|
|
mkdir -p $dev
|
|
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
|
|
|
|
# Fix up library paths for split outputs
|
|
substituteInPlace $dev/share/cmake/Torch/TorchConfig.cmake \
|
|
--replace \''${TORCH_INSTALL_PREFIX}/lib "$out/lib" \
|
|
|
|
substituteInPlace \
|
|
$dev/share/cmake/Caffe2/Caffe2Targets-release.cmake \
|
|
--replace \''${_IMPORT_PREFIX}/lib "$out/lib"
|
|
'';
|
|
|
|
postFixup =
|
|
let
|
|
rpath = lib.makeLibraryPath (
|
|
[ stdenv.cc.cc ]
|
|
++ lib.optionals cudaSupport [
|
|
cudaPackages_13.cuda_cudart # libcuda.so
|
|
cudaPackages_13.libcufft
|
|
cudaPackages_13.libcurand
|
|
cudaPackages_13.libcusolver
|
|
cudaPackages_13.libcusparse
|
|
libz
|
|
]
|
|
);
|
|
in
|
|
lib.optionalString stdenv.hostPlatform.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 ''
|
|
addDriverRunpath "$lib"
|
|
''}
|
|
done
|
|
''
|
|
+ lib.optionalString stdenv.hostPlatform.isDarwin ''
|
|
for f in $out/lib/*.dylib; do
|
|
otool -L $f
|
|
done
|
|
for f in $out/lib/*.dylib; do
|
|
if otool -L $f | grep "@rpath/libomp.dylib" >& /dev/null; then
|
|
install_name_tool -change "@rpath/libomp.dylib" ${llvmPackages.openmp}/lib/libomp.dylib $f
|
|
fi
|
|
install_name_tool -id $out/lib/$(basename $f) $f || true
|
|
for rpath in $(otool -L $f | grep rpath | awk '{print $1}');do
|
|
install_name_tool -change $rpath $out/lib/$(basename $rpath) $f
|
|
done
|
|
done
|
|
for f in $out/lib/*.dylib; do
|
|
otool -L $f
|
|
done
|
|
'';
|
|
|
|
outputs = [
|
|
"out"
|
|
"dev"
|
|
];
|
|
|
|
passthru.tests.cmake = callPackage ./test {
|
|
inherit cudaSupport;
|
|
};
|
|
|
|
meta = {
|
|
description = "C++ API of the PyTorch machine learning framework";
|
|
homepage = "https://pytorch.org/";
|
|
sourceProvenance = with lib.sourceTypes; [ binaryNativeCode ];
|
|
# Includes CUDA and Intel MKL, but redistributions of the binary are not limited.
|
|
# https://docs.nvidia.com/cuda/eula/index.html
|
|
# https://www.intel.com/content/www/us/en/developer/articles/license/onemkl-license-faq.html
|
|
license = lib.licenses.bsd3;
|
|
maintainers = with lib.maintainers; [ junjihashimoto ];
|
|
platforms = [
|
|
"aarch64-darwin"
|
|
"x86_64-linux"
|
|
];
|
|
};
|
|
}
|