@@ -58,23 +58,40 @@ void Functions<S,D>::gw_oro_src(
5858 const Real hdsp = 2 * sgh; // Surface streamline displacement height (2*sgh).
5959
6060 Int k = pver-1 ;
61- src_level = k;
6261
6362 // Averages over source region.
64- Real rsrc = pmid (k+ 1 )/(C::Rair*t (k)) * dpm (k); // Density.
63+ Real rsrc = pmid (k)/(C::Rair*t (k)) * dpm (k); // Density.
6564 Real usrc = u (k) * dpm (k); // Zonal wind.
6665 Real vsrc = v (k) * dpm (k); // Meridional wind.
6766 Real nsrc = nm (k)* dpm (k); // B-V frequency.
6867
69- for (k = pver - 2 ; k >= pver/2 -1 ; --k) {
70- if (hdsp > std::sqrt (zm (k)*zm (k+1 ))) {
71- src_level = k;
72- rsrc = rsrc + pmid (k+1 ) / (C::Rair*t (k))* dpm (k);
73- usrc = usrc + u (k) * dpm (k);
74- vsrc = vsrc + v (k) * dpm (k);
75- nsrc = nsrc + nm (k)* dpm (k);
76- }
77- }
68+ Real rsrc_sum, usrc_sum, vsrc_sum, nsrc_sum;
69+ Kokkos::parallel_reduce (
70+ Kokkos::TeamVectorRange (team, pver/2 - 1 , pver-1 ), [&] (const int k, Real& lrsrc, Real& lusrc, Real& lvsrc, Real& lnsrc) {
71+ if (hdsp > std::sqrt (zm (k)*zm (k+1 ))) {
72+ lrsrc += pmid (k) / (C::Rair*t (k))* dpm (k);
73+ lusrc += u (k) * dpm (k);
74+ lvsrc += v (k) * dpm (k);
75+ lnsrc += (nm (k)* dpm (k));
76+ }
77+ }, rsrc_sum, usrc_sum, vsrc_sum, nsrc_sum);
78+
79+ Kokkos::parallel_reduce (
80+ Kokkos::TeamVectorRange (team, pver/2 - 1 , pver-1 ), [&] (const int k, Int& lmin) {
81+ if (lmin > pver - 1 ) {
82+ lmin = pver - 1 ;
83+ }
84+ if (hdsp > std::sqrt (zm (k)*zm (k+1 )) && lmin > k) {
85+ lmin = k;
86+ }
87+ }, Kokkos::Min<Int>(src_level));
88+
89+ team.team_barrier ();
90+
91+ rsrc += rsrc_sum;
92+ usrc += usrc_sum;
93+ vsrc += vsrc_sum;
94+ nsrc += nsrc_sum;
7895
7996 // Difference in interface pressure across source region.
8097 const Real dpsrc = pint (pver) - pint (src_level);
@@ -88,16 +105,21 @@ void Functions<S,D>::gw_oro_src(
88105 get_unit_vector (usrc, vsrc, xv, yv, ubi (pver));
89106
90107 // Project the local wind at midpoints onto the source wind.
91- for (k = 0 ; k < pver; ++k) {
92- ubm (k) = dot_2d (u (k), v (k), xv, yv);
93- }
108+ Kokkos::parallel_for (
109+ Kokkos::TeamVectorRange (team, pver), [&] (const int k) {
110+ ubm (k) = dot_2d (u (k), v (k), xv, yv);
111+ });
112+
113+ team.team_barrier ();
94114
95115 // Compute the interface wind projection by averaging the midpoint winds.
96116 // Use the top level wind at the top interface.
97117 ubi (0 ) = ubm (0 );
98118
99119 midpoint_interp (team, ubm, ekat::subview (ubi, Kokkos::pair<int , int >{1 , pver}));
100120
121+ team.team_barrier ();
122+
101123 // Determine the orographic c=0 source term following McFarlane (1987).
102124 // Set the source top interface index to pver, if the orographic term is
103125 // zero.
@@ -115,18 +137,24 @@ void Functions<S,D>::gw_oro_src(
115137 // Set the phase speeds and wave numbers in the direction of the source
116138 // wind. Set the source stress magnitude (positive only, note that the
117139 // sign of the stress is the same as (c-u).
118- for (k = pver; k >= src_level; --k) {
119- tau (0 , k) = tauoro;
120- }
140+ Kokkos::parallel_for (
141+ Kokkos::TeamVectorRange (team, src_level, pver+1 ), [&] (const int k) {
142+ tau (pgwv, k) = tauoro;
143+ });
121144
122145 // Allow wind tendencies all the way to the model bottom.
123146 tend_level = pver - 1 ;
124- --src_level;
147+
148+ // adjust to c indexing. Up to this point, src_level was used to index into 0:pver arrays
149+ Kokkos::single (Kokkos::PerTeam (team), [&] {
150+ --src_level;
151+ });
125152
126153 // No spectrum; phase speed is just 0.
127- for (k = 0 ; k < (Int)c.size (); ++k) {
128- c (k) = 0 ;
129- }
154+ Kokkos::parallel_for (
155+ Kokkos::TeamVectorRange (team, c.size ()), [&] (const int k) {
156+ c (k) = 0 ;
157+ });
130158}
131159
132160} // namespace gw
0 commit comments