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] 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 {