From 2b54132b9ac5462e867dd964a016f26c88499239 Mon Sep 17 00:00:00 2001
From: Harmen Stoppels <harmenstoppels@gmail.com>
Date: Tue, 19 Oct 2021 03:06:14 +0200
Subject: cosma: add new versions and improve package (#24136)

* cosma: add new versions and improve package

* Move method below depends_on's
---
 var/spack/repos/builtin/packages/cosma/package.py | 91 +++++++++++++++--------
 1 file changed, 61 insertions(+), 30 deletions(-)

diff --git a/var/spack/repos/builtin/packages/cosma/package.py b/var/spack/repos/builtin/packages/cosma/package.py
index ea65b41c26..4aa7b479b8 100644
--- a/var/spack/repos/builtin/packages/cosma/package.py
+++ b/var/spack/repos/builtin/packages/cosma/package.py
@@ -7,61 +7,92 @@
 from spack import *
 
 
-class Cosma(CMakePackage, CudaPackage):
+class Cosma(CMakePackage):
     """
     Distributed Communication-Optimal Matrix-Matrix Multiplication Library
     """
 
     maintainers = ['haampie', 'kabicm', 'teonnik']
     homepage = 'https://github.com/eth-cscs/COSMA'
-    url = 'https://github.com/eth-cscs/COSMA/releases/download/v2.2.0/cosma.tar.gz'
+    url = 'https://github.com/eth-cscs/COSMA/releases/download/v2.5.1/COSMA-v2.5.1.tar.gz'
     git = 'https://github.com/eth-cscs/COSMA.git'
 
     # note: The default archives produced with github do not have the archives
     #       of the submodules.
     version('master', branch='master', submodules=True)
+    version('2.5.1', sha256='085b7787597374244bbb1eb89bc69bf58c35f6c85be805e881e1c0b25166c3ce')
+    version('2.5.0', sha256='7f68bb0ee5c80f9b8df858afcbd017ad4ed87ac09439d13d7d890844dbdd3d54')
+    version('2.4.0', sha256='5714315ce06d48037f86cfee2d7f19340643fee95e9d7f1e92dc1b623b67e395')
     version('2.3.0', sha256='0c01c2deb5a0cd177952178350188a62c42ce55e604d7948ac472f55bf0d4815')
     version('2.2.0', sha256='1eb92a98110df595070a12193b9221eecf9d103ced8836c960f6c79a2bd553ca')
     version('2.0.7', sha256='8d70bfcbda6239b6a8fbeaca138790bbe58c0c3aa576879480d2632d4936cf7e')
     version('2.0.2', sha256='4f3354828bc718f3eef2f0098c3bdca3499297497a220da32db1acd57920c68d')
 
-    variant('scalapack', default=False,
-            description='Build with ScaLAPACK API.')
+    # We just need the libraries of cuda and rocm, so no need to extend
+    # CudaPackage or ROCmPackage.
+    variant('cuda', default=False, description='Build with cuBLAS support')
+    variant('rocm', default=False, description='Build with rocBLAS support')
+    variant('scalapack', default=False, description='Build with ScaLAPACK API')
 
     depends_on('cmake@3.12:', type='build')
     depends_on('mpi@3:')
-    depends_on('blas', when='~cuda')
+    depends_on('blas', when='~cuda ~rocm')
     depends_on('scalapack', when='+scalapack')
+    depends_on('cuda', when='+cuda')
+    depends_on('rocblas', when='+rocm')
+
+    def url_for_version(self, version):
+        if version <= Version('2.3.0'):
+            return 'https://github.com/eth-cscs/COSMA/releases/download/v{0}/cosma.tar.gz'.format(version)
+
+        return 'https://github.com/eth-cscs/COSMA/releases/download/v{0}/COSMA-v{1}.tar.gz'.format(version, version)
 
     def setup_build_environment(self, env):
         if '+cuda' in self.spec:
             env.set('CUDA_PATH', self.spec['cuda'].prefix)
 
-    def cmake_args(self):
+    def cosma_blas_cmake_arg(self):
+        query_to_cmake_arg = [
+            ('^intel-mkl', 'MKL'),
+            ('^intel-oneapi-mkl', 'MKL'),
+            ('^cray-libsci', 'CRAY_LIBSCI'),
+            ('^netlib-lapack', 'CUSTOM'),
+            ('^openblas', 'OPENBLAS'),
+            ('+cuda', 'CUDA'),
+            ('+rocm', 'ROCM')
+        ]
+
+        if self.version >= Version('2.4.0'):
+            query_to_cmake_arg.extend([
+                ('^blis', 'BLIS'),
+                ('^amdblis', 'BLIS'),
+                ('^atlas', 'ATLAS')
+            ])
+
+        for (query, cmake_arg) in query_to_cmake_arg:
+            if query in self.spec:
+                return cmake_arg
+
+        return 'CUSTOM'
+
+    def cosma_scalapack_cmake_arg(self):
         spec = self.spec
-        args = ['-DCOSMA_WITH_TESTS=OFF',
-                '-DCOSMA_WITH_APPS=OFF',
-                '-DCOSMA_WITH_PROFILING=OFF',
-                '-DCOSMA_WITH_BENCHMARKS=OFF']
 
-        if '^mkl' in spec:
-            args += ['-DCOSMA_BLAS=MKL']
+        if '~scalapack' in spec:
+            return 'OFF'
+        elif '^intel-mkl' in spec or '^intel-oneapi-mkl' in spec:
+            return 'MKL'
         elif '^cray-libsci' in spec:
-            args += ['-DCOSMA_BLAS=CRAY_LIBSCI']
-        elif '^netlib-lapack' in spec:
-            args += ['-DCOSMA_BLAS=CUSTOM']
-        elif '^openblas' in spec:
-            args += ['-DCOSMA_BLAS=OPENBLAS']
-        elif '+cuda' in spec:
-            args += ['-DCOSMA_BLAS=CUDA']
-        else:  # TODO '^rocm' in spec:
-            args += ['-DCOSMA_BLAS=ROCM']
-
-        if '+scalapack' and '^mkl' in spec:
-            args += ['-DCOSMA_SCALAPACK=MKL']
-        elif '+scalapack' and '^cray-libsci' in spec:
-            args += ['-DCOSMA_SCALAPACK=CRAY_LIBSCI']
-        elif '+scalapack' and '^netlib-scalapack' in spec:
-            args += ['-DCOSMA_SCALAPACK=CUSTOM']
-
-        return args
+            return 'CRAY_LIBSCI'
+
+        return 'CUSTOM'
+
+    def cmake_args(self):
+        return [
+            self.define('COSMA_WITH_TESTS', 'OFF'),
+            self.define('COSMA_WITH_APPS', 'OFF'),
+            self.define('COSMA_WITH_PROFILING', 'OFF'),
+            self.define('COSMA_WITH_BENCHMARKS', 'OFF'),
+            self.define('COSMA_BLAS', self.cosma_blas_cmake_arg()),
+            self.define('COSMA_SCALAPACK', self.cosma_scalapack_cmake_arg())
+        ]
-- 
cgit v1.2.3-70-g09d2