hypre: add rocm support (#29147)

Co-authored-by: Sarah Osborn <30503782+osborn9@users.noreply.github.com>
This commit is contained in:
Jon Rood 2022-03-09 01:39:59 -07:00 committed by GitHub
parent 9d1d52d51f
commit 68f0973339
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -9,7 +9,7 @@
from spack import * from spack import *
class Hypre(AutotoolsPackage, CudaPackage): class Hypre(AutotoolsPackage, CudaPackage, ROCmPackage):
"""Hypre is a library of high performance preconditioners that """Hypre is a library of high performance preconditioners that
features parallel multigrid methods for both structured and features parallel multigrid methods for both structured and
unstructured grid problems.""" unstructured grid problems."""
@ -93,7 +93,9 @@ class Hypre(AutotoolsPackage, CudaPackage):
depends_on('superlu-dist', when='+superlu-dist+mpi') depends_on('superlu-dist', when='+superlu-dist+mpi')
conflicts('+cuda', when='+int64') 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') conflicts('+gptune', when='~mpi')
# Patch to build shared libraries on Darwin does not apply to # Patch to build shared libraries on Darwin does not apply to
@ -140,6 +142,10 @@ def configure_args(self):
if '+fortran' in spec: if '+fortran' in spec:
os.environ['F77'] = spec['mpi'].mpif77 os.environ['F77'] = spec['mpi'].mpif77
configure_args.append('--with-MPI') 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: else:
configure_args.append('--without-MPI') configure_args.append('--without-MPI')
@ -176,7 +182,13 @@ def configure_args(self):
configure_args.extend([ configure_args.extend([
'--with-cuda', '--with-cuda',
'--enable-curand', '--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 # New in 2.21.0: replaces --enable-cub
if '@2.21.0:' in spec: if '@2.21.0:' in spec:
configure_args.append('--enable-device-memory-pool') configure_args.append('--enable-device-memory-pool')
@ -188,10 +200,29 @@ def configure_args(self):
configure_args.extend([ configure_args.extend([
'--without-cuda', '--without-cuda',
'--disable-curand', '--disable-curand',
'--disable-cusparse',
]) ])
if '@:2.20.99' in spec: if '@:2.20.99' in spec:
configure_args.append('--disable-cub') 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: if '+unified-memory' in spec:
configure_args.append('--enable-unified-memory') configure_args.append('--enable-unified-memory')
@ -210,10 +241,6 @@ def setup_build_environment(self, env):
if '+cuda' in spec: if '+cuda' in spec:
env.set('CUDA_HOME', spec['cuda'].prefix) env.set('CUDA_HOME', spec['cuda'].prefix)
env.set('CUDA_PATH', 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 # In CUDA builds hypre currently doesn't handle flags correctly
env.append_flags( env.append_flags(
'CXXFLAGS', '-O2' if '~debug' in spec else '-g') 'CXXFLAGS', '-O2' if '~debug' in spec else '-g')