NixOS · samuela · Feb 27, 2023 · Feb 21, 2023
diff --git a/pkgs/development/libraries/science/math/magma/default.nix b/pkgs/development/libraries/science/math/magma/default.nix
@@ -1,83 +1,53 @@
-{ lib
-, stdenv
-, fetchurl
-, cmake
-, ninja
-, gfortran
-, libpthreadstubs
-, lapack
-, blas
-, cudaPackages
-, hip
-, hipblas
-, hipsparse
-, openmp
-, useCUDA ? true
-, useROCM ? false
-, gpuTargets ? [ ]
+args@{ callPackage
+, lib
+, ...
 }:
 
-let
-  inherit (cudaPackages) cudatoolkit cudaFlags;
-in stdenv.mkDerivation (finalAttrs: {
-  pname = "magma";
-  version = "2.6.2";
+# Type aliases
+# Release = {
+#  version: String
+#  hash: String
+#  supportedGpuTargets: List String
+# }
 
-  src = fetchurl {
-    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=";
+let
+  inherit (lib) lists strings trivial;
+
+  computeName = version: "magma_${strings.replaceStrings [ "." ] [ "_" ] version}";
+
+  # buildMagmaPackage :: Release -> Derivation
+  buildMagmaPackage = magmaRelease: callPackage ./generic.nix (
+    (builtins.removeAttrs args [ "callPackage" ]) // {
+      inherit magmaRelease;
+    }
+  );
+
+  # Reverse the list to have the latest release first
+  # magmaReleases :: List Release
+  magmaReleases = lists.reverseList (builtins.import ./releases.nix);
+
+  # The latest release is the first element of the list and will be our default choice
+  # latestReleaseName :: String
+  latestReleaseName = computeName (builtins.head magmaReleases).version;
+
+  # Function to transform our releases into build attributes
+  # toBuildAttrs :: Release -> { name: String, value: Derivation }
+  toBuildAttrs = release: {
+    name = computeName release.version;
+    value = buildMagmaPackage release;
   };
 
-  nativeBuildInputs = [
-    cmake
-    ninja
-    gfortran
-  ];
-
-  buildInputs = [
-    libpthreadstubs
-    lapack
-    blas
-  ] ++ lib.optionals useCUDA [
-    cudatoolkit
-  ] ++ lib.optionals useROCM [
-    hip
-    hipblas
-    hipsparse
-    openmp
-  ];
+  # Add all supported builds as attributes
+  # allBuilds :: AttrSet String Derivation
+  allBuilds = builtins.listToAttrs (lists.map toBuildAttrs magmaReleases);
 
-  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.cudaRealArches}"
-  ] ++ 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)}"
-  ];
+  # The latest release will be our default build
+  # defaultBuild :: AttrSet String Derivation
+  defaultBuild.magma = allBuilds.${latestReleaseName};
 
-  buildFlags = [
-    "magma"
-    "magma_sparse"
-  ];
+  # builds :: AttrSet String Derivation
+  builds = allBuilds // defaultBuild;
+in
 
-  doCheck = false;
+builds
 
