Skip to content

Commit 2a8c37e

Browse files
whannah1whannah1
committed
enable host mirroring of ZM variables
1 parent 55c11d3 commit 2a8c37e

File tree

3 files changed

+163
-43
lines changed

3 files changed

+163
-43
lines changed

‎components/eamxx/src/physics/zm/eamxx_zm_process_interface.cpp‎

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,45 @@ void ZMDeepConvection::initialize_impl (const RunType)
102102

103103
// initialize variables on the fortran side
104104
zm::zm_eamxx_bridge_init(m_nlev);
105+
106+
//----------------------------------------------------------------------------
107+
// initialize host mirror variables
108+
zm_input.h_z_mid = Kokkos::create_mirror_view(zm_input.f_z_mid);
109+
zm_input.h_p_mid = Kokkos::create_mirror_view(zm_input.f_p_mid);
110+
zm_input.h_p_del = Kokkos::create_mirror_view(zm_input.f_p_del);
111+
zm_input.h_T_mid = Kokkos::create_mirror_view(zm_input.f_T_mid);
112+
zm_input.h_qv = Kokkos::create_mirror_view(zm_input.f_qv);
113+
zm_input.h_uwind = Kokkos::create_mirror_view(zm_input.f_uwind);
114+
zm_input.h_vwind = Kokkos::create_mirror_view(zm_input.f_vwind);
115+
zm_input.h_omega = Kokkos::create_mirror_view(zm_input.f_omega);
116+
zm_input.h_cldfrac = Kokkos::create_mirror_view(zm_input.f_cldfrac);
117+
zm_input.h_z_int = Kokkos::create_mirror_view(zm_input.f_z_int);
118+
zm_input.h_p_int = Kokkos::create_mirror_view(zm_input.f_p_int);
119+
zm_input.h_tpert = Kokkos::create_mirror_view(zm_input.tpert);
120+
121+
zm_output.h_tend_s = Kokkos::create_mirror_view(zm_output.f_tend_s);
122+
zm_output.h_tend_qv = Kokkos::create_mirror_view(zm_output.f_tend_qv);
123+
zm_output.h_tend_u = Kokkos::create_mirror_view(zm_output.f_tend_u);
124+
zm_output.h_tend_v = Kokkos::create_mirror_view(zm_output.f_tend_v);
125+
zm_output.h_rain_prod = Kokkos::create_mirror_view(zm_output.f_rain_prod);
126+
zm_output.h_snow_prod = Kokkos::create_mirror_view(zm_output.f_snow_prod);
127+
zm_output.h_prec_flux = Kokkos::create_mirror_view(zm_output.f_prec_flux);
128+
zm_output.h_snow_flux = Kokkos::create_mirror_view(zm_output.f_snow_flux);
129+
zm_output.h_mass_flux = Kokkos::create_mirror_view(zm_output.f_mass_flux);
130+
zm_output.h_prec = Kokkos::create_mirror_view(zm_output.prec);
131+
zm_output.h_snow = Kokkos::create_mirror_view(zm_output.snow);
132+
zm_output.h_cape = Kokkos::create_mirror_view(zm_output.cape);
133+
zm_output.h_activity = Kokkos::create_mirror_view(zm_output.activity);
134+
135+
//----------------------------------------------------------------------------
136+
// initialize host mirror variables for managed views
137+
const auto& phis = get_field_in("phis") .get_view<const Real*>();
138+
const auto& pblh = get_field_in("pbl_height") .get_view<const Real*>();
139+
const auto& landfrac = get_field_in("landfrac") .get_view<const Real*>();
140+
zm_input.h_phis = Kokkos::create_mirror_view(phis);
141+
zm_input.h_pblh = Kokkos::create_mirror_view(pblh);
142+
zm_input.h_landfrac = Kokkos::create_mirror_view(landfrac);
143+
105144
}
106145

