Skip to content

Commit d977eaf

Browse files
odiazibodiazib
committed
MAM4xx: Adding nested parallels in chm_diags.
1 parent 453fc0a commit d977eaf

File tree

3 files changed

+89
-83
lines changed

3 files changed

+89
-83
lines changed

‎src/mam4xx/mo_chm_diags.hpp‎

Lines changed: 88 additions & 78 deletions
Original file line numberDiff line numberDiff line change
@@ -153,7 +153,6 @@ void chm_diags(
153153
// ... local variables
154154
//--------------------------------------------------------------------
155155

156-
Real wgt;
157156
// Real pointer :: fldcw(:,:) //working pointer to extract data from pbuf for
158157
// sum of mass for aerosol classes
159158

@@ -170,7 +169,8 @@ void chm_diags(
170169
//--------------------------------------------------------------------
171170
// ... "diagnostic" groups
172171
//--------------------------------------------------------------------
173-
for (int kk = 0; kk < pver; kk++) {
172+
constexpr int pver_local = pver;
173+
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pver_local), [&](int kk) {
174174
vmr_nox(kk) = 0;
175175
vmr_noy(kk) = 0;
176176
vmr_clox(kk) = 0;
@@ -180,7 +180,7 @@ void chm_diags(
180180
mmr_noy(kk) = 0;
181181
mmr_sox(kk) = 0;
182182
mmr_nhx(kk) = 0;
183-
}
183+
});
184184
df_noy(0) = 0;
185185
df_sox(0) = 0;
186186
df_nhx(0) = 0;
@@ -199,99 +199,116 @@ void chm_diags(
199199
// initialize the mass arrays
200200
// if (history_aerosol .and. .not. history_verbose) then // what was this used
201201
// for?
202-
for (int kk = 0; kk < pver; kk++) {
202+
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pver_local), [&](int kk) {
203203
mass_bc(kk) = 0;
204204
mass_dst(kk) = 0;
205205
mass_mom(kk) = 0;
206206
mass_ncl(kk) = 0;
207207
mass_pom(kk) = 0;
208208
mass_so4(kk) = 0;
209209
mass_soa(kk) = 0;
210-
}
210+
});
211+
212+
Kokkos::single(Kokkos::PerTeam(team),
213+
[=]() { area(0) *= haero::square(rearth); });
211214

212-
area(0) *= haero::square(rearth);
215+
team.team_barrier();
213216

214-
for (int kk = 0; kk < pver; kk++) {
217+
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pver_local), [&](int kk) {
215218
mass(kk) = pdel(kk) * area(0) * rgrav;
216219
drymass(kk) = pdeldry(kk) * area(0) * rgrav;
217-
}
220+
});
218221

219222
// convert ozone from mol/mol (w.r.t. dry air mass) to DU
220-
for (int kk = 0; kk < pver; kk++) {
223+
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pver_local), [&](int kk) {
221224
ozone_layer(kk) =
222225
pdeldry(kk) * vmr[id_o3](kk) * avogadro * rgrav / mwdry / DUfac * 1e3;
223-
}
226+
});
224227
// total column ozone
225228
ozone_col(0) = 0;
226229
ozone_trop(0) = 0;
227230
ozone_strat(0) = 0;
228-
229-
for (int kk = 0; kk < pver; kk++) {
230-
ozone_col(0) += ozone_layer(kk);
231-
if (kk <= ltrop) {
232-
// stratospheric column ozone
233-
ozone_strat(0) += ozone_layer(kk);
234-
} else {
235-
// tropospheric column ozone
236-
ozone_trop(0) += ozone_layer(kk);
237-
}
238-
}
231+
team.team_barrier();
232+
233+
Kokkos::parallel_reduce(
234+
Kokkos::TeamVectorRange(team, pver_local),
235+
[&](int kk, Real &update) { update += ozone_layer(kk); }, ozone_col(0));
236+
237+
// stratospheric column ozone
238+
Kokkos::parallel_reduce(
239+
Kokkos::TeamVectorRange(team, pver_local),
240+
[&](int kk, Real &update) {
241+
if (kk <= ltrop)
242+
update += ozone_layer(kk);
243+
},
244+
ozone_strat(0));
245+
246+
// tropospheric column ozone
247+
Kokkos::parallel_reduce(
248+
Kokkos::TeamVectorRange(team, pver_local),
249+
[&](int kk, Real &update) {
250+
if (kk > ltrop)
251+
update += ozone_layer(kk);
252+
},
253+
ozone_trop(0));
239254

