Merge pull request #113210 from danieldk/libtorch-bin-fixes

libtorch-bin: remove dependency on nvidia_x11, fix passthru test

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

@@ -8,10 +8,16 @@
 , fixDarwinDylibNames
 
 , cudaSupport
-, nvidia_x11
+, cudatoolkit_10_2
+, cudnn_cudatoolkit_10_2
 }:
 
 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 = "1.7.1";
   device = if cudaSupport then "cuda" else "cpu";
   srcs = import ./binary-hashes.nix version;
@@ -24,12 +30,7 @@ in stdenv.mkDerivation {
 
   nativeBuildInputs =
     if stdenv.isDarwin then [ fixDarwinDylibNames ]
-    else [ addOpenGLRunpath patchelf ]
-      ++ lib.optionals cudaSupport [ addOpenGLRunpath ];
-
-  buildInputs = [
-    stdenv.cc.cc
-  ] ++ lib.optionals cudaSupport [ nvidia_x11 ];
+    else [ patchelf ] ++ lib.optionals cudaSupport [ addOpenGLRunpath ];
 
   dontBuild = true;
   dontConfigure = true;
@@ -56,9 +57,7 @@ in stdenv.mkDerivation {
   '';
 
   postFixup = let
-    libPaths = [ stdenv.cc.cc.lib ]
-      ++ lib.optionals cudaSupport [ nvidia_x11 ];
-    rpath = lib.makeLibraryPath libPaths;
+    rpath = lib.makeLibraryPath [ stdenv.cc.cc.lib ];
   in lib.optionalString stdenv.isLinux ''
     find $out/lib -type f \( -name '*.so' -or -name '*.so.*' \) | while read lib; do
       echo "setting rpath for $lib..."
@@ -108,12 +107,17 @@ in stdenv.mkDerivation {
 
   outputs = [ "out" "dev" ];
 
-  passthru.tests.cmake = callPackage ./test { };
+  passthru.tests.cmake = callPackage ./test {
+    inherit cudaSupport;
+    cudatoolkit = cudatoolkit_10_2;
+    cudnn = cudnn_cudatoolkit_10_2;
+  };
 
   meta = with lib; {
     description = "C++ API of the PyTorch machine learning framework";
     homepage = "https://pytorch.org/";
     license = licenses.unfree; # Includes CUDA and Intel MKL.
+    maintainers = with maintainers; [ danieldk ];
     platforms = with platforms; linux ++ darwin;
   };
 }

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

@@ -1,6 +1,28 @@
-{ stdenv, cmake, libtorch-bin, symlinkJoin }:
+{ lib
+, stdenv
+, cmake
+, libtorch-bin
+, linkFarm
+, symlinkJoin
 
-stdenv.mkDerivation {
+, cudaSupport
+, cudatoolkit
+, cudnn
+}:
+let
+  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;
 
@@ -8,7 +30,11 @@ stdenv.mkDerivation {
 
   nativeBuildInputs = [ cmake ];
 
-  buildInputs = [ libtorch-bin ];
+  buildInputs = [ libtorch-bin ] ++
+    lib.optionals cudaSupport [ cudnn ];
+
+  cmakeFlags = lib.optionals cudaSupport
+    [ "-DCUDA_TOOLKIT_ROOT_DIR=${cudatoolkit_joined}" ];
 
   doCheck = true;
 
@@ -17,6 +43,7 @@ stdenv.mkDerivation {
   '';
 
   checkPhase = ''
-    ./test
+    LD_LIBRARY_PATH=${cudaStub}''${LD_LIBRARY_PATH:+:}$LD_LIBRARY_PATH \
+      ./test
   '';
 }

pkgs/top-level/all-packages.nix

@@ -1677,7 +1677,6 @@ in
     else libtensorflow-bin;
 
   libtorch-bin = callPackage ../development/libraries/science/math/libtorch/bin.nix {
-    inherit (linuxPackages) nvidia_x11;
     cudaSupport = config.cudaSupport or false;
   };