More RRTMGP performance work

components/eamxx/CMakeLists.txt

@@ -212,8 +212,8 @@ endif()
 #     #cmakedefine RRTMGP_EXPENSIVE_CHECKS
 option (SCREAM_RRTMGP_DEBUG "Turn on extra debug checks in RRTMGP" ${SCREAM_DEBUG})
 
-option(SCREAM_RRTMGP_ENABLE_YAKL "Use YAKL under rrtmgp" TRUE)
-option(SCREAM_RRTMGP_ENABLE_KOKKOS "Use Kokkos under rrtmgp" FALSE)
+option(SCREAM_RRTMGP_ENABLE_YAKL "Use YAKL under rrtmgp" FALSE)
+option(SCREAM_RRTMGP_ENABLE_KOKKOS "Use Kokkos under rrtmgp" TRUE)
 if (SCREAM_RRTMGP_ENABLE_YAKL)
   add_definitions("-DRRTMGP_ENABLE_YAKL")
 endif()

components/eamxx/cime_config/namelist_defaults_scream.xml

@@ -531,6 +531,7 @@ be lost if SCREAM_HACK_XML is not enabled.
       <do_subcol_sampling type="logical" doc="Flag to turn on/off subcolumn sampling of optical properties; if false treat cells as either completely clear or cloudy">
           true
       </do_subcol_sampling>
+      <pool_size_multiplier type="real">1.0</pool_size_multiplier>
     </rrtmgp>
 
     <mac_aero_mic inherit="atm_proc_group">

components/eamxx/cime_config/testdefs/testmods_dirs/scream/kokkos_rrtmgp/shell_commands

@@ -0,0 +1,3 @@
+./xmlchange --append SCREAM_CMAKE_OPTIONS='SCREAM_RRTMGP_ENABLE_YAKL Off'
+./xmlchange --append SCREAM_CMAKE_OPTIONS='SCREAM_RRTMGP_ENABLE_KOKKOS On'
+

components/eamxx/cime_config/testdefs/testmods_dirs/scream/yakl_rrtmgp/shell_commands

@@ -0,0 +1,2 @@
+./xmlchange --append SCREAM_CMAKE_OPTIONS='SCREAM_RRTMGP_ENABLE_YAKL On'
+./xmlchange --append SCREAM_CMAKE_OPTIONS='SCREAM_RRTMGP_ENABLE_KOKKOS Off'

components/eamxx/src/physics/rrtmgp/CMakeLists.txt

@@ -51,59 +51,60 @@ endmacro()
 ##################################
 
 # RRTMGP++ requires YAKL
+if (SCREAM_RRTMGP_ENABLE_YAKL)
+  string(TOLOWER "${CMAKE_BUILD_TYPE}" CMAKE_BUILD_TYPE_ci)
+  if (TARGET yakl)
+    # Other E3SM components are building YAKL...
+    message ("It appears some other part of E3SM is building YAKL.\n"
+      "We will reuse that, but if this is a debug build we will\n"
+      "add the --fmad=false flag to the cuda flags used by YAKL\n")
+  else ()
+    # Prepare CUDA/HIP flags for YAKL
+    if (CUDA_BUILD)
+      string(REPLACE ";" " " KOKKOS_CUDA_OPTIONS_STR "${KOKKOS_CUDA_OPTIONS}")
+      set(YAKL_ARCH "CUDA")
+      set(YAKL_CUDA_FLAGS "-DYAKL_ARCH_CUDA ${KOKKOS_CUDA_OPTIONS_STR} --expt-relaxed-constexpr -ccbin ${CMAKE_CXX_COMPILER}")
+      string (REPLACE " " ";" YAKL_CUDA_FLAGS_LIST ${YAKL_CUDA_FLAGS})
+    endif()
+    if (HIP_BUILD)
+      set(YAKL_ARCH "HIP")
+      set(YAKL_HIP_FLAGS "-DYAKL_ARCH_HIP -O3 -D__HIP_ROCclr__ -D__HIP_ARCH_GFX90A__=1 --rocm-path=${ROCM_PATH} --offload-arch=gfx90a -x hip")
+      string (REPLACE " " ";" YAKL_HIP_FLAGS_LIST ${YAKL_HIP_FLAGS})
+    endif()
 
