From 66d4709b4d8f1556c7558b2719a0a2ce96adb5d7 Mon Sep 17 00:00:00 2001
From: James Foucar <jgfouca@sandia.gov>
Date: Thu, 13 Jun 2024 12:43:07 -0600
Subject: [PATCH] Clean up kokkos/tril situation

If Trilinos is being used, we should use the kokkos that comes
with it. This removes the need to explicitly set Kokkos_ROOT in
config_machines.xml to the Trilinos one.
---
 cime_config/machines/config_machines.xml |  2 --
 components/cmake/find_dep_packages.cmake | 10 +++++++---
 share/build/buildlib.kokkos              |  8 +++++++-
 3 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/cime_config/machines/config_machines.xml b/cime_config/machines/config_machines.xml
index 37691bb669da..978eb1ff470d 100644
--- a/cime_config/machines/config_machines.xml
+++ b/cime_config/machines/config_machines.xml
@@ -277,7 +277,6 @@
       <env name="BLA_VENDOR">Generic</env>
       <env name="Albany_ROOT">$SHELL{if [ -z "$Albany_ROOT" ]; then echo /global/common/software/e3sm/mali_tpls/albany-e3sm-serial-release-gcc-cmake-fix; else echo "$Albany_ROOT"; fi}</env>
       <env name="Trilinos_ROOT">$SHELL{if [ -z "$Trilinos_ROOT" ]; then echo /global/common/software/e3sm/mali_tpls/trilinos-e3sm-serial-release-gcc; else echo "$Trilinos_ROOT"; fi}</env>
-      <env name="Kokkos_ROOT">$SHELL{if [ -z "$Kokkos_ROOT" ]; then echo /global/common/software/e3sm/mali_tpls/trilinos-e3sm-serial-release-gcc; else echo "$Kokkos_ROOT"; fi}</env>
     </environment_variables>
     <environment_variables compiler="nvidia" mpilib="mpich">
       <env name="ADIOS2_ROOT">$SHELL{if [ -z "$ADIOS2_ROOT" ]; then echo /global/cfs/cdirs/e3sm/3rdparty/adios2/2.9.1/cray-mpich-8.1.25/nvidia-22.7; else echo "$ADIOS2_ROOT"; fi}</env>
@@ -2656,7 +2655,6 @@
     <environment_variables compiler="gnu">
       <env name="Albany_ROOT">$SHELL{if [ -z "$Albany_ROOT" ]; then echo /lcrc/group/e3sm/ac.jwatkins/LandIce/AlbanyBuilds/build-gcc-sfad12-e3sm/install; else echo "$Albany_ROOT"; fi}</env>
       <env name="Trilinos_ROOT">$SHELL{if [ -z "$Trilinos_ROOT" ]; then echo /lcrc/group/e3sm/ac.jwatkins/LandIce/TrilinosBuilds/build-gcc-e3sm/install; else echo "$Trilinos_ROOT"; fi}</env>
-      <env name="Kokkos_ROOT">$SHELL{if [ -z "$Kokkos_ROOT" ]; then echo /lcrc/group/e3sm/ac.jwatkins/LandIce/TrilinosBuilds/build-gcc-e3sm/install; else echo "$Kokkos_ROOT"; fi}</env>
     </environment_variables>
   </machine>
 
diff --git a/components/cmake/find_dep_packages.cmake b/components/cmake/find_dep_packages.cmake
index 9d5628bb0bd4..a963fd494330 100644
--- a/components/cmake/find_dep_packages.cmake
+++ b/components/cmake/find_dep_packages.cmake
@@ -16,10 +16,14 @@
 # errors.
 if (USE_KOKKOS)
 
-  # Kokkos will be built in the sharedlibs if Kokkos_ROOT is
-  # unset.
   if (NOT DEFINED ENV{Kokkos_ROOT})
-    set(ENV{Kokkos_ROOT} ${INSTALL_SHAREDPATH})
+    # We should use the kokkos from Trilinos if we are using Trilinos
+    if (USE_TRILINOS)
+      set (ENV{Kokkos_ROOT} ${Trilinos_ROOT})
+    else()
+      # Kokkos will be built in the sharedlibs if Kokkos_ROOT is unset.
+      set(ENV{Kokkos_ROOT} ${INSTALL_SHAREDPATH})
+    endif()
   endif()
 
   find_package(Kokkos REQUIRED)
diff --git a/share/build/buildlib.kokkos b/share/build/buildlib.kokkos
index 035e86b0df2a..8184997615a6 100755
--- a/share/build/buildlib.kokkos
+++ b/share/build/buildlib.kokkos
@@ -50,6 +50,11 @@ OR
 ###############################################################################
 def buildlib(bldroot, installpath, case):
 ###############################################################################
+    # If we are using Trilinos, we must use the Kokkos that comes with it
+    if case.get_value("USE_TRILINOS"):
+        print ("You case is using Trilinos and will use its Kokkos")
+        return
+
     installed_kokkos_dir = os.environ.get("Kokkos_ROOT")
     if installed_kokkos_dir is not None:
         # We are trying to use a pre-installed kokkos. Look for the relevant folders/libs,
@@ -65,9 +70,10 @@ def buildlib(bldroot, installpath, case):
         # often arch dependent (e.g., $prefix/lib or $prefix/lib64)... The best thing
         # would be to run a small cmake script, that calls find_package. But E3SM's
         # cmake build system will do that soon enough, so any error will be caught there.
+        print (f"Using pre-installed Kokkos_ROOT: {kokkos_root}")
         return
     else:
-        print ("no value foudn in env for Kokkos_ROOT. building from scratch")
+        print ("no value found in env for Kokkos_ROOT. building from scratch")
 
     srcroot = case.get_value("SRCROOT")
     ekat_dir = os.path.join(srcroot, "externals", "ekat")