cudaPackages.cuda_nvcc: patch nvcc.profile

let nvcc know about cudart and the stub drivers; this way
we do not need to search for cudart in the setupCudaHook

3 files changed, 25 insertions, 21 deletions

diff --git a/pkgs/development/cuda-modules/cuda/overrides.nix b/pkgs/development/cuda-modules/cuda/overrides.nix
index fd32978bfb59c..ff462b4cc970b 100644
--- a/pkgs/development/cuda-modules/cuda/overrides.nix
+++ b/pkgs/development/cuda-modules/cuda/overrides.nix
@@ -61,7 +61,30 @@ attrsets.filterAttrs (attr: _: (builtins.hasAttr attr prev)) {
 
   cuda_nvcc = prev.cuda_nvcc.overrideAttrs (
     oldAttrs: {
-      propagatedBuildInputs = [final.setupCudaHook];
+
+      outputs = oldAttrs.outputs ++ [ "lib" ];
+
+      postPatch =
+        (oldAttrs.postPatch or "")
+        + ''
+          substituteInPlace bin/nvcc.profile \
+            --replace \
+              '$(TOP)/lib' \
+              "''${!outputLib}/lib"
+
+          cat << EOF >> bin/nvcc.profile
+          LIBRARIES +=  "-L''${!outputBin}/nvvm/lib" "-L$static/lib" "-L${final.cuda_cudart.lib}/lib" "-L${final.cuda_cudart.lib}/lib/stubs"
+          INCLUDES += "-I''${!outputDev}/include" "-I''${!outputBin}/nvvm/include" "-I${final.cuda_cudart.dev}/include"
+          EOF
+        '';
+
+      propagatedBuildInputs = [ final.setupCudaHook ];
+
+      postInstall =
+        (oldAttrs.postInstall or "")
+        + ''
+          moveToOutput "nvvm" "''${!outputBin}"
+        '';
 
       meta = (oldAttrs.meta or {}) // {
         mainProgram = "nvcc";
diff --git a/pkgs/development/cuda-modules/saxpy/default.nix b/pkgs/development/cuda-modules/saxpy/default.nix
index fff52801e3cca..73e17b28757b0 100644
--- a/pkgs/development/cuda-modules/saxpy/default.nix
+++ b/pkgs/development/cuda-modules/saxpy/default.nix
@@ -14,6 +14,7 @@ let
     cudaVersion
     flags
     libcublas
+    setupCudaHook
     ;
 in
 backendStdenv.mkDerivation {
diff --git a/pkgs/development/cuda-modules/setup-hooks/setup-cuda-hook.sh b/pkgs/development/cuda-modules/setup-hooks/setup-cuda-hook.sh
index 7b7b3bdde80e3..a4a444fcd2417 100644
--- a/pkgs/development/cuda-modules/setup-hooks/setup-cuda-hook.sh
+++ b/pkgs/development/cuda-modules/setup-hooks/setup-cuda-hook.sh
@@ -93,26 +93,6 @@ setupCUDAToolkitCompilers() {
     if [[ -z "${dontCompressFatbin-}" ]]; then
         export NVCC_PREPEND_FLAGS+=" -Xfatbin=-compress-all"
     fi
-
-    # CMake's enable_language(CUDA) runs a compiler test and it doesn't account for
-    # CUDAToolkit_ROOT. We have to help it locate libcudart
-    if [[ -z "${nvccDontPrependCudartFlags-}" ]] ; then
-        if [[ ! -v cudaOutputToPath["cuda_cudart-out"] ]] ; then
-            echo "setupCUDAToolkitCompilers: missing cudaPackages.cuda_cudart. This may become an an error in the future" >&2
-            # exit 1
-        fi
-        for pkg in "${!cudaOutputToPath[@]}" ; do
-            [[ ! "$pkg" = cuda_cudart* ]] && continue
-
-            local path="${cudaOutputToPath[$pkg]}"
-            if [[ -d "$path/include" ]] ; then
-                export NVCC_PREPEND_FLAGS+=" -I$path/include"
-            fi
-            if [[ -d "$path/lib" ]] ; then
-                export NVCC_PREPEND_FLAGS+=" -L$path/lib"
-            fi
-        done
-    fi
 }
 preConfigureHooks+=(setupCUDAToolkitCompilers)