Merge pull request #2965 from E3SM-Project/mahf708/small-kernels

options to build sk p3 and shoc separately

cime_config/tests.py

@@ -668,6 +668,9 @@
             "SMS_D_Ln9.ne4_ne4.F2010-SCREAMv1-noAero.scream-output-preset-3",
             "ERP_Ln22.ne4pg2_ne4pg2.F2010-SCREAMv1.scream-output-preset-4",
             "ERS_D_Ln22.ne4pg2_ne4pg2.F2010-SCREAMv1.scream-rad_frequency_2--scream-output-preset-5",
+            "ERS_Ln22.ne4pg2_ne4pg2.F2010-SCREAMv1.scream-small_kernels--scream-output-preset-5",
+            "ERS_Ln22.ne4pg2_ne4pg2.F2010-SCREAMv1.scream-small_kernels_p3--scream-output-preset-5",
+            "ERS_Ln22.ne4pg2_ne4pg2.F2010-SCREAMv1.scream-small_kernels_shoc--scream-output-preset-5",
             )
     },
 

components/eamxx/CMakeLists.txt

@@ -198,8 +198,10 @@ set(NetCDF_C_PATH ${DEFAULT_NetCDF_C_PATH} CACHE FILEPATH "Path to netcdf C inst
 set(SCREAM_MACHINE ${DEFAULT_SCREAM_MACHINE} CACHE STRING "The CIME/SCREAM name for the current machine")
 option(SCREAM_MPI_ON_DEVICE "Whether to use device pointers for MPI calls" ON)
 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")
-if (NOT SCREAM_SMALL_KERNELS)
+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()
 

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

@@ -0,0 +1,2 @@
+./xmlchange --append SCREAM_CMAKE_OPTIONS='SCREAM_P3_SMALL_KERNELS On'
+$CIMEROOT/../components/eamxx/scripts/atmchange --all internal_diagnostics_level=1 atmosphere_processes::internal_diagnostics_level=0 -b

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

@@ -0,0 +1,2 @@
+./xmlchange --append SCREAM_CMAKE_OPTIONS='SCREAM_SHOC_SMALL_KERNELS On'
+$CIMEROOT/../components/eamxx/scripts/atmchange --all internal_diagnostics_level=1 atmosphere_processes::internal_diagnostics_level=0 -b

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

@@ -74,14 +74,14 @@ set(P3_SK_SRCS
     )
 
 set(P3_LIBS "p3")
-if (SCREAM_SMALL_KERNELS)
+if (SCREAM_P3_SMALL_KERNELS)
   add_library(p3 ${P3_SRCS} ${P3_SK_SRCS})
 else()
   add_library(p3 ${P3_SRCS})
   if (NOT SCREAM_LIBS_ONLY AND NOT SCREAM_ONLY_GENERATE_BASELINES)
     add_library(p3_sk ${P3_SRCS} ${P3_SK_SRCS})
-    # Always build shoc_sk with SCREAM_SMALL_KERNELS on
-    target_compile_definitions(p3_sk PUBLIC "SCREAM_SMALL_KERNELS")
+    # Always build p3_sk with SCREAM_P3_SMALL_KERNELS on
+    target_compile_definitions(p3_sk PUBLIC "SCREAM_P3_SMALL_KERNELS")
     list(APPEND P3_LIBS "p3_sk")
   endif()
 endif()

components/eamxx/src/physics/p3/impl/p3_main_impl.hpp

@@ -349,7 +349,7 @@ ::p3_main(
   Int nk,
   const physics::P3_Constants<S> & p3constants)
 {
-#ifndef SCREAM_SMALL_KERNELS
+#ifndef SCREAM_P3_SMALL_KERNELS
   return p3_main_internal(runtime_options,
                          prognostic_state,
                          diagnostic_inputs,

components/eamxx/src/physics/p3/p3_functions.hpp

@@ -420,7 +420,7 @@ struct Functions
     const uview_1d<Spack>& nc_tend,
     Scalar& precip_liq_surf);
 
-#ifdef SCREAM_SMALL_KERNELS
+#ifdef SCREAM_P3_SMALL_KERNELS
   static void cloud_sedimentation_disp(
     const uview_2d<Spack>& qc_incld,
     const uview_2d<const Spack>& rho,
@@ -468,7 +468,7 @@ struct Functions
     Scalar& precip_liq_surf,
     const physics::P3_Constants<ScalarT> & p3constants);
 
-#ifdef SCREAM_SMALL_KERNELS
+#ifdef SCREAM_P3_SMALL_KERNELS
   static void rain_sedimentation_disp(
     const uview_2d<const Spack>& rho,
     const uview_2d<const Spack>& inv_rho,
@@ -518,7 +518,7 @@ struct Functions
     Scalar& precip_ice_surf,
     const physics::P3_Constants<ScalarT> & p3constants);
 
-#ifdef SCREAM_SMALL_KERNELS
+#ifdef SCREAM_P3_SMALL_KERNELS
   static void ice_sedimentation_disp(
     const uview_2d<const Spack>& rho,
     const uview_2d<const Spack>& inv_rho,
@@ -562,7 +562,7 @@ struct Functions
     const uview_1d<Spack>& bm,
     const uview_1d<Spack>& th_atm);
 
-#ifdef SCREAM_SMALL_KERNELS
+#ifdef SCREAM_P3_SMALL_KERNELS
   static void homogeneous_freezing_disp(
     const uview_2d<const Spack>& T_atm,
     const uview_2d<const Spack>& inv_exner,
@@ -868,7 +868,7 @@ struct Functions
                            const Int& timestepcount, const bool& force_abort, const Int& source_ind, const MemberType& team,
                            const uview_1d<const Scalar>& col_loc);
 
-#ifdef SCREAM_SMALL_KERNELS
+#ifdef SCREAM_P3_SMALL_KERNELS
   static void check_values_disp(const uview_2d<const Spack>& qv, const uview_2d<const Spack>& temp, const Int& ktop, const Int& kbot,
                            const Int& timestepcount, const bool& force_abort, const Int& source_ind,
                            const uview_2d<const Scalar>& col_loc, const Int& nj, const Int& nk);
@@ -914,7 +914,7 @@ struct Functions
     Scalar& precip_ice_surf,
     view_1d_ptr_array<Spack, 36>& zero_init);
 
-#ifdef SCREAM_SMALL_KERNELS
+#ifdef SCREAM_P3_SMALL_KERNELS
   static void p3_main_init_disp(
     const Int& nj,const Int& nk_pack,
     const uview_2d<const Spack>& cld_frac_i, const uview_2d<const Spack>& cld_frac_l,
@@ -988,7 +988,7 @@ struct Functions
     bool& is_hydromet_present,
     const physics::P3_Constants<ScalarT> & p3constants);
 
-#ifdef SCREAM_SMALL_KERNELS
+#ifdef SCREAM_P3_SMALL_KERNELS
   static void p3_main_part1_disp(
     const Int& nj,
     const Int& nk,
@@ -1121,7 +1121,7 @@ struct Functions
     const Int& nk,
     const physics::P3_Constants<ScalarT> & p3constants);
 
-#ifdef SCREAM_SMALL_KERNELS
+#ifdef SCREAM_P3_SMALL_KERNELS
   static void p3_main_part2_disp(
     const Int& nj,
     const Int& nk,
@@ -1246,7 +1246,7 @@ struct Functions
     const uview_1d<Spack>& diag_eff_radius_qr,
     const physics::P3_Constants<ScalarT> & p3constants);
 
-#ifdef SCREAM_SMALL_KERNELS
+#ifdef SCREAM_P3_SMALL_KERNELS
   static void p3_main_part3_disp(
     const Int& nj,
     const Int& nk_pack,
@@ -1319,7 +1319,7 @@ struct Functions
     Int nk, // number of vertical cells per column
     const physics::P3_Constants<ScalarT> & p3constants);
 
-#ifdef SCREAM_SMALL_KERNELS
+#ifdef SCREAM_P3_SMALL_KERNELS
   static Int p3_main_internal_disp(
     const P3Runtime& runtime_options,
     const P3PrognosticState& prognostic_state,

components/eamxx/src/physics/p3/tests/CMakeLists.txt

@@ -60,7 +60,7 @@ if (NOT SCREAM_ONLY_GENERATE_BASELINES)
       LABELS "p3;physics;fail"
       ${FORCE_RUN_DIFF_FAILS})
 
-  if (NOT SCREAM_SMALL_KERNELS)
+  if (NOT SCREAM_P3_SMALL_KERNELS)
      CreateUnitTest(p3_sk_tests "${P3_TESTS_SRCS}"
         LIBS p3_sk
         THREADS 1 ${SCREAM_TEST_MAX_THREADS} ${SCREAM_TEST_THREAD_INC}

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

@@ -104,14 +104,14 @@ if (HIP_BUILD)
 endif()
 
 set(SHOC_LIBS "shoc")
-if (SCREAM_SMALL_KERNELS)
+if (SCREAM_SHOC_SMALL_KERNELS)
   add_library(shoc ${SHOC_SRCS} ${SHOC_SK_SRCS})
 else()
   add_library(shoc ${SHOC_SRCS})
   if (NOT SCREAM_LIBS_ONLY AND NOT SCREAM_ONLY_GENERATE_BASELINES)
     add_library(shoc_sk ${SHOC_SRCS} ${SHOC_SK_SRCS})
-    # Always build shoc_sk with SCREAM_SMALL_KERNELS on
-    target_compile_definitions(shoc_sk PUBLIC "SCREAM_SMALL_KERNELS")
+    # Always build shoc_sk with SCREAM_SHOC_SMALL_KERNELS on
+    target_compile_definitions(shoc_sk PUBLIC "SCREAM_SHOC_SMALL_KERNELS")
     list(APPEND SHOC_LIBS "shoc_sk")
   endif()
 endif()

components/eamxx/src/physics/shoc/eamxx_shoc_process_interface.cpp

@@ -156,7 +156,7 @@ void SHOCMacrophysics::init_buffers(const ATMBufferManager &buffer_manager)
   using scalar_view_t = decltype(m_buffer.wpthlp_sfc);
   scalar_view_t* _1d_scalar_view_ptrs[Buffer::num_1d_scalar_ncol] =
     {&m_buffer.wpthlp_sfc, &m_buffer.wprtp_sfc, &m_buffer.upwp_sfc, &m_buffer.vpwp_sfc
-#ifdef SCREAM_SMALL_KERNELS
+#ifdef SCREAM_SHOC_SMALL_KERNELS
      , &m_buffer.se_b, &m_buffer.ke_b, &m_buffer.wv_b, &m_buffer.wl_b
      , &m_buffer.se_a, &m_buffer.ke_a, &m_buffer.wv_a, &m_buffer.wl_a
      , &m_buffer.kbfs, &m_buffer.ustar2, &m_buffer.wstar
@@ -182,7 +182,7 @@ void SHOCMacrophysics::init_buffers(const ATMBufferManager &buffer_manager)
     &m_buffer.z_mid, &m_buffer.rrho, &m_buffer.thv, &m_buffer.dz, &m_buffer.zt_grid, &m_buffer.wm_zt,
     &m_buffer.inv_exner, &m_buffer.thlm, &m_buffer.qw, &m_buffer.dse, &m_buffer.tke_copy, &m_buffer.qc_copy,
     &m_buffer.shoc_ql2, &m_buffer.shoc_mix, &m_buffer.isotropy, &m_buffer.w_sec, &m_buffer.wqls_sec, &m_buffer.brunt
-#ifdef SCREAM_SMALL_KERNELS
+#ifdef SCREAM_SHOC_SMALL_KERNELS
     , &m_buffer.rho_zt, &m_buffer.shoc_qv, &m_buffer.tabs, &m_buffer.dz_zt
 #endif
   };
@@ -191,7 +191,7 @@ void SHOCMacrophysics::init_buffers(const ATMBufferManager &buffer_manager)
     &m_buffer.z_int, &m_buffer.rrho_i, &m_buffer.zi_grid, &m_buffer.thl_sec, &m_buffer.qw_sec,
     &m_buffer.qwthl_sec, &m_buffer.wthl_sec, &m_buffer.wqw_sec, &m_buffer.wtke_sec, &m_buffer.uw_sec,
     &m_buffer.vw_sec, &m_buffer.w3
-#ifdef SCREAM_SMALL_KERNELS
+#ifdef SCREAM_SHOC_SMALL_KERNELS
     , &m_buffer.dz_zi
 #endif
   };
@@ -353,7 +353,7 @@ void SHOCMacrophysics::initialize_impl (const RunType run_type)
   history_output.wqls_sec  = m_buffer.wqls_sec;
   history_output.brunt     = m_buffer.brunt;
 
-#ifdef SCREAM_SMALL_KERNELS
+#ifdef SCREAM_SHOC_SMALL_KERNELS
   temporaries.se_b = m_buffer.se_b;
   temporaries.ke_b = m_buffer.ke_b;
   temporaries.wv_b = m_buffer.wv_b;
@@ -485,7 +485,7 @@ void SHOCMacrophysics::run_impl (const double dt)
   // Run shoc main
   SHF::shoc_main(m_num_cols, m_num_levs, m_num_levs+1, m_npbl, m_nadv, m_num_tracers, dt,
                  workspace_mgr,runtime_options,input,input_output,output,history_output
-#ifdef SCREAM_SMALL_KERNELS
+#ifdef SCREAM_SHOC_SMALL_KERNELS
                  , temporaries
 #endif
                  );

components/eamxx/src/physics/shoc/eamxx_shoc_process_interface.hpp

@@ -382,13 +382,13 @@ class SHOCMacrophysics : public scream::AtmosphereProcess
 
   // Structure for storing local variables initialized using the ATMBufferManager
   struct Buffer {
-#ifndef SCREAM_SMALL_KERNELS
+#ifndef SCREAM_SHOC_SMALL_KERNELS
     static constexpr int num_1d_scalar_ncol = 4;
 #else
     static constexpr int num_1d_scalar_ncol = 15;
 #endif
     static constexpr int num_1d_scalar_nlev = 1;
-#ifndef SCREAM_SMALL_KERNELS
+#ifndef SCREAM_SHOC_SMALL_KERNELS
     static constexpr int num_2d_vector_mid  = 18;
     static constexpr int num_2d_vector_int  = 12;
 #else
@@ -401,7 +401,7 @@ class SHOCMacrophysics : public scream::AtmosphereProcess
     uview_1d<Real> wprtp_sfc;
     uview_1d<Real> upwp_sfc;
     uview_1d<Real> vpwp_sfc;
-#ifdef SCREAM_SMALL_KERNELS
+#ifdef SCREAM_SHOC_SMALL_KERNELS
     uview_1d<Real> se_b;
     uview_1d<Real> ke_b;
     uview_1d<Real> wv_b;
@@ -448,7 +448,7 @@ class SHOCMacrophysics : public scream::AtmosphereProcess
     uview_2d<Spack> w3;
     uview_2d<Spack> wqls_sec;
     uview_2d<Spack> brunt;
-#ifdef SCREAM_SMALL_KERNELS
+#ifdef SCREAM_SHOC_SMALL_KERNELS
     uview_2d<Spack> rho_zt;
     uview_2d<Spack> shoc_qv;
     uview_2d<Spack> tabs;
@@ -506,7 +506,7 @@ class SHOCMacrophysics : public scream::AtmosphereProcess
   SHF::SHOCOutput output;
   SHF::SHOCHistoryOutput history_output;
   SHF::SHOCRuntime runtime_options;
-#ifdef SCREAM_SMALL_KERNELS
+#ifdef SCREAM_SHOC_SMALL_KERNELS
   SHF::SHOCTemporaries temporaries;
 #endif
 

components/eamxx/src/physics/shoc/impl/shoc_main_impl.hpp

@@ -61,7 +61,7 @@ Int Functions<S,D>::shoc_init(
   return host_view(0);
 }
 
-#ifndef SCREAM_SMALL_KERNELS
+#ifndef SCREAM_SHOC_SMALL_KERNELS
 template<typename S, typename D>
 KOKKOS_FUNCTION
 void Functions<S,D>::shoc_main_internal(
@@ -586,7 +586,7 @@ Int Functions<S,D>::shoc_main(
   const SHOCInputOutput&   shoc_input_output,   // Input/Output
   const SHOCOutput&        shoc_output,         // Output
   const SHOCHistoryOutput& shoc_history_output  // Output (diagnostic)
-#ifdef SCREAM_SMALL_KERNELS
+#ifdef SCREAM_SHOC_SMALL_KERNELS
   , const SHOCTemporaries& shoc_temporaries     // Temporaries for small kernels
 #endif
                               )
@@ -608,7 +608,7 @@ Int Functions<S,D>::shoc_main(
   const Scalar Ckh           = shoc_runtime.Ckh;
   const Scalar Ckm           = shoc_runtime.Ckm;
 
-#ifndef SCREAM_SMALL_KERNELS
+#ifndef SCREAM_SHOC_SMALL_KERNELS
   using ExeSpace = typename KT::ExeSpace;
 
   // SHOC main loop