Port eamxx rrtmgp to kokkos

components/eamxx/CMakeLists.txt

@@ -203,13 +203,24 @@ if (NOT SCREAM_SMALL_KERNELS)
   set(EKAT_DISABLE_WORKSPACE_SHARING TRUE CACHE STRING "")
 endif()
 
-# Add RRTMGP debug checks. Note, we might consider also adding RRTMGP_EXPENSIVE_CHECKS
+# Add RRTMGP settings. Note, we might consider also adding RRTMGP_EXPENSIVE_CHECKS
 # to turn on the RRTMGP internal checks here as well, via
 #     option (RRTMGP_EXPENSIVE_CHECKS "Turn on internal RRTMGP error checking" ${SCREAM_DEBUG})
 # and then adding to scream_config.h:
 #     #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)
+if (SCREAM_RRTMGP_ENABLE_YAKL)
+  add_definitions("-DRRTMGP_ENABLE_YAKL")
+endif()
+
+if (SCREAM_RRTMGP_ENABLE_KOKKOS)
+  add_definitions("-DRRTMGP_ENABLE_KOKKOS")
+endif()
+
+
 set(SCREAM_DOUBLE_PRECISION TRUE CACHE BOOL "Set to double precision (default True)")
 
 # For now, only used in share/grid/remap/refining_remapper_rma.*pp
@@ -618,4 +629,3 @@ print_var(SCREAM_TEST_THREAD_INC)
 print_var(SCREAM_TEST_MAX_RANKS)
 
 message ("**************************************************")
-

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

@@ -116,6 +116,9 @@ set(EXTERNAL_SRC
   ${EAM_RRTMGP_DIR}/external/cpp/rte/kernels/mo_optical_props_kernels.cpp
   ${EAM_RRTMGP_DIR}/external/cpp/rte/kernels/mo_rte_solver_kernels.cpp
   ${EAM_RRTMGP_DIR}/external/cpp/extensions/fluxes_byband/mo_fluxes_byband_kernels.cpp
+  ${EAM_RRTMGP_DIR}/external/cpp/examples/all-sky/mo_garand_atmos_io.cpp
+  ${EAM_RRTMGP_DIR}/external/cpp/examples/all-sky/mo_load_cloud_coefficients.cpp
+  ${EAM_RRTMGP_DIR}/external/cpp/examples/mo_load_coefficients.cpp
 )
 add_library(rrtmgp ${EXTERNAL_SRC})
 target_compile_definitions(rrtmgp PUBLIC EAMXX_HAS_RRTMGP)
@@ -129,10 +132,14 @@ target_link_libraries(rrtmgp yakl Kokkos::kokkos)
 target_include_directories(rrtmgp PUBLIC
     ${SCREAM_BASE_DIR}/../../externals/YAKL
     ${EAM_RRTMGP_DIR}/external/cpp
+    ${EAM_RRTMGP_DIR}/external/cpp/extensions/cloud_optics
+    ${EAM_RRTMGP_DIR}/external/cpp/examples
+    ${EAM_RRTMGP_DIR}/external/cpp/examples/all-sky
     ${EAM_RRTMGP_DIR}/external/cpp/rte
     ${EAM_RRTMGP_DIR}/external/cpp/rte/kernels
     ${EAM_RRTMGP_DIR}/external/cpp/rrtmgp
     ${EAM_RRTMGP_DIR}/external/cpp/rrtmgp/kernels
+    ${NetCDF_C_PATH}/include
 )
 
 # Build RRTMGP interface; note that we separate the SCREAM-specific RRTMGP interface
