Add driver library path to some packages to find CUDA libraries.

This is to avoid relying on LD_LIBRARY_PATH for finding the CUDA driver libraries.

pkgs/development/libraries/science/math/cudnn/generic.nix

@@ -7,6 +7,7 @@
 , lib
 , cudatoolkit
 , fetchurl
+, addOpenGLRunpath
 }:
 
 stdenv.mkDerivation rec {
@@ -19,6 +20,8 @@ stdenv.mkDerivation rec {
     inherit sha256;
   };
 
+  nativeBuildInputs = [ addOpenGLRunpath ];
+
   installPhase = ''
     function fixRunPath {
       p=$(patchelf --print-rpath $1)
@@ -31,6 +34,12 @@ stdenv.mkDerivation rec {
     cp -a lib64 $out/lib64
   '';
 
+  # Set RUNPATH so that libcuda in /run/opengl-driver(-32)/lib can be found.
+  # See the explanation in addOpenGLRunpath.
+  postFixup = ''
+    addOpenGLRunpath $out/lib/lib*.so
+  '';
+
   propagatedBuildInputs = [
     cudatoolkit
   ];

pkgs/development/libraries/science/math/nccl/default.nix

@@ -1,4 +1,4 @@
-{ stdenv, fetchFromGitHub, which, cudatoolkit }:
+{ stdenv, fetchFromGitHub, which, cudatoolkit, addOpenGLRunpath }:
 
 stdenv.mkDerivation rec {
   name = "nccl-${version}-cuda-${cudatoolkit.majorVersion}";
@@ -13,7 +13,7 @@ stdenv.mkDerivation rec {
 
   outputs = [ "out" "dev" ];
 
-  nativeBuildInputs = [ which ];
+  nativeBuildInputs = [ which addOpenGLRunpath ];
 
   buildInputs = [ cudatoolkit ];
 
@@ -28,6 +28,10 @@ stdenv.mkDerivation rec {
 
   postFixup = ''
     moveToOutput lib/libnccl_static.a $dev
+
+    # Set RUNPATH so that libnvidia-ml in /run/opengl-driver(-32)/lib can be found.
+    # See the explanation in addOpenGLRunpath.
+    addOpenGLRunpath $out/lib/lib*.so
   '';
 
   NIX_CFLAGS_COMPILE = [ "-Wno-unused-function" ];