From c2c335d15e22174f1c08ccea7c98e04f91eeadbb Mon Sep 17 00:00:00 2001
From: Connor Baker <ConnorBaker01@gmail.com>
Date: Tue, 13 May 2025 17:30:40 +0000
Subject: [PATCH] _cuda.extensions: make overriding all CUDA package sets
 easier

Signed-off-by: Connor Baker <ConnorBaker01@gmail.com>
(cherry picked from commit 36d409bc3a3f3a6967848ca3bb969d747cdf2c86)
---
 doc/languages-frameworks/cuda.section.md        | 16 ++++++++++++++++
 doc/redirects.json                              |  3 +++
 pkgs/development/cuda-modules/_cuda/default.nix |  1 +
 pkgs/top-level/cuda-packages.nix                |  1 +
 4 files changed, 21 insertions(+)

diff --git a/doc/languages-frameworks/cuda.section.md b/doc/languages-frameworks/cuda.section.md
index 01d92a322d72..b4475614748f 100644
--- a/doc/languages-frameworks/cuda.section.md
+++ b/doc/languages-frameworks/cuda.section.md
@@ -95,6 +95,22 @@ final: prev: {
 }
 ```
 
+## Extending CUDA package sets {#cuda-extending-cuda-package-sets}
+
+CUDA package sets are scopes, so they provide the usual `overrideScope` attribute for overriding package attributes (see the note about `cudaMajorMinorVersion` and `_cuda` in [Configuring CUDA package sets](#cuda-configuring-cuda-package-sets)).
+
+Inspired by `pythonPackagesExtensions`, the `_cuda.extensions` attribute is a list of extensions applied to every version of the CUDA package set, allowing modification of all versions of the CUDA package set without having to know what they are or find a way to enumerate and modify them explicitly. As an example, disabling `cuda_compat` across all CUDA package sets can be accomplished with this overlay:
+
+```nix
+final: prev: {
+  _cuda = prev._cuda.extend (
+    _: prevAttrs: {
+      extensions = prevAttrs.extensions ++ [ (_: _: { cuda_compat = null; }) ];
+    }
+  );
+}
+```
+
 ## Using cudaPackages {#cuda-using-cudapackages}
 
 ::: {.caution}
diff --git a/doc/redirects.json b/doc/redirects.json
index 89aea52d2943..2ffb1a1f05d3 100644
--- a/doc/redirects.json
+++ b/doc/redirects.json
@@ -32,6 +32,9 @@
   "cuda-docker-podman": [
     "index.html#cuda-docker-podman"
   ],
+  "cuda-extending-cuda-package-sets": [
+    "index.html#cuda-extending-cuda-package-sets"
+  ],
   "cuda-package-set-maintenance": [
     "index.html#cuda-package-set-maintenance",
     "index.html#adding-a-new-cuda-release"
diff --git a/pkgs/development/cuda-modules/_cuda/default.nix b/pkgs/development/cuda-modules/_cuda/default.nix
index 0f6f80506616..fdbac3c8fbd7 100644
--- a/pkgs/development/cuda-modules/_cuda/default.nix
+++ b/pkgs/development/cuda-modules/_cuda/default.nix
@@ -22,6 +22,7 @@ lib.fixedPoints.makeExtensible (final: {
     inherit (final) bootstrapData db;
     inherit lib;
   };
+  extensions = [ ]; # Extensions applied to every CUDA package set.
   fixups = import ./fixups { inherit lib; };
   lib = import ./lib {
     _cuda = final;
diff --git a/pkgs/top-level/cuda-packages.nix b/pkgs/top-level/cuda-packages.nix
index 53dd4b6029ea..72fea9de9ec4 100644
--- a/pkgs/top-level/cuda-packages.nix
+++ b/pkgs/top-level/cuda-packages.nix
@@ -205,6 +205,7 @@ let
     ++ lib.optionals config.allowAliases [
       (import ../development/cuda-modules/aliases.nix { inherit lib; })
     ]
+    ++ _cuda.extensions
   );
 
   cudaPackages = customisation.makeScope newScope (