Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add cuda-related variants for Albany and Trilinos #31

Merged
merged 3 commits into from
Sep 10, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 29 additions & 2 deletions var/spack/repos/builtin/packages/albany/package.py
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,10 @@ class Albany(CMakePackage):
description="Enable SFad build")
variant("sandybridge", default=False,
description="Compile Trilinos used by Albany for Sandybridge architecture")
variant("zen3", default=False,
description="Compile Trilinos used by Albany for Zen3 architecture")
variant("ampere80", default=False,
description="Compile Trilinos used by Albany for Ampere80 architecture")
variant("mpas", default=False,
description="Enable MPAS interface in build")
variant("sfadsize", default="4", values=("4", "6", "8", "12", "24"), multi=False,
Expand All @@ -53,6 +57,12 @@ class Albany(CMakePackage):
description="Enable packages needed for Albany optimization capabilities (SuperLU, FROSCh, ROL)"),
variant("omegah", default=False,
description="Enable Omega_h in build")
variant("cuda", default=False,
description="Enable CUDA in build")
variant("uvm", default=False, when="+cuda",
description="Enable UVM for CUDA builds")
variant("slfad", default=False,
description="Enable SLFad build")

# Add dependencies
depends_on("mpi")
Expand All @@ -68,6 +78,11 @@ class Albany(CMakePackage):
depends_on("trilinos-for-albany~superlu-dist+exodus+chaco~isorropia+tempus+teko~intrepid+intrepid2+minitensor+phalanx+pnetcdf+nox+piro+shards+stk+amesos2~amesos~hypre+ifpack2~mumps~suite-sparse~epetra~ifpack~ml+muelu~aztec+superlu+rol+frosch gotype=long_long", when="~sandybridge+optimization")
depends_on("trilinos-for-albany~superlu-dist+exodus+chaco~isorropia+tempus+teko~intrepid+intrepid2+minitensor+phalanx+pnetcdf+nox+piro+shards+stk+amesos2~amesos~hypre+ifpack2~mumps~suite-sparse+sandybridge~epetra~ifpack~ml+muelu~aztec+superlu+rol+frosch gotype=long_long", when="+sandybridge+optimization")

depends_on("trilinos-for-albany+cuda+uvm", when="+uvm")
depends_on("trilinos-for-albany+cuda", when="+cuda~uvm")
depends_on("trilinos-for-albany+zen3", when="+zen3")
depends_on("trilinos-for-albany+ampere80", when="+ampere80")

