From 88de99f496b510090522eb26a11fe1327f97ee79 Mon Sep 17 00:00:00 2001
From: Xylar Asay-Davis <xylarstorm@gmail.com>
Date: Sat, 26 Oct 2024 10:00:13 -0600
Subject: [PATCH] Update trilinos and albany variants on pm-gpu

Other GPU machines will need to set variants in a similar way
---
 compass/machines/pm-gpu.cfg |  4 ++++
 conda/bootstrap.py          | 15 ++++-----------
 conda/default.cfg           |  3 +++
 3 files changed, 11 insertions(+), 11 deletions(-)

diff --git a/compass/machines/pm-gpu.cfg b/compass/machines/pm-gpu.cfg
index b82faf518..020e32c35 100644
--- a/compass/machines/pm-gpu.cfg
+++ b/compass/machines/pm-gpu.cfg
@@ -31,6 +31,10 @@ spack = /global/cfs/cdirs/e3sm/software/compass/pm-gpu/spack
 # pnetcdf as E3SM (spack modules are used otherwise)
 use_e3sm_hdf5_netcdf = True
 
+# spack variants for Albany and Trilinos
+albany_variants = +mpas~py+unit_tests+cuda+uvm+sfad sfadsize=12
+trilinos_variants = +cuda+uvm+ampere80+zen3
+
 # The parallel section describes options related to running jobs in parallel.
 # Most options in this section come from mache so here we just add or override
 # some defaults
diff --git a/conda/bootstrap.py b/conda/bootstrap.py
index 1e42a3b5d..a5682ba2c 100755
--- a/conda/bootstrap.py
+++ b/conda/bootstrap.py
@@ -470,6 +470,8 @@ def build_spack_env(config, update_spack, machine, compiler, mpi,  # noqa: C901
                     tmpdir, logger):
 
     albany = config.get('deploy', 'albany')
+    albany_variants = config.get('deploy', 'albany_variants')
+    trilinos_variants = config.get('deploy', 'trilinos_variants')
     cmake = config.get('deploy', 'cmake')
     esmf = config.get('deploy', 'esmf')
     lapack = config.get('deploy', 'lapack')
@@ -479,9 +481,6 @@ def build_spack_env(config, update_spack, machine, compiler, mpi,  # noqa: C901
     scorpio = config.get('deploy', 'scorpio')
     parallelio = config.get('deploy', 'parallelio')
 
-    # for now, we'll assume Cuda is needed anytime GPUs are present
-    with_cuda = config.has_option('parallel', 'gpus_per_node')
-
     if config.has_option('deploy', 'spack_mirror'):
         spack_mirror = config.get('deploy', 'spack_mirror')
     else:
@@ -543,14 +542,8 @@ def build_spack_env(config, update_spack, machine, compiler, mpi,  # noqa: C901
             f'@{parallelio}+pnetcdf~timing"')
 
     if albany != 'None':
-        if with_cuda:
-            albany_cuda = '+cuda+uvm+sfad sfadsize=12'
-            trilinos_cuda = '+cuda+uvm'
-        else:
-            albany_cuda = ''
-            trilinos_cuda = ''
-        specs.append(f'"trilinos-for-albany@{albany}{trilinos_cuda}"')
-        specs.append(f'"albany@{albany}+mpas~py+unit_tests{albany_cuda}"')
+        specs.append(f'"trilinos-for-albany@{albany}{trilinos_variants}"')
+        specs.append(f'"albany@{albany}{albany_variants}"')
 
     yaml_template = f'{spack_template_path}/{machine}_{compiler}_{mpi}.yaml'
     if not os.path.exists(yaml_template):
diff --git a/conda/default.cfg b/conda/default.cfg
index 6cae25589..ba4c1bfe4 100644
--- a/conda/default.cfg
+++ b/conda/default.cfg
@@ -21,6 +21,9 @@ mpi = nompi
 
 # the version of various packages to include if using spack
 albany = compass-2024-03-13
+# spack variants for Albany and Trilinos
+albany_variants = +mpas~py+unit_tests
+trilinos_variants =
 # cmake newer than 3.23.0 needed for Trilinos
 cmake = 3.23.0:
 esmf = 8.6.1