240255
for (int mm = 0; mm < gas_pcnst; mm++) {
241256
// other options of species are not used, only use weight=1
242-
wgt = 1;
257+
const Real wgt = 1;
243258

244259
for (int i = 0; i < 3; i++) { // FIXME: bad constant (len of sox species)
245260
if (sox_species[i] == mm) {
246-
for (int kk = 0; kk < pver; kk++) {
247-
mmr_sox(kk) = mmr_sox(kk) + wgt * mmr[mm](kk);
248-
}
261+
262+
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pver_local),
263+
[&](int kk) { mmr_sox(kk) += wgt * mmr[mm](kk); });
249264
}
250265
}
251266

252267
if (aer_species[mm] == mm) {
253268
const char *symbol = solsym[mm];
254269
if (name_matches(symbol, "bc_a")) {
255-
for (int kk = 0; kk < pver; kk++) {
256-
mass_bc(kk) += mmr[mm](kk);
257-
}
270+
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pver_local),
271+
[&](int kk) { mass_bc(kk) += mmr[mm](kk); });
272+
258273
} else if (name_matches(symbol, "dst_a")) {
259-
for (int kk = 0; kk < pver; kk++) {
260-
mass_dst(kk) += mmr[mm](kk);
261-
}
274+
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pver_local),
275+
[&](int kk) { mass_dst(kk) += mmr[mm](kk); });
262276
} else if (name_matches(symbol, "mom_a")) {
263-
for (int kk = 0; kk < pver; kk++) {
264-
mass_mom(kk) += mmr[mm](kk);
265-
}
277+
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pver_local),
278+
[&](int kk) { mass_mom(kk) += mmr[mm](kk); });
279+
266280
} else if (name_matches(symbol, "ncl_a")) {
267-
for (int kk = 0; kk < pver; kk++) {
268-
mass_ncl(kk) += mmr[mm](kk);
269-
}
281+
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pver_local),
282+
[&](int kk) { mass_ncl(kk) += mmr[mm](kk); });
270283
} else if (name_matches(symbol, "pom_a")) {
271-
for (int kk = 0; kk < pver; kk++) {
272-
mass_pom(kk) += mmr[mm](kk);
273-
}
284+
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pver_local),
285+
[&](int kk) { mass_pom(kk) += mmr[mm](kk); });
274286
} else if (name_matches(symbol, "so4_a")) {
275-
for (int kk = 0; kk < pver; kk++) {
276-
mass_so4(kk) += mmr[mm](kk);
277-
}
287+
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pver_local),
288+
[&](int kk) { mass_so4(kk) += mmr[mm](kk); });
278289
} else if (name_matches(symbol, "soa_a")) {
279-
for (int kk = 0; kk < pver; kk++) {
280-
mass_soa(kk) += mmr[mm](kk);
281-
}
282-
}
283-
}
284-
285-
for (int i = 0; i < 3; i++) { // FIXME: bad constant (len of sox species)
286-
if (sox_species[i] == mm) {
287-
df_sox(0) += wgt * depflx[mm] * S_molwgt / adv_mass[mm];
290+
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pver_local),
291+
[&](int kk) { mass_soa(kk) += mmr[mm](kk); });
288292
}
289293
}
290-
291-
for (int kk = 0; kk < pver; kk++) {
292-
net_chem(kk) = mmr_tend[mm](kk) * mass(kk);
293-
}
294+
Kokkos::parallel_for(
295+
Kokkos::TeamVectorRange(team, pver_local),
296+
[&](int kk) { net_chem(kk) = mmr_tend[mm](kk) * mass(kk); });
294297
}
298+
constexpr int gas_pcnst_local = gas_pcnst;
299+
Kokkos::parallel_reduce(
300+
Kokkos::TeamVectorRange(team, gas_pcnst_local),
301+
[&](int mm, Real &update) {
302+
// other options of species are not used, only use weight=1
303+
const Real wgt = 1;
304+
for (int i = 0; i < 3;
305+
i++) { // FIXME: bad constant (len of sox species)
306+
if (sox_species[i] == mm) {
307+
update += wgt * depflx[mm] * S_molwgt / adv_mass[mm];
308+
}
309+
}
310+
},
311+
df_sox(0));
295312