-string(TOLOWER "${CMAKE_BUILD_TYPE}" CMAKE_BUILD_TYPE_ci)
-if (TARGET yakl)
-  # Other E3SM components are building YAKL...
-  message ("It appears some other part of E3SM is building YAKL.\n"
-           "We will reuse that, but if this is a debug build we will\n"
-           "add the --fmad=false flag to the cuda flags used by YAKL\n")
-else ()
-  # Prepare CUDA/HIP flags for YAKL
-  if (CUDA_BUILD)
-    string(REPLACE ";" " " KOKKOS_CUDA_OPTIONS_STR "${KOKKOS_CUDA_OPTIONS}")
-    set(YAKL_ARCH "CUDA")
-    set(YAKL_CUDA_FLAGS "-DYAKL_ARCH_CUDA ${KOKKOS_CUDA_OPTIONS_STR} --expt-relaxed-constexpr -ccbin ${CMAKE_CXX_COMPILER}")
-    string (REPLACE " " ";" YAKL_CUDA_FLAGS_LIST ${YAKL_CUDA_FLAGS})
-  endif()
-  if (HIP_BUILD)
-    set(YAKL_ARCH "HIP")
-    set(YAKL_HIP_FLAGS "-DYAKL_ARCH_HIP -O3 -D__HIP_ROCclr__ -D__HIP_ARCH_GFX90A__=1 --rocm-path=${ROCM_PATH} --offload-arch=gfx90a -x hip")
-    string (REPLACE " " ";" YAKL_HIP_FLAGS_LIST ${YAKL_HIP_FLAGS})
-  endif()
+    set (YAKL_SOURCE_DIR ${SCREAM_BASE_DIR}/../../externals/YAKL)
+    add_subdirectory(${YAKL_SOURCE_DIR} ${CMAKE_BINARY_DIR}/externals/YAKL)
 
-  set (YAKL_SOURCE_DIR ${SCREAM_BASE_DIR}/../../externals/YAKL)
-  add_subdirectory(${YAKL_SOURCE_DIR} ${CMAKE_BINARY_DIR}/externals/YAKL)
+    # Set some additional flag/cpp option on the yakl target
 
-  # Set some additional flag/cpp option on the yakl target
+    cmake_policy (SET CMP0079 NEW) # Allow to link to a tgt from a different directory
 
-  cmake_policy (SET CMP0079 NEW) # Allow to link to a tgt from a different directory
+    # EAMxx *requires* MPI, so simply look for it, then link against it
+    find_package(MPI REQUIRED COMPONENTS C)
+    target_link_libraries (yakl INTERFACE MPI::MPI_C)
 
-  # EAMxx *requires* MPI, so simply look for it, then link against it
-  find_package(MPI REQUIRED COMPONENTS C)
-  target_link_libraries (yakl INTERFACE MPI::MPI_C)
+    # For debug builds, set -DYAKL_DEBUG
+    if (CMAKE_BUILD_TYPE_ci STREQUAL "debug")
+      target_compile_definitions(yakl INTERFACE YAKL_DEBUG)
+    endif()
+  endif()
 
-  # For debug builds, set -DYAKL_DEBUG
-  if (CMAKE_BUILD_TYPE_ci STREQUAL "debug")
-    target_compile_definitions(yakl INTERFACE YAKL_DEBUG)
+  # See eamxx/src/dynamics/homme/CMakeLists.txt for an explanation of this
+  # workaround.
+  if ((SCREAM_MACHINE STREQUAL "ascent" OR SCREAM_MACHINE STREQUAL "pm-gpu") AND CMAKE_BUILD_TYPE_ci STREQUAL "debug")
+    SetCudaFlagsYakl(yakl CUDA_LANG FLAGS -UNDEBUG)
+  else()
+    SetCudaFlagsYakl(yakl CUDA_LANG)
   endif()
-endif()
 
-# See eamxx/src/dynamics/homme/CMakeLists.txt for an explanation of this
-# workaround.
-if ((SCREAM_MACHINE STREQUAL "ascent" OR SCREAM_MACHINE STREQUAL "pm-gpu") AND CMAKE_BUILD_TYPE_ci STREQUAL "debug")
-  SetCudaFlagsYakl(yakl CUDA_LANG FLAGS -UNDEBUG)
-else()
-  SetCudaFlagsYakl(yakl CUDA_LANG)
+  list(APPEND CMAKE_MODULE_PATH ${YAKL_SOURCE_DIR})
+  include (yakl_utils)
 endif()
 
 ##################################
 #           RRTMGP               #
 ##################################
 
-list(APPEND CMAKE_MODULE_PATH ${YAKL_SOURCE_DIR})
-include (yakl_utils)
-
 set(EAM_RRTMGP_DIR ${SCREAM_BASE_DIR}/../eam/src/physics/rrtmgp)
 # Build RRTMGP library; this builds the core RRTMGP external source as a library named "rrtmgp"
 # NOTE: The external RRTMGP build needs some fixes to work with CUDA in a library build, so for now we will build these ourselves
