Merge branch 'jgfouca/refactor_ekat_build' into next (PR #7634)

jgfouca · jgfouca · commit deee180c0fa5 · 2025-08-29T09:30:45.000-06:00
Refactor EKAT/Kokkos integration with E3SM

Change list:
1) Fix DEBUG propagation to EKAT build, this will fix CI
2) Add new setting KOKKOS_USE_EKAT_MACH_FILE to E3SM cmake macros. This
   will allow machines to opt-in to using an EKAT machine file to configure
   kokkos. This should hopefully reduce duplication and tedium of Kokkos
   cmake configuration.
3) Make several machine take advantage of this feature.
4) Make eamxx standalone use the ekat machine files where possible
5) Remove unused EKAT_DISABLE_WORKSPACE_SHARING
6) Allow users to turn on kokkos view bounds checking via the env variable E3SM_ENABLE_KOKKOS_BOUNDS_CHECKING
7) ALWAYS turn off openmp for non-threaded cases

Guide to configuring Kokkos/EKAT:
1) If the machine/compiler macro does nothing, it will use the universal
   Kokkos settings (Serial + (OpenMP IF compile_threaded).
2) The macro can manually set configuration by setting KOKKOS_OPTIONS.
3) The macro can just follow EKAT machine files by setting KOKKOS_USE_EKAT_MACH_FILE to TRUE.
   This only works if ekat/cmake/machine-files has a file ${MACH}.cmake.

[BFB]
diff --git a/cime_config/machines/cmake_macros/chrysalis.cmake b/cime_config/machines/cmake_macros/chrysalis.cmake
@@ -0,0 +1 @@
+set(KOKKOS_USE_EKAT_MACH_FILE "TRUE")
diff --git a/cime_config/machines/cmake_macros/compy.cmake b/cime_config/machines/cmake_macros/compy.cmake
@@ -0,0 +1 @@
+set(KOKKOS_USE_EKAT_MACH_FILE "TRUE")
diff --git a/cime_config/machines/cmake_macros/gnu_ghci-snl-cpu.cmake b/cime_config/machines/cmake_macros/gnu_ghci-snl-cpu.cmake
@@ -0,0 +1 @@
+set(KOKKOS_USE_EKAT_MACH_FILE "TRUE")
diff --git a/cime_config/machines/cmake_macros/gnu_ghci-snl-gpu.cmake b/cime_config/machines/cmake_macros/gnu_ghci-snl-gpu.cmake
@@ -0,0 +1 @@
+set(KOKKOS_USE_EKAT_MACH_FILE "TRUE")
diff --git a/cime_config/machines/cmake_macros/gnu_mappy.cmake b/cime_config/machines/cmake_macros/gnu_mappy.cmake
@@ -6,3 +6,5 @@ string(APPEND CMAKE_Fortran_FLAGS_RELEASE " -O2")
 if (MPILIB STREQUAL mpi-serial AND NOT compile_threaded)
   set(PFUNIT_PATH "$ENV{SEMS_PFUNIT_ROOT}")
 endif()
+
+set(KOKKOS_USE_EKAT_MACH_FILE "TRUE")
diff --git a/cime_config/machines/cmake_macros/pm-cpu.cmake b/cime_config/machines/cmake_macros/pm-cpu.cmake
@@ -0,0 +1 @@
+set(KOKKOS_USE_EKAT_MACH_FILE "TRUE")
diff --git a/cime_config/machines/cmake_macros/pm-gpu.cmake b/cime_config/machines/cmake_macros/pm-gpu.cmake
@@ -0,0 +1 @@
+set(KOKKOS_USE_EKAT_MACH_FILE "TRUE")
diff --git a/cime_config/machines/cmake_macros/universal.cmake b/cime_config/machines/cmake_macros/universal.cmake
@@ -7,6 +7,7 @@ set(KOKKOS_OPTIONS "-DKokkos_ENABLE_SERIAL=On")
 if (compile_threaded)
   string(APPEND KOKKOS_OPTIONS " -DKokkos_ENABLE_OPENMP=On")
 endif()
