From 2639fa16b9361aef298b02d753a75881da273b19 Mon Sep 17 00:00:00 2001
From: Connor Baker <connor.baker@tweag.io>
Date: Wed, 3 Apr 2024 16:23:20 +0000
Subject: [PATCH 01/15] cudaPackages.backendStdenv: prune and sort callPackage
 args

---
 pkgs/development/cuda-modules/backend-stdenv.nix | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/pkgs/development/cuda-modules/backend-stdenv.nix b/pkgs/development/cuda-modules/backend-stdenv.nix
index 5d1c0c735806..57219ad607c0 100644
--- a/pkgs/development/cuda-modules/backend-stdenv.nix
+++ b/pkgs/development/cuda-modules/backend-stdenv.nix
@@ -1,11 +1,9 @@
 {
+  cudaVersion,
   lib,
   nvccCompatibilities,
-  cudaVersion,
   pkgs,
-  overrideCC,
   stdenv,
-  wrapCCWith,
   stdenvAdapters,
 }:
 

From 14635b40927383531ad9b6b5335911fdab7d6ac9 Mon Sep 17 00:00:00 2001
From: Connor Baker <connor.baker@tweag.io>
Date: Wed, 3 Apr 2024 16:32:02 +0000
Subject: [PATCH 02/15] cuda-modules: don't use *Platform attributes from pkgs

For more information about *why* this is desirable, see
https://github.com/NixOS/nixpkgs/pull/45717
and
https://github.com/NixOS/nixpkgs/issues/27069
---
 .../cuda-modules/cuda-library-samples/extension.nix      | 9 ++++++++-
 pkgs/development/cuda-modules/cuda-samples/extension.nix | 6 +++++-
 pkgs/development/cuda-modules/cutensor/extension.nix     | 6 +++++-
 pkgs/development/cuda-modules/flags.nix                  | 6 +++++-
 .../cuda-modules/generic-builders/manifest.nix           | 5 +++--
 .../cuda-modules/generic-builders/multiplex.nix          | 6 +++++-
 pkgs/development/cuda-modules/tensorrt/fixup.nix         | 3 ++-
 7 files changed, 33 insertions(+), 8 deletions(-)

diff --git a/pkgs/development/cuda-modules/cuda-library-samples/extension.nix b/pkgs/development/cuda-modules/cuda-library-samples/extension.nix
index 456ab8168a45..eb291cd4da09 100644
--- a/pkgs/development/cuda-modules/cuda-library-samples/extension.nix
+++ b/pkgs/development/cuda-modules/cuda-library-samples/extension.nix
@@ -1,5 +1,12 @@
-{ hostPlatform, lib }:
+{
+  lib,
+  # NOTE: We would use backendStdenv to get the hostPlatform, but it's not available in the callPackage context
+  # we're using to call this file.
+  stdenv,
+}:
 let
+  inherit (stdenv) hostPlatform;
+
   # Samples are built around the CUDA Toolkit, which is not available for
   # aarch64. Check for both CUDA version and platform.
   platformIsSupported = hostPlatform.isx86_64 && hostPlatform.isLinux;
diff --git a/pkgs/development/cuda-modules/cuda-samples/extension.nix b/pkgs/development/cuda-modules/cuda-samples/extension.nix
index d41da90cd5d0..e4a02e9a5e0c 100644
--- a/pkgs/development/cuda-modules/cuda-samples/extension.nix
+++ b/pkgs/development/cuda-modules/cuda-samples/extension.nix
@@ -1,7 +1,9 @@
 {
   cudaVersion,
-  hostPlatform,
   lib,
+  # NOTE: We would use backendStdenv to get the hostPlatform, but it's not available in the callPackage context
+  # we're using to call this file.
+  stdenv,
 }:
 let
   cudaVersionToHash = {
@@ -23,6 +25,8 @@ let
     "12.3" = "sha256-fjVp0G6uRCWxsfe+gOwWTN+esZfk0O5uxS623u0REAk=";
   };
 
+  inherit (stdenv) hostPlatform;
+
   # Samples are built around the CUDA Toolkit, which is not available for
   # aarch64. Check for both CUDA version and platform.
   cudaVersionIsSupported = cudaVersionToHash ? ${cudaVersion};
diff --git a/pkgs/development/cuda-modules/cutensor/extension.nix b/pkgs/development/cuda-modules/cutensor/extension.nix
index c41113939ca2..ee9b31972e18 100644
--- a/pkgs/development/cuda-modules/cutensor/extension.nix
+++ b/pkgs/development/cuda-modules/cutensor/extension.nix
@@ -15,9 +15,11 @@
 {
   cudaVersion,
   flags,
-  hostPlatform,
   lib,
   mkVersionedPackageName,
+  # NOTE: We would use backendStdenv to get the hostPlatform, but it's not available in the callPackage context
+  # we're using to call this file.
+  stdenv,
 }:
 let
   inherit (lib)
@@ -29,6 +31,8 @@ let
     trivial
     ;
 
+  inherit (stdenv) hostPlatform;
+
   redistName = "cutensor";
   pname = "libcutensor";
 
diff --git a/pkgs/development/cuda-modules/flags.nix b/pkgs/development/cuda-modules/flags.nix
index 196b6b9f8f99..3f5be88423b7 100644
--- a/pkgs/development/cuda-modules/flags.nix
+++ b/pkgs/development/cuda-modules/flags.nix
@@ -7,7 +7,9 @@
   cudaForwardCompat ? (config.cudaForwardCompat or true),
   lib,
   cudaVersion,
-  hostPlatform,
+  # NOTE: We cannot refer to backendStdenv as flags exists outside the fixed-point;
+  # otherwise, we would use that to get the hostPlatform.
+  stdenv,
   # gpus :: List Gpu
   gpus,
 }:
@@ -20,6 +22,8 @@ let
     trivial
     ;
 
+  inherit (stdenv) hostPlatform;
+
   # Flags are determined based on your CUDA toolkit by default.  You may benefit
   # from improved performance, reduced file size, or greater hardware support by
   # passing a configuration based on your specific GPU environment.
diff --git a/pkgs/development/cuda-modules/generic-builders/manifest.nix b/pkgs/development/cuda-modules/generic-builders/manifest.nix
index 73c34b0c86ee..e5831ef99dd0 100644
--- a/pkgs/development/cuda-modules/generic-builders/manifest.nix
+++ b/pkgs/development/cuda-modules/generic-builders/manifest.nix
@@ -10,7 +10,6 @@
   markForCudatoolkitRootHook,
   flags,
   stdenv,
-  hostPlatform,
   # Builder-specific arguments
   # Short package name (e.g., "cuda_cccl")
   # pname : String
@@ -40,6 +39,8 @@ let
     sourceTypes
     ;
 
+  inherit (backendStdenv) hostPlatform;
+
   # Get the redist architectures for which package provides distributables.
   # These are used by meta.platforms.
   supportedRedistArchs = builtins.attrNames featureRelease;
@@ -48,7 +49,7 @@ let
   # It is `"unsupported"` if the redistributable is not supported on the target platform.
   redistArch = flags.getRedistArch hostPlatform.system;
 
