65 lines
1.2 KiB
Nix
65 lines
1.2 KiB
Nix
|
{
|
||
|
|
lib,
|
||
|
|
stdenv,
|
||
|
|
cmake,
|
||
|
|
libtorch-bin,
|
||
|
|
linkFarm,
|
||
|
|
symlinkJoin,
|
||
|
|
|
||
|
|
cudaSupport,
|
||
|
|
cudaPackages ? { },
|
||
|
|
}:
|
||
|
|
let
|
||
|
|
inherit (cudaPackages) cudatoolkit cudnn;
|
||
|
|
|
||
|
|
cudatoolkit_joined = symlinkJoin {
|
||
|
|
name = "${cudatoolkit.name}-unsplit";
|
||
|
|
paths = [
|
||
|
|
cudatoolkit.out
|
||
|
|
cudatoolkit.lib
|
||
|
|
];
|
||
|
|
};
|
||
|
|
|
||
|
|
# We do not have access to /run/opengl-driver/lib in the sandbox,
|
||
|
|
# so use a stub instead.
|
||
|
|
cudaStub = linkFarm "cuda-stub" [
|
||
|
|
{
|
||
|
|
name = "libcuda.so.1";
|
||
|
|
path = "${cudatoolkit}/lib/stubs/libcuda.so";
|
||
|
|
}
|
||
|
|
];
|
||
|
|
|
||
|
|
in
|
||
|
|
stdenv.mkDerivation {
|
||
|
|
pname = "libtorch-test";
|
||
|
|
version = libtorch-bin.version;
|
||
|
|
|
||
|
|
src = lib.fileset.toSource {
|
||
|
|
root = ./.;
|
||
|
|
fileset = lib.fileset.unions [
|
||
|
|
./CMakeLists.txt
|
||
|
|
./test.cpp
|
||
|
|
];
|
||
|
|
};
|
||
|
|
|
||
|
|
nativeBuildInputs = [ cmake ];
|
||
|
|
|
||
|
|
buildInputs = [ libtorch-bin ] ++ lib.optionals cudaSupport [ cudnn ];
|
||
|
|
|
||
|
|
cmakeFlags = lib.optionals cudaSupport [ "-DCUDA_TOOLKIT_ROOT_DIR=${cudatoolkit_joined}" ];
|
||
|
|
|
||
|
|
doCheck = true;
|
||
|
|
|
||
|
|
installPhase = ''
|
||
|
|
touch $out
|
||
|
|
'';
|
||
|
|
|
||
|
|
checkPhase =
|
||
|
|
lib.optionalString cudaSupport ''
|
||
|
|
LD_LIBRARY_PATH=${cudaStub}''${LD_LIBRARY_PATH:+:}$LD_LIBRARY_PATH \
|
||
|
|
''
|
||
|
|
+ ''
|
||
|
|
./test
|
||
|
|
'';
|
||
|
|
}
|