+set(KOKKOS_USE_EKAT_MACH_FILE "FALSE")
 
 # Unless told otherwise, set has_contiguous to FALSE
 set(HAS_F2008_CONTIGUOUS "FALSE")
diff --git a/components/eamxx/CMakeLists.txt b/components/eamxx/CMakeLists.txt
@@ -71,20 +71,22 @@ endif()
 #                  Kokkos related settings                         #
 ####################################################################
 
+# Note: this only does anything for standalone eamxx. CIME cases will
+# use the sharedlib EKAT/kokkos
 if (Kokkos_ENABLE_CUDA)
   if (Kokkos_ENABLE_CUDA_RELOCATABLE_DEVICE_CODE)
     if (Kokkos_ENABLE_DEBUG_BOUNDS_CHECK)
       string (CONCAT msg
-          "Kokkos_ENALBE_CUDA_RELOCATABLE_DEVICE_CODE=ON, and Kokkos_ENALBE_DEBUG_BOUNDS_CHECK=ON.\n"
+          "Kokkos_ENABLE_CUDA_RELOCATABLE_DEVICE_CODE=ON, and Kokkos_ENABLE_DEBUG_BOUNDS_CHECK=ON.\n"
           "   -> Disabling bounds checks, to prevent internal compiler errors.")
       message(WARNING "${msg}")
-      set (Kokkos_ENABLE_DEBUG_BOUNDS_CHECK OFF CACHE BOOL "" FORCE)
+      set(Kokkos_ENABLE_DEBUG_BOUNDS_CHECK OFF CACHE BOOL "" FORCE)
     else()
       string (CONCAT msg
-          "Kokkos_ENALBE_CUDA_RELOCATABLE_DEVICE_CODE=ON.\n"
+          "Kokkos_ENABLE_CUDA_RELOCATABLE_DEVICE_CODE=ON.\n"
           "   -> Disabling bounds checks, to prevent internal compiler errors.")
       message(STATUS "${msg}")
-      set (Kokkos_ENABLE_DEBUG_BOUNDS_CHECK OFF CACHE BOOL "" FORCE)
+      set(Kokkos_ENABLE_DEBUG_BOUNDS_CHECK OFF CACHE BOOL "" FORCE)
     endif()
   endif()
 endif()
@@ -204,9 +206,6 @@ option(SCREAM_ENABLE_MAM "Whether to enable MAM aerosol support" ON)
 set(SCREAM_SMALL_KERNELS ${DEFAULT_SMALL_KERNELS} CACHE STRING "Use small, non-monolothic kokkos kernels for ALL components that support them")
 set(SCREAM_P3_SMALL_KERNELS ${SCREAM_SMALL_KERNELS} CACHE STRING "Use small, non-monolothic kokkos kernels for P3 only")
 set(SCREAM_SHOC_SMALL_KERNELS ${SCREAM_SMALL_KERNELS} CACHE STRING "Use small, non-monolothic kokkos kernels for SHOC only")
-if (NOT SCREAM_P3_SMALL_KERNELS AND NOT SCREAM_SHOC_SMALL_KERNELS)
-  set(EKAT_DISABLE_WORKSPACE_SHARING TRUE CACHE STRING "")
-endif()
 
 # Add RRTMGP settings. Note, we might consider also adding RRTMGP_EXPENSIVE_CHECKS
 # to turn on the RRTMGP internal checks here as well, via
diff --git a/components/eamxx/cmake/machine-files/chrysalis.cmake b/components/eamxx/cmake/machine-files/chrysalis.cmake
@@ -1,15 +1,2 @@
 include(${CMAKE_CURRENT_LIST_DIR}/common.cmake)
 common_setup()