@@ -122,15 +123,25 @@ set(EXTERNAL_SRC
 add_library(rrtmgp ${EXTERNAL_SRC})
 target_compile_definitions(rrtmgp PUBLIC EAMXX_HAS_RRTMGP)
 EkatDisableAllWarning(rrtmgp)
-yakl_process_target(rrtmgp)
+if (SCREAM_RRTMGP_ENABLE_YAKL)
+  yakl_process_target(rrtmgp)
+else()
+  if (CUDA_BUILD)
+    target_compile_options(rrtmgp PUBLIC $<$<COMPILE_LANGUAGE:CXX>:--expt-relaxed-constexpr>)
+  endif()
+endif()
 
 # NOTE: cannot use 'PUBLIC' in target_link_libraries,
 #       since yakl_process_target already used it
 #       with the "plain" signature
 if (NOT TARGET Kokkos::kokkos)
   find_package(Kokkos REQUIRED)
 endif ()
-target_link_libraries(rrtmgp yakl Kokkos::kokkos)
+if (SCREAM_RRTMGP_ENABLE_YAKL)
+  target_link_libraries(rrtmgp yakl Kokkos::kokkos)
+else()
+  target_link_libraries(rrtmgp Kokkos::kokkos)
+endif()
 target_include_directories(rrtmgp PUBLIC
     ${SCREAM_BASE_DIR}/../../externals/YAKL
     ${EAM_RRTMGP_DIR}/external/cpp
@@ -158,21 +169,23 @@ target_include_directories(rrtmgp PUBLIC
 #       SCREAM_RRTMGP_YAKL       #
 ##################################
 
-set(SCREAM_RRTMGP_SOURCES_YAKL
+set(SCREAM_RRTMGP_SOURCES_INTERFACE
   scream_rrtmgp_interface.cpp
 )
 
-add_library(scream_rrtmgp_yakl ${SCREAM_RRTMGP_SOURCES_YAKL})
-yakl_process_target(scream_rrtmgp_yakl)
+add_library(scream_rrtmgp_interface ${SCREAM_RRTMGP_SOURCES_INTERFACE})
+if (SCREAM_RRTMGP_ENABLE_YAKL)
+  yakl_process_target(scream_rrtmgp_interface)
+endif()
 
 # NOTE: cannot use 'PUBLIC' in target_link_libraries,
 #       since yakl_process_target already used it
 #       with the "plain" signature
 find_library(NETCDF_C netcdf HINTS ${NetCDF_C_PATH}/lib)
-target_link_libraries(scream_rrtmgp_yakl ${NETCDF_C} rrtmgp scream_share Kokkos::kokkos)
-target_include_directories(scream_rrtmgp_yakl PUBLIC
+target_link_libraries(scream_rrtmgp_interface ${NETCDF_C} rrtmgp scream_share Kokkos::kokkos)
+target_include_directories(scream_rrtmgp_interface PUBLIC
     ${CMAKE_CURRENT_SOURCE_DIR})
-target_include_directories(scream_rrtmgp_yakl SYSTEM PUBLIC
+target_include_directories(scream_rrtmgp_interface SYSTEM PUBLIC
     ${NetCDF_C_PATH}/include
     ${EAM_RRTMGP_DIR}/external)
 
@@ -186,7 +199,7 @@ set(SCREAM_RRTMGP_SOURCES
 )
 
 add_library(scream_rrtmgp ${SCREAM_RRTMGP_SOURCES})
-target_link_libraries(scream_rrtmgp PUBLIC scream_share physics_share csm_share scream_rrtmgp_yakl Kokkos::kokkos)
+target_link_libraries(scream_rrtmgp PUBLIC scream_share physics_share csm_share scream_rrtmgp_interface Kokkos::kokkos)
 set_target_properties(scream_rrtmgp PROPERTIES
   Fortran_MODULE_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/modules
 )
@@ -199,9 +212,11 @@ target_include_directories(scream_rrtmgp PUBLIC
 # ${YAKL_${YAKL_ARCH}_FLAGS} flags to the CXX flags of scream_rrtmgp.
 # In particular, this will ensure that all the yakl macros
 # are correctly defined in YAKL headers, depending on the backend
-if (YAKL_ARCH)
-  target_compile_options(scream_rrtmgp PUBLIC
-        "$<$<COMPILE_LANGUAGE:CXX>:${YAKL_${YAKL_ARCH}_FLAGS_LIST}>")
+if (SCREAM_RRTMGP_ENABLE_YAKL)
+  if (YAKL_ARCH)
+    target_compile_options(scream_rrtmgp PUBLIC
+      "$<$<COMPILE_LANGUAGE:CXX>:${YAKL_${YAKL_ARCH}_FLAGS_LIST}>")
+  endif()
 endif()