nixpkgs/pkgs/development/libraries/science/math/libtorch/test/default.nix
Frederik Rietdijk 1d63f89caa cudaPackages: overhaul of how we package cuda packages
There are many different versions of the `cudatoolkit` and related
cuda packages, and it can be tricky to ensure they remain compatible.

- `cudaPackages` is now a package set with `cudatoolkit`, `cudnn`, `cutensor`, `nccl`, as well as `cudatoolkit` split into smaller packages ("redist");
- expressions should now use `cudaPackages` as parameter instead of the individual cuda packages;
- `makeScope` is now used, so it is possible to use `.overrideScope'` to set e.g. a different `cudnn` version;
- `release-cuda.nix` is introduced to easily evaluate cuda packages using hydra.
2022-04-09 08:50:22 +02:00

52 lines
1.0 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 = ./.;
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
'';
}