-
-set(EKAT_MACH_FILES_PATH ${CMAKE_CURRENT_LIST_DIR}/../../../../externals/ekat/cmake/machine-files)
-
-# Get AMD arch settings
-include(${EKAT_MACH_FILES_PATH}/kokkos/amd-zen2.cmake)
-
-# Add OpenMP settings in standalone mode OR e3sm with compile_threaded=ON
-if (NOT "${PROJECT_NAME}" STREQUAL "E3SM" OR compile_threaded)
-  include(${EKAT_MACH_FILES_PATH}/kokkos/openmp.cmake)
-endif()
-
-# Use srun for standalone testing
-include(${EKAT_MACH_FILES_PATH}/mpi/srun.cmake)
diff --git a/components/eamxx/cmake/machine-files/compy.cmake b/components/eamxx/cmake/machine-files/compy.cmake
@@ -1,10 +1,2 @@
 include(${CMAKE_CURRENT_LIST_DIR}/common.cmake)
 common_setup()
-
-# Load all kokkos settings from Ekat's mach file
-include (${EKAT_MACH_FILES_PATH}/kokkos/intel-skx.cmake)
-include (${EKAT_MACH_FILES_PATH}/kokkos/openmp.cmake)
-include (${EKAT_MACH_FILES_PATH}/mpi/srun.cmake)
-
-#Compy SLURM specific settings
-set(EKAT_MPI_NP_FLAG "-p short --mpi=pmi2 -n" CACHE STRING "" FORCE)
diff --git a/components/eamxx/cmake/machine-files/pm-cpu.cmake b/components/eamxx/cmake/machine-files/pm-cpu.cmake
@@ -1,9 +1,6 @@
 include(${CMAKE_CURRENT_LIST_DIR}/common.cmake)
 common_setup()
 
-include (${EKAT_MACH_FILES_PATH}/kokkos/amd-zen3.cmake)
-include (${EKAT_MACH_FILES_PATH}/kokkos/openmp.cmake)
-
 set(CMAKE_CXX_FLAGS "-DTHRUST_IGNORE_CUB_VERSION_CHECK" CACHE STRING "" FORCE)
 
 #message(STATUS "pm-cpu CMAKE_CXX_COMPILER_ID=${CMAKE_CXX_COMPILER_ID} CMAKE_Fortran_COMPILER_VERSION=${CMAKE_Fortran_COMPILER_VERSION}")
diff --git a/components/eamxx/cmake/machine-files/pm-gpu.cmake b/components/eamxx/cmake/machine-files/pm-gpu.cmake
@@ -1,23 +1,6 @@
 include(${CMAKE_CURRENT_LIST_DIR}/common.cmake)
 common_setup()
 
-#message(STATUS "pm-gpu PROJECT_NAME=${PROJECT_NAME} USE_CUDA=${USE_CUDA} KOKKOS_ENABLE_CUDA=${KOKKOS_ENABLE_CUDA}")
-if ("${PROJECT_NAME}" STREQUAL "E3SM")
-  if (USE_CUDA)
-    include (${EKAT_MACH_FILES_PATH}/kokkos/nvidia-a100.cmake)
-    include (${EKAT_MACH_FILES_PATH}/kokkos/cuda.cmake)
-  else()
-    include (${EKAT_MACH_FILES_PATH}/kokkos/amd-zen3.cmake)
-    include (${EKAT_MACH_FILES_PATH}/kokkos/openmp.cmake)
-    #include (${EKAT_MACH_FILES_PATH}/kokkos/serial.cmake)
-  endif()
-else()
-  include (${EKAT_MACH_FILES_PATH}/kokkos/nvidia-a100.cmake)
-  include (${EKAT_MACH_FILES_PATH}/kokkos/cuda.cmake)
-endif()
-
-include (${EKAT_MACH_FILES_PATH}/mpi/srun.cmake)
-
 set(EKAT_MPI_EXTRA_ARGS "${EKAT_MPI_EXTRA_ARGS} --gpus-per-task=1" CACHE STRING "" FORCE)
 
 #option(Kokkos_ARCH_AMPERE80 "" ON)
diff --git a/externals/ekat b/externals/ekat
@@ -1 +1 @@
-Subproject commit 10ab9d34ceba4eead542702a29676adf77ad91af
+Subproject commit 092d6e16f5c665d1d7dc7887fb4c60f0cfe1b689
diff --git a/share/build/buildlib.ekat b/share/build/buildlib.ekat
@@ -4,9 +4,8 @@ import os, sys, argparse, logging, shutil
 
 from standard_script_setup import *
 from CIME import utils
