For top-level beres, need team-shared locals

Author: jgfouca

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

@@ -544,6 +544,7 @@ struct Functions
   static void gw_beres_src(
     // Inputs
     const MemberType& team,
+    const Workspace& workspace,
     const GwCommonInit& init,
     const GwConvectInit& cinit,
     const Int& pver,

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

@@ -16,6 +16,7 @@ KOKKOS_FUNCTION
 void Functions<S,D>::gw_beres_src(
   // Inputs
   const MemberType& team,
+  const Workspace& workspace,
   const GwCommonInit& init,
   const GwConvectInit& cinit,
   const Int& pver,
@@ -42,20 +43,26 @@ void Functions<S,D>::gw_beres_src(
   Real& hdepth,
   Real& maxq0_out)
 {
+  // Note all locals must be shared per team so that we have team-consistent
+  // views of these values.
+  auto local_storage = workspace.take("local_storage");
+
   // Maximum heating rate.
-  Real maxq0(0);
+  Real& maxq0 = local_storage(0); maxq0 = 0;
 
   // Bottom/top heating range index.
-  Int mini(0), maxi(0);
+  Int& mini = reinterpret_cast<Int&>(local_storage(1)); mini = 0;
+  Int& maxi = reinterpret_cast<Int&>(local_storage(2)); maxi = 0;
 
   // Mean wind in heating region.
-  Real uh(0);
+  Real& uh = local_storage(3); uh = 0;
 
   // Min/max projected wind value in each column.
-  Real Umin(0), Umax(0);
+  Real& Umin = local_storage(4); Umin = 0;
+  Real& Umax = local_storage(5); Umax = 0;
 
   // Speed of convective cells relative to storm.
-  Int storm_speed(0);
+  Int& storm_speed = reinterpret_cast<Int&>(local_storage(6)); storm_speed = 0;
 
   // note: the heating_altitude_max is probably not needed because there is
   // rarely any convective heating above this level and the performance impact
@@ -109,6 +116,8 @@ void Functions<S,D>::gw_beres_src(
     Kokkos::TeamVectorRange(team, init.cref.size()), [&] (const int l) {
       c(l) = init.cref(l);
     });
+
+  workspace.release(local_storage);
 }
 
 } // namespace gw

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

@@ -49,7 +49,7 @@ void Functions<S,D>::gw_convect_gw_sources(
     //------------------------------------------------------------------
 
     // Shift spectrum so that it is relative to the ground.
-    const Int shift = -static_cast<Int>(std::round(storm_speed/init.dc));
+    const Int shift = -static_cast<Int>(std::round(storm_speed/init.dc)) % num_pgwv;
 
     // Adjust for critical level filtering.
     const Int Umini = ekat::impl::max(static_cast<Int>(std::round(umin/init.dc)), -pgwv) + pgwv;

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

@@ -35,10 +35,6 @@ void Functions<S,D>::gw_heating_depth(
   // Find indices for the top and bottom of the heating range.
   EKAT_KERNEL_REQUIRE_MSG(!use_gw_convect_old, "The old gw convect algorithm is not supported");
 
-  // The type that holds both min and max results
-  using ResultType = Kokkos::MinMax<Int>::value_type;
-  ResultType min_max_result;
-
   //---------------------------------------------------------------------
   // cleaner version that addresses bug in original where heating max and
   // depth were too low whenever heating <=0 occurred in the middle of

components/eamxx/src/physics/gw/tests/infra/gw_test_data.cpp

@@ -1384,6 +1384,8 @@ void gw_beres_src(GwBeresSrcData& d)
 
   const auto policy = ekat::TeamPolicyFactory<ExeSpace>::get_default_team_policy(d.ncol, d.init.init.pver);
 
+  WSM wsm(16, 1, policy);
+
   GWF::GwCommonInit init_cp = GWF::s_common_init;
   GWF::GwConvectInit cinit_cp = GWF::s_convect_init;
 
@@ -1412,6 +1414,7 @@ void gw_beres_src(GwBeresSrcData& d)
 
     GWF::gw_beres_src(
       team,
+      wsm.get_workspace(team),
       init_cp,
       cinit_cp,
       pver,