From ef3daffc8807a360ec29f84e02a1a36acf936832 Mon Sep 17 00:00:00 2001
From: Madoura <madouura@gmail.com>
Date: Fri, 23 Dec 2022 17:21:59 -0600
Subject: [PATCH] magma: add rocm support

---
 .../libraries/science/math/magma/default.nix  | 83 +++++++++++++------
 pkgs/top-level/all-packages.nix               | 15 +++-
 2 files changed, 73 insertions(+), 25 deletions(-)

diff --git a/pkgs/development/libraries/science/math/magma/default.nix b/pkgs/development/libraries/science/math/magma/default.nix
index 06b4e12d04e7..b7223690f435 100644
--- a/pkgs/development/libraries/science/math/magma/default.nix
+++ b/pkgs/development/libraries/science/math/magma/default.nix
@@ -1,40 +1,75 @@
-{ lib, stdenv, fetchurl, cmake, gfortran, ninja, cudaPackages, libpthreadstubs, lapack, blas }:
+{ lib
+, stdenv
+, fetchurl
+, cmake
+, ninja
+, gfortran
+, libpthreadstubs
+, lapack
+, blas
+, cudaPackages
+, hip
+, hipblas
+, hipsparse
+, openmp
+, useCUDA ? true
+, useROCM ? false
+, gpuTargets ? [ ]
+}:
 
 let
   inherit (cudaPackages) cudatoolkit cudaFlags;
-in
-
-assert let majorIs = lib.versions.major cudatoolkit.version;
-       in majorIs == "9" || majorIs == "10" || majorIs == "11";
-
-let
+in stdenv.mkDerivation (finalAttrs: {
+  pname = "magma";
   version = "2.6.2";
 
-in stdenv.mkDerivation {
-  pname = "magma";
-  inherit version;
   src = fetchurl {
-    url = "https://icl.cs.utk.edu/projectsfiles/magma/downloads/magma-${version}.tar.gz";
+    name = "magma-${finalAttrs.version}.tar.gz";
+    url = "https://icl.cs.utk.edu/projectsfiles/magma/downloads/magma-${finalAttrs.version}.tar.gz";
     hash = "sha256-dbVU2rAJA+LRC5cskT5Q5/iMvGLzrkMrWghsfk7aCnE=";
-    name = "magma-${version}.tar.gz";
   };
 
-  nativeBuildInputs = [ gfortran cmake ninja ];
+  nativeBuildInputs = [
+    cmake
+    ninja
+    gfortran
+  ];
 
-  buildInputs = [ cudatoolkit libpthreadstubs lapack blas ];
+  buildInputs = [
+    libpthreadstubs
+    lapack
+    blas
+  ] ++ lib.optionals useCUDA [
+    cudatoolkit
+  ] ++ lib.optionals useROCM [
+    hip
+    hipblas
+    hipsparse
+    openmp
+  ];
 
-  cmakeFlags = [
+  cmakeFlags = lib.optionals useCUDA [
+    "-DCMAKE_C_COMPILER=${cudatoolkit.cc}/bin/gcc"
+    "-DCMAKE_CXX_COMPILER=${cudatoolkit.cc}/bin/g++"
+    "-DMAGMA_ENABLE_CUDA=ON"
     "-DGPU_TARGET=${builtins.concatStringsSep "," cudaFlags.cudaRealArchs}"
+  ] ++ lib.optionals useROCM [
+    "-DCMAKE_C_COMPILER=${hip}/bin/hipcc"
+    "-DCMAKE_CXX_COMPILER=${hip}/bin/hipcc"
+    "-DMAGMA_ENABLE_HIP=ON"
+    "-DGPU_TARGET=${builtins.concatStringsSep "," (if gpuTargets == [ ] then hip.gpuTargets else gpuTargets)}"
+  ];
+
+  buildFlags = [
+    "magma"
+    "magma_sparse"
   ];
 
   doCheck = false;
 
-  preConfigure = ''
-    export CC=${cudatoolkit.cc}/bin/gcc CXX=${cudatoolkit.cc}/bin/g++
-  '';
-
-  enableParallelBuilding=true;
-  buildFlags = [ "magma" "magma_sparse" ];
+  passthru = {
+    inherit cudatoolkit;
+  };
 
   meta = with lib; {
     description = "Matrix Algebra on GPU and Multicore Architectures";
@@ -42,7 +77,7 @@ in stdenv.mkDerivation {
     homepage = "http://icl.cs.utk.edu/magma/index.html";
     platforms = platforms.unix;
     maintainers = with maintainers; [ tbenst ];
+    # CUDA and ROCm are mutually exclusive
+    broken = useCUDA && useROCM || useCUDA && versionOlder cudatoolkit.version "9";
   };
-
-  passthru.cudatoolkit = cudatoolkit;
-}
+})
diff --git a/pkgs/top-level/all-packages.nix b/pkgs/top-level/all-packages.nix
index 39b518b743e4..fe595d410ffb 100644
--- a/pkgs/top-level/all-packages.nix
+++ b/pkgs/top-level/all-packages.nix
@@ -36119,7 +36119,20 @@ with pkgs;
 
   lie = callPackage ../applications/science/math/LiE { };
 
-  magma = callPackage ../development/libraries/science/math/magma { };
+  magma = callPackage ../development/libraries/science/math/magma {
+    inherit (llvmPackages_rocm) openmp;
+  };
+
+  magma-cuda = magma.override {
+    useCUDA = true;
+    useROCM = false;
+  };
+
+  magma-hip = magma.override {
+    useCUDA = false;
+    useROCM = true;
+  };
+
   clmagma = callPackage ../development/libraries/science/math/clmagma { };
 
   mathematica = callPackage ../applications/science/math/mathematica { };