-  passthru = {
-    inherit cudatoolkit;
-  };
-
-  meta = with lib; {
-    description = "Matrix Algebra on GPU and Multicore Architectures";
-    license = licenses.bsd3;
-    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";
-  };
-})
diff --git a/pkgs/development/libraries/science/math/magma/generic.nix b/pkgs/development/libraries/science/math/magma/generic.nix
@@ -0,0 +1,160 @@
+# Type aliases
+# Release = {
+#  version: String
+#  hash: String
+#  supportedGpuTargets: List String
+# }
+
+{ blas
+, cmake
+, cudaPackages
+, cudaSupport ? true
+, fetchurl
+, gfortran
+, gpuTargets ? [ ]
+, hip
+, hipblas
+, hipsparse
+, lapack
+, lib
+, libpthreadstubs
+, magmaRelease
+, ninja
+, openmp
+, rocmSupport ? false
+, stdenv
+, symlinkJoin
+}:
+
+
+let
+  inherit (lib) lists strings trivial;
+  inherit (cudaPackages) cudatoolkit cudaFlags cudaVersion;
+  inherit (magmaRelease) version hash supportedGpuTargets;
+
+  # NOTE: The lists.subtractLists function is perhaps a bit unintuitive. It subtracts the elements
+  #   of the first list *from* the second list. That means:
+  #   lists.subtractLists a b = b - a
+
+  # For CUDA
+  supportedCudaSmArches = lists.intersectLists cudaFlags.cudaRealArches supportedGpuTargets;
+  # Subtract the supported SM architectures from the real SM architectures to get the unsupported
+  # SM architectures.
+  unsupportedCudaSmArches = lists.subtractLists supportedCudaSmArches cudaFlags.cudaRealArches;
+
+  # For ROCm
+  # NOTE: The hip.gpuTargets are prefixed with "gfx" instead of "sm" like cudaFlags.cudaRealArches.
+  #   For some reason, Magma's CMakeLists.txt file does not handle the "gfx" prefix, so we must
+  #   remove it.
+  rocmArches = lists.map (x: strings.removePrefix "gfx" x) hip.gpuTargets;
+  supportedRocmArches = lists.intersectLists rocmArches supportedGpuTargets;
+  unsupportedRocmArches = lists.subtractLists supportedRocmArches rocmArches;
+
+  supportedCustomGpuTargets = lists.intersectLists gpuTargets supportedGpuTargets;
+  unsupportedCustomGpuTargets = lists.subtractLists supportedCustomGpuTargets gpuTargets;
+
+  # Use trivial.warnIf to print a warning if any unsupported GPU targets are specified.
+  gpuArchWarner = supported: unsupported:
+    trivial.throwIf (supported == [ ])
+      (
+        "No supported GPU targets specified. Requested GPU targets: "
+        + strings.concatStringsSep ", " unsupported
+      )
+      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
+      throw "No GPU targets specified"
+  );
+
+  cuda_joined = symlinkJoin {
+    name = "cuda-redist-${cudaVersion}";
+    paths = with cudaPackages; [
+      cuda_nvcc
+      cuda_cudart # cuda_runtime.h
+      libcublas
+      libcusparse
+      cuda_nvprof # <cuda_profiler_api.h>
+    ];
+  };
+in
+
+stdenv.mkDerivation {
+  pname = "magma";
+  inherit version;
+
+  src = fetchurl {
+    name = "magma-${version}.tar.gz";
+    url = "https://icl.cs.utk.edu/projectsfiles/magma/downloads/magma-${version}.tar.gz";
+    inherit hash;
+  };
+
+  nativeBuildInputs = [
+    cmake
+    ninja
+    gfortran
+  ];
+
+  buildInputs = [
+    libpthreadstubs
+    lapack
+    blas
+  ] ++ lists.optionals cudaSupport [
+    cuda_joined
+  ] ++ lists.optionals rocmSupport [
+    hip
+    hipblas
+    hipsparse
+    openmp
+  ];
+
+  cmakeFlags = lists.optionals cudaSupport [
+    "-DCMAKE_C_COMPILER=${cudatoolkit.cc}/bin/cc"
+    "-DCMAKE_CXX_COMPILER=${cudatoolkit.cc}/bin/c++"
+    "-DMAGMA_ENABLE_CUDA=ON"
+  ] ++ lists.optionals rocmSupport [
+    "-DCMAKE_C_COMPILER=${hip}/bin/hipcc"
+    "-DCMAKE_CXX_COMPILER=${hip}/bin/hipcc"
+    "-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 ''
+    export CUDAHOSTCXX=${cudatoolkit.cc}/bin/c++
+  '';
+
+  buildFlags = [
+    "magma"
+    "magma_sparse"
+  ];
+
+  doCheck = false;
+
+  passthru = {
+    inherit cudaPackages cudaSupport;
+  };
+
+  meta = with lib; {
+    description = "Matrix Algebra on GPU and Multicore Architectures";
+    license = licenses.bsd3;
+    homepage = "http://icl.cs.utk.edu/magma/index.html";
+    platforms = platforms.unix;
+    maintainers = with maintainers; [ tbenst ];
+    # CUDA and ROCm are mutually exclusive
+    broken = cudaSupport && rocmSupport || cudaSupport && strings.versionOlder cudaVersion "9";
+  };
+}
diff --git a/pkgs/development/libraries/science/math/magma/releases.nix b/pkgs/development/libraries/science/math/magma/releases.nix
@@ -0,0 +1,98 @@
+# 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
+[
+  {
+    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"
+      "703"
+      "704"
+      "705"
+      "801"
+      "802"
+      "803"
+      "805"
+      "810"
+      "900"
+      "902"
+      "904"
+      "906"
+      "908"
+      "909"
+      "90c"
+      "1010"
+      "1011"
+      "1012"
+      "1030"
+      "1031"
+      "1032"
+      "1033"
+    ];
+  }
+  {
+    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"
+      "703"
+      "704"
+      "705"
+      "801"
+      "802"
+      "803"
+      "805"
+      "810"
+      "900"
+      "902"
+      "904"
+      "906"
+      "908"
+      "909"
+      "90c"
+      "1010"
+      "1011"
+      "1012"
+      "1030"
+      "1031"
+      "1032"
+      "1033"
+    ];
+  }
+]
diff --git a/pkgs/development/python-modules/torch/default.nix b/pkgs/development/python-modules/torch/default.nix
@@ -51,7 +51,7 @@ assert !cudaSupport || (let majorIs = lib.versions.major cudatoolkit.version;
 
 # confirm that cudatoolkits are sync'd across dependencies
 assert !(MPISupport && cudaSupport) || mpi.cudatoolkit == cudatoolkit;
-assert !cudaSupport || magma.cudatoolkit == cudatoolkit;
+assert !cudaSupport || magma.cudaPackages.cudatoolkit == cudatoolkit;
 
 let
   setBool = v: if v then "1" else "0";