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


1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51

{ 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
  '';
}