@@ -151,9 +158,6 @@ target_include_directories(rrtmgp PUBLIC
 
 set(SCREAM_RRTMGP_SOURCES_YAKL
   scream_rrtmgp_interface.cpp
-  mo_garand_atmos_io.cpp
-  mo_load_coefficients.cpp
-  mo_load_cloud_coefficients.cpp
 )
 
 add_library(scream_rrtmgp_yakl ${SCREAM_RRTMGP_SOURCES_YAKL})

components/eamxx/src/physics/rrtmgp/eamxx_rrtmgp_process_interface.hpp

@@ -77,7 +77,7 @@ class RRTMGPRadiation : public AtmosphereProcess {
   // This is only used if a positive value is supplied
   Real m_fixed_solar_zenith_angle;
 
-  // Need to hard-code some dimension sizes for now. 
+  // Need to hard-code some dimension sizes for now.
   // TODO: find a better way of configuring this
   const int m_nswbands = 14;
   const int m_nlwbands = 16;
@@ -88,7 +88,12 @@ class RRTMGPRadiation : public AtmosphereProcess {
   int m_ngas;
   std::vector<ci_string>   m_gas_names;
   view_1d_real             m_gas_mol_weights;
+#ifdef RRTMGP_ENABLE_YAKL
   GasConcs                 m_gas_concs;
+#endif
+#ifdef RRTMGP_ENABLE_KOKKOS
+  GasConcsK                m_gas_concs_k;
+#endif
 
   // Prescribed greenhouse gas surface concentrations in moles / moles air
   Real m_co2vmr;
@@ -119,18 +124,33 @@ class RRTMGPRadiation : public AtmosphereProcess {
     static constexpr int num_3d_nlay_nlwgpts = 1;
 
     // 1d size (ncol)
+    uview_1d<Real> cosine_zenith;
+#ifdef RRTMGP_ENABLE_YAKL
     real1d mu0;
     real1d sfc_alb_dir_vis;
     real1d sfc_alb_dir_nir;
     real1d sfc_alb_dif_vis;
     real1d sfc_alb_dif_nir;
-    uview_1d<Real> cosine_zenith;
     real1d sfc_flux_dir_vis;
     real1d sfc_flux_dir_nir;
     real1d sfc_flux_dif_vis;
     real1d sfc_flux_dif_nir;
+#endif
+#ifdef RRTMGP_ENABLE_KOKKOS
+    real1dk mu0_k;
+    real1dk sfc_alb_dir_vis_k;
+    real1dk sfc_alb_dir_nir_k;
+    real1dk sfc_alb_dif_vis_k;
+    real1dk sfc_alb_dif_nir_k;
+    real1dk sfc_flux_dir_vis_k;
+    real1dk sfc_flux_dir_nir_k;
+    real1dk sfc_flux_dif_vis_k;
+    real1dk sfc_flux_dif_nir_k;
+#endif
 
     // 2d size (ncol, nlay)
+    uview_2d<Real> d_dz;
+#ifdef RRTMGP_ENABLE_YAKL
     real2d p_lay;
     real2d t_lay;
     real2d z_del;
@@ -146,9 +166,28 @@ class RRTMGPRadiation : public AtmosphereProcess {
     real2d iwp;
     real2d sw_heating;
     real2d lw_heating;
-    uview_2d<Real> d_dz;
+#endif
+#ifdef RRTMGP_ENABLE_KOKKOS
+    real2dk p_lay_k;
+    real2dk t_lay_k;
+    real2dk z_del_k;
+    real2dk p_del_k;
+    real2dk qc_k;
+    real2dk nc_k;
+    real2dk qi_k;
+    real2dk cldfrac_tot_k;
+    real2dk eff_radius_qc_k;
+    real2dk eff_radius_qi_k;
+    real2dk tmp2d_k;
+    real2dk lwp_k;
+    real2dk iwp_k;
+    real2dk sw_heating_k;
+    real2dk lw_heating_k;
+#endif
 
     // 2d size (ncol, nlay+1)
+    uview_2d<Real> d_tint;
+#ifdef RRTMGP_ENABLE_YAKL
     real2d p_lev;
     real2d t_lev;
     real2d sw_flux_up;
@@ -171,35 +210,99 @@ class RRTMGPRadiation : public AtmosphereProcess {
     real2d lw_clrsky_flux_dn;
     real2d lw_clnsky_flux_up;
     real2d lw_clnsky_flux_dn;
-    uview_2d<Real> d_tint;
+#endif
+#ifdef RRTMGP_ENABLE_KOKKOS
+    real2dk p_lev_k;
+    real2dk t_lev_k;
+    real2dk sw_flux_up_k;
+    real2dk sw_flux_dn_k;
+    real2dk sw_flux_dn_dir_k;
+    real2dk lw_flux_up_k;
+    real2dk lw_flux_dn_k;
+    real2dk sw_clnclrsky_flux_up_k;
+    real2dk sw_clnclrsky_flux_dn_k;
+    real2dk sw_clnclrsky_flux_dn_dir_k;
+    real2dk sw_clrsky_flux_up_k;
+    real2dk sw_clrsky_flux_dn_k;
+    real2dk sw_clrsky_flux_dn_dir_k;
+    real2dk sw_clnsky_flux_up_k;
+    real2dk sw_clnsky_flux_dn_k;
+    real2dk sw_clnsky_flux_dn_dir_k;
+    real2dk lw_clnclrsky_flux_up_k;
+    real2dk lw_clnclrsky_flux_dn_k;
+    real2dk lw_clrsky_flux_up_k;
+    real2dk lw_clrsky_flux_dn_k;
+    real2dk lw_clnsky_flux_up_k;
+    real2dk lw_clnsky_flux_dn_k;
+#endif
 
     // 3d size (ncol, nlay+1, nswbands)
+#ifdef RRTMGP_ENABLE_YAKL
     real3d sw_bnd_flux_up;
     real3d sw_bnd_flux_dn;
     real3d sw_bnd_flux_dir;
     real3d sw_bnd_flux_dif;
+#endif
+#ifdef RRTMGP_ENABLE_KOKKOS
+    real3dk sw_bnd_flux_up_k;
+    real3dk sw_bnd_flux_dn_k;
+    real3dk sw_bnd_flux_dir_k;
+    real3dk sw_bnd_flux_dif_k;
+#endif
 
     // 3d size (ncol, nlay+1, nlwbands)
+#ifdef RRTMGP_ENABLE_YAKL
     real3d lw_bnd_flux_up;
     real3d lw_bnd_flux_dn;
+#endif
+#ifdef RRTMGP_ENABLE_KOKKOS
+    real3dk lw_bnd_flux_up_k;
+    real3dk lw_bnd_flux_dn_k;
+#endif
 
     // 2d size (ncol, nswbands)
+#ifdef RRTMGP_ENABLE_YAKL
     real2d sfc_alb_dir;
     real2d sfc_alb_dif;
+#endif
+#ifdef RRTMGP_ENABLE_KOKKOS
+    real2dk sfc_alb_dir_k;
+    real2dk sfc_alb_dif_k;
+#endif
 
     // 3d size (ncol, nlay, n[sw,lw]bands)
+#ifdef RRTMGP_ENABLE_YAKL
     real3d aero_tau_sw;
     real3d aero_ssa_sw;
     real3d aero_g_sw;
     real3d aero_tau_lw;
+#endif
+#ifdef RRTMGP_ENABLE_KOKKOS
+    real3dk aero_tau_sw_k;
+    real3dk aero_ssa_sw_k;
+    real3dk aero_g_sw_k;
+    real3dk aero_tau_lw_k;
+#endif
 
     // 3d size (ncol, nlay, n[sw,lw]bnds)
+#ifdef RRTMGP_ENABLE_YAKL
     real3d cld_tau_sw_bnd;
     real3d cld_tau_lw_bnd;
+#endif
+#ifdef RRTMGP_ENABLE_KOKKOS
+    real3dk cld_tau_sw_bnd_k;
+    real3dk cld_tau_lw_bnd_k;
+#endif
 
     // 3d size (ncol, nlay, n[sw,lw]gpts)
+#ifdef RRTMGP_ENABLE_YAKL
     real3d cld_tau_sw_gpt;
     real3d cld_tau_lw_gpt;
+#endif
+#ifdef RRTMGP_ENABLE_KOKKOS
+    real3dk cld_tau_sw_gpt_k;
+    real3dk cld_tau_lw_gpt_k;
+#endif
 
   };