296313
// diagnostics for cloud-borne aerosols, then add to corresponding mass
297314
// accumulators
@@ -320,36 +337,29 @@ void chm_diags(
320337
}
321338
}
322339
if (name_matches(symbol_cw, "bc_c")) {
323-
for (int kk = 0; kk < pver; kk++) {
324-
mass_bc(kk) += fldcw[nn](kk);
325-
}
340+
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pver_local),
341+
[&](int kk) { mass_bc(kk) += fldcw[nn](kk); });
326342
} else if (name_matches(symbol_cw, "dst_c")) {
327-
for (int kk = 0; kk < pver; kk++) {
328-
mass_dst(kk) += fldcw[nn](kk);
329-
}
343+
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pver_local),
344+
[&](int kk) { mass_dst(kk) += fldcw[nn](kk); });
330345
} else if (name_matches(symbol_cw, "mom_c")) {
331-
for (int kk = 0; kk < pver; kk++) {
332-
mass_mom(kk) += fldcw[nn](kk);
333-
}
346+
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pver_local),
347+
[&](int kk) { mass_mom(kk) += fldcw[nn](kk); });
334348
} else if (name_matches(symbol_cw, "ncl_c")) {
335-
for (int kk = 0; kk < pver; kk++) {
336-
mass_ncl(kk) += fldcw[nn](kk);
337-
}
349+
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pver_local),
350+
[&](int kk) { mass_ncl(kk) += fldcw[nn](kk); });
338351
} else if (name_matches(symbol_cw, "pom_c")) {
339-
for (int kk = 0; kk < pver; kk++) {
340-
mass_pom(kk) += fldcw[nn](kk);
341-
}
352+
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pver_local),
353+
[&](int kk) { mass_pom(kk) += fldcw[nn](kk); });
342354
} else if (name_matches(symbol_cw, "so4_c")) {
343-
for (int kk = 0; kk < pver; kk++) {
344-
mass_so4(kk) += fldcw[nn](kk);
345-
}
355+
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pver_local),
356+
[&](int kk) { mass_so4(kk) += fldcw[nn](kk); });
346357
} else if (name_matches(symbol_cw, "soa_c")) {
347-
for (int kk = 0; kk < pver; kk++) {
348-
mass_soa(kk) += fldcw[nn](kk);
349-
}
358+
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pver_local),
359+
[&](int kk) { mass_soa(kk) += fldcw[nn](kk); });
350360
}
351361
}
352362
} // chm_diags
353363
} // namespace mo_chm_diags
354364
} // namespace mam4
355-
#endif
365+
#endif

‎src/validation/mo_chm_diags/CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ set(TEST_LIST
3636
set(DEFAULT_TOL 1e-13)
3737
set(ERROR_THRESHOLDS
3838
9e-2
39-
9e-2
39+
1.3e-5
4040
)
4141
foreach(input tol IN ZIP_LISTS TEST_LIST ERROR_THRESHOLDS)
4242
# copy the baseline file into place.

‎src/validation/mo_chm_diags/chm_diags.cpp‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -78,10 +78,6 @@ void chm_diags(Ensemble *ensemble) {
7878
"bc_a4", "mom_a4", "num_a4"};
7979

8080
//=========init views==========
81-
82-
for (int kk = 0; kk < pver; ++kk) {
83-
}
84-
8581
ColumnView mmr[gas_pcnst];
8682
ColumnView vmr[gas_pcnst];
8783
ColumnView mmr_tend[gas_pcnst];

0 commit comments

Comments
 (0)