-  sourceMatchesHost = flags.getNixSystem redistArch == stdenv.hostPlatform.system;
+  sourceMatchesHost = flags.getNixSystem redistArch == hostPlatform.system;
 in
 backendStdenv.mkDerivation (finalAttrs: {
   # NOTE: Even though there's no actual buildPhase going on here, the derivations of the
diff --git a/pkgs/development/cuda-modules/generic-builders/multiplex.nix b/pkgs/development/cuda-modules/generic-builders/multiplex.nix
index 0b523e56b8c4..fd06d19e60f0 100644
--- a/pkgs/development/cuda-modules/generic-builders/multiplex.nix
+++ b/pkgs/development/cuda-modules/generic-builders/multiplex.nix
@@ -3,7 +3,9 @@
   lib,
   cudaVersion,
   flags,
-  hostPlatform,
+  # NOTE: We would use backendStdenv to get the hostPlatform, but it's not available in the callPackage context
+  # we're using to call this file (cudaPackages.callPackage).
+  stdenv,
   # Expected to be passed by the caller
   mkVersionedPackageName,
   # pname :: String
@@ -40,6 +42,8 @@ let
     strings
     ;
 
+  inherit (stdenv) hostPlatform;
+
   evaluatedModules = modules.evalModules {
     modules = [
       ../modules
diff --git a/pkgs/development/cuda-modules/tensorrt/fixup.nix b/pkgs/development/cuda-modules/tensorrt/fixup.nix
index 3615284fb080..fb843aa4e1e1 100644
--- a/pkgs/development/cuda-modules/tensorrt/fixup.nix
+++ b/pkgs/development/cuda-modules/tensorrt/fixup.nix
@@ -1,7 +1,7 @@
 {
+  backendStdenv,
   cudaVersion,
   final,
-  hostPlatform,
   lib,
   mkVersionedPackageName,
   package,
@@ -17,6 +17,7 @@ let
     strings
     versions
     ;
+  inherit (backendStdenv) hostPlatform;
   # targetArch :: String
   targetArch = attrsets.attrByPath [ hostPlatform.system ] "unsupported" {
     x86_64-linux = "x86_64-linux-gnu";

From d22f03ede96f524c2215fb1a94970714123f39d8 Mon Sep 17 00:00:00 2001
From: Connor Baker <connor.baker@tweag.io>
Date: Wed, 3 Apr 2024 16:38:52 +0000
Subject: [PATCH 03/15] cuda-modules: update uses of substituteInPlace to
 --replace-fail to fix deprecation warning

---
 .../cuda-modules/cuda-library-samples/generic.nix           | 2 +-
 pkgs/development/cuda-modules/cuda/overrides.nix            | 6 +++---
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/pkgs/development/cuda-modules/cuda-library-samples/generic.nix b/pkgs/development/cuda-modules/cuda-library-samples/generic.nix
index 4797871731b8..64131ab59b82 100644
--- a/pkgs/development/cuda-modules/cuda-library-samples/generic.nix
+++ b/pkgs/development/cuda-modules/cuda-library-samples/generic.nix
@@ -76,7 +76,7 @@ in
       # CUTENSOR_ROOT is double escaped
       postPatch = ''
         substituteInPlace CMakeLists.txt \
-          --replace "\''${CUTENSOR_ROOT}/include" "${cutensor.dev}/include"
+          --replace-fail "\''${CUTENSOR_ROOT}/include" "${cutensor.dev}/include"
       '';
 
       CUTENSOR_ROOT = cutensor;
diff --git a/pkgs/development/cuda-modules/cuda/overrides.nix b/pkgs/development/cuda-modules/cuda/overrides.nix
index 9a2360d7f7c1..c170e3d39c91 100644
--- a/pkgs/development/cuda-modules/cuda/overrides.nix
+++ b/pkgs/development/cuda-modules/cuda/overrides.nix
@@ -134,13 +134,13 @@ attrsets.filterAttrs (attr: _: (builtins.hasAttr attr prev)) {
         (oldAttrs.postPatch or "")
         + ''
           substituteInPlace bin/nvcc.profile \
-            --replace \
+            --replace-fail \
               '$(TOP)/lib' \
               "''${!outputLib}/lib" \
-            --replace \
+            --replace-fail \
               '$(TOP)/$(_NVVM_BRANCH_)' \
               "''${!outputBin}/nvvm" \
-            --replace \
+            --replace-fail \
               '$(TOP)/$(_TARGET_DIR_)/include' \
               "''${!outputDev}/include"
 

From ad2a4e862cc760386ed6f0b86dc0f2774114838d Mon Sep 17 00:00:00 2001
From: Connor Baker <connor.baker@tweag.io>
Date: Mon, 8 Apr 2024 21:27:02 +0000
Subject: [PATCH 04/15] cudaPackages.markForCudatoolkitRootHook: always set
 return explicitly

---
 .../cuda-modules/setup-hooks/mark-for-cudatoolkit-root-hook.sh  | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pkgs/development/cuda-modules/setup-hooks/mark-for-cudatoolkit-root-hook.sh b/pkgs/development/cuda-modules/setup-hooks/mark-for-cudatoolkit-root-hook.sh
index ba04c2e0806a..cea82d817a12 100644
--- a/pkgs/development/cuda-modules/setup-hooks/mark-for-cudatoolkit-root-hook.sh
+++ b/pkgs/development/cuda-modules/setup-hooks/mark-for-cudatoolkit-root-hook.sh
@@ -7,7 +7,7 @@ echo "Sourcing mark-for-cudatoolkit-root-hook" >&2
 
 markForCUDAToolkit_ROOT() {
     mkdir -p "${prefix}/nix-support"
-    [[ -f "${prefix}/nix-support/include-in-cudatoolkit-root" ]] && return
+    [[ -f "${prefix}/nix-support/include-in-cudatoolkit-root" ]] && return 0
     echo "$pname-$output" > "${prefix}/nix-support/include-in-cudatoolkit-root"
 }
 

From 9230c2cc8d9b98f302fda10436b08185f5cb9eea Mon Sep 17 00:00:00 2001
From: Connor Baker <connor.baker@tweag.io>
Date: Mon, 8 Apr 2024 21:33:48 +0000
Subject: [PATCH 05/15] cudaPackages.markForCudatoolkitRootHook: fix bug with
 strictDeps

The setupCudaHook always checks for the existence of `$prefix/nix-support/include-in-cudatoolkit-root`, so we need to be sure
it always exists. It isn't populated when `strictDeps` is set, but it must exist.
---
 .../mark-for-cudatoolkit-root-hook.sh         | 21 ++++++++++++++-----
 1 file changed, 16 insertions(+), 5 deletions(-)

diff --git a/pkgs/development/cuda-modules/setup-hooks/mark-for-cudatoolkit-root-hook.sh b/pkgs/development/cuda-modules/setup-hooks/mark-for-cudatoolkit-root-hook.sh
index cea82d817a12..0abd651005c6 100644
--- a/pkgs/development/cuda-modules/setup-hooks/mark-for-cudatoolkit-root-hook.sh
+++ b/pkgs/development/cuda-modules/setup-hooks/mark-for-cudatoolkit-root-hook.sh
@@ -1,14 +1,25 @@
 # shellcheck shell=bash
 
-# Should we mimick cc-wrapper's "hygiene"?
-[[ -z ${strictDeps-} ]] || (( "$hostOffset" < 0 )) || return 0
+(( ${hostOffset:?} == -1 && ${targetOffset:?} == 0)) || return 0
 
 echo "Sourcing mark-for-cudatoolkit-root-hook" >&2
 
 markForCUDAToolkit_ROOT() {
-    mkdir -p "${prefix}/nix-support"
-    [[ -f "${prefix}/nix-support/include-in-cudatoolkit-root" ]] && return 0
-    echo "$pname-$output" > "${prefix}/nix-support/include-in-cudatoolkit-root"
+    mkdir -p "${prefix:?}/nix-support"
+    local markerPath="$prefix/nix-support/include-in-cudatoolkit-root"
+
+    # Return early if the file already exists.
+    [[ -f "$markerPath" ]] && return 0
+
+    # Always create the file, even if it's empty, since setup-cuda-hook relies on its existence.
+    # However, only populate it if strictDeps is not set.
+    touch "$markerPath"
+
+    # Return early if strictDeps is set.
+    [[ -n "${strictDeps-}" ]] && return 0
+
+    # Populate the file with the package name and output.
+    echo "${pname:?}-${output:?}" > "$markerPath"
 }
 
 fixupOutputHooks+=(markForCUDAToolkit_ROOT)

From a31b0d7f95292293b4cddfc381c9242895c53c18 Mon Sep 17 00:00:00 2001
From: Connor Baker <connor.baker@tweag.io>
Date: Wed, 3 Apr 2024 17:09:07 +0000
Subject: [PATCH 06/15] cudaPackages.setupCudaHook: sourcing messages should
 match hook name

---
 pkgs/development/cuda-modules/setup-hooks/setup-cuda-hook.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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 a4a444fcd241..95f5fd884e42 100644
--- a/pkgs/development/cuda-modules/setup-hooks/setup-cuda-hook.sh
+++ b/pkgs/development/cuda-modules/setup-hooks/setup-cuda-hook.sh
@@ -9,7 +9,7 @@ reason=
 [[ -n ${cudaSetupHookOnce-} ]] && guard=Skipping && reason=" because the hook has been propagated more than once"
 
 if (( "${NIX_DEBUG:-0}" >= 1 )) ; then
-    echo "$guard hostOffset=$hostOffset targetOffset=$targetOffset setupCudaHook$reason" >&2
+    echo "$guard hostOffset=$hostOffset targetOffset=$targetOffset setup-cuda-hook$reason" >&2
 else
     echo "$guard setup-cuda-hook$reason" >&2
 fi

From 02966afc5481730c3ff8e09cb2c7a79d96b63be0 Mon Sep 17 00:00:00 2001
From: Connor Baker <connor.baker@tweag.io>
Date: Wed, 3 Apr 2024 18:50:44 +0000
Subject: [PATCH 07/15] cudaPackages.setupCudaHook: always set return
 explicitly

On NixOS, the the return status is that of the last command executed within the function or script.

When we're doing tests immediately before the return, the value you end up returning might not be what you want.
---
 .../cuda-modules/setup-hooks/setup-cuda-hook.sh        | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

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 95f5fd884e42..dad78cfebded 100644
--- a/pkgs/development/cuda-modules/setup-hooks/setup-cuda-hook.sh
+++ b/pkgs/development/cuda-modules/setup-hooks/setup-cuda-hook.sh
@@ -24,8 +24,8 @@ extendcudaHostPathsSeen() {
     (( "${NIX_DEBUG:-0}" >= 1 )) && echo "extendcudaHostPathsSeen $1" >&2
 
     local markerPath="$1/nix-support/include-in-cudatoolkit-root"
-    [[ ! -f "${markerPath}" ]] && return
-    [[ -v cudaHostPathsSeen[$1] ]] && return
+    [[ ! -f "${markerPath}" ]] && return 0
+    [[ -v cudaHostPathsSeen[$1] ]] && return 0
 
     cudaHostPathsSeen["$1"]=1
 
@@ -33,7 +33,7 @@ extendcudaHostPathsSeen() {
     local cudaOutputName
     read -r cudaOutputName < "$markerPath"
 
-    [[ -z "$cudaOutputName" ]] && return
+    [[ -z "$cudaOutputName" ]] && return 0
 
     local oldPath="${cudaOutputToPath[$cudaOutputName]-}"
     [[ -n "$oldPath" ]] && echo "extendcudaHostPathsSeen: warning: overwriting $cudaOutputName from $oldPath to $1" >&2
@@ -59,7 +59,7 @@ setupCUDAToolkitCompilers() {
     echo Executing setupCUDAToolkitCompilers >&2
 
     if [[ -n "${dontSetupCUDAToolkitCompilers-}" ]] ; then
-        return
+        return 0
     fi
 
     # Point NVCC at a compatible compiler
@@ -99,7 +99,7 @@ preConfigureHooks+=(setupCUDAToolkitCompilers)
 propagateCudaLibraries() {
     (( "${NIX_DEBUG:-0}" >= 1 )) && echo "propagateCudaLibraries: cudaPropagateToOutput=$cudaPropagateToOutput cudaHostPathsSeen=${!cudaHostPathsSeen[*]}" >&2
 
-    [[ -z "${cudaPropagateToOutput-}" ]] && return
+    [[ -z "${cudaPropagateToOutput-}" ]] && return 0
 
     mkdir -p "${!cudaPropagateToOutput}/nix-support"
     # One'd expect this should be propagated-bulid-build-deps, but that doesn't seem to work

From 5ee7bfc443cc021354e020a3a2dc352227500f0e Mon Sep 17 00:00:00 2001
From: Connor Baker <connor.baker@tweag.io>
Date: Wed, 3 Apr 2024 21:20:51 +0000
Subject: [PATCH 08/15] cudaPackages.setupCudaHook: fix error when reading
 empty marker

---
 pkgs/development/cuda-modules/setup-hooks/setup-cuda-hook.sh | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

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 dad78cfebded..6e57c7b1072e 100644
--- a/pkgs/development/cuda-modules/setup-hooks/setup-cuda-hook.sh
+++ b/pkgs/development/cuda-modules/setup-hooks/setup-cuda-hook.sh
@@ -31,7 +31,10 @@ extendcudaHostPathsSeen() {
 
     # E.g. cuda_cudart-lib
     local cudaOutputName
-    read -r cudaOutputName < "$markerPath"
+    # Fail gracefully if the file is empty.
+    # One reason the file may be empty: the package was built with strictDeps set, but the current build does not have
+    # strictDeps set.
+    read -r cudaOutputName < "$markerPath" || return 0
 
     [[ -z "$cudaOutputName" ]] && return 0
 

From e77b24b15926bf690a6e69d2fb805edce0f4f4b1 Mon Sep 17 00:00:00 2001
From: Connor Baker <connor.baker@tweag.io>
Date: Wed, 3 Apr 2024 22:35:46 +0000
Subject: [PATCH 09/15] cudaPackages.flags: add cmakeCudaArchitectures and
 cmakeCudaArchitecturesString

---
 pkgs/development/cuda-modules/flags.nix | 22 ++++++++++++++++++++++
 1 file changed, 22 insertions(+)

diff --git a/pkgs/development/cuda-modules/flags.nix b/pkgs/development/cuda-modules/flags.nix
index 3f5be88423b7..7bba27a2f3e3 100644
--- a/pkgs/development/cuda-modules/flags.nix
+++ b/pkgs/development/cuda-modules/flags.nix
@@ -211,6 +211,16 @@ let
       # E.g. "-gencode=arch=compute_75,code=sm_75 ... -gencode=arch=compute_86,code=compute_86"
       gencodeString = strings.concatStringsSep " " gencode;
 
+      # cmakeCudaArchitectures :: List String
+      # A list of CUDA capabilities without dots, suitable for passing to CMake.
+      # E.g. [ "75" "86" ]
+      cmakeCudaArchitectures = lists.map dropDot cudaCapabilities;
+
+      # cmakeCudaArchitecturesString :: String
+      # A semicolon-separated string of CUDA capabilities without dots, suitable for passing to CMake.
+      # E.g. "75;86"
+      cmakeCudaArchitecturesString = strings.concatStringsSep ";" cmakeCudaArchitectures;
+
       # Jetson devices cannot be targeted by the same binaries which target non-Jetson devices. While
       # NVIDIA provides both `linux-aarch64` and `linux-sbsa` packages, which both target `aarch64`,
       # they are built with different settings and cannot be mixed.
@@ -274,6 +284,12 @@ assert
       ];
       gencodeString = "-gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_86,code=sm_86 -gencode=arch=compute_86,code=compute_86";
 
+      cmakeCudaArchitectures = [
+        "75"
+        "86"
+      ];
+      cmakeCudaArchitecturesString = "75;86";
+
       isJetsonBuild = false;
     };
     actual = formatCapabilities {
@@ -343,6 +359,12 @@ assert
       ];
       gencodeString = "-gencode=arch=compute_62,code=sm_62 -gencode=arch=compute_72,code=sm_72 -gencode=arch=compute_72,code=compute_72";
 
+      cmakeCudaArchitectures = [
+        "62"
+        "72"
+      ];
+      cmakeCudaArchitecturesString = "62;72";
+
       isJetsonBuild = true;
     };
     actual = formatCapabilities {

From 5ed9f23d218223ce5ea280e43bdcf6739d8ace07 Mon Sep 17 00:00:00 2001
From: Connor Baker <connor.baker@tweag.io>
Date: Wed, 3 Apr 2024 22:16:24 +0000
Subject: [PATCH 10/15] cudaPackages.saxpy: switch to cudaAtLeast, cudaOlder,
 __structuredAttrs, and enable on Jetson post-11.4

---
 .../cuda-modules/saxpy/default.nix            | 22 +++++++++----------
 1 file changed, 11 insertions(+), 11 deletions(-)

diff --git a/pkgs/development/cuda-modules/saxpy/default.nix b/pkgs/development/cuda-modules/saxpy/default.nix
index 2a2eedbcb1db..a36cec3e692b 100644
--- a/pkgs/development/cuda-modules/saxpy/default.nix
+++ b/pkgs/development/cuda-modules/saxpy/default.nix
@@ -10,8 +10,9 @@ let
     cuda_cccl
     cuda_cudart
     cuda_nvcc
+    cudaAtLeast
+    cudaOlder
     cudatoolkit
-    cudaVersion
     flags
     libcublas
     setupCudaHook
@@ -24,6 +25,7 @@ backendStdenv.mkDerivation {
 
   src = ./.;
 
+  __structuredAttrs = true;
   strictDeps = true;
 
   nativeBuildInputs =
@@ -31,24 +33,22 @@ backendStdenv.mkDerivation {
       cmake
       autoAddDriverRunpath
     ]
-    ++ lib.optionals (lib.versionOlder cudaVersion "11.4") [ cudatoolkit ]
-    ++ lib.optionals (lib.versionAtLeast cudaVersion "11.4") [ cuda_nvcc ];
+    ++ lib.optionals (cudaOlder "11.4") [ cudatoolkit ]
+    ++ lib.optionals (cudaAtLeast "11.4") [ cuda_nvcc ];
 
   buildInputs =
-    lib.optionals (lib.versionOlder cudaVersion "11.4") [ cudatoolkit ]
-    ++ lib.optionals (lib.versionAtLeast cudaVersion "11.4") [
+    lib.optionals (cudaOlder "11.4") [ cudatoolkit ]
+    ++ lib.optionals (cudaAtLeast "11.4") [
       (getDev libcublas)
       (getLib libcublas)
       (getOutput "static" libcublas)
       cuda_cudart
     ]
-    ++ lib.optionals (lib.versionAtLeast cudaVersion "12.0") [ cuda_cccl ];
+    ++ lib.optionals (cudaAtLeast "12.0") [ cuda_cccl ];
 
-  cmakeFlags = [
+  cmakeFlagsArray = [
     (lib.cmakeBool "CMAKE_VERBOSE_MAKEFILE" true)
-    (lib.cmakeFeature "CMAKE_CUDA_ARCHITECTURES" (
-      with flags; lib.concatStringsSep ";" (lib.lists.map dropDot cudaCapabilities)
-    ))
+    (lib.cmakeFeature "CMAKE_CUDA_ARCHITECTURES" flags.cmakeCudaArchitecturesString)
   ];
 
   meta = rec {
@@ -56,6 +56,6 @@ backendStdenv.mkDerivation {
     license = lib.licenses.mit;
     maintainers = lib.teams.cuda.members;
     platforms = lib.platforms.unix;
-    badPlatforms = lib.optionals flags.isJetsonBuild platforms;
+    badPlatforms = lib.optionals (flags.isJetsonBuild && cudaOlder "11.4") platforms;
   };
 }

From 0494330fad2dde171bbb3f09795e4e6347f50ed8 Mon Sep 17 00:00:00 2001
From: Connor Baker <connor.baker@tweag.io>
Date: Wed, 3 Apr 2024 22:27:03 +0000
Subject: [PATCH 11/15] cudaPackages.nccl: switch to cudaAtLeast, cudaOlder,
 and __structuredAttrs

---
 .../development/cuda-modules/nccl/default.nix | 30 ++++++++++---------
 1 file changed, 16 insertions(+), 14 deletions(-)

diff --git a/pkgs/development/cuda-modules/nccl/default.nix b/pkgs/development/cuda-modules/nccl/default.nix
index 9db08c722acd..ec84b8dfb906 100644
--- a/pkgs/development/cuda-modules/nccl/default.nix
+++ b/pkgs/development/cuda-modules/nccl/default.nix
@@ -17,9 +17,10 @@ let
     cuda_cccl
     cuda_cudart
     cuda_nvcc
+    cudaAtLeast
     cudaFlags
+    cudaOlder
     cudatoolkit
-    cudaVersion
     ;
 in
 backendStdenv.mkDerivation (finalAttrs: {
@@ -33,6 +34,7 @@ backendStdenv.mkDerivation (finalAttrs: {
     hash = "sha256-ModIjD6RaRD/57a/PA1oTgYhZsAQPrrvhl5sNVXnO6c=";
   };
 
+  __structuredAttrs = true;
   strictDeps = true;
 
   outputs = [
@@ -46,12 +48,12 @@ backendStdenv.mkDerivation (finalAttrs: {
       autoAddDriverRunpath
       python3
     ]
-    ++ lib.optionals (lib.versionOlder cudaVersion "11.4") [ cudatoolkit ]
-    ++ lib.optionals (lib.versionAtLeast cudaVersion "11.4") [ cuda_nvcc ];
+    ++ lib.optionals (cudaOlder "11.4") [ cudatoolkit ]
+    ++ lib.optionals (cudaAtLeast "11.4") [ cuda_nvcc ];
 
   buildInputs =
-    lib.optionals (lib.versionOlder cudaVersion "11.4") [ cudatoolkit ]
-    ++ lib.optionals (lib.versionAtLeast cudaVersion "11.4") [
+    lib.optionals (cudaOlder "11.4") [ cudatoolkit ]
+    ++ lib.optionals (cudaAtLeast "11.4") [
       cuda_nvcc.dev # crt/host_config.h
       cuda_cudart
     ]
@@ -59,25 +61,25 @@ backendStdenv.mkDerivation (finalAttrs: {
     # against other version, like below, it's important that we use the same format. Otherwise,
     # we'll get incorrect results.
     # For example, lib.versionAtLeast "12.0" "12.0.0" == false.
-    ++ lib.optionals (lib.versionAtLeast cudaVersion "12.0") [ cuda_cccl ];
+    ++ lib.optionals (cudaAtLeast "12.0") [ cuda_cccl ];
 
   env.NIX_CFLAGS_COMPILE = toString [ "-Wno-unused-function" ];
 
-  preConfigure = ''
+  postPatch = ''
     patchShebangs ./src/device/generate.py
-    makeFlagsArray+=(
-      "NVCC_GENCODE=${lib.concatStringsSep " " cudaFlags.gencode}"
-    )
   '';
 
-  makeFlags =
-    [ "PREFIX=$(out)" ]
-    ++ lib.optionals (lib.versionOlder cudaVersion "11.4") [
+  makeFlagsArray =
+    [
+      "PREFIX=$(out)"
+      "NVCC_GENCODE=${cudaFlags.gencodeString}"
+    ]
+    ++ lib.optionals (cudaOlder "11.4") [
       "CUDA_HOME=${cudatoolkit}"
       "CUDA_LIB=${lib.getLib cudatoolkit}/lib"
       "CUDA_INC=${lib.getDev cudatoolkit}/include"
     ]
-    ++ lib.optionals (lib.versionAtLeast cudaVersion "11.4") [
+    ++ lib.optionals (cudaAtLeast "11.4") [
       "CUDA_HOME=${cuda_nvcc}"
       "CUDA_LIB=${lib.getLib cuda_cudart}/lib"
       "CUDA_INC=${lib.getDev cuda_cudart}/include"

From d5cbe889f0476c97a7acfd8826ba0e7b88319205 Mon Sep 17 00:00:00 2001
From: Connor Baker <connor.baker@tweag.io>
Date: Wed, 3 Apr 2024 23:25:54 +0000
Subject: [PATCH 12/15] cuda-modules/generic-builders/manifest: add to
 brokenConditions, simplify src, break out comments in postPatch

---
 .../generic-builders/manifest.nix             | 90 ++++++++++---------
 1 file changed, 50 insertions(+), 40 deletions(-)

diff --git a/pkgs/development/cuda-modules/generic-builders/manifest.nix b/pkgs/development/cuda-modules/generic-builders/manifest.nix
index e5831ef99dd0..ce45ada68fc2 100644
--- a/pkgs/development/cuda-modules/generic-builders/manifest.nix
+++ b/pkgs/development/cuda-modules/generic-builders/manifest.nix
@@ -128,7 +128,18 @@ backendStdenv.mkDerivation (finalAttrs: {
   # brokenConditions :: AttrSet Bool
   # Sets `meta.broken = true` if any of the conditions are true.
   # Example: Broken on a specific version of CUDA or when a dependency has a specific version.
-  brokenConditions = { };
+  brokenConditions = {
+    # Unclear how this is handled by Nix internals.
+    "Duplicate entries in outputs" = finalAttrs.outputs != lists.unique finalAttrs.outputs;
+    # Typically this results in the static output being empty, as all libraries are moved
+    # back to the lib output.
+    "lib output follows static output" =
+      let
+        libIndex = lists.findFirstIndex (x: x == "lib") null finalAttrs.outputs;
+        staticIndex = lists.findFirstIndex (x: x == "static") null finalAttrs.outputs;
+      in
+      libIndex != null && staticIndex != null && libIndex > staticIndex;
+  };
 
   # badPlatformsConditions :: AttrSet Bool
   # Sets `meta.badPlatforms = meta.platforms` if any of the conditions are true.
@@ -138,44 +149,43 @@ backendStdenv.mkDerivation (finalAttrs: {
   };
 
   # src :: Optional Derivation
-  src = trivial.pipe redistArch [
-    # If redistArch doesn't exist in redistribRelease, return null.
-    (redistArch: redistribRelease.${redistArch} or null)
-    # If the release is non-null, fetch the source; otherwise, return null.
-    (trivial.mapNullable (
-      { relative_path, sha256, ... }:
-      fetchurl {
-        url = "https://developer.download.nvidia.com/compute/${redistName}/redist/${relative_path}";
-        inherit sha256;
-      }
-    ))
-  ];
-
-  # Handle the pkg-config files:
-  # 1. No FHS
-  # 2. Location expected by the pkg-config wrapper
-  # 3. Generate unversioned names too
-  postPatch = ''
-    for path in pkg-config pkgconfig ; do
-      [[ -d "$path" ]] || continue
-      mkdir -p share/pkgconfig
-      mv "$path"/* share/pkgconfig/
-      rmdir "$path"
-    done
-
-    for pc in share/pkgconfig/*.pc ; do
-      sed -i \
-        -e "s|^cudaroot\s*=.*\$|cudaroot=''${!outputDev}|" \
-        -e "s|^libdir\s*=.*/lib\$|libdir=''${!outputLib}/lib|" \
-        -e "s|^includedir\s*=.*/include\$|includedir=''${!outputDev}/include|" \
-        "$pc"
-    done
+  # If redistArch doesn't exist in redistribRelease, return null.
+  src = trivial.mapNullable (
+    { relative_path, sha256, ... }:
+    fetchurl {
+      url = "https://developer.download.nvidia.com/compute/${redistName}/redist/${relative_path}";
+      inherit sha256;
+    }
+  ) (redistribRelease.${redistArch} or null);
 
+  postPatch =
+    # Pkg-config's setup hook expects configuration files in $out/share/pkgconfig
+    ''
+      for path in pkg-config pkgconfig; do
+        [[ -d "$path" ]] || continue
+        mkdir -p share/pkgconfig
+        mv "$path"/* share/pkgconfig/
+        rmdir "$path"
+      done
+    ''
+    # Rewrite FHS paths with store paths
+    # NOTE: output* fall back to out if the corresponding output isn't defined.
+    + ''
+      for pc in share/pkgconfig/*.pc; do
+        sed -i \
+          -e "s|^cudaroot\s*=.*\$|cudaroot=''${!outputDev}|" \
+          -e "s|^libdir\s*=.*/lib\$|libdir=''${!outputLib}/lib|" \
+          -e "s|^includedir\s*=.*/include\$|includedir=''${!outputDev}/include|" \
+          "$pc"
+      done
+    ''
+    # Generate unversioned names.
     # E.g. cuda-11.8.pc -> cuda.pc
-    for pc in share/pkgconfig/*-"$majorMinorVersion.pc" ; do
-      ln -s "$(basename "$pc")" "''${pc%-$majorMinorVersion.pc}".pc
-    done
-  '';
+    + ''
+      for pc in share/pkgconfig/*-"$majorMinorVersion.pc"; do
+        ln -s "$(basename "$pc")" "''${pc%-$majorMinorVersion.pc}".pc
+      done
+    '';
 
   env.majorMinorVersion = cudaMajorMinorVersion;
 
@@ -234,7 +244,7 @@ backendStdenv.mkDerivation (finalAttrs: {
     # Handle the existence of libPath, which requires us to re-arrange the lib directory
     + strings.optionalString (libPath != null) ''
       full_lib_path="lib/${libPath}"
-      if [[ ! -d "$full_lib_path" ]] ; then
+      if [[ ! -d "$full_lib_path" ]]; then
         echo "${finalAttrs.pname}: '$full_lib_path' does not exist, only found:" >&2
         find lib/ -mindepth 1 -maxdepth 1 >&2
         echo "This release might not support your CUDA version" >&2
@@ -265,9 +275,9 @@ backendStdenv.mkDerivation (finalAttrs: {
   postInstallCheck = ''
     echo "Executing postInstallCheck"
 
-    if [[ -z "''${allowFHSReferences-}" ]] ; then
+    if [[ -z "''${allowFHSReferences-}" ]]; then
       mapfile -t outputPaths < <(for o in $(getAllOutputNames); do echo "''${!o}"; done)
-      if grep --max-count=5 --recursive --exclude=LICENSE /usr/ "''${outputPaths[@]}" ; then
+      if grep --max-count=5 --recursive --exclude=LICENSE /usr/ "''${outputPaths[@]}"; then
         echo "Detected references to /usr" >&2
         exit 1
       fi

From 843b91a0eee065972051c4be0945dd0ad8c628d3 Mon Sep 17 00:00:00 2001
From: Connor Baker <connor.baker@tweag.io>
Date: Thu, 4 Apr 2024 02:03:20 +0000
Subject: [PATCH 13/15] cuda-modules/cuda/overrides: refactor

---
 .../cuda-modules/cuda/overrides.nix           | 452 +++++++++++-------
 pkgs/top-level/cuda-packages.nix              |   2 +-
 2 files changed, 270 insertions(+), 184 deletions(-)

diff --git a/pkgs/development/cuda-modules/cuda/overrides.nix b/pkgs/development/cuda-modules/cuda/overrides.nix
index c170e3d39c91..5d23d8f7f2a1 100644
--- a/pkgs/development/cuda-modules/cuda/overrides.nix
+++ b/pkgs/development/cuda-modules/cuda/overrides.nix
@@ -1,122 +1,178 @@
-{
-  cudaVersion,
-  lib,
-  addDriverRunpath,
-}:
 let
-  inherit (lib) attrsets lists strings;
-  # cudaVersionOlder : Version -> Boolean
-  cudaVersionOlder = strings.versionOlder cudaVersion;
-  # cudaVersionAtLeast : Version -> Boolean
-  cudaVersionAtLeast = strings.versionAtLeast cudaVersion;
+  filterAndCreateOverrides =
+    createOverrideAttrs: final: prev:
+    let
+      # It is imperative that we use `final.callPackage` to perform overrides,
+      # so the final package set is available to the override functions.
+      inherit (final) callPackage;
 
-  addBuildInputs =
-    drv: buildInputs:
-    drv.overrideAttrs (prevAttrs: {
-      buildInputs = prevAttrs.buildInputs ++ buildInputs;
-    });
-in
-# NOTE: Filter out attributes that are not present in the previous version of
-# the package set. This is necessary to prevent the appearance of attributes
-# like `cuda_nvcc` in `cudaPackages_10_0, which predates redistributables.
-final: prev:
-attrsets.filterAttrs (attr: _: (builtins.hasAttr attr prev)) {
-  libcufile = prev.libcufile.overrideAttrs (prevAttrs: {
-    buildInputs = prevAttrs.buildInputs ++ [
-      final.libcublas.lib
-      final.pkgs.numactl
-      final.pkgs.rdma-core
+      # NOTE(@connorbaker): We MUST use `lib` from `prev` because the attribute
+      # names CAN NOT depend on `final`.
+      inherit (prev.lib.attrsets) filterAttrs mapAttrs;
+      inherit (prev.lib.trivial) pipe;
+
+      # NOTE: Filter out attributes that are not present in the previous version of
+      # the package set. This is necessary to prevent the appearance of attributes
+      # like `cuda_nvcc` in `cudaPackages_10_0, which predates redistributables.
+      filterOutNewAttrs = filterAttrs (name: _: prev ? ${name});
+
+      # Apply callPackage to each attribute value, yielding a value to be passed
+      # to overrideAttrs.
+      callPackageThenOverrideAttrs = mapAttrs (
+        name: value: prev.${name}.overrideAttrs (callPackage value { })
+      );
+    in
+    pipe createOverrideAttrs [
+      filterOutNewAttrs
+      callPackageThenOverrideAttrs
     ];
-    # Before 11.7 libcufile depends on itself for some reason.
-    autoPatchelfIgnoreMissingDeps =
-      prevAttrs.autoPatchelfIgnoreMissingDeps
-      ++ lists.optionals (cudaVersionOlder "11.7") [ "libcufile.so.0" ];
-  });
+in
+# Each attribute name is the name of an existing package in the previous version
+# of the package set.
+# The value is a function (to be provided to callPackage), which yields a value
+# to be provided to overrideAttrs. This allows us to override the attributes of
+# a package without losing access to the fixed point of the package set --
+# especially useful given that some packages may depend on each other!
+filterAndCreateOverrides {
+  libcufile =
+    {
+      cudaOlder,
+      lib,
+      libcublas,
+      numactl,
+      rdma-core,
+    }:
+    prevAttrs: {
+      buildInputs = prevAttrs.buildInputs ++ [
+        libcublas.lib
+        numactl
+        rdma-core
+      ];
+      # Before 11.7 libcufile depends on itself for some reason.
+      autoPatchelfIgnoreMissingDeps =
+        prevAttrs.autoPatchelfIgnoreMissingDeps
+        ++ lib.lists.optionals (cudaOlder "11.7") [ "libcufile.so.0" ];
+    };
 
-  libcusolver = addBuildInputs prev.libcusolver (
-    # Always depends on this
-    [ final.libcublas.lib ]
-    # Dependency from 12.0 and on
-    ++ lists.optionals (cudaVersionAtLeast "12.0") [ final.libnvjitlink.lib ]
-    # Dependency from 12.1 and on
-    ++ lists.optionals (cudaVersionAtLeast "12.1") [ final.libcusparse.lib ]
-  );
+  libcusolver =
+    {
+      cudaAtLeast,
+      lib,
+      libcublas,
+      libcusparse ? null,
+      libnvjitlink ? null,
+    }:
+    prevAttrs: {
+      buildInputs =
+        prevAttrs.buildInputs
+        # Always depends on this
+        ++ [ libcublas.lib ]
+        # Dependency from 12.0 and on
+        ++ lib.lists.optionals (cudaAtLeast "12.0") [ libnvjitlink.lib ]
+        # Dependency from 12.1 and on
+        ++ lib.lists.optionals (cudaAtLeast "12.1") [ libcusparse.lib ];
 
-  libcusparse = addBuildInputs prev.libcusparse (
-    lists.optionals (cudaVersionAtLeast "12.0") [ final.libnvjitlink.lib ]
-  );
+      brokenConditions = prevAttrs.brokenConditions // {
+        "libnvjitlink missing (CUDA >= 12.0)" =
+          !(cudaAtLeast "12.0" -> (libnvjitlink != null && libnvjitlink.lib != null));
+        "libcusparse missing (CUDA >= 12.1)" =
+          !(cudaAtLeast "12.1" -> (libcusparse != null && libcusparse.lib != null));
+      };
+    };
 
-  cuda_cudart = prev.cuda_cudart.overrideAttrs (prevAttrs: {
-    # Remove once cuda-find-redist-features has a special case for libcuda
-    outputs =
-      prevAttrs.outputs
-      ++ lists.optionals (!(builtins.elem "stubs" prevAttrs.outputs)) [ "stubs" ];
+  libcusparse =
+    {
+      cudaAtLeast,
+      lib,
+      libnvjitlink ? null,
+    }:
+    prevAttrs: {
+      buildInputs =
+        prevAttrs.buildInputs
+        # Dependency from 12.0 and on
+        ++ lib.lists.optionals (cudaAtLeast "12.0") [ libnvjitlink.lib ];
 
-    allowFHSReferences = false;
+      brokenConditions = prevAttrs.brokenConditions // {
+        "libnvjitlink missing (CUDA >= 12.0)" =
+          !(cudaAtLeast "12.0" -> (libnvjitlink != null && libnvjitlink.lib != null));
+      };
+    };
 
-    # The libcuda stub's pkg-config doesn't follow the general pattern:
-    postPatch =
-      prevAttrs.postPatch or ""
-      + ''
-        while IFS= read -r -d $'\0' path ; do
-          sed -i \
-            -e "s|^libdir\s*=.*/lib\$|libdir=''${!outputLib}/lib/stubs|" \
-            -e "s|^Libs\s*:\(.*\)\$|Libs: \1 -Wl,-rpath,${addDriverRunpath.driverLink}/lib|" \
-            "$path"
-        done < <(find -iname 'cuda-*.pc' -print0)
-      ''
-      + ''
+  # TODO(@connorbaker): cuda_cudart.dev depends on crt/host_config.h, which is from
+  # cuda_nvcc.dev. It would be nice to be able to encode that.
+  cuda_cudart =
+    { addDriverRunpath, lib }:
+    prevAttrs: {
+      # Remove once cuda-find-redist-features has a special case for libcuda
+      outputs =
+        prevAttrs.outputs
+        ++ lib.lists.optionals (!(builtins.elem "stubs" prevAttrs.outputs)) [ "stubs" ];
+
+      allowFHSReferences = false;
+
+      # The libcuda stub's pkg-config doesn't follow the general pattern:
+      postPatch =
+        prevAttrs.postPatch or ""
+        + ''
+          while IFS= read -r -d $'\0' path; do
+            sed -i \
+              -e "s|^libdir\s*=.*/lib\$|libdir=''${!outputLib}/lib/stubs|" \
+              -e "s|^Libs\s*:\(.*\)\$|Libs: \1 -Wl,-rpath,${addDriverRunpath.driverLink}/lib|" \
+              "$path"
+          done < <(find -iname 'cuda-*.pc' -print0)
+        ''
         # Namelink may not be enough, add a soname.
         # Cf. https://gitlab.kitware.com/cmake/cmake/-/issues/25536
-        if [[ -f lib/stubs/libcuda.so && ! -f lib/stubs/libcuda.so.1 ]] ; then
-          ln -s libcuda.so lib/stubs/libcuda.so.1
-        fi
-      '';
+        + ''
+          if [[ -f lib/stubs/libcuda.so && ! -f lib/stubs/libcuda.so.1 ]]; then
+            ln -s libcuda.so lib/stubs/libcuda.so.1
+          fi
+        '';
 
-    postFixup =
-      prevAttrs.postFixup or ""
-      + ''
-        moveToOutput lib/stubs "$stubs"
-        ln -s "$stubs"/lib/stubs/* "$stubs"/lib/
-        ln -s "$stubs"/lib/stubs "''${!outputLib}/lib/stubs"
-      '';
-  });
-
-  cuda_compat = prev.cuda_compat.overrideAttrs (prevAttrs: {
-    autoPatchelfIgnoreMissingDeps = prevAttrs.autoPatchelfIgnoreMissingDeps ++ [
-      "libnvrm_gpu.so"
-      "libnvrm_mem.so"
-      "libnvdla_runtime.so"
-    ];
-    # `cuda_compat` only works on aarch64-linux, and only when building for Jetson devices.
-    badPlatformsConditions = prevAttrs.badPlatformsConditions // {
-      "Trying to use cuda_compat on aarch64-linux targeting non-Jetson devices" =
-        !final.flags.isJetsonBuild;
+      postFixup =
+        prevAttrs.postFixup or ""
+        + ''
+          moveToOutput lib/stubs "$stubs"
+          ln -s "$stubs"/lib/stubs/* "$stubs"/lib/
+          ln -s "$stubs"/lib/stubs "''${!outputLib}/lib/stubs"
+        '';
     };
-  });
 
-  cuda_gdb = addBuildInputs prev.cuda_gdb (
-    # x86_64 only needs gmp from 12.0 and on
-    lists.optionals (cudaVersionAtLeast "12.0") [ final.pkgs.gmp ]
-  );
-
-  cuda_nvcc = prev.cuda_nvcc.overrideAttrs (
-    oldAttrs:
-    let
-      # This replicates the logic in stdenvAdapters.useLibsFrom, except we use
-      # gcc from pkgsHostTarget and not from buildPackages.
-      ccForLibs-wrapper = final.pkgs.stdenv.cc;
-      gccMajorVersion = final.nvccCompatibilities.${cudaVersion}.gccMaxMajorVersion;
-      cc = final.pkgs.wrapCCWith {
-        cc = final.pkgs."gcc${gccMajorVersion}".cc;
-        useCcForLibs = true;
-        gccForLibs = ccForLibs-wrapper.cc;
+  cuda_compat =
+    { flags, lib }:
+    prevAttrs: {
+      autoPatchelfIgnoreMissingDeps = prevAttrs.autoPatchelfIgnoreMissingDeps ++ [
+        "libnvrm_gpu.so"
+        "libnvrm_mem.so"
+        "libnvdla_runtime.so"
+      ];
+      # `cuda_compat` only works on aarch64-linux, and only when building for Jetson devices.
+      badPlatformsConditions = prevAttrs.badPlatformsConditions // {
+        "Trying to use cuda_compat on aarch64-linux targeting non-Jetson devices" = !flags.isJetsonBuild;
       };
-    in
+    };
+
+  cuda_gdb =
     {
+      cudaAtLeast,
+      gmp,
+      lib,
+    }:
+    prevAttrs: {
+      buildInputs =
+        prevAttrs.buildInputs
+        # x86_64 only needs gmp from 12.0 and on
+        ++ lib.lists.optionals (cudaAtLeast "12.0") [ gmp ];
+    };
 
-      outputs = oldAttrs.outputs ++ lists.optionals (!(builtins.elem "lib" oldAttrs.outputs)) [ "lib" ];
-
+  cuda_nvcc =
+    {
+      backendStdenv,
+      cuda_cudart,
+      lib,
+      setupCudaHook,
+    }:
+    prevAttrs: {
       # Patch the nvcc.profile.
       # Syntax:
       # - `=` for assignment,
@@ -131,12 +187,9 @@ attrsets.filterAttrs (attr: _: (builtins.hasAttr attr prev)) {
       # backend-stdenv.nix
 
       postPatch =
-        (oldAttrs.postPatch or "")
+        (prevAttrs.postPatch or "")
         + ''
           substituteInPlace bin/nvcc.profile \
-            --replace-fail \
-              '$(TOP)/lib' \
-              "''${!outputLib}/lib" \
             --replace-fail \
               '$(TOP)/$(_NVVM_BRANCH_)' \
               "''${!outputBin}/nvvm" \
@@ -147,22 +200,24 @@ attrsets.filterAttrs (attr: _: (builtins.hasAttr attr prev)) {
           cat << EOF >> bin/nvcc.profile
 
           # Fix a compatible backend compiler
-          PATH += ${lib.getBin cc}/bin:
+          PATH += "${backendStdenv.cc}/bin":
 
           # Expose the split-out nvvm
-          LIBRARIES =+ -L''${!outputBin}/nvvm/lib
-          INCLUDES =+ -I''${!outputBin}/nvvm/include
-
-          # Expose cudart and the libcuda stubs
-          LIBRARIES =+ -L$static/lib" "-L${final.cuda_cudart.lib}/lib -L${final.cuda_cudart.lib}/lib/stubs
-          INCLUDES =+ -I${final.cuda_cudart.dev}/include
+          LIBRARIES =+ "-L''${!outputBin}/nvvm/lib"
+          INCLUDES =+ "-I''${!outputBin}/nvvm/include"
           EOF
         '';
 
-      propagatedBuildInputs = [ final.setupCudaHook ];
+      # NOTE(@connorbaker):
+      # Though it might seem odd or counter-intuitive to add the setup hook to `propagatedBuildInputs` instead of
+      # `propagatedNativeBuildInputs`, it is necessary! If you move the setup hook from `propagatedBuildInputs` to
+      # `propagatedNativeBuildInputs`, it stops being propagated to downstream packages during their build because
+      # setup hooks in `propagatedNativeBuildInputs` are not designed to affect the runtime or build environment of
+      # dependencies; they are only meant to affect the build environment of the package that directly includes them.
+      propagatedBuildInputs = (prevAttrs.propagatedBuildInputs or [ ]) ++ [ setupCudaHook ];
 
       postInstall =
-        (oldAttrs.postInstall or "")
+        (prevAttrs.postInstall or "")
         + ''
           moveToOutput "nvvm" "''${!outputBin}"
         '';
@@ -170,48 +225,77 @@ attrsets.filterAttrs (attr: _: (builtins.hasAttr attr prev)) {
       # The nvcc and cicc binaries contain hard-coded references to /usr
       allowFHSReferences = true;
 
-      meta = (oldAttrs.meta or { }) // {
+      meta = (prevAttrs.meta or { }) // {
         mainProgram = "nvcc";
       };
-    }
-  );
+    };
 
-  cuda_nvprof = prev.cuda_nvprof.overrideAttrs (prevAttrs: {
-    buildInputs = prevAttrs.buildInputs ++ [ final.cuda_cupti.lib ];
-  });
+  cuda_nvprof =
+    { cuda_cupti }: prevAttrs: { buildInputs = prevAttrs.buildInputs ++ [ cuda_cupti.lib ]; };
 
-  cuda_demo_suite = addBuildInputs prev.cuda_demo_suite [
-    final.pkgs.freeglut
-    final.pkgs.libGLU
-    final.pkgs.libglvnd
-    final.pkgs.mesa
-    final.libcufft.lib
-    final.libcurand.lib
-  ];
+  cuda_demo_suite =
+    {
+      freeglut,
+      libcufft,
+      libcurand,
+      libGLU,
+      libglvnd,
+      mesa,
+    }:
+    prevAttrs: {
+      buildInputs = prevAttrs.buildInputs ++ [
+        freeglut
+        libcufft.lib
+        libcurand.lib
+        libGLU
+        libglvnd
+        mesa
+      ];
+    };
 
-  nsight_compute = prev.nsight_compute.overrideAttrs (prevAttrs: {
-    nativeBuildInputs =
-      prevAttrs.nativeBuildInputs
-      ++ (
-        if (strings.versionOlder prev.nsight_compute.version "2022.2.0") then
-          [ final.pkgs.qt5.wrapQtAppsHook ]
-        else
-          [ final.pkgs.qt6.wrapQtAppsHook ]
-      );
-    buildInputs =
-      prevAttrs.buildInputs
-      ++ (
-        if (strings.versionOlder prev.nsight_compute.version "2022.2.0") then
-          [ final.pkgs.qt5.qtwebview ]
-        else
-          [ final.pkgs.qt6.qtwebview ]
-      );
-  });
-
-  nsight_systems = prev.nsight_systems.overrideAttrs (
+  nsight_compute =
+    {
+      lib,
+      qt5 ? null,
+      qt6 ? null,
+    }:
     prevAttrs:
     let
-      qt = if lib.versionOlder prevAttrs.version "2022.4.2.1" then final.pkgs.qt5 else final.pkgs.qt6;
+      inherit (lib.strings) versionOlder versionAtLeast;
+      inherit (prevAttrs) version;
+      qt = if versionOlder version "2022.2.0" then qt5 else qt6;
+      inherit (qt) wrapQtAppsHook qtwebview;
+    in
+    {
+      nativeBuildInputs = prevAttrs.nativeBuildInputs ++ [ wrapQtAppsHook ];
+      buildInputs = prevAttrs.buildInputs ++ [ qtwebview ];
+      brokenConditions = prevAttrs.brokenConditions // {
+        "Qt 5 missing (<2022.2.0)" = !(versionOlder version "2022.2.0" -> qt5 != null);
+        "Qt 6 missing (>=2022.2.0)" = !(versionAtLeast version "2022.2.0" -> qt6 != null);
+      };
+    };
+
+  nsight_systems =
+    {
+      cuda_cudart,
+      cudaOlder,
+      gst_all_1,
+      lib,
+      nss,
+      numactl,
+      pulseaudio,
+      qt5 ? null,
+      qt6 ? null,
+      rdma-core,
+      ucx,
+      wayland,
+      xorg,
+    }:
+    prevAttrs:
+    let
+      inherit (lib.strings) versionOlder versionAtLeast;
+      inherit (prevAttrs) version;
+      qt = if lib.strings.versionOlder prevAttrs.version "2022.4.2.1" then qt5 else qt6;
       qtwayland =
         if lib.versions.major qt.qtbase.version == "5" then
           lib.getBin qt.qtwayland
@@ -223,55 +307,57 @@ attrsets.filterAttrs (attr: _: (builtins.hasAttr attr prev)) {
       # An ad hoc replacement for
       # https://github.com/ConnorBaker/cuda-redist-find-features/issues/11
       env.rmPatterns = toString [
+        "nsight-systems/*/*/lib{arrow,jpeg}*"
+        "nsight-systems/*/*/lib{ssl,ssh,crypto}*"
+        "nsight-systems/*/*/libboost*"
+        "nsight-systems/*/*/libexec"
         "nsight-systems/*/*/libQt*"
         "nsight-systems/*/*/libstdc*"
-        "nsight-systems/*/*/libboost*"
-        "nsight-systems/*/*/lib{ssl,ssh,crypto}*"
-        "nsight-systems/*/*/lib{arrow,jpeg}*"
         "nsight-systems/*/*/Mesa"
-        "nsight-systems/*/*/python/bin/python"
-        "nsight-systems/*/*/libexec"
         "nsight-systems/*/*/Plugins"
+        "nsight-systems/*/*/python/bin/python"
       ];
       postPatch =
         prevAttrs.postPatch or ""
         + ''
-          for path in $rmPatterns ; do
+          for path in $rmPatterns; do
             rm -r "$path"
           done
         '';
       nativeBuildInputs = prevAttrs.nativeBuildInputs ++ [ qt.wrapQtAppsHook ];
       buildInputs = prevAttrs.buildInputs ++ [
-        final.cuda_cudart.stubs
-        final.pkgs.alsa-lib
-        final.pkgs.boost178
-        final.pkgs.e2fsprogs
-        final.pkgs.gst_all_1.gst-plugins-base
-        final.pkgs.gst_all_1.gstreamer
-        final.pkgs.nss
-        final.pkgs.numactl
-        final.pkgs.pulseaudio
-        final.pkgs.rdma-core
-        final.pkgs.ucx
-        final.pkgs.wayland
-        final.pkgs.xorg.libXcursor
-        final.pkgs.xorg.libXdamage
-        final.pkgs.xorg.libXrandr
-        final.pkgs.xorg.libXtst
-        qt.qtbase
         (qt.qtdeclarative or qt.full)
         (qt.qtsvg or qt.full)
+        cuda_cudart.stubs
+        gst_all_1.gst-plugins-base
+        gst_all_1.gstreamer
+        nss
+        numactl
+        pulseaudio
+        qt.qtbase
         qtWaylandPlugins
+        rdma-core
+        ucx
+        wayland
+        xorg.libXcursor
+        xorg.libXdamage
+        xorg.libXrandr
+        xorg.libXtst
       ];
 
-      # Older releases require boost 1.70 deprecated in Nixpkgs
-      meta.broken = prevAttrs.meta.broken or false || lib.versionOlder final.cudaVersion "11.8";
-    }
-  );
+      brokenConditions = prevAttrs.brokenConditions // {
+        # Older releases require boost 1.70, which is deprecated in Nixpkgs
+        "CUDA too old (<11.8)" = cudaOlder "11.8";
+        "Qt 5 missing (<2022.4.2.1)" = !(versionOlder version "2022.4.2.1" -> qt5 != null);
+        "Qt 6 missing (>=2022.4.2.1)" = !(versionAtLeast version "2022.4.2.1" -> qt6 != null);
+      };
+    };
 
-  nvidia_driver = prev.nvidia_driver.overrideAttrs {
-    # No need to support this package as we have drivers already
-    # in linuxPackages.
-    meta.broken = true;
-  };
+  nvidia_driver =
+    { }:
+    prevAttrs: {
+      brokenConditions = prevAttrs.brokenConditions // {
+        "Package is not supported; use drivers from linuxPackages" = true;
+      };
+    };
 }
diff --git a/pkgs/top-level/cuda-packages.nix b/pkgs/top-level/cuda-packages.nix
index f573d2358147..c1952d821d19 100644
--- a/pkgs/top-level/cuda-packages.nix
+++ b/pkgs/top-level/cuda-packages.nix
@@ -86,7 +86,7 @@ let
     [
       (import ../development/cuda-modules/setup-hooks/extension.nix)
       (callPackage ../development/cuda-modules/cuda/extension.nix { inherit cudaVersion; })
-      (callPackage ../development/cuda-modules/cuda/overrides.nix { inherit cudaVersion; })
+      (import ../development/cuda-modules/cuda/overrides.nix)
       (callPackage ../development/cuda-modules/generic-builders/multiplex.nix {
         inherit cudaVersion flags mkVersionedPackageName;
         pname = "cudnn";

From 7a66dcf83f948c506472dec40f73330c4a9a37d1 Mon Sep 17 00:00:00 2001
From: Connor Baker <connor.baker@tweag.io>
Date: Thu, 18 Apr 2024 15:46:47 +0000
Subject: [PATCH 14/15] cuda-modules: use stdenv instead of backendStdenv for
 *Platform access

---
 .../cuda-modules/cuda-library-samples/extension.nix        | 7 +------
 pkgs/development/cuda-modules/cuda-samples/extension.nix   | 2 --
 pkgs/development/cuda-modules/cuda-samples/generic.nix     | 3 ++-
 pkgs/development/cuda-modules/cutensor/extension.nix       | 2 --
 pkgs/development/cuda-modules/flags.nix                    | 2 --
 .../development/cuda-modules/generic-builders/manifest.nix | 2 +-
 .../cuda-modules/generic-builders/multiplex.nix            | 2 --
 pkgs/development/cuda-modules/tensorrt/fixup.nix           | 4 ++--
 8 files changed, 6 insertions(+), 18 deletions(-)

diff --git a/pkgs/development/cuda-modules/cuda-library-samples/extension.nix b/pkgs/development/cuda-modules/cuda-library-samples/extension.nix
index eb291cd4da09..1184547c7f93 100644
--- a/pkgs/development/cuda-modules/cuda-library-samples/extension.nix
+++ b/pkgs/development/cuda-modules/cuda-library-samples/extension.nix
@@ -1,9 +1,4 @@
-{
-  lib,
-  # NOTE: We would use backendStdenv to get the hostPlatform, but it's not available in the callPackage context
-  # we're using to call this file.
-  stdenv,
-}:
+{ lib, stdenv }:
 let
   inherit (stdenv) hostPlatform;
 
diff --git a/pkgs/development/cuda-modules/cuda-samples/extension.nix b/pkgs/development/cuda-modules/cuda-samples/extension.nix
index e4a02e9a5e0c..0a8a3f9ff6ea 100644
--- a/pkgs/development/cuda-modules/cuda-samples/extension.nix
+++ b/pkgs/development/cuda-modules/cuda-samples/extension.nix
@@ -1,8 +1,6 @@
 {
   cudaVersion,
   lib,
-  # NOTE: We would use backendStdenv to get the hostPlatform, but it's not available in the callPackage context
-  # we're using to call this file.
   stdenv,
 }:
 let
diff --git a/pkgs/development/cuda-modules/cuda-samples/generic.nix b/pkgs/development/cuda-modules/cuda-samples/generic.nix
index a6a382c8a219..7b3a46acea92 100644
--- a/pkgs/development/cuda-modules/cuda-samples/generic.nix
+++ b/pkgs/development/cuda-modules/cuda-samples/generic.nix
@@ -11,6 +11,7 @@
   hash,
   lib,
   pkg-config,
+  stdenv,
 }:
 let
   inherit (lib) lists strings;
@@ -63,7 +64,7 @@ backendStdenv.mkDerivation (finalAttrs: {
   installPhase = ''
     runHook preInstall
 
-    install -Dm755 -t $out/bin bin/${backendStdenv.hostPlatform.parsed.cpu.name}/${backendStdenv.hostPlatform.parsed.kernel.name}/release/*
+    install -Dm755 -t $out/bin bin/${stdenv.hostPlatform.parsed.cpu.name}/${stdenv.hostPlatform.parsed.kernel.name}/release/*
 
     runHook postInstall
   '';
diff --git a/pkgs/development/cuda-modules/cutensor/extension.nix b/pkgs/development/cuda-modules/cutensor/extension.nix
index ee9b31972e18..5fdf356df916 100644
--- a/pkgs/development/cuda-modules/cutensor/extension.nix
+++ b/pkgs/development/cuda-modules/cutensor/extension.nix
@@ -17,8 +17,6 @@
   flags,
   lib,
   mkVersionedPackageName,
-  # NOTE: We would use backendStdenv to get the hostPlatform, but it's not available in the callPackage context
-  # we're using to call this file.
   stdenv,
 }:
 let
diff --git a/pkgs/development/cuda-modules/flags.nix b/pkgs/development/cuda-modules/flags.nix
index 7bba27a2f3e3..3a6bafdd8246 100644
--- a/pkgs/development/cuda-modules/flags.nix
+++ b/pkgs/development/cuda-modules/flags.nix
@@ -7,8 +7,6 @@
   cudaForwardCompat ? (config.cudaForwardCompat or true),
   lib,
   cudaVersion,
-  # NOTE: We cannot refer to backendStdenv as flags exists outside the fixed-point;
-  # otherwise, we would use that to get the hostPlatform.
   stdenv,
   # gpus :: List Gpu
   gpus,
diff --git a/pkgs/development/cuda-modules/generic-builders/manifest.nix b/pkgs/development/cuda-modules/generic-builders/manifest.nix
index ce45ada68fc2..006abb456cdc 100644
--- a/pkgs/development/cuda-modules/generic-builders/manifest.nix
+++ b/pkgs/development/cuda-modules/generic-builders/manifest.nix
@@ -39,7 +39,7 @@ let
     sourceTypes
     ;
 
-  inherit (backendStdenv) hostPlatform;
+  inherit (stdenv) hostPlatform;
 
   # Get the redist architectures for which package provides distributables.
   # These are used by meta.platforms.
diff --git a/pkgs/development/cuda-modules/generic-builders/multiplex.nix b/pkgs/development/cuda-modules/generic-builders/multiplex.nix
index fd06d19e60f0..ce50653bb5d4 100644
--- a/pkgs/development/cuda-modules/generic-builders/multiplex.nix
+++ b/pkgs/development/cuda-modules/generic-builders/multiplex.nix
@@ -3,8 +3,6 @@
   lib,
   cudaVersion,
   flags,
-  # NOTE: We would use backendStdenv to get the hostPlatform, but it's not available in the callPackage context
-  # we're using to call this file (cudaPackages.callPackage).
   stdenv,
   # Expected to be passed by the caller
   mkVersionedPackageName,
diff --git a/pkgs/development/cuda-modules/tensorrt/fixup.nix b/pkgs/development/cuda-modules/tensorrt/fixup.nix
index fb843aa4e1e1..42359aedac11 100644
--- a/pkgs/development/cuda-modules/tensorrt/fixup.nix
+++ b/pkgs/development/cuda-modules/tensorrt/fixup.nix
@@ -1,5 +1,4 @@
 {
-  backendStdenv,
   cudaVersion,
   final,
   lib,
@@ -7,6 +6,7 @@
   package,
   patchelf,
   requireFile,
+  stdenv,
   ...
 }:
 let
@@ -17,7 +17,7 @@ let
     strings
     versions
     ;
-  inherit (backendStdenv) hostPlatform;
+  inherit (stdenv) hostPlatform;
   # targetArch :: String
   targetArch = attrsets.attrByPath [ hostPlatform.system ] "unsupported" {
     x86_64-linux = "x86_64-linux-gnu";

From 6208368a822cf27adfc7b743e3be6d96ebe91095 Mon Sep 17 00:00:00 2001
From: Connor Baker <connor.baker@tweag.io>
Date: Thu, 18 Apr 2024 15:51:03 +0000
Subject: [PATCH 15/15] cudaPackages.cudaFlags: drop unused
 cmakeCudaArchitectures attribute

---
 pkgs/development/cuda-modules/flags.nix | 15 +--------------
 1 file changed, 1 insertion(+), 14 deletions(-)

diff --git a/pkgs/development/cuda-modules/flags.nix b/pkgs/development/cuda-modules/flags.nix
index 3a6bafdd8246..93952a66216b 100644
--- a/pkgs/development/cuda-modules/flags.nix
+++ b/pkgs/development/cuda-modules/flags.nix
@@ -209,15 +209,10 @@ let
       # E.g. "-gencode=arch=compute_75,code=sm_75 ... -gencode=arch=compute_86,code=compute_86"
       gencodeString = strings.concatStringsSep " " gencode;
 
-      # cmakeCudaArchitectures :: List String
-      # A list of CUDA capabilities without dots, suitable for passing to CMake.
-      # E.g. [ "75" "86" ]
-      cmakeCudaArchitectures = lists.map dropDot cudaCapabilities;
-
       # cmakeCudaArchitecturesString :: String
       # A semicolon-separated string of CUDA capabilities without dots, suitable for passing to CMake.
       # E.g. "75;86"
-      cmakeCudaArchitecturesString = strings.concatStringsSep ";" cmakeCudaArchitectures;
+      cmakeCudaArchitecturesString = strings.concatMapStringsSep ";" dropDot cudaCapabilities;
 
       # Jetson devices cannot be targeted by the same binaries which target non-Jetson devices. While
       # NVIDIA provides both `linux-aarch64` and `linux-sbsa` packages, which both target `aarch64`,
@@ -282,10 +277,6 @@ assert
       ];
       gencodeString = "-gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_86,code=sm_86 -gencode=arch=compute_86,code=compute_86";
 
-      cmakeCudaArchitectures = [
-        "75"
-        "86"
-      ];
       cmakeCudaArchitecturesString = "75;86";
 
       isJetsonBuild = false;
@@ -357,10 +348,6 @@ assert
       ];
       gencodeString = "-gencode=arch=compute_62,code=sm_62 -gencode=arch=compute_72,code=sm_72 -gencode=arch=compute_72,code=compute_72";
 
-      cmakeCudaArchitectures = [
-        "62"
-        "72"
-      ];
       cmakeCudaArchitecturesString = "62;72";
 
       isJetsonBuild = true;