cudaPackages: move config expressions to cuda-modules

6 files changed, 1199 insertions, 0 deletions

diff --git a/pkgs/development/cuda-modules/cudatoolkit/releases.nix b/pkgs/development/cuda-modules/cudatoolkit/releases.nix
new file mode 100644
index 000000000000..b66f60060702
--- /dev/null
+++ b/pkgs/development/cuda-modules/cudatoolkit/releases.nix
@@ -0,0 +1,102 @@
+# Type Aliases
+# CudaVersion = String (two-component version, e.g. "10.0")
+# Release = {
+#   version: String
+#     - The version of CUDA.
+#   url: String
+#     - The URL to download the CUDA installer from.
+#   sha256: String
+#     - The SHA256 checksum of the CUDA installer.
+# }
+# Releases = AttrSet CudaVersion Release
+{
+  "10.0" = {
+    version = "10.0.130";
+    url = "https://developer.nvidia.com/compute/cuda/10.0/Prod/local_installers/cuda_10.0.130_410.48_linux";
+    sha256 = "16p3bv1lwmyqpxil8r951h385sy9asc578afrc7lssa68c71ydcj";
+  };
+
+  "10.1" = {
+    version = "10.1.243";
+    url = "https://developer.download.nvidia.com/compute/cuda/10.1/Prod/local_installers/cuda_10.1.243_418.87.00_linux.run";
+    sha256 = "0caxhlv2bdq863dfp6wj7nad66ml81vasq2ayf11psvq2b12vhp7";
+  };
+
+  "10.2" = {
+    version = "10.2.89";
+    url = "http://developer.download.nvidia.com/compute/cuda/10.2/Prod/local_installers/cuda_10.2.89_440.33.01_linux.run";
+    sha256 = "04fasl9sjkb1jvchvqgaqxprnprcz7a8r52249zp2ijarzyhf3an";
+  };
+
+  "11.0" = {
+    version = "11.0.3";
+    url = "https://developer.download.nvidia.com/compute/cuda/11.0.3/local_installers/cuda_11.0.3_450.51.06_linux.run";
+    sha256 = "1h4c69nfrgm09jzv8xjnjcvpq8n4gnlii17v3wzqry5d13jc8ydh";
+  };
+
+  "11.1" = {
+    version = "11.1.1";
+    url = "https://developer.download.nvidia.com/compute/cuda/11.1.1/local_installers/cuda_11.1.1_455.32.00_linux.run";
+    sha256 = "13yxv2fgvdnqqbwh1zb80x4xhyfkbajfkwyfpdg9493010kngbiy";
+  };
+
+  "11.2" = {
+    version = "11.2.1";
+    url = "https://developer.download.nvidia.com/compute/cuda/11.2.1/local_installers/cuda_11.2.1_460.32.03_linux.run";
+    sha256 = "sha256-HamMuJfMX1inRFpKZspPaSaGdwbLOvWKZpzc2Nw9F8g=";
+  };
+
+  "11.3" = {
+    version = "11.3.1";
+    url = "https://developer.download.nvidia.com/compute/cuda/11.3.1/local_installers/cuda_11.3.1_465.19.01_linux.run";
+    sha256 = "0d19pwcqin76scbw1s5kgj8n0z1p4v1hyfldqmamilyfxycfm4xd";
+  };
+
+  "11.4" = {
+    version = "11.4.2";
+    url = "https://developer.download.nvidia.com/compute/cuda/11.4.2/local_installers/cuda_11.4.2_470.57.02_linux.run";
+    sha256 = "sha256-u9h8oOkT+DdFSnljZ0c1E83e9VUILk2G7Zo4ZZzIHwo=";
+  };
+
+  "11.5" = {
+    version = "11.5.0";
+    url = "https://developer.download.nvidia.com/compute/cuda/11.5.0/local_installers/cuda_11.5.0_495.29.05_linux.run";
+    sha256 = "sha256-rgoWk9lJfPPYHmlIlD43lGNpANtxyY1Y7v2sr38aHkw=";
+  };
+
+  "11.6" = {
+    version = "11.6.1";
+    url = "https://developer.download.nvidia.com/compute/cuda/11.6.1/local_installers/cuda_11.6.1_510.47.03_linux.run";
+    sha256 = "sha256-qyGa/OALdCABEyaYZvv/derQN7z8I1UagzjCaEyYTX4=";
+  };
+
+  "11.7" = {
+    version = "11.7.0";
+    url = "https://developer.download.nvidia.com/compute/cuda/11.7.0/local_installers/cuda_11.7.0_515.43.04_linux.run";
+    sha256 = "sha256-CH/fy7ofeVQ7H3jkOo39rF9tskLQQt3oIOFtwYWJLyY=";
+  };
+
+  "11.8" = {
+    version = "11.8.0";
+    url = "https://developer.download.nvidia.com/compute/cuda/11.8.0/local_installers/cuda_11.8.0_520.61.05_linux.run";
+    sha256 = "sha256-kiPErzrr5Ke77Zq9mxY7A6GzS4VfvCtKDRtwasCaWhY=";
+  };
+
+  "12.0" = {
+    version = "12.0.1";
+    url = "https://developer.download.nvidia.com/compute/cuda/12.0.1/local_installers/cuda_12.0.1_525.85.12_linux.run";
+    sha256 = "sha256-GyBaBicvFGP0dydv2rkD8/ZmkXwGjlIHOAAeacehh1s=";
+  };
+
+  "12.1" = {
+    version = "12.1.1";
+    url = "https://developer.download.nvidia.com/compute/cuda/12.1.1/local_installers/cuda_12.1.1_530.30.02_linux.run";
+    sha256 = "sha256-10Ai1B2AEFMZ36Ib7qObd6W5kZU5wEh6BcqvJEbWpw4=";
+  };
+
+  "12.2" = {
+    version = "12.2.0";
+    url = "https://developer.download.nvidia.com/compute/cuda/12.2.0/local_installers/cuda_12.2.0_535.54.03_linux.run";
+    sha256 = "sha256-7PPSr63LrAKfD0UFeFgQ1S0AbkuHunn/P5hDNqK79Rg=";
+  };
+}
diff --git a/pkgs/development/cuda-modules/cudnn/releases.nix b/pkgs/development/cuda-modules/cudnn/releases.nix
new file mode 100644
index 000000000000..701cb4c20d36
--- /dev/null
+++ b/pkgs/development/cuda-modules/cudnn/releases.nix
@@ -0,0 +1,262 @@
+# NOTE: Check https://docs.nvidia.com/deeplearning/cudnn/archives/index.html for support matrices.
+# Version policy is to keep the latest minor release for each major release.
+{
+  cudnn.releases = {
+    # jetson
+    linux-aarch64 = [
+      {
+        version = "8.9.5.30";
+        minCudaVersion = "12.0";
+        maxCudaVersion = "12.2";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-aarch64/cudnn-linux-aarch64-8.9.5.30_cuda12-archive.tar.xz";
+        hash = "sha256-BJH3sC9VwiB362eL8xTB+RdSS9UHz1tlgjm/mKRyM6E=";
+      }
+    ];
+    # powerpc
+    linux-ppc64le = [];
+    # server-grade arm
+    linux-sbsa = [
+      {
+        version = "8.4.1.50";
+        minCudaVersion = "11.0";
+        maxCudaVersion = "11.7";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-sbsa/cudnn-linux-sbsa-8.4.1.50_cuda11.6-archive.tar.xz";
+        hash = "sha256-CxufrFt4l04v2qp0hD2xj2Ns6PPZmdYv8qYVuZePw2A=";
+      }
+      {
+        version = "8.5.0.96";
+        minCudaVersion = "11.0";
+        maxCudaVersion = "11.7";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-sbsa/cudnn-linux-sbsa-8.5.0.96_cuda11-archive.tar.xz";
+        hash = "sha256-hngKu+zUY05zY/rR0ACuI7eQWl+Dg73b9zMsaTR5Hd4=";
+      }
+      {
+        version = "8.6.0.163";
+        minCudaVersion = "11.0";
+        maxCudaVersion = "11.8";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-sbsa/cudnn-linux-sbsa-8.6.0.163_cuda11-archive.tar.xz";
+        hash = "sha256-oCAieNPL1POtw/eBa/9gcWIcsEKwkDaYtHesrIkorAY=";
+      }
+      {
+        version = "8.7.0.84";
+        minCudaVersion = "11.0";
+        maxCudaVersion = "11.8";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-sbsa/cudnn-linux-sbsa-8.7.0.84_cuda11-archive.tar.xz";
+        hash = "sha256-z5Z/eNv2wHUkPMg6oYdZ43DbN1SqFbEqChTov2ejqdQ=";
+      }
+      {
+        version = "8.8.1.3";
+        minCudaVersion = "11.0";
+        maxCudaVersion = "11.8";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-sbsa/cudnn-linux-sbsa-8.8.1.3_cuda11-archive.tar.xz";
+        hash = "sha256-OzWq+aQkmIbZONmWSYyFoZzem3RldoXyJy7GVT6GM1k=";
+      }
+      {
+        version = "8.8.1.3";
+        minCudaVersion = "12.0";
+        maxCudaVersion = "12.0";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-sbsa/cudnn-linux-sbsa-8.8.1.3_cuda12-archive.tar.xz";
+        hash = "sha256-njl3qhudBuuGC1gqyJM2MGdaAkMCnCWb/sW7VpmGfSA=";
+      }
+      {
+        version = "8.9.6.50";
+        minCudaVersion = "11.0";
+        maxCudaVersion = "11.8";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-sbsa/cudnn-linux-sbsa-8.9.6.50_cuda11-archive.tar.xz";
+        hash = "sha256-nlQWYSOJWci7o3wFGIuxrkoo8d3ddg4F2hU/qJySvBE=";
+      }
+      {
+        version = "8.9.6.50";
+        minCudaVersion = "12.0";
+        maxCudaVersion = "12.2";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-sbsa/cudnn-linux-sbsa-8.9.6.50_cuda12-archive.tar.xz";
+        hash = "sha256-L20O26RelmeynVfjohEADW3Vaj3VbFS2dTUadTKlXdg=";
+      }
+    ];
+    # x86_64
+    linux-x86_64 = [
+      {
+        version = "7.4.2.24";
+        minCudaVersion = "10.0";
+        maxCudaVersion = "10.0";
+        url = "https://developer.download.nvidia.com/compute/redist/cudnn/v7.4.2/cudnn-10.0-linux-x64-v7.4.2.24.tgz";
+        hash = "sha256-Lt/IagK1DRfojEeJVaMy5qHoF05+U6NFi06lH68C2qM=";
+      }
+      {
+        version = "7.6.5.32";
+        minCudaVersion = "10.0";
+        maxCudaVersion = "10.0";
+        url = "https://developer.download.nvidia.com/compute/redist/cudnn/v7.6.5/cudnn-10.0-linux-x64-v7.6.5.32.tgz";
+        hash = "sha256-KDVeOV8LK5OsLIO2E2CzW6bNA3fkTni+GXtrYbS0kro=";
+      }
+      {
+        version = "7.6.5.32";
+        minCudaVersion = "10.1";
+        maxCudaVersion = "10.1";
+        url = "https://developer.download.nvidia.com/compute/redist/cudnn/v7.6.5/cudnn-10.1-linux-x64-v7.6.5.32.tgz";
+        hash = "sha256-fq7IA5osMKsLx1jTA1iHZ2k972v0myJIWiwAvy4TbLM=";
+      }
+      {
+        version = "7.6.5.32";
+        minCudaVersion = "10.2";
+        maxCudaVersion = "10.2";
+        url = "https://developer.download.nvidia.com/compute/redist/cudnn/v7.6.5/cudnn-10.2-linux-x64-v7.6.5.32.tgz";
+        hash = "sha256-fq7IA5osMKsLx1jTA1iHZ2k972v0myJIWiwAvy4TbLN=";
+      }
+      {
+        version = "8.0.5.39";
+        minCudaVersion = "10.1";
+        maxCudaVersion = "10.1";
+        url = "https://developer.download.nvidia.com/compute/redist/cudnn/v8.0.5/cudnn-10.1-linux-x64-v8.0.5.39.tgz";
+        hash = "sha256-kJCElSmIlrM6qVBjo0cfk8NmJ9esAcF9w211xl7qSgA=";
+      }
+      {
+        version = "8.0.5.39";
+        minCudaVersion = "10.2";
+        maxCudaVersion = "10.2";
+        url = "https://developer.download.nvidia.com/compute/redist/cudnn/v8.0.5/cudnn-10.2-linux-x64-v8.0.5.39.tgz";
+        hash = "sha256-IfhMBcZ78eyFnnfDjM1b8VSWT6HDCPRJlZvkw1bjgvM=";
+      }
+      {
+        version = "8.0.5.39";
+        minCudaVersion = "11.0";
+        maxCudaVersion = "11.0";
+        url = "https://developer.download.nvidia.com/compute/redist/cudnn/v8.0.5/cudnn-11.0-linux-x64-v8.0.5.39.tgz";
+        hash = "sha256-ThbueJXetKixwZS4ErpJWG730mkCBRQB03F1EYmKm3M=";
+      }
+      {
+        version = "8.0.5.39";
+        minCudaVersion = "11.1";
+        maxCudaVersion = "11.1";
+        url = "https://developer.download.nvidia.com/compute/redist/cudnn/v8.0.5/cudnn-11.1-linux-x64-v8.0.5.39.tgz";
+        hash = "sha256-HQRr+nk5navMb2yxUHkYdUQ5RC6gyp4Pvs3URvmwDM4=";
+      }
+      {
+        version = "8.1.1.33";
+        minCudaVersion = "10.2";
+        maxCudaVersion = "10.2";
+        url = "https://developer.download.nvidia.com/compute/redist/cudnn/v8.1.1/cudnn-10.2-linux-x64-v8.1.1.33.tgz";
+        hash = "sha256-Kkp7mabpv6aQ6xm7QeSVU/KnpJGls6v8rpAOFmxbbr0=";
+      }
+      {
+        version = "8.1.1.33";
+        minCudaVersion = "11.0";
+        maxCudaVersion = "11.2";
+        url = "https://developer.download.nvidia.com/compute/redist/cudnn/v8.1.1/cudnn-11.2-linux-x64-v8.1.1.33.tgz";
+        hash = "sha256-mKh4TpKGLyABjSDCgbMNSgzZUfk2lPZDPM9K6cUCumo=";
+      }
+      {
+        version = "8.2.4.15";
+        minCudaVersion = "10.2";
+        maxCudaVersion = "10.2";
+        url = "https://developer.download.nvidia.com/compute/redist/cudnn/v8.2.4/cudnn-10.2-linux-x64-v8.2.4.15.tgz";
+        hash = "sha256-0jyUoxFaHHcRamwSfZF1+/WfcjNkN08mo0aZB18yIvE=";
+      }
+      {
+        version = "8.2.4.15";
+        minCudaVersion = "11.0";
+        maxCudaVersion = "11.4";
+        url = "https://developer.download.nvidia.com/compute/redist/cudnn/v8.2.4/cudnn-11.4-linux-x64-v8.2.4.15.tgz";
+        hash = "sha256-Dl0t+JC5ln76ZhnaQhMQ2XMjVlp58FoajLm3Fluq0Nc=";
+      }
+      {
+        version = "8.3.3.40";
+        minCudaVersion = "10.2";
+        maxCudaVersion = "10.2";
+        url = "https://developer.download.nvidia.com/compute/redist/cudnn/v8.3.3/local_installers/10.2/cudnn-linux-x86_64-8.3.3.40_cuda10.2-archive.tar.xz";
+        hash = "sha256-2FVPKzLmKV1fyPOsJeaPlAWLAYyAHaucFD42gS+JJqs=";
+      }
+      {
+        version = "8.3.3.40";
+        minCudaVersion = "11.0";
+        maxCudaVersion = "11.6";
+        url = "https://developer.download.nvidia.com/compute/redist/cudnn/v8.3.3/local_installers/11.5/cudnn-linux-x86_64-8.3.3.40_cuda11.5-archive.tar.xz";
+        hash = "sha256-6r6Wx1zwPqT1N5iU2RTx+K4UzqsSGYnoSwg22Sf7dzE=";
+      }
+      {
+        version = "8.4.1.50";
+        minCudaVersion = "10.2";
+        maxCudaVersion = "10.2";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-8.4.1.50_cuda10.2-archive.tar.xz";
+        hash = "sha256-I88qMmU6lIiLVmaPuX7TTbisgTav839mssxUo3lQNjg=";
+      }
+      {
+        version = "8.4.1.50";
+        minCudaVersion = "11.0";
+        maxCudaVersion = "11.7";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-8.4.1.50_cuda11.6-archive.tar.xz";
+        hash = "sha256-7JbSN22B/KQr3T1MPXBambKaBlurV/kgVhx2PinGfQE=";
+      }
+      {
+        version = "8.5.0.96";
+        minCudaVersion = "10.2";
+        maxCudaVersion = "10.2";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-8.5.0.96_cuda10-archive.tar.xz";
+        hash = "sha256-1mzhbbzR40WKkHnQLtJHhg0vYgf7G8a0OBcCwIOkJjM=";
+      }
+      {
+        version = "8.5.0.96";
+        minCudaVersion = "11.0";
+        maxCudaVersion = "11.7";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-8.5.0.96_cuda11-archive.tar.xz";
+        hash = "sha256-VFSm/ZTwCHKMqumtrZk8ToXvNjAuJrzkO+p9RYpee20=";
+      }
+      {
+        version = "8.6.0.163";
+        minCudaVersion = "10.2";
+        maxCudaVersion = "10.2";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-8.6.0.163_cuda10-archive.tar.xz";
+        hash = "sha256-t4sr/GrFqqdxu2VhaJQk5K1Xm/0lU4chXG8hVL09R9k=";
+      }
+      {
+        version = "8.6.0.163";
+        minCudaVersion = "11.0";
+        maxCudaVersion = "11.8";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-8.6.0.163_cuda11-archive.tar.xz";
+        hash = "sha256-u8OW30cpTGV+3AnGAGdNYIyxv8gLgtz0VHBgwhcRFZ4=";
+      }
+      {
+        version = "8.7.0.84";
+        minCudaVersion = "10.2";
+        maxCudaVersion = "10.2";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-8.7.0.84_cuda10-archive.tar.xz";
+        hash = "sha256-bZhaqc8+GbPV2FQvvbbufd8VnEJgvfkICc2N3/gitRg=";
+      }
+      {
+        version = "8.7.0.84";
+        minCudaVersion = "11.0";
+        maxCudaVersion = "11.8";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-8.7.0.84_cuda11-archive.tar.xz";
+        hash = "sha256-l2xMunIzyXrnQAavq1Fyl2MAukD1slCiH4z3H1nJ920=";
+      }
+      {
+        version = "8.8.1.3";
+        minCudaVersion = "11.0";
+        maxCudaVersion = "11.8";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-8.8.1.3_cuda11-archive.tar.xz";
+        hash = "sha256-r3WEyuDMVSS1kT7wjCm6YVQRPGDrCjegWQqRtRWoqPk=";
+      }
+      {
+        version = "8.8.1.3";
+        minCudaVersion = "12.0";
+        maxCudaVersion = "12.0";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-8.8.1.3_cuda12-archive.tar.xz";
+        hash = "sha256-edd6dpx+cXWrx7XC7VxJQUjAYYqGQThyLIh/lcYjd3w=";
+      }
+      {
+        version = "8.9.6.50";
+        minCudaVersion = "11.0";
+        maxCudaVersion = "11.8";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-8.9.6.50_cuda11-archive.tar.xz";
+        hash = "sha256-oOLvVemfTNZH99HaqlqkUE/6M1ujAYbVwyiPL0ffBX4=";
+      }
+      {
+        version = "8.9.6.50";
+        minCudaVersion = "12.0";
+        maxCudaVersion = "12.2";
+        url = "https://developer.download.nvidia.com/compute/cudnn/redist/cudnn/linux-x86_64/cudnn-linux-x86_64-8.9.6.50_cuda12-archive.tar.xz";
+        hash = "sha256-FyIlnblSZbs4E0OKWhxuzZed6JrkU2YDkEBC4STTAtU=";
+      }
+    ];
+  };
+}
diff --git a/pkgs/development/cuda-modules/flags.nix b/pkgs/development/cuda-modules/flags.nix
new file mode 100644
index 000000000000..7b7922085ff6
--- /dev/null
+++ b/pkgs/development/cuda-modules/flags.nix
@@ -0,0 +1,387 @@
+# Type aliases
+# Gpu :: AttrSet
+#   - See the documentation in ./gpus.nix.
+{
+  config,
+  cudaCapabilities ? (config.cudaCapabilities or []),
+  cudaForwardCompat ? (config.cudaForwardCompat or true),
+  lib,
+  cudaVersion,
+  hostPlatform,
+  # gpus :: List Gpu
+  gpus,
+}:
+let
+  inherit (lib)
+    asserts
+    attrsets
+    lists
+    strings
+    trivial
+    ;
+
+  # 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.
+  #
+  # cudaCapabilities :: List Capability
+  # List of hardware generations to build.
+  # E.g. [ "8.0" ]
+  # Currently, the last item is considered the optional forward-compatibility arch,
+  # but this may change in the future.
+  #
+  # cudaForwardCompat :: Bool
+  # Whether to include the forward compatibility gencode (+PTX)
+  # to support future GPU generations.
+  # E.g. true
+  #
+  # Please see the accompanying documentation or https://github.com/NixOS/nixpkgs/pull/205351
+
+  # isSupported :: Gpu -> Bool
+  isSupported =
+    gpu:
+    let
+      inherit (gpu) minCudaVersion maxCudaVersion;
+      lowerBoundSatisfied = strings.versionAtLeast cudaVersion minCudaVersion;
+      upperBoundSatisfied =
+        (maxCudaVersion == null) || !(strings.versionOlder maxCudaVersion cudaVersion);
+    in
+    lowerBoundSatisfied && upperBoundSatisfied;
+
+  # NOTE: Jetson is never built by default.
+  # isDefault :: Gpu -> Bool
+  isDefault =
+    gpu:
+    let
+      inherit (gpu) dontDefaultAfter;
+      newGpu = dontDefaultAfter == null;
+      recentGpu = newGpu || strings.versionAtLeast dontDefaultAfter cudaVersion;
+    in
+    recentGpu;
+
+  # supportedGpus :: List Gpu
+  # GPUs which are supported by the provided CUDA version.
+  supportedGpus = builtins.filter isSupported gpus;
+
+  # defaultGpus :: List Gpu
+  # GPUs which are supported by the provided CUDA version and we want to build for by default.
+  defaultGpus = builtins.filter isDefault supportedGpus;
+
+  # supportedCapabilities :: List Capability
+  supportedCapabilities = lists.map (gpu: gpu.computeCapability) supportedGpus;
+
+  # defaultCapabilities :: List Capability
+  # The default capabilities to target, if not overridden by the user.
+  defaultCapabilities = lists.map (gpu: gpu.computeCapability) defaultGpus;
+
+  # cudaArchNameToVersions :: AttrSet String (List String)
+  # Maps the name of a GPU architecture to different versions of that architecture.
+  # For example, "Ampere" maps to [ "8.0" "8.6" "8.7" ].
+  cudaArchNameToVersions =
+    lists.groupBy' (versions: gpu: versions ++ [gpu.computeCapability]) [] (gpu: gpu.archName)
+      supportedGpus;
+
+  # cudaComputeCapabilityToName :: AttrSet String String
+  # Maps the version of a GPU architecture to the name of that architecture.
+  # For example, "8.0" maps to "Ampere".
+  cudaComputeCapabilityToName = builtins.listToAttrs (
+    lists.map (gpu: attrsets.nameValuePair gpu.computeCapability gpu.archName) supportedGpus
+  );
+
+  # cudaComputeCapabilityToIsJetson :: AttrSet String Boolean
+  cudaComputeCapabilityToIsJetson = builtins.listToAttrs (
+    lists.map (attrs: attrsets.nameValuePair attrs.computeCapability attrs.isJetson) supportedGpus
+  );
+
+  # jetsonComputeCapabilities :: List String
+  jetsonComputeCapabilities = trivial.pipe cudaComputeCapabilityToIsJetson [
+    (attrsets.filterAttrs (_: isJetson: isJetson))
+    builtins.attrNames
+  ];
+
+  # Find the intersection with the user-specified list of cudaCapabilities.
+  # NOTE: Jetson devices are never built by default because they cannot be targeted along
+  # non-Jetson devices and require an aarch64 host platform. As such, if they're present anywhere,
+  # they must be in the user-specified cudaCapabilities.
+  # NOTE: We don't need to worry about mixes of Jetson and non-Jetson devices here -- there's
+  # sanity-checking for all that in cudaFlags.
+  jetsonTargets = lists.intersectLists jetsonComputeCapabilities cudaCapabilities;
+
+  # dropDot :: String -> String
+  dropDot = ver: builtins.replaceStrings ["."] [""] ver;
+
+  # archMapper :: String -> List String -> List String
+  # Maps a feature across a list of architecture versions to produce a list of architectures.
+  # For example, "sm" and [ "8.0" "8.6" "8.7" ] produces [ "sm_80" "sm_86" "sm_87" ].
+  archMapper = feat: lists.map (computeCapability: "${feat}_${dropDot computeCapability}");
+
+  # gencodeMapper :: String -> List String -> List String
+  # Maps a feature across a list of architecture versions to produce a list of gencode arguments.
+  # For example, "sm" and [ "8.0" "8.6" "8.7" ] produces [ "-gencode=arch=compute_80,code=sm_80"
+  # "-gencode=arch=compute_86,code=sm_86" "-gencode=arch=compute_87,code=sm_87" ].
+  gencodeMapper =
+    feat:
+    lists.map (
+      computeCapability:
+      "-gencode=arch=compute_${dropDot computeCapability},code=${feat}_${dropDot computeCapability}"
+    );
+
+  # Maps Nix system to NVIDIA redist arch.
+  # NOTE: We swap out the default `linux-sbsa` redist (for server-grade ARM chips) with the
+  # `linux-aarch64` redist (which is for Jetson devices) if we're building any Jetson devices.
+  # Since both are based on aarch64, we can only have one or the other, otherwise there's an
+  # ambiguity as to which should be used.
+  # getRedistArch :: String -> String
+  getRedistArch =
+    nixSystem:
+    if nixSystem == "aarch64-linux" then
+      if jetsonTargets != [] then "linux-aarch64" else "linux-sbsa"
+    else if nixSystem == "x86_64-linux" then
+      "linux-x86_64"
+    else if nixSystem == "ppc64le-linux" then
+      "linux-ppc64le"
+    else if nixSystem == "x86_64-windows" then
+      "windows-x86_64"
+    else
+      builtins.throw "Unsupported Nix system: ${nixSystem}";
+
+  # Maps NVIDIA redist arch to Nix system.
+  getNixSystem =
+    redistArch:
+    if
+      lists.elem redistArch [
+        "linux-aarch64"
+        "linux-sbsa"
+      ]
+    then
+      "aarch64-linux"
+    else if redistArch == "linux-x86_64" then
+      "x86_64-linux"
+    else if redistArch == "linux-ppc64le" then
+      "ppc64le-linux"
+    else if redistArch == "windows-x86_64" then
+      "x86_64-windows"
+    else
+      builtins.throw "Unsupported NVIDIA redist arch: ${redistArch}";
+
+  formatCapabilities =
+    {
+      cudaCapabilities,
+      enableForwardCompat ? true,
+    }:
+    rec {
+      inherit cudaCapabilities enableForwardCompat;
+
+      # archNames :: List String
+      # E.g. [ "Turing" "Ampere" ]
+      archNames = lists.unique (
+        lists.map (cap: cudaComputeCapabilityToName.${cap} or (throw "missing cuda compute capability"))
+          cudaCapabilities
+      );
+
+      # realArches :: List String
+      # The real architectures are physical architectures supported by the CUDA version.
+      # E.g. [ "sm_75" "sm_86" ]
+      realArches = archMapper "sm" cudaCapabilities;
+
+      # virtualArches :: List String
+      # The virtual architectures are typically used for forward compatibility, when trying to support
+      # an architecture newer than the CUDA version allows.
+      # E.g. [ "compute_75" "compute_86" ]
+      virtualArches = archMapper "compute" cudaCapabilities;
+
+      # arches :: List String
+      # By default, build for all supported architectures and forward compatibility via a virtual
+      # architecture for the newest supported architecture.
+      # E.g. [ "sm_75" "sm_86" "compute_86" ]
+      arches = realArches ++ lists.optional enableForwardCompat (lists.last virtualArches);
+
+      # gencode :: List String
+      # A list of CUDA gencode arguments to pass to NVCC.
+      # E.g. [ "-gencode=arch=compute_75,code=sm_75" ... "-gencode=arch=compute_86,code=compute_86" ]
+      gencode =
+        let
+          base = gencodeMapper "sm" cudaCapabilities;
+          forward = gencodeMapper "compute" [(lists.last cudaCapabilities)];
+        in
+        base ++ lib.optionals enableForwardCompat forward;
+
+      # gencodeString :: String
+      # A space-separated string of CUDA gencode arguments to pass to NVCC.
+      # E.g. "-gencode=arch=compute_75,code=sm_75 ... -gencode=arch=compute_86,code=compute_86"
+      gencodeString = strings.concatStringsSep " " gencode;
+
+      # 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.
+      # isJetsonBuild :: Boolean
+      isJetsonBuild =
+        let
+          # List of booleans representing whether any of the currently targeted capabilities are
+          # Jetson devices.
+          # isJetsons :: List Boolean
+          isJetsons =
+            lists.map (trivial.flip builtins.getAttr cudaComputeCapabilityToIsJetson)
+              cudaCapabilities;
+          anyJetsons = lists.any (trivial.id) isJetsons;
+          allJetsons = lists.all (trivial.id) isJetsons;
+          hostIsAarch64 = hostPlatform.isAarch64;
+        in
+        trivial.throwIfNot (anyJetsons -> (allJetsons && hostIsAarch64))
+          ''
+            Jetson devices cannot be targeted with non-Jetson devices. Additionally, they require hostPlatform to be aarch64.
+            You requested ${builtins.toJSON cudaCapabilities} for host platform ${hostPlatform.system}.
+            Exactly one of the following must be true:
+            - All CUDA capabilities belong to Jetson devices (${trivial.boolToString allJetsons}) and the hostPlatform is aarch64 (${trivial.boolToString hostIsAarch64}).
+            - No CUDA capabilities belong to Jetson devices (${trivial.boolToString (!anyJetsons)}).
+            See ${./gpus.nix} for a list of architectures supported by this version of Nixpkgs.
+          ''
+          allJetsons;
+    };
+in
+# When changing names or formats: pause, validate, and update the assert
+assert let
+  expected = {
+    cudaCapabilities = [
+      "7.5"
+      "8.6"
+    ];
+    enableForwardCompat = true;
+
+    archNames = [
+      "Turing"
+      "Ampere"
+    ];
+    realArches = [
+      "sm_75"
+      "sm_86"
+    ];
+    virtualArches = [
+      "compute_75"
+      "compute_86"
+    ];
+    arches = [
+      "sm_75"
+      "sm_86"
+      "compute_86"
+    ];
+
+    gencode = [
+      "-gencode=arch=compute_75,code=sm_75"
+      "-gencode=arch=compute_86,code=sm_86"
+      "-gencode=arch=compute_86,code=compute_86"
+    ];
+    gencodeString = "-gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_86,code=sm_86 -gencode=arch=compute_86,code=compute_86";
+
+    isJetsonBuild = false;
+  };
+  actual = formatCapabilities {
+    cudaCapabilities = [
+      "7.5"
+      "8.6"
+    ];
+  };
+  actualWrapped = (builtins.tryEval (builtins.deepSeq actual actual)).value;
+in
+asserts.assertMsg (expected == actualWrapped) ''
+  This test should only fail when using a version of CUDA older than 11.2, the first to support
+  8.6.
+  Expected: ${builtins.toJSON expected}
+  Actual: ${builtins.toJSON actualWrapped}
+'';
+# Check mixed Jetson and non-Jetson devices
+assert let
+  expected = false;
+  actual = formatCapabilities {
+    cudaCapabilities = [
+      "7.2"
+      "7.5"
+    ];
+  };
+  actualWrapped = (builtins.tryEval (builtins.deepSeq actual actual)).value;
+in
+asserts.assertMsg (expected == actualWrapped) ''
+  Jetson devices capabilities cannot be mixed with non-jetson devices.
+  Capability 7.5 is non-Jetson and should not be allowed with Jetson 7.2.
+  Expected: ${builtins.toJSON expected}
+  Actual: ${builtins.toJSON actualWrapped}
+'';
+# Check Jetson-only
+assert let
+  expected = {
+    cudaCapabilities = [
+      "6.2"
+      "7.2"
+    ];
+    enableForwardCompat = true;
+
+    archNames = [
+      "Pascal"
+      "Volta"
+    ];
+    realArches = [
+      "sm_62"
+      "sm_72"
+    ];
+    virtualArches = [
+      "compute_62"
+      "compute_72"
+    ];
+    arches = [
+      "sm_62"
+      "sm_72"
+      "compute_72"
+    ];
+
+    gencode = [
+      "-gencode=arch=compute_62,code=sm_62"
+      "-gencode=arch=compute_72,code=sm_72"
+      "-gencode=arch=compute_72,code=compute_72"
+    ];
+    gencodeString = "-gencode=arch=compute_62,code=sm_62 -gencode=arch=compute_72,code=sm_72 -gencode=arch=compute_72,code=compute_72";
+
+    isJetsonBuild = true;
+  };
+  actual = formatCapabilities {
+    cudaCapabilities = [
+      "6.2"
+      "7.2"
+    ];
+  };
+  actualWrapped = (builtins.tryEval (builtins.deepSeq actual actual)).value;
+in
+asserts.assertMsg
+  # We can't do this test unless we're targeting aarch64
+  (hostPlatform.isAarch64 -> (expected == actualWrapped))
+  ''
+    Jetson devices can only be built with other Jetson devices.
+    Both 6.2 and 7.2 are Jetson devices.
+    Expected: ${builtins.toJSON expected}
+    Actual: ${builtins.toJSON actualWrapped}
+  '';
+{
+  # formatCapabilities :: { cudaCapabilities: List Capability, enableForwardCompat: Boolean } ->  { ... }
+  inherit formatCapabilities;
+
+  # cudaArchNameToVersions :: String => String
+  inherit cudaArchNameToVersions;
+
+  # cudaComputeCapabilityToName :: String => String
+  inherit cudaComputeCapabilityToName;
+
+  # dropDot :: String -> String
+  inherit dropDot;
+
+  inherit
+    defaultCapabilities
+    supportedCapabilities
+    jetsonComputeCapabilities
+    jetsonTargets
+    getNixSystem
+    getRedistArch
+    ;
+}
+// formatCapabilities {
+  cudaCapabilities = if cudaCapabilities == [] then defaultCapabilities else cudaCapabilities;
+  enableForwardCompat = cudaForwardCompat;
+}
diff --git a/pkgs/development/cuda-modules/gpus.nix b/pkgs/development/cuda-modules/gpus.nix
new file mode 100644
index 000000000000..d17073936138
--- /dev/null
+++ b/pkgs/development/cuda-modules/gpus.nix
@@ -0,0 +1,204 @@
+# Type aliases
+#
+# Gpu = {
+#   archName: String
+#     - The name of the microarchitecture.
+#   computeCapability: String
+#     - The compute capability of the GPU.
+#   isJetson: Boolean
+#     - Whether a GPU is part of NVIDIA's line of Jetson embedded computers. This field is
+#       notable because it tells us what architecture to build for (as Jetson devices are
+#       aarch64).
+#       More on Jetson devices here:
+#       https://www.nvidia.com/en-us/autonomous-machines/embedded-systems/
+#       NOTE: These architectures are only built upon request.
+#   minCudaVersion: String
+#     - The minimum (inclusive) CUDA version that supports this GPU.
+#   dontDefaultAfter: null | String
+#     - The CUDA version after which to exclude this GPU from the list of default capabilities
+#       we build. null means we always include this GPU in the default capabilities if it is
+#       supported.
+#   maxCudaVersion: null | String
+#     - The maximum (exclusive) CUDA version that supports this GPU. null means there is no
+#       maximum.
+# }
+#
+# Many thanks to Arnon Shimoni for maintaining a list of these architectures and capabilities.
+# Without your work, this would have been much more difficult.
+# https://arnon.dk/matching-sm-architectures-arch-and-gencode-for-various-nvidia-cards/
+[
+  {
+    # GeForce 700, GT-730
+    archName = "Kepler";
+    computeCapability = "3.0";
+    isJetson = false;
+    minCudaVersion = "10.0";
+    dontDefaultAfter = "10.2";
+    maxCudaVersion = "10.2";
+  }
+  {
+    archName = "Kepler";
+    computeCapability = "3.2";
+    isJetson = false;
+    minCudaVersion = "10.0";
+    dontDefaultAfter = "10.2";
+    maxCudaVersion = "10.2";
+  }
+  {
+    # Tesla K40
+    archName = "Kepler";
+    computeCapability = "3.5";
+    isJetson = false;
+    minCudaVersion = "10.0";
+    dontDefaultAfter = "11.0";
+    maxCudaVersion = "11.8";
+  }
+  {
+    # Tesla K80
+    archName = "Kepler";
+    computeCapability = "3.7";
+    isJetson = false;
+    minCudaVersion = "10.0";
+    dontDefaultAfter = "11.0";
+    maxCudaVersion = "11.8";
+  }
+  {
+    # Tesla/Quadro M series
+    archName = "Maxwell";
+    computeCapability = "5.0";
+    isJetson = false;
+    minCudaVersion = "10.0";
+    dontDefaultAfter = "11.0";
+    maxCudaVersion = null;
+  }
+  {
+    # Quadro M6000 , GeForce 900, GTX-970, GTX-980, GTX Titan X
+    archName = "Maxwell";
+    computeCapability = "5.2";
+    isJetson = false;
+    minCudaVersion = "10.0";
+    dontDefaultAfter = "11.0";
+    maxCudaVersion = null;
+  }
+  {
+    # Tegra (Jetson) TX1 / Tegra X1, Drive CX, Drive PX, Jetson Nano
+    archName = "Maxwell";
+    computeCapability = "5.3";
+    isJetson = true;
+    minCudaVersion = "10.0";
+    dontDefaultAfter = "11.0";
+    maxCudaVersion = null;
+  }
+  {
+    # Quadro GP100, Tesla P100, DGX-1 (Generic Pascal)
+    archName = "Pascal";
+    computeCapability = "6.0";
+    isJetson = false;
+    minCudaVersion = "10.0";
+    dontDefaultAfter = null;
+    maxCudaVersion = null;
+  }
+  {
+    # GTX 1080, GTX 1070, GTX 1060, GTX 1050, GTX 1030 (GP108), GT 1010 (GP108) Titan Xp, Tesla
+    # P40, Tesla P4, Discrete GPU on the NVIDIA Drive PX2
+    archName = "Pascal";
+    computeCapability = "6.1";
+    isJetson = false;
+    minCudaVersion = "10.0";
+    dontDefaultAfter = null;
+    maxCudaVersion = null;
+  }
+  {
+    # Integrated GPU on the NVIDIA Drive PX2, Tegra (Jetson) TX2
+    archName = "Pascal";
+    computeCapability = "6.2";
+    isJetson = true;
+    minCudaVersion = "10.0";
+    dontDefaultAfter = null;
+    maxCud