Add ROCmPackage to trilinos (#28424)

* Add ROCmPackage to trilinos.

* Simplify Trilinos ROCm support. Add MI200 to Kokkos AMD GPU arch map.
This commit is contained in:
Jon Rood 2022-01-18 04:24:32 -07:00 committed by GitHub
parent 64913af04a
commit 7ae8a5d55b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 31 additions and 2 deletions

View File

@ -140,7 +140,8 @@ class Kokkos(CMakePackage, CudaPackage, ROCmPackage):
amdgpu_arch_map = { amdgpu_arch_map = {
'gfx900': 'vega900', 'gfx900': 'vega900',
'gfx906': 'vega906', 'gfx906': 'vega906',
'gfx908': 'vega908' 'gfx908': 'vega908',
'gfx90a': 'vega90A'
} }
amd_support_conflict_msg = ( amd_support_conflict_msg = (
'{0} is not supported; ' '{0} is not supported; '

View File

@ -22,7 +22,7 @@
# https://github.com/trilinos/Trilinos/issues/175 # https://github.com/trilinos/Trilinos/issues/175
class Trilinos(CMakePackage, CudaPackage): class Trilinos(CMakePackage, CudaPackage, ROCmPackage):
"""The Trilinos Project is an effort to develop algorithms and enabling """The Trilinos Project is an effort to develop algorithms and enabling
technologies within an object-oriented software framework for the solution technologies within an object-oriented software framework for the solution
of large-scale, complex multi-physics engineering and scientific problems. of large-scale, complex multi-physics engineering and scientific problems.
@ -192,6 +192,7 @@ class Trilinos(CMakePackage, CudaPackage):
# Tpetra packages # Tpetra packages
with when('~kokkos'): with when('~kokkos'):
conflicts('+cuda') conflicts('+cuda')
conflicts('+rocm')
conflicts('+tpetra') conflicts('+tpetra')
conflicts('+intrepid2') conflicts('+intrepid2')
conflicts('+phalanx') conflicts('+phalanx')
@ -411,6 +412,17 @@ def setup_build_environment(self, env):
else: else:
env.set('CXX', spec["kokkos-nvcc-wrapper"].kokkos_cxx) env.set('CXX', spec["kokkos-nvcc-wrapper"].kokkos_cxx)
if '+rocm' in spec:
if '+mpi' in spec:
env.set('OMPI_CXX', self.spec['hip'].hipcc)
env.set('MPICH_CXX', self.spec['hip'].hipcc)
env.set('MPICXX_CXX', self.spec['hip'].hipcc)
else:
env.set('CXX', self.spec['hip'].hipcc)
if '+stk' in spec:
# Using CXXFLAGS for hipcc which doesn't use flags in the spack wrappers
env.set('CXXFLAGS', '-DSTK_NO_BOOST_STACKTRACE')
def cmake_args(self): def cmake_args(self):
options = [] options = []
@ -717,6 +729,22 @@ def define_tpl(trilinos_name, spack_name, have_dep):
for arch in spec.variants['cuda_arch'].value for arch in spec.variants['cuda_arch'].value
) )
if '+rocm' in spec:
options.extend([
define_kok_enable('ROCM', False),
define_kok_enable('HIP', True)
])
if '+tpetra' in spec:
options.append(define('Tpetra_INST_HIP', True))
amdgpu_arch_map = Kokkos.amdgpu_arch_map
for amd_target in spec.variants['amdgpu_target'].value:
try:
arch = amdgpu_arch_map[amd_target]
except KeyError:
pass
else:
options.append(define("Kokkos_ARCH_" + arch.upper(), True))
# ################# System-specific ###################### # ################# System-specific ######################
# Fortran lib (assumes clang is built with gfortran!) # Fortran lib (assumes clang is built with gfortran!)