depends_on("trilinos-for-albany@develop", when="@develop")
depends_on("trilinos-for-albany@compass-2023-08-03", when="@compass-2023-08-03")
depends_on("trilinos-for-albany@compass-2024-03-13", when="@compass-2024-03-13")
Expand Down Expand Up @@ -111,8 +126,6 @@ def cmake_args(self):
"ON" if "+perf" in spec else "OFF"),
"-DENABLE_64BIT_INT:BOOL=%s" % (
"ON" if "+64bit" in spec else "OFF"),
"-DENABLE_FAD_TYPE:STRING=%s" % (
"SFad" if "+sfad" in spec else "DFad"),
"-DENABLE_MPAS_INTERFACE:BOOL=%s" % (
"ON" if "+mpas" in spec else "OFF"),
"-DENABLE_ALBANY_PYTHON:BOOL=%s" % (
Expand All @@ -129,8 +142,22 @@ def cmake_args(self):

if "+sfad" in spec:
options.extend([
"-DENABLE_FAD_TYPE:STRING=SFad",
"-DALBANY_SFAD_SIZE=%d" % int(spec.variants["sfadsize"].value)
])
elif "+slfad" in spec:
options.extend([
"-DENABLE_FAD_TYPE:STRING=SLFad",
"-DENABLE_TAN_FAD_TYPE:STRING=SLFad",
"-DENABLE_HES_VEC_FAD_TYPE:STRING=SLFad",
"-DALBANY_SLFAD_SIZE=90",
"-DALBANY_TAN_SLFAD_SIZE=100",
"-DALBANY_HES_VEC_SLFAD_SIZE=100"
])

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Note: this is needed for Albany tests but it might not be needed for mali/e3sm tests. in fact, I think one of the e3sm tests fails with this because of oom issues. sfad12 might be sufficient for all mali/e3sm tests.

else:
options.extend([
"-DENABLE_FAD_TYPE:STRING=DFad"])

if "+debug" in spec:
options.extend([
"-DCMAKE_BUILD_TYPE:STRING=DEBUG"
Expand Down
63 changes: 63 additions & 0 deletions var/spack/repos/builtin/packages/trilinos-for-albany/package.py
Original file line number Diff line number Diff line change
Expand Up @@ -79,6 +79,14 @@ class TrilinosForAlbany(CMakePackage):
description='Compile with Fortran support')
variant('openmp', default=False,
description='Enable OpenMP')
variant('cuda', default=False,
description='Enable CUDA')
variant("uvm", default=False, when="+cuda",
description="Turn on UVM for CUDA build")
variant("aware", default=False, when="+cuda+tpetra",
description="Turn on Tpetra GPU-aware MPI for CUDA build")
variant("wrapper", default=False,
description="Use nvcc-wrapper for CUDA build")
variant('shared', default=True,
description='Enables the build of shared libraries')
variant('debug', default=False,
Expand All @@ -87,6 +95,10 @@ class TrilinosForAlbany(CMakePackage):
description='Compile using the default xSDK configuration')
variant('sandybridge', default=False,
description='Compile Trilinos for Sandybridge architecture')
variant('zen3', default=False,
description='Compile Trilinos for Zen3 architecture')
variant('ampere80', default=False,
description='Compile Trilinos for Ampere80 architecture')

# TPLs (alphabet order)
variant('boost', default=True,
Expand Down Expand Up @@ -314,6 +326,20 @@ class TrilinosForAlbany(CMakePackage):
# ADIOS2 was only added after v12.14.1
conflicts('+adios2', when='@:12.14.1')
conflicts('+adios2', when='@xsdk-0.2.0')

# CUDA without wrapper requires clang
requires(
"%clang",
when="+cuda~wrapper",
msg="trilinos~wrapper+cuda can only be built with the Clang compiler",
)
conflicts("+wrapper", when="~cuda")
conflicts("+wrapper", when="%clang")
# Old trilinos fails with new CUDA (see #27180)
conflicts("@:13.0.1 +cuda", when="^cuda@11:")
# Cuda UVM must be enabled prior to 13.2
# See https://github.com/spack/spack/issues/28869
conflicts("~uvm", when="@:13.1 +cuda")
# ###################### Dependencies ##########################

# Everything should be compiled position independent (-fpic)
Expand Down Expand Up @@ -363,6 +389,7 @@ class TrilinosForAlbany(CMakePackage):
depends_on('python', when='+python')
depends_on('py-numpy', when='+python', type=('build', 'run'))
depends_on('swig', when='+python')
depends_on("kokkos-nvcc-wrapper", when="+wrapper")

patch('umfpack_from_suitesparse.patch', when='@11.14.1:12.8.1')
patch('xlf_seacas.patch', when='@12.10.1:12.12.1 %xl')
Expand All @@ -380,6 +407,22 @@ def setup_run_environment(self, env):
if "+exodus" in self.spec:
env.prepend_path("PYTHONPATH", self.prefix.lib)

def setup_dependent_package(self, module, dependent_spec):
if "+wrapper" in self.spec:
self.spec.kokkos_cxx = self.spec["kokkos-nvcc-wrapper"].kokkos_cxx
else:
self.spec.kokkos_cxx = spack_cxx

def setup_build_environment(self, env):
spec = self.spec
if "+cuda" in spec and "+wrapper" in spec:
if "+mpi" in spec:
env.set("OMPI_CXX", spec["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("MPICH_CXX", spec["kokkos-nvcc-wrapper"].kokkos_cxx)
env.set("MPICXX_CXX", spec["kokkos-nvcc-wrapper"].kokkos_cxx)
else:
env.set("CXX", spec["kokkos-nvcc-wrapper"].kokkos_cxx)

def cmake_args(self):
spec = self.spec

Expand All @@ -400,6 +443,10 @@ def cmake_args(self):
'ON' if '+debug' in spec else 'OFF'),
'-DKokkos_ARCH_SNB:BOOL=%s' % (
'ON' if '+sandybridge' in spec else 'OFF'),
'-DKokkos_ARCH_ZEN3:BOOL=%s' % (
'ON' if '+zen3' in spec else 'OFF'),
'-DKokkos_ARCH_AMPERE80:BOOL=%s' % (
'ON' if '+ampere80' in spec else 'OFF'),

# The following can cause problems on systems that don't have
# static libraries available for things like dl and pthreads
Expand Down Expand Up @@ -759,6 +806,22 @@ def cmake_args(self):
'-DTpetra_INST_OPENMP:BOOL=ON'
])

# CUDA
if '+cuda' in spec:
options.extend([
'-DKokkos_ENABLE_CUDA:BOOL=ON',
'-DKokkos_ENABLE_CUDA_LAMBDA:BOOL=ON',
'-DKokkos_ENABLE_CUDA_UVM:BOOL=%s' % 'ON' if '+uvm' in spec else 'OFF',
'-DKokkos_ENABLE_IMPL_CUDA_MALLOC_ASYNC:BOOL=OFF',
'-DTPL_ENABLE_CUDA:BOOL=ON',
'-DTPL_ENABLE_CUSPARSE:BOOL=OFF',
jewatkins marked this conversation as resolved.
Show resolved Hide resolved
'-DTPL_ENABLE_THRUST:BOOL=ON'
])
if '+tpetra' in spec:
options.extend([
'-DTpetra_ASSUME_GPU_AWARE_MPI:BOOL=%s' % 'ON' if '+aware' in spec else 'OFF'
])

# Fortran lib
if '+fortran' in spec:
if spec.satisfies('%gcc') or spec.satisfies('%clang'):
Expand Down
Loading