From 83b3634174c5532e9c020f84facdf56feef49b81 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Dani=C3=ABl=20de=20Kok?= <me@danieldk.eu>
Date: Wed, 10 Feb 2021 10:19:40 +0100
Subject: [PATCH] magma: fix CUDA 11 build

By default, MAGMA builds against compute capability 3.0 (among other
capabilities). However, support for this capability was dropped in
CUDA 11. This change fixes CUDA 11 support by excluding 3.0.

While at it, this change also adds support for newer compute
capabilities that are not enabled by MAGMA by default (to support
older CUDA versions).
---
 .../libraries/science/math/magma/default.nix  | 38 ++++++++++++++++++-
 1 file changed, 36 insertions(+), 2 deletions(-)

diff --git a/pkgs/development/libraries/science/math/magma/default.nix b/pkgs/development/libraries/science/math/magma/default.nix
index b8e3999ffe8..c5960dddf51 100644
--- a/pkgs/development/libraries/science/math/magma/default.nix
+++ b/pkgs/development/libraries/science/math/magma/default.nix
@@ -1,8 +1,40 @@
 { lib, stdenv, fetchurl, cmake, gfortran, ninja, cudatoolkit, libpthreadstubs, lapack, blas }:
 
-with lib;
+assert let majorIs = lib.versions.major cudatoolkit.version;
+       in majorIs == "9" || majorIs == "10" || majorIs == "11";
 
-let version = "2.5.4";
+let
+  version = "2.5.4";
+
+  # We define a specific set of CUDA compute capabilities here,
+  # because CUDA 11 does not support compute capability 3.0. Also,
+  # we use it to enable newer capabilities that are not enabled
+  # by magma by default. The list of supported architectures
+  # can be found in magma's top-level CMakeLists.txt.
+  cudaCapabilities = rec {
+    cuda9 = [
+      "Kepler"  # 3.0, 3.5
+      "Maxwell" # 5.0
+      "Pascal"  # 6.0
+      "Volta"   # 7.0
+    ];
+
+    cuda10 = [
+      "Turing"  # 7.5
+    ] ++ cuda9;
+
+    cuda11 = [
+      "sm_35"   # sm_30 is not supported by CUDA 11
+      "Maxwell" # 5.0
+      "Pascal"  # 6.0
+      "Volta"   # 7.0
+      "Turing"  # 7.5
+      "Ampere"  # 8.0
+    ];
+  };
+
+  capabilityString = lib.strings.concatStringsSep ","
+    cudaCapabilities."cuda${lib.versions.major cudatoolkit.version}";
 
 in stdenv.mkDerivation {
   pname = "magma";
@@ -17,6 +49,8 @@ in stdenv.mkDerivation {
 
   buildInputs = [ cudatoolkit libpthreadstubs lapack blas ];
 
+  cmakeFlags = [ "-DGPU_TARGET=${capabilityString}" ];
+
   doCheck = false;
 
   preConfigure = ''