From 36d409bc3a3f3a6967848ca3bb969d747cdf2c86 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>
---
 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 10319679abd1..181c5f17aa61 100644
--- a/doc/redirects.json
+++ b/doc/redirects.json
@@ -35,6 +35,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 53cc23af8a9d..52c357c83209 100644
--- a/pkgs/top-level/cuda-packages.nix
+++ b/pkgs/top-level/cuda-packages.nix
@@ -209,6 +209,7 @@ let
     ++ lib.optionals config.allowAliases [
       (import ../development/cuda-modules/aliases.nix { inherit lib; })
     ]
+    ++ _cuda.extensions
   );
 
   cudaPackages = customisation.makeScope newScope (