magma: use CMAKE_CUDA_ARCHITECTURES directly

2 files changed, 27 insertions, 46 deletions

diff --git a/pkgs/development/libraries/science/math/magma/generic.nix b/pkgs/development/libraries/science/math/magma/generic.nix
index e27107ca15d80..f61f1877019b2 100644
--- a/pkgs/development/libraries/science/math/magma/generic.nix
+++ b/pkgs/development/libraries/science/math/magma/generic.nix
@@ -11,7 +11,8 @@
 , cudaSupport ? true
 , fetchurl
 , gfortran
-, gpuTargets ? [ ]
+, cudaCapabilities ? cudaPackages.cudaFlags.cudaCapabilities
+, gpuTargets ? [ ] # Non-CUDA targets, that is HIP
 , hip
 , hipblas
 , hipsparse
@@ -36,12 +37,6 @@ let
   #   of the first list *from* the second list. That means:
   #   lists.subtractLists a b = b - a
 
-  # For CUDA
-  supportedCudaSmArches = lists.intersectLists cudaFlags.realArches supportedGpuTargets;
-  # Subtract the supported SM architectures from the real SM architectures to get the unsupported
-  # SM architectures.
-  unsupportedCudaSmArches = lists.subtractLists supportedCudaSmArches cudaFlags.realArches;
-
   # For ROCm
   # NOTE: The hip.gpuTargets are prefixed with "gfx" instead of "sm" like cudaFlags.realArches.
   #   For some reason, Magma's CMakeLists.txt file does not handle the "gfx" prefix, so we must
@@ -62,19 +57,32 @@ let
       )
       supported;
 
-  # Create the gpuTargetString.
   gpuTargetString = strings.concatStringsSep "," (
     if gpuTargets != [ ] then
     # If gpuTargets is specified, it always takes priority.
       gpuArchWarner supportedCustomGpuTargets unsupportedCustomGpuTargets
-    else if cudaSupport then
-      gpuArchWarner supportedCudaSmArches unsupportedCudaSmArches
     else if rocmSupport then
       gpuArchWarner supportedRocmArches unsupportedRocmArches
+    else if cudaSupport then
+      [ ] # It's important we pass explicit -DGPU_TARGET to reset magma's defaults
     else
       throw "No GPU targets specified"
   );
 
+  # E.g. [ "80" "86" "90" ]
+  cudaArchitectures = (builtins.map cudaFlags.dropDot cudaCapabilities);
+
+  cudaArchitecturesString = strings.concatStringsSep ";" cudaArchitectures;
+  minArch =
+    let
+      minArch' = builtins.head (builtins.sort builtins.lessThan cudaArchitectures);
+    in
+    # If this fails some day, something must've changed and we should re-validate our assumptions
+    assert builtins.stringLength minArch' == 2;
+    # "75" -> "750"  Cf. https://bitbucket.org/icl/magma/src/f4ec79e2c13a2347eff8a77a3be6f83bc2daec20/CMakeLists.txt#lines-273
+    "${minArch'}0";
+
+
   cuda_joined = symlinkJoin {
     name = "cuda-redist-${cudaVersion}";
     paths = with cudaPackages; [
@@ -87,6 +95,8 @@ let
   };
 in
 
+assert (builtins.match "[^[:space:]]*" gpuTargetString) != null;
+
 stdenv.mkDerivation {
   pname = "magma";
   inherit version;
@@ -116,7 +126,11 @@ stdenv.mkDerivation {
     openmp
   ];
 
-  cmakeFlags = lists.optionals cudaSupport [
+  cmakeFlags = [
+    "-DGPU_TARGET=${gpuTargetString}"
+  ] ++ lists.optionals cudaSupport [
+    "-DCMAKE_CUDA_ARCHITECTURES=${cudaArchitecturesString}"
+    "-DMIN_ARCH=${minArch}" # Disarms magma's asserts
     "-DCMAKE_C_COMPILER=${cudatoolkit.cc}/bin/cc"
     "-DCMAKE_CXX_COMPILER=${cudatoolkit.cc}/bin/c++"
     "-DMAGMA_ENABLE_CUDA=ON"
@@ -126,14 +140,10 @@ stdenv.mkDerivation {
     "-DMAGMA_ENABLE_HIP=ON"
   ];
 
-  # NOTE: We must set GPU_TARGET in preConfigure in this way because it may contain spaces.
-  preConfigure = ''
-    cmakeFlagsArray+=("-DGPU_TARGET=${gpuTargetString}")
-  ''
   # NOTE: The stdenv's CXX is used when compiling the CMake test to determine the version of
   #   CUDA available. This isn't necessarily the same as cudatoolkit.cc, so we must set
   #   CUDAHOSTCXX.
-  + strings.optionalString cudaSupport ''
+  preConfigure = strings.optionalString cudaSupport ''
     export CUDAHOSTCXX=${cudatoolkit.cc}/bin/c++
   '';
 
diff --git a/pkgs/development/libraries/science/math/magma/releases.nix b/pkgs/development/libraries/science/math/magma/releases.nix
index 3d08aa95d4d18..029f418edce3c 100644
--- a/pkgs/development/libraries/science/math/magma/releases.nix
+++ b/pkgs/development/libraries/science/math/magma/releases.nix
@@ -1,27 +1,13 @@
 # NOTE: Order matters! Put the oldest version first, and the newest version last.
 # NOTE: Make sure the supportedGpuTargets are in order of oldest to newest.
 #   You can update the supportedGpuTargets by looking at the CMakeLists.txt file.
-#   CUDA starts here: https://bitbucket.org/icl/magma/src/f4ec79e2c13a2347eff8a77a3be6f83bc2daec20/CMakeLists.txt#lines-175
 #   HIP is here: https://bitbucket.org/icl/magma/src/f4ec79e2c13a2347eff8a77a3be6f83bc2daec20/CMakeLists.txt#lines-386
+#   CUDA works around magma's wrappers and uses FindCUDAToolkit directly
 [
   {
     version = "2.6.2";
     hash = "sha256-dbVU2rAJA+LRC5cskT5Q5/iMvGLzrkMrWghsfk7aCnE=";
     supportedGpuTargets = [
-      "sm_20"
-      "sm_30"
-      "sm_35"
-      "sm_37"
-      "sm_50"
-      "sm_52"
-      "sm_53"
-      "sm_60"
-      "sm_61"
-      "sm_62"
-      "sm_70"
-      "sm_71"
-      "sm_75"
-      "sm_80"
       "700"
       "701"
       "702"
@@ -53,21 +39,6 @@
     version = "2.7.1";
     hash = "sha256-2chxHAR6OMrhbv3nS+4uszMyF/0nEeHpuGBsu7SuGlA=";
     supportedGpuTargets = [
-      "sm_20"
-      "sm_30"
-      "sm_35"
-      "sm_37"
-      "sm_50"
-      "sm_52"
-      "sm_53"
-      "sm_60"
-      "sm_61"
-      "sm_62"
-      "sm_70"
-      "sm_71"
-      "sm_75"
-      "sm_80"
-      "sm_90"
       "700"
       "701"
       "702"