Merge branch 'jgfouca/port_gwd_compute_tendencies_from_stress_divergence' into master (PR E3SM-Project#7525)

Change list:
1) Main change: Fill out CXX impl of gwd_compute_tendencies_from_stress_divergence
2) gw_common_init: Allow it to be called multiple times by deallocating cref and alpha if
   they had been previously allocated.
3) Begin to fill in constants needed by GW
4) Create struct for containing common init data in C++
5) Add C++ version of gw_common_init and gw_common_finalize
6) Don't just assume all C++ data is packed. Just have it all unpacked for now
7) Fix randomized test data for gwd_compute_tendencies_from_stress_divergence so
   that some actual interesting computations happen.
8) For now, have a f90 and cxx bridge available for this function and use the f90
   one for baseline generation. Once we are sure we have things working, the f90
   calls can be removed.

[BFB] (for everything outside of eamxx/gw)

Author: jgfouca

components/eam/src/physics/cam/gw/gw_common.F90

@@ -141,6 +141,7 @@ subroutine gw_common_init(pver_in, pgwv_in, dc_in, cref_in, &
   pver = pver_in
   pgwv = pgwv_in
   dc = dc_in
+  if (allocated(cref)) deallocate(cref)
   allocate(cref(-pgwv:pgwv), stat=ierr, errmsg=errstring)
   if (ierr /= 0) return
   cref = cref_in
@@ -154,6 +155,7 @@ subroutine gw_common_init(pver_in, pgwv_in, dc_in, cref_in, &
   kwv = kwv_in
   gravit = gravit_in
   rair = rair_in
+  if (allocated(alpha)) deallocate(alpha)
   allocate(alpha(0:pver), stat=ierr, errmsg=errstring)
   if (ierr /= 0) return
   alpha = alpha_in

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

@@ -31,6 +31,7 @@ if (NOT EAMXX_ENABLE_GPU OR Kokkos_ENABLE_CUDA_RELOCATABLE_DEVICE_CODE OR Kokkos
     eti/gw_gw_ediff.cpp
     eti/gw_gw_diff_tend.cpp
     eti/gw_gw_oro_src.cpp
+    eti/gw_gw_common_init.cpp
   ) # GW ETI SRCS
 endif()
 

components/eamxx/src/physics/gw/eti/gw_gw_common_init.cpp

@@ -0,0 +1,14 @@
+#include "impl/gw_gw_common_init_impl.hpp"
+
+namespace scream {
+namespace gw {
+
+/*
+ * Explicit instantiation for doing gw_common_init on Reals using the
+ * default device.
+ */
+
+template struct Functions<Real,DefaultDevice>;
+
+} // namespace gw
+} // namespace scream

components/eamxx/src/physics/gw/gw_functions.hpp

@@ -21,12 +21,6 @@ namespace gw {
 template <typename ScalarT, typename DeviceT>
 struct Functions
 {
-  //
-  // ---------- GW constants ---------
-  //
-  struct GWC {
-  };
-
   //
   // ------- Types --------
   //
@@ -54,47 +48,159 @@ struct Functions
   using view_1d = typename KT::template view_1d<S>;
   template <typename S>
   using view_2d = typename KT::template view_2d<S>;
+  template <typename S>
+  using view_3d = typename KT::template view_3d<S>;
 
   template <typename S>
   using uview_1d = typename ekat::template Unmanaged<view_1d<S> >;
   template <typename S>
   using uview_2d = typename ekat::template Unmanaged<view_2d<S> >;
+  template <typename S>
+  using uview_3d = typename ekat::template Unmanaged<view_3d<S> >;
 
   using MemberType = typename KT::MemberType;
 
-  using WorkspaceManager = typename ekat::WorkspaceManager<Spack, Device>;
+  using WorkspaceManager = typename ekat::WorkspaceManager<Scalar, Device>;
   using Workspace        = typename WorkspaceManager::Workspace;
 
+  //
+  // ---------- GW constants ---------
+  //
+  struct GWC {
+    // Index the cardinal directions.
+    static inline constexpr int west = 0;
+    static inline constexpr int east = 1;
+    static inline constexpr int south = 2;
+    static inline constexpr int north = 3;
+
+    // rair/gravit
+    static inline constexpr Real rog = C::Rair / C::gravit;
+
+    // Background diffusivity.
+    static inline constexpr Real dback = 0.05;
+
+    // Minimum non-zero stress.
+    static inline constexpr Real taumin = 1.e-10;
+
+    // Maximum allowed change in u-c (before efficiency applied).
+    static inline constexpr Real umcfac = 0.5;
+
+    // Minimum value of (u-c)**2.
+    static inline constexpr Real ubmc2mn = 0.01;
+  };
+
+  //
+  // --------- Structs for organizing data ---------
+  //
+
+  struct GwCommonInit {
+    GwCommonInit() : initialized(false) {}
+
+    // Tell us if initialize has been called
+    bool initialized;
+
+    // This flag preserves answers for vanilla CAM by making a few changes (e.g.
+    // order of operations) when only orographic waves are on.
+    bool orographic_only; // = .false.
+
+    // Number of levels in the atmosphere.
+    int pver;
+
+    // Maximum number of waves allowed (i.e. wavenumbers are -pgwv:pgwv).
+    int pgwv;
+
+    // Bin width for spectrum.
+    Real dc; // = huge(1._r8)
+
+    // Reference speeds for the spectrum.
+    view_1d<Real> cref;
+
+    // Whether or not molecular diffusion is being done, and bottom level where
+    // it is done.
+    bool do_molec_diff; // = .false.
+    int nbot_molec; // = huge(1)
+
+    // Whether or not to enforce an upper boundary condition of tau = 0.
+    bool tau_0_ubc; // = .false.
+
+    // Critical Froude number.
+    Real fcrit2; // = huge(1._r8)
+
+    // Effective horizontal wave number.
+    Real kwv; // = huge(1._r8)
+
+    // Interface levels for gravity wave sources.
+    int ktop; // = huge(1)
+    int kbotbg; // = huge(1)
+
+    // Effective wavenumber.
+    Real effkwv; // = huge(1._r8)
+
+    // Newtonian cooling coefficients.
+    view_1d<Real> alpha;
+
+    // Maximum wind tendency from stress divergence (before efficiency applied).
+    Real tndmax; // = huge(1._r8)
+  };
+
+  //
+  // --------- Init/Finalize Functions ---------
+  //
+  static void gw_common_init(
+    // Inputs
+    const Int& pver_in,
+    const Int& pgwv_in,
+    const Real& dc_in,
+    const uview_1d<const Real>& cref_in,
+    const bool& orographic_only_in,
+    const bool& do_molec_diff_in,
+    const bool& tau_0_ubc_in,
+    const Int& nbot_molec_in,
+    const Int& ktop_in,
+    const Int& kbotbg_in,
+    const Real& fcrit2_in,
+    const Real& kwv_in,
+    const uview_1d<const Real>& alpha_in);
+
+  static void gw_common_finalize()
+  {
+    s_common_init.cref  = decltype(s_common_init.cref)();
+    s_common_init.alpha = decltype(s_common_init.alpha)();
+  }
+
   //
   // --------- Functions ---------
   //
 
   KOKKOS_FUNCTION
   static void gwd_compute_tendencies_from_stress_divergence(
     // Inputs
-    const Int& ncol,
+    const MemberType& team,
+    const Workspace& workspace,
+    const GwCommonInit& init,
     const Int& pver,
     const Int& pgwv,
-    const Int& ngwv,
     const bool& do_taper,
-    const Spack& dt,
-    const Spack& effgw,
-    const uview_1d<const Int>& tend_level,
-    const uview_1d<const Spack>& lat,
-    const uview_1d<const Spack>& dpm,
-    const uview_1d<const Spack>& rdpm,
-    const uview_1d<const Spack>& c,
-    const uview_1d<const Spack>& ubm,
-    const uview_1d<const Spack>& t,
-    const uview_1d<const Spack>& nm,
-    const uview_1d<const Spack>& xv,
-    const uview_1d<const Spack>& yv,
+    const Real& dt,
+    const Real& effgw,
+    const Int& tend_level,
+    const Int& max_level,
+    const Real& lat,
+    const uview_1d<const Real>& dpm,
+    const uview_1d<const Real>& rdpm,
+    const uview_1d<const Real>& c,
+    const uview_1d<const Real>& ubm,
+    const uview_1d<const Real>& t,
+    const uview_1d<const Real>& nm,
+    const Real& xv,
+    const Real& yv,
     // Inputs/Outputs
-    const uview_1d<Spack>& tau,
+    const uview_2d<Real>& tau,
+    const uview_2d<Real>& work,
     // Outputs
-    const uview_1d<Spack>& gwut,
-    const uview_1d<Spack>& utgw,
-    const uview_1d<Spack>& vtgw);
+    const uview_2d<Real>& gwut,
+    const uview_1d<Real>& utgw,
+    const uview_1d<Real>& vtgw);
 
   KOKKOS_FUNCTION
   static void gw_prof(
@@ -137,7 +243,6 @@ struct Functions
     const Int& pver,
     const Int& pgwv,
     const Int& ncol,
-    const Int& ngwv,
     const uview_1d<const Int>& src_level,
     const uview_1d<const Spack>& ubi,
     const uview_1d<const Spack>& c,
@@ -156,7 +261,6 @@ struct Functions
     const Int& pver,
     const Int& pgwv,
     const Int& ncol,
-    const Int& ngwv,
     const uview_1d<const Int>& tend_level,
     const uview_1d<const Spack>& tau,
     const uview_1d<const Spack>& ubi,
@@ -172,7 +276,6 @@ struct Functions
     const Int& pver,
     const Int& pgwv,
     const Int& ncol,
-    const Int& ngwv,
     const Spack& dt,
     const uview_1d<const Int>& tend_level,
     const uview_1d<const Spack>& pmid,
@@ -198,7 +301,6 @@ struct Functions
     const Int& pver,
     const Int& pgwv,
     const Int& ncol,
-    const Int& ngwv,
     const uview_1d<const Int>& src_level,
     const uview_1d<const Int>& tend_level,
     const bool& do_taper,
@@ -256,7 +358,6 @@ struct Functions
     const Int& pver,
     const Int& pgwv,
     const Int& ncol,
-    const Int& ngwv,
     const Int& kbot,
     const uview_1d<const Spack>& frontgf,
     // Outputs
@@ -268,7 +369,6 @@ struct Functions
     const Int& pver,
     const Int& pgwv,
     const Int& ncol,
-    const Int& ngwv,
     const Int& kbot,
     const uview_1d<const Spack>& u,
     const uview_1d<const Spack>& v,
@@ -334,7 +434,6 @@ struct Functions
     const Int& pver,
     const Int& pgwv,
     const Int& ncol,
-    const Int& ngwv,
     const uview_1d<const Spack>& lat,
     const Spack& hdepth_min,
     const uview_1d<const Spack>& hdepth,
@@ -355,7 +454,6 @@ struct Functions
     const Int& pver,
     const Int& pgwv,
     const Int& ncol,
-    const Int& ngwv,
     const uview_1d<const Spack>& lat,
     const uview_1d<const Spack>& u,
     const uview_1d<const Spack>& v,
@@ -384,7 +482,6 @@ struct Functions
     // Inputs
     const Int& ncol,
     const Int& pver,
-    const Int& ngwv,
     const Int& kbot,
     const Int& ktop,
     const uview_1d<const Int>& tend_level,
@@ -436,6 +533,12 @@ struct Functions
     const uview_1d<Spack>& xv,
     const uview_1d<Spack>& yv,
     const uview_1d<Spack>& c);
+
+  //
+  // --------- Members ---------
+  //
+  inline static GwCommonInit s_common_init;
+
 }; // struct Functions
 
 } // namespace gw
@@ -463,5 +566,6 @@ struct Functions
 # include "impl/gw_gw_ediff_impl.hpp"
 # include "impl/gw_gw_diff_tend_impl.hpp"
 # include "impl/gw_gw_oro_src_impl.hpp"
+# include "impl/gw_gw_common_init_impl.hpp"
 #endif // GPU && !KOKKOS_ENABLE_*_RELOCATABLE_DEVICE_CODE
 #endif // P3_FUNCTIONS_HPP

components/eamxx/src/physics/gw/impl/gw_gw_beres_src_impl.hpp

@@ -18,7 +18,6 @@ void Functions<S,D>::gw_beres_src(
 const Int& pver,
 const Int& pgwv,
 const Int& ncol,
-const Int& ngwv,
 const uview_1d<const Spack>& lat,
 const uview_1d<const Spack>& u,
 const uview_1d<const Spack>& v,

components/eamxx/src/physics/gw/impl/gw_gw_cm_src_impl.hpp

@@ -18,7 +18,6 @@ void Functions<S,D>::gw_cm_src(
 const Int& pver,
 const Int& pgwv,
 const Int& ncol,
-const Int& ngwv,
 const Int& kbot,
 const uview_1d<const Spack>& u,
 const uview_1d<const Spack>& v,

components/eamxx/src/physics/gw/impl/gw_gw_common_init_impl.hpp

@@ -0,0 +1,54 @@
+#ifndef GW_GW_COMMON_INIT_IMPL_HPP
+#define GW_GW_COMMON_INIT_IMPL_HPP
+
+#include "gw_functions.hpp" // for ETI only but harmless for GPU
+
+namespace scream {
+namespace gw {
+
+/*
+ * Implementation of gw gw_common_init. Clients should NOT
+ * #include this file, but include gw_functions.hpp instead.
+ */
+
+template<typename S, typename D>
+void Functions<S,D>::gw_common_init(
+// Inputs
+  const Int& pver_in,
+  const Int& pgwv_in,
+  const Real& dc_in,
+  const uview_1d<const Real>& cref_in,
+  const bool& orographic_only_in,
+  const bool& do_molec_diff_in,
+  const bool& tau_0_ubc_in,
+  const Int& nbot_molec_in,
+  const Int& ktop_in,
+  const Int& kbotbg_in,
+  const Real& fcrit2_in,
+  const Real& kwv_in,
+  const uview_1d<const Real>& alpha_in)
+{
+  s_common_init.initialized = true;
+  s_common_init.pver = pver_in;
+  s_common_init.pgwv = pgwv_in;
+  s_common_init.dc = dc_in;
+  s_common_init.cref = view_1d<Real>("cref", cref_in.size());
+  Kokkos::deep_copy(s_common_init.cref, cref_in);
+  s_common_init.orographic_only = orographic_only_in;
+  s_common_init.do_molec_diff = do_molec_diff_in;
+  s_common_init.tau_0_ubc = tau_0_ubc_in;
+  s_common_init.nbot_molec = nbot_molec_in;
+  s_common_init.ktop = ktop_in;
+  s_common_init.kbotbg = kbotbg_in;
+  s_common_init.fcrit2 = fcrit2_in;
+  s_common_init.kwv = kwv_in;
+  s_common_init.alpha = view_1d<Real>("alpha", alpha_in.size());
+  Kokkos::deep_copy(s_common_init.alpha, alpha_in);
+  s_common_init.effkwv = kwv_in * fcrit2_in;
+  s_common_init.tndmax = orographic_only_in ? 500. / 86400. : 400. / 86400.;
+}
+
+} // namespace gw
+} // namespace scream
+
+#endif