-from CIME.utils import expect, run_bld_cmd_ensure_logging, run_cmd_no_fail, run_cmd
+from CIME.utils import expect, run_bld_cmd_ensure_logging
 from CIME.case import Case
-from CIME.build import get_standard_makefile_args
 
 sys.path.append(os.path.dirname(__file__))
 from buildlib_util import extract_from_macros
@@ -78,19 +77,35 @@ def buildlib(bldroot, installpath, case):
 
     srcroot = case.get_value("SRCROOT")
     ekat_dir = os.path.join(srcroot, "externals", "ekat")
+    ekat_cmake_files = os.path.join(ekat_dir, "cmake", "machine-files")
+    mach = case.get_value("MACH")
+    ekat_cmake_mach_file = os.path.join(ekat_cmake_files, f"{mach}.cmake")
     expect(os.path.isdir(ekat_dir), "Missing ekat submodule")
 
     # We want to get the compiler and kokkos_options from Macros.make
     # (generated from cmake_macros), but we want to otherwise
     # let kokkos control flags
     mpilib = case.get_value("MPILIB")
-    macro_vals = extract_from_macros(case, "kokkos", extra_vars=("KOKKOS_OPTIONS",))
-    cc, cxx, kokkos_options_raw = macro_vals[4], macro_vals[5], macro_vals[-1]
+    macro_vals = extract_from_macros(case, "kokkos",
+                                     extra_vars=("KOKKOS_OPTIONS", "KOKKOS_USE_EKAT_MACH_FILE"))
+    cc, cxx, kokkos_options_raw, use_ekat_mach_file = \
+        macro_vals[4], macro_vals[5], macro_vals[-2], macro_vals[-1]
+    use_ekat_mach_file = True if use_ekat_mach_file.upper() == "TRUE" else False
     if mpilib == "mpi-serial":
         # Use serial compilers
         cc, cxx = macro_vals[1], macro_vals[2]
 
-    kokkos_options = kokkos_options_raw.split(":=")[-1].strip()
+    # Use the ekat mach file if requested
+    if use_ekat_mach_file:
+        expect(os.path.exists(ekat_cmake_mach_file),
+               f"No KOKKOS_OPTIONS provided and no ekat machine file exists for {mach}")
+        kokkos_options = f"-C {ekat_cmake_mach_file}"
+    else:
+        kokkos_options = kokkos_options_raw.split(":=")[-1].strip()
+
+    # Turn off threading for non-threaded cases
+    if not case.get_build_threaded():
+        kokkos_options += " -DKokkos_ENABLE_OPENMP=Off"
 
     # We will take care of setting CMAKE_CXX_COMPILER unless user has explictly
     # set one in their KOKKOS_OPTIONS.
@@ -113,9 +128,15 @@ def buildlib(bldroot, installpath, case):
     gmake_j = case.get_value("GMAKE_J")
 
     # Just enable all ekat packages?
-    ekat_cmake_options = "-DEKAT_ENABLE_ALL_PACKAGES=On -DEKAT_ENABLE_TESTS=Off"
+    debug = case.get_value("DEBUG")
+    ekat_cmake_options = "-DEKAT_ENABLE_ALL_PACKAGES=On -DEKAT_ENABLE_TESTS=Off -DKokkos_ENABLE_SERIAL=On"
     if mpilib == "mpi-serial":
         ekat_cmake_options += " -DEKAT_ENABLE_MPI=Off"
+    if debug:
+        ekat_cmake_options += " -DCMAKE_BUILD_TYPE=Debug"
+        # Kokkos bounds checking is too expensive to always turn on
+        if "E3SM_ENABLE_KOKKOS_BOUNDS_CHECKING" in os.environ:
+            ekat_cmake_options += " -DKokkos_ENABLE_DEBUG_BOUNDS_CHECK=On"
 
     gen_makefile_cmd = f"cmake {kokkos_options} {cc} {cxx} -DCMAKE_CXX_STANDARD=17 {ekat_cmake_options} -DCMAKE_INSTALL_PREFIX={installpath} {ekat_dir}"
 
