hypre: add rocm support (#29147)

Co-authored-by: Sarah Osborn <30503782+osborn9@users.noreply.github.com>

1 files changed, 33 insertions, 6 deletions

diff --git a/var/spack/repos/builtin/packages/hypre/package.py b/var/spack/repos/builtin/packages/hypre/package.py
index 293eec5cb9..cc665b7d5c 100644
--- a/var/spack/repos/builtin/packages/hypre/package.py
+++ b/var/spack/repos/builtin/packages/hypre/package.py
@@ -9,7 +9,7 @@ import sys
 from spack import *
 
 
-class Hypre(AutotoolsPackage, CudaPackage):
+class Hypre(AutotoolsPackage, CudaPackage, ROCmPackage):
     """Hypre is a library of high performance preconditioners that
        features parallel multigrid methods for both structured and
        unstructured grid problems."""
@@ -93,7 +93,9 @@ class Hypre(AutotoolsPackage, CudaPackage):
     depends_on('superlu-dist', when='+superlu-dist+mpi')
 
     conflicts('+cuda', when='+int64')
-    conflicts('+unified-memory', when='~cuda')
+    conflicts('+rocm', when='+int64')
+    conflicts('+rocm', when='@:2.20')
+    conflicts('+unified-memory', when='~cuda~rocm')
     conflicts('+gptune', when='~mpi')
 
     # Patch to build shared libraries on Darwin does not apply to
@@ -140,6 +142,10 @@ class Hypre(AutotoolsPackage, CudaPackage):
             if '+fortran' in spec:
                 os.environ['F77'] = spec['mpi'].mpif77
             configure_args.append('--with-MPI')
+            configure_args.append('--with-MPI-lib-dirs={0}'.format(
+                spec['mpi'].prefix.lib))
+            configure_args.append('--with-MPI-include={0}'.format(
+                spec['mpi'].prefix.include))
         else:
             configure_args.append('--without-MPI')
 
@@ -176,7 +182,13 @@ class Hypre(AutotoolsPackage, CudaPackage):
             configure_args.extend([
                 '--with-cuda',
                 '--enable-curand',
+                '--enable-cusparse',
             ])
+            cuda_arch_vals = spec.variants['cuda_arch'].value
+            if cuda_arch_vals:
+                cuda_arch_sorted = list(sorted(cuda_arch_vals, reverse=True))
+                cuda_arch = cuda_arch_sorted[0]
+                configure_args.append('--with-gpu-arch={0}'.format(cuda_arch))
             # New in 2.21.0: replaces --enable-cub
             if '@2.21.0:' in spec:
                 configure_args.append('--enable-device-memory-pool')
@@ -188,10 +200,29 @@ class Hypre(AutotoolsPackage, CudaPackage):
             configure_args.extend([
                 '--without-cuda',
                 '--disable-curand',
+                '--disable-cusparse',
             ])
             if '@:2.20.99' in spec:
                 configure_args.append('--disable-cub')
 
+        if '+rocm' in spec:
+            configure_args.extend([
+                '--with-hip',
+                '--enable-rocrand',
+                '--enable-rocsparse',
+            ])
+            rocm_arch_vals = spec.variants['amdgpu_target'].value
+            if rocm_arch_vals:
+                rocm_arch_sorted = list(sorted(rocm_arch_vals, reverse=True))
+                rocm_arch = rocm_arch_sorted[0]
+                configure_args.append('--with-gpu-arch={0}'.format(rocm_arch))
+        else:
+            configure_args.extend([
+                '--without-hip',
+                '--disable-rocrand',
+                '--disable-rocsparse',
+            ])
+
         if '+unified-memory' in spec:
             configure_args.append('--enable-unified-memory')
 
@@ -210,10 +241,6 @@ class Hypre(AutotoolsPackage, CudaPackage):
         if '+cuda' in spec:
             env.set('CUDA_HOME', spec['cuda'].prefix)
             env.set('CUDA_PATH', spec['cuda'].prefix)
-            cuda_arch = spec.variants['cuda_arch'].value
-            if cuda_arch:
-                arch_sorted = list(sorted(cuda_arch, reverse=True))
-                env.set('HYPRE_CUDA_SM', arch_sorted[0])
             # In CUDA builds hypre currently doesn't handle flags correctly
             env.append_flags(
                 'CXXFLAGS', '-O2' if '~debug' in spec else '-g')