107146
/*------------------------------------------------------------------------------------------------*/
@@ -428,6 +467,7 @@ void ZMDeepConvection::init_buffers(const ATMBufferManager &buffer_manager)
428467
size_t used_mem = (reinterpret_cast<Real*>(total_mem) - buffer_manager.get_memory())*sizeof(Real);
429468
auto mem_chk = ( used_mem == requested_buffer_size_in_bytes() );
430469
EKAT_REQUIRE_MSG(mem_chk,"Error! Used memory != requested memory for ZMDeepConvection.");
470+
//----------------------------------------------------------------------------
431471
}
432472

433473
/*------------------------------------------------------------------------------------------------*/

‎components/eamxx/src/physics/zm/fortran_bridge/zm_eamxx_bridge.cpp‎

Lines changed: 28 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -61,34 +61,34 @@ void zm_eamxx_bridge_run( Int ncol, Int pver,
6161
zm_eamxx_bridge_run_c( ncol, // 01
6262
zm_input.dtime, // 02
6363
zm_input.is_first_step, // 03
64-
zm_input.phis .data(), // 04
65-
zm_input.f_z_mid .data(), // 05
66-
zm_input.f_z_int .data(), // 06
67-
zm_input.f_p_mid .data(), // 07
68-
zm_input.f_p_int .data(), // 08
69-
zm_input.f_p_del .data(), // 09
70-
zm_input.f_T_mid .data(), // 10
71-
zm_input.f_qv .data(), // 11
72-
zm_input.f_uwind .data(), // 12
73-
zm_input.f_vwind .data(), // 13
74-
zm_input.f_omega .data(), // 14
75-
zm_input.f_cldfrac .data(), // 15
76-
zm_input.pblh .data(), // 16
77-
zm_input.tpert .data(), // 17
78-
zm_input.landfrac .data(), // 18
79-
zm_output.prec .data(), // 19
80-
zm_output.snow .data(), // 20
81-
zm_output.cape .data(), // 21
82-
zm_output.activity .data(), // 22
83-
zm_output.f_tend_s .data(), // 23
84-
zm_output.f_tend_qv .data(), // 24
85-
zm_output.f_tend_u .data(), // 25
86-
zm_output.f_tend_v .data(), // 26
87-
zm_output.f_rain_prod .data(), // 27
88-
zm_output.f_snow_prod .data(), // 28
89-
zm_output.f_prec_flux .data(), // 29
90-
zm_output.f_snow_flux .data(), // 30
91-
zm_output.f_mass_flux .data() // 31
64+
zm_input.h_phis .data(), // 04
65+
zm_input.h_z_mid .data(), // 05
66+
zm_input.h_z_int .data(), // 06
67+
zm_input.h_p_mid .data(), // 07
68+
zm_input.h_p_int .data(), // 08
69+
zm_input.h_p_del .data(), // 09
70+
zm_input.h_T_mid .data(), // 10
71+
zm_input.h_qv .data(), // 11
72+
zm_input.h_uwind .data(), // 12
73+
zm_input.h_vwind .data(), // 13
74+
zm_input.h_omega .data(), // 14
75+
zm_input.h_cldfrac .data(), // 15
76+
zm_input.h_pblh .data(), // 16
77+
zm_input.h_tpert .data(), // 17
78+
zm_input.h_landfrac .data(), // 18
79+
zm_output.h_prec .data(), // 19
80+
zm_output.h_snow .data(), // 20
81+
zm_output.h_cape .data(), // 21
82+
zm_output.h_activity .data(), // 22
83+
zm_output.h_tend_s .data(), // 23
84+
zm_output.h_tend_qv .data(), // 24
85+
zm_output.h_tend_u .data(), // 25
86+
zm_output.h_tend_v .data(), // 26
87+
zm_output.h_rain_prod .data(), // 27
88+
zm_output.h_snow_prod .data(), // 28
89+
zm_output.h_prec_flux .data(), // 29
90+
zm_output.h_snow_flux .data(), // 30
91+
zm_output.h_mass_flux .data() // 31
9292
);
9393

9494
zm_input.transpose<ekat::TransposeDirection::f2c>(ncol,pver);

‎components/eamxx/src/physics/zm/zm_functions.hpp‎

Lines changed: 95 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,8 @@ struct Functions {
4141
template <typename S> using uview_2d = typename ekat::template Unmanaged<view_2d<S> >;
4242
template <typename S> using uview_2dl = typename ekat::template Unmanaged<view_2dl<S> >;
4343
template <typename S> using uview_2d_strided = typename ekat::template Unmanaged<view_2d_strided<S> >;
44+
template <typename S> using uview_2dh = typename ekat::template Unmanaged<view_2dl<S>>::HostMirror;
45+
template <typename S> using uview_1dh = typename ekat::template Unmanaged<view_1d<S>>::HostMirror;
4446

4547
// ---------------------------------------------------------------------------
4648
// Structs
@@ -89,18 +91,35 @@ struct Functions {
8991
view_1d<const Scalar> landfrac; // land area fraction
9092

9193
// unmanaged LayoutLeft views for fortran bridging
92-
uview_2dl<Real> f_z_mid;
93-
uview_2dl<Real> f_p_mid;
94-
uview_2dl<Real> f_p_del;
95-
uview_2dl<Real> f_T_mid;
96-
uview_2dl<Real> f_qv;
97-
uview_2dl<Real> f_uwind;
98-
uview_2dl<Real> f_vwind;
99-
uview_2dl<Real> f_omega;
100-
uview_2dl<Real> f_cldfrac;
101-
102-
uview_2dl<Real> f_z_int;
103-
uview_2dl<Real> f_p_int;
94+
uview_2dl<Real> f_z_mid;
95+
uview_2dl<Real> f_p_mid;
96+
uview_2dl<Real> f_p_del;
97+
uview_2dl<Real> f_T_mid;
98+
uview_2dl<Real> f_qv;
99+
uview_2dl<Real> f_uwind;
100+
uview_2dl<Real> f_vwind;
101+
uview_2dl<Real> f_omega;
102+
uview_2dl<Real> f_cldfrac;
103+
uview_2dl<Real> f_z_int;
104+
uview_2dl<Real> f_p_int;
105+
106+
// host mirror versions of ZM interface variables
107+
uview_2dh<Real> h_z_mid;
108+
uview_2dh<Real> h_p_mid;
109+
uview_2dh<Real> h_p_del;
110+
uview_2dh<Real> h_T_mid;
111+
uview_2dh<Real> h_qv;
112+
uview_2dh<Real> h_uwind;
113+
uview_2dh<Real> h_vwind;
114+
uview_2dh<Real> h_omega;
115+
uview_2dh<Real> h_cldfrac;
116+
uview_2dh<Real> h_z_int;
117+
uview_2dh<Real> h_p_int;
118+
119+
uview_1dh<Scalar> h_phis;
120+
uview_1dh<Scalar> h_pblh;
121+
uview_1dh<Scalar> h_tpert;
122+
uview_1dh<Scalar> h_landfrac;
104123

105124
// -------------------------------------------------------------------------
106125
// transpose method for fortran bridging
@@ -125,6 +144,23 @@ struct Functions {
125144
f_p_int (i,j) = p_int (i,j/Spack::n)[j%Spack::n];
126145
}
127146
}
147+
//----------------------------------------------------------------------
148+
// copy to host mirrors
149+
Kokkos::deep_copy(h_z_mid, f_z_mid);
150+
Kokkos::deep_copy(h_p_mid, f_p_mid);
151+
Kokkos::deep_copy(h_p_del, f_p_del);
152+
Kokkos::deep_copy(h_T_mid, f_T_mid);
153+
Kokkos::deep_copy(h_qv, f_qv);
154+
Kokkos::deep_copy(h_uwind, f_uwind);
155+
Kokkos::deep_copy(h_vwind, f_vwind);
156+
Kokkos::deep_copy(h_omega, f_omega);
157+
Kokkos::deep_copy(h_cldfrac, f_cldfrac);
158+
Kokkos::deep_copy(h_z_int, f_z_int);
159+
Kokkos::deep_copy(h_p_int, f_p_int);
160+
Kokkos::deep_copy(h_phis, phis);
161+
Kokkos::deep_copy(h_pblh, pblh);
162+
Kokkos::deep_copy(h_tpert, tpert);
163+
Kokkos::deep_copy(h_landfrac, landfrac);
128164
}
129165
}
130166
// -------------------------------------------------------------------------
@@ -164,11 +200,26 @@ struct Functions {
164200
uview_2dl<Real> f_tend_v;
165201
uview_2dl<Real> f_rain_prod;
166202
uview_2dl<Real> f_snow_prod;
167-
168203
uview_2dl<Real> f_prec_flux;
169204
uview_2dl<Real> f_snow_flux;
170205
uview_2dl<Real> f_mass_flux;
171206

207+
// host mirror versions of ZM interface variables
208+
uview_2dh<Real> h_tend_s;
209+
uview_2dh<Real> h_tend_qv;
210+
uview_2dh<Real> h_tend_u;
211+
uview_2dh<Real> h_tend_v;
212+
uview_2dh<Real> h_rain_prod;
213+
uview_2dh<Real> h_snow_prod;
214+
uview_2dh<Real> h_prec_flux;
215+
uview_2dh<Real> h_snow_flux;
216+
uview_2dh<Real> h_mass_flux;
217+
218+
uview_1dh<Scalar> h_prec;
219+
uview_1dh<Scalar> h_snow;
220+
uview_1dh<Scalar> h_cape;
221+
uview_1dh<Int> h_activity;
222+
172223
// -------------------------------------------------------------------------
173224
// transpose method for fortran bridging
174225
template <ekat::TransposeDirection::Enum D>
@@ -192,10 +243,39 @@ struct Functions {
192243
f_mass_flux(i,j) = mass_flux(i,j/Spack::n)[j%Spack::n];
193244
}
194245
}
195-
// sync_to_host here?
246+
//----------------------------------------------------------------------
247+
// copy to host mirror
248+
Kokkos::deep_copy(h_tend_s, f_tend_s);
249+
Kokkos::deep_copy(h_tend_qv, f_tend_qv);
250+
Kokkos::deep_copy(h_tend_u, f_tend_u);
251+
Kokkos::deep_copy(h_tend_v, f_tend_v);
252+
Kokkos::deep_copy(h_rain_prod,f_rain_prod);
253+
Kokkos::deep_copy(h_snow_prod,f_snow_prod);
254+
Kokkos::deep_copy(h_prec_flux,f_prec_flux);
255+
Kokkos::deep_copy(h_snow_flux,f_snow_flux);
256+
Kokkos::deep_copy(h_mass_flux,f_mass_flux);
257+
Kokkos::deep_copy(h_prec, prec);
258+
Kokkos::deep_copy(h_snow, snow);
259+
Kokkos::deep_copy(h_cape, cape);
260+
Kokkos::deep_copy(h_activity, activity);
196261
}
197262
if (D == ekat::TransposeDirection::f2c) {
198-
// sync_to_device?
263+
//----------------------------------------------------------------------
264+
// copy to host mirror
265+
Kokkos::deep_copy(f_tend_s, h_tend_s);
266+
Kokkos::deep_copy(f_tend_qv, h_tend_qv);
267+
Kokkos::deep_copy(f_tend_u, h_tend_u);
268+
Kokkos::deep_copy(f_tend_v, h_tend_v);
269+
Kokkos::deep_copy(f_rain_prod,h_rain_prod);
270+
Kokkos::deep_copy(f_snow_prod,h_snow_prod);
271+
Kokkos::deep_copy(f_prec_flux,h_prec_flux);
272+
Kokkos::deep_copy(f_snow_flux,h_snow_flux);
273+
Kokkos::deep_copy(f_mass_flux,h_mass_flux);
274+
Kokkos::deep_copy(prec, h_prec);
275+
Kokkos::deep_copy(snow, h_snow);
276+
Kokkos::deep_copy(cape, h_cape);
277+
Kokkos::deep_copy(activity, h_activity);
278+
//----------------------------------------------------------------------
199279
for (int i=0; i<ncol_in; ++i) {
200280
// mid-point level variables
201281
for (int j=0; j<pver_in; ++j) {

0 commit comments

Comments
 (0)