Skip to content

Commit

Permalink
Merge Pull Request #2980 from E3SM-Project/scream/tcclevenger/view_al…
Browse files Browse the repository at this point in the history 
…locs_in_p3

Automatically Merged using E3SM Pull Request AutoTester
PR Title: View alloc fixes for P3
PR Author: tcclevenger
PR LABELS: p3, AT: AUTOMERGE
  • Loading branch information
@E3SM-Bot
E3SM-Bot authored Sep 10, 2024
2 parents 474edd8 + 3a85062 commit e65666f
Show file tree
Hide file tree
Showing 36 changed files with 807 additions and 678 deletions.
13 changes: 6 additions & 7 deletions components/eamxx/src/physics/p3/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,7 +45,6 @@ if (NOT EAMXX_ENABLE_GPU OR Kokkos_ENABLE_CUDA_RELOCATABLE_DEVICE_CODE OR Kokkos
eti/p3_ice_relaxation_timescale.cpp
eti/p3_ice_nucleation.cpp
eti/p3_ice_cldliq_wet_growth.cpp
eti/p3_get_latent_heat.cpp
eti/p3_check_values.cpp
eti/p3_incloud_mixingratios.cpp
eti/p3_subgrid_variance_scaling.cpp
Expand All @@ -63,13 +62,13 @@ endif()

# List of dispatch source files if monolithic kernels are off
set(P3_SK_SRCS
disp/p3_check_values_impl_disp.cpp
disp/p3_ice_sed_impl_disp.cpp
disp/p3_main_impl_part1_disp.cpp
disp/p3_check_values_impl_disp.cpp
disp/p3_ice_sed_impl_disp.cpp
disp/p3_main_impl_part1_disp.cpp
disp/p3_main_impl_part3_disp.cpp
disp/p3_cloud_sed_impl_disp.cpp
disp/p3_main_impl_disp.cpp
disp/p3_main_impl_part2_disp.cpp
disp/p3_cloud_sed_impl_disp.cpp
disp/p3_main_impl_disp.cpp
disp/p3_main_impl_part2_disp.cpp
disp/p3_rain_sed_impl_disp.cpp
)

Expand Down
7 changes: 3 additions & 4 deletions components/eamxx/src/physics/p3/disp/p3_ice_sed_impl_disp.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -46,8 +46,8 @@ ::ice_sedimentation_disp(

// Ice sedimentation: (adaptive substepping)
ice_sedimentation(
ekat::subview(rho, i), ekat::subview(inv_rho, i), ekat::subview(rhofaci, i), ekat::subview(cld_frac_i, i),
ekat::subview(inv_dz, i), team, workspace, nk, ktop, kbot, kdir, dt, inv_dt,
ekat::subview(rho, i), ekat::subview(inv_rho, i), ekat::subview(rhofaci, i), ekat::subview(cld_frac_i, i),
ekat::subview(inv_dz, i), team, workspace, nk, ktop, kbot, kdir, dt, inv_dt,
ekat::subview(qi, i), ekat::subview(qi_incld, i), ekat::subview(ni, i), ekat::subview(ni_incld, i),
ekat::subview(qm, i), ekat::subview(qm_incld, i), ekat::subview(bm, i), ekat::subview(bm_incld, i), ekat::subview(qi_tend, i), ekat::subview(ni_tend, i),
ice_table_vals, precip_ice_surf(i), p3constants);
Expand All @@ -60,7 +60,6 @@ void Functions<Real,DefaultDevice>
::homogeneous_freezing_disp(
const uview_2d<const Spack>& T_atm,
const uview_2d<const Spack>& inv_exner,
const uview_2d<const Spack>& latent_heat_fusion,
const Int& nj, const Int& nk, const Int& ktop, const Int& kbot, const Int& kdir,
const uview_2d<Spack>& qc,
const uview_2d<Spack>& nc,
Expand Down Expand Up @@ -89,7 +88,7 @@ ::homogeneous_freezing_disp(

// homogeneous freezing of cloud and rain
homogeneous_freezing(
ekat::subview(T_atm, i), ekat::subview(inv_exner, i), ekat::subview(latent_heat_fusion, i), team, nk, ktop, kbot, kdir,
ekat::subview(T_atm, i), ekat::subview(inv_exner, i), team, nk, ktop, kbot, kdir,
ekat::subview(qc, i), ekat::subview(nc, i), ekat::subview(qr, i), ekat::subview(nr, i), ekat::subview(qi, i),
ekat::subview(ni, i), ekat::subview(qm, i), ekat::subview(bm, i), ekat::subview(th_atm, i));

Expand Down
198 changes: 103 additions & 95 deletions components/eamxx/src/physics/p3/disp/p3_main_impl_disp.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,16 +32,16 @@ ::p3_main_init_disp(
const uview_2d<Spack>& qv_supersat_i, const uview_2d<Spack>& qtend_ignore, const uview_2d<Spack>& ntend_ignore, const uview_2d<Spack>& mu_c,
const uview_2d<Spack>& lamc, const uview_2d<Spack>& rho_qi, const uview_2d<Spack>& qv2qi_depos_tend, const uview_2d<Spack>& precip_total_tend,
const uview_2d<Spack>& nevapr, const uview_2d<Spack>& precip_liq_flux, const uview_2d<Spack>& precip_ice_flux)
{
{
using ExeSpace = typename KT::ExeSpace;
const auto policy = ekat::ExeSpaceUtils<ExeSpace>::get_default_team_policy(nj, nk_pack);

Kokkos::parallel_for("p3_main_init",
policy, KOKKOS_LAMBDA(const MemberType& team) {
const Int i = team.league_rank();
precip_liq_surf(i) = 0;
precip_ice_surf(i) = 0;

const Int i = team.league_rank();
precip_liq_surf(i) = 0;
precip_ice_surf(i) = 0;

Kokkos::parallel_for(
Kokkos::TeamVectorRange(team, nk_pack), [&] (Int k) {
Expand Down Expand Up @@ -107,17 +107,14 @@ ::p3_main_internal_disp(
const P3Infrastructure& infrastructure,
const P3HistoryOnly& history_only,
const P3LookupTables& lookup_tables,
const P3Temporaries& temporaries,
const WorkspaceManager& workspace_mgr,
Int nj,
Int nk,
const physics::P3_Constants<Real> & p3constants)
{
using ExeSpace = typename KT::ExeSpace;

view_2d<Spack> latent_heat_sublim("latent_heat_sublim", nj, nk), latent_heat_vapor("latent_heat_vapor", nj, nk), latent_heat_fusion("latent_heat_fusion", nj, nk);

get_latent_heat(nj, nk, latent_heat_vapor, latent_heat_sublim, latent_heat_fusion);

const Int nk_pack = ekat::npack<Spack>(nk);

// load constants into local vars
Expand All @@ -131,87 +128,99 @@ ::p3_main_internal_disp(
view_1d<bool> nucleationPossible("nucleationPossible", nj);
view_1d<bool> hydrometeorsPresent("hydrometeorsPresent", nj);

//
// Create temporary variables needed for p3
//
view_2d<Spack>
mu_r("mu_r", nj, nk_pack), // shape parameter of rain
T_atm("T_atm", nj, nk_pack), // temperature at the beginning of the microphysics step [K]

// 2D size distribution and fallspeed parameters
lamr("lamr", nj, nk_pack), logn0r("logn0r", nj, nk_pack), nu("nu", nj, nk_pack),
cdist("cdist", nj, nk_pack), cdist1("cdist1", nj, nk_pack), cdistr("cdistr", nj, nk_pack),

// Variables needed for in-cloud calculations
// Inverse cloud fractions (1/cld)
inv_cld_frac_i("inv_cld_frac_i", nj, nk_pack), inv_cld_frac_l("inv_cld_frac_l", nj, nk_pack), inv_cld_frac_r("inv_cld_frac_r", nj, nk_pack),
// In cloud mass-mixing ratios
qc_incld("qc_incld", nj, nk_pack), qr_incld("qr_incld", nj, nk_pack), qi_incld("qi_incld", nj, nk_pack), qm_incld("qm_incld", nj, nk_pack),
// In cloud number concentrations
nc_incld("nc_incld", nj, nk_pack), nr_incld("nr_incld", nj, nk_pack), ni_incld("ni_incld", nj, nk_pack), bm_incld("bm_incld", nj, nk_pack),

// Other
inv_dz("inv_dz", nj, nk_pack), inv_rho("inv_rho", nj, nk_pack), ze_ice("ze_ice", nj, nk_pack), ze_rain("ze_rain", nj, nk_pack),
prec("prec", nj, nk_pack), rho("rho", nj, nk_pack), rhofacr("rhofacr", nj, nk_pack), rhofaci("rhofaci", nj, nk_pack),
acn("acn", nj, nk_pack), qv_sat_l("qv_sat", nj, nk_pack), qv_sat_i("qv_sat_i", nj, nk_pack), sup("sup", nj, nk_pack),
qv_supersat_i("qv_supersat", nj, nk_pack), tmparr2("tmparr2", nj, nk_pack), exner("exner", nj, nk_pack),
diag_equiv_reflectivity("diag_equiv_ref", nj, nk_pack), diag_vm_qi("diag_vm_qi", nj, nk_pack), diag_diam_qi("diag_diam_qi", nj, nk_pack),
pratot("pratot", nj, nk_pack), prctot("prctot", nj, nk_pack),

// p3_tend_out, may not need these
qtend_ignore("qtend_ignore", nj, nk_pack), ntend_ignore("ntend_ignore", nj, nk_pack),

// Variables still used in F90 but removed from C++ interface
mu_c("mu_c", nj, nk_pack), lamc("lamc", nj, nk_pack),
qr_evap_tend("qr_evap_tend", nj, nk_pack),

// cloud sedimentation
v_qc("v_qc", nj, nk_pack), v_nc("v_nc", nj, nk_pack), flux_qx("flux_qx", nj, nk_pack), flux_nx("flux_nx", nj, nk_pack),

// ice sedimentation
v_qit("v_qit", nj, nk_pack), v_nit("v_nit", nj, nk_pack), flux_nit("flux_nit", nj, nk_pack), flux_bir("flux_bir", nj, nk_pack),
flux_qir("flux_qir", nj, nk_pack), flux_qit("flux_qit", nj, nk_pack),

// rain sedimentation
v_qr("v_qr", nj, nk_pack), v_nr("v_nr", nj, nk_pack);

// Get views of all inputs
auto pres = diagnostic_inputs.pres;
auto dz = diagnostic_inputs.dz;
auto nc_nuceat_tend = diagnostic_inputs.nc_nuceat_tend;
auto nccn_prescribed = diagnostic_inputs.nccn;
auto ni_activated = diagnostic_inputs.ni_activated;
auto inv_qc_relvar = diagnostic_inputs.inv_qc_relvar;
auto dpres = diagnostic_inputs.dpres;
auto inv_exner = diagnostic_inputs.inv_exner;
auto cld_frac_i = diagnostic_inputs.cld_frac_i;
auto cld_frac_l = diagnostic_inputs.cld_frac_l;
auto cld_frac_r = diagnostic_inputs.cld_frac_r;
auto col_location = infrastructure.col_location;
auto qc = prognostic_state.qc;
auto nc = prognostic_state.nc;
auto qr = prognostic_state.qr;
auto nr = prognostic_state.nr;
auto qi = prognostic_state.qi;
auto qm = prognostic_state.qm;
auto ni = prognostic_state.ni;
auto bm = prognostic_state.bm;
auto qv = prognostic_state.qv;
auto th = prognostic_state.th;
auto diag_eff_radius_qc = diagnostic_outputs.diag_eff_radius_qc;
auto diag_eff_radius_qi = diagnostic_outputs.diag_eff_radius_qi;
auto diag_eff_radius_qr = diagnostic_outputs.diag_eff_radius_qr;
auto qv2qi_depos_tend = diagnostic_outputs.qv2qi_depos_tend;
auto rho_qi = diagnostic_outputs.rho_qi;
auto precip_liq_flux = diagnostic_outputs.precip_liq_flux;
auto precip_ice_flux = diagnostic_outputs.precip_ice_flux;
auto precip_total_tend = diagnostic_outputs.precip_total_tend;
auto nevapr = diagnostic_outputs.nevapr;
auto qv_prev = diagnostic_inputs.qv_prev;
auto t_prev = diagnostic_inputs.t_prev;
auto liq_ice_exchange = history_only.liq_ice_exchange;
auto vap_liq_exchange = history_only.vap_liq_exchange;
auto vap_ice_exchange = history_only.vap_ice_exchange;
auto pres = diagnostic_inputs.pres;
auto dz = diagnostic_inputs.dz;
auto nc_nuceat_tend = diagnostic_inputs.nc_nuceat_tend;
auto nccn_prescribed = diagnostic_inputs.nccn;
auto ni_activated = diagnostic_inputs.ni_activated;
auto inv_qc_relvar = diagnostic_inputs.inv_qc_relvar;
auto dpres = diagnostic_inputs.dpres;
auto inv_exner = diagnostic_inputs.inv_exner;
auto cld_frac_i = diagnostic_inputs.cld_frac_i;
auto cld_frac_l = diagnostic_inputs.cld_frac_l;
auto cld_frac_r = diagnostic_inputs.cld_frac_r;
auto col_location = infrastructure.col_location;
auto qc = prognostic_state.qc;
auto nc = prognostic_state.nc;
auto qr = prognostic_state.qr;
auto nr = prognostic_state.nr;
auto qi = prognostic_state.qi;
auto qm = prognostic_state.qm;
auto ni = prognostic_state.ni;
auto bm = prognostic_state.bm;
auto qv = prognostic_state.qv;
auto th = prognostic_state.th;
auto diag_eff_radius_qc = diagnostic_outputs.diag_eff_radius_qc;
auto diag_eff_radius_qi = diagnostic_outputs.diag_eff_radius_qi;
auto diag_eff_radius_qr = diagnostic_outputs.diag_eff_radius_qr;
auto qv2qi_depos_tend = diagnostic_outputs.qv2qi_depos_tend;
auto rho_qi = diagnostic_outputs.rho_qi;
auto precip_liq_flux = diagnostic_outputs.precip_liq_flux;
auto precip_ice_flux = diagnostic_outputs.precip_ice_flux;
auto precip_total_tend = diagnostic_outputs.precip_total_tend;
auto nevapr = diagnostic_outputs.nevapr;
auto qv_prev = diagnostic_inputs.qv_prev;
auto t_prev = diagnostic_inputs.t_prev;
auto liq_ice_exchange = history_only.liq_ice_exchange;
auto vap_liq_exchange = history_only.vap_liq_exchange;
auto vap_ice_exchange = history_only.vap_ice_exchange;
auto mu_r = temporaries.mu_r;
auto T_atm = temporaries.T_atm;
auto lamr = temporaries.lamr;
auto logn0r = temporaries.logn0r;
auto nu = temporaries.nu;
auto cdist = temporaries.cdist;
auto cdist1 = temporaries.cdist1;
auto cdistr = temporaries.cdistr;
auto inv_cld_frac_i = temporaries.inv_cld_frac_i;
auto inv_cld_frac_l = temporaries.inv_cld_frac_l;
auto inv_cld_frac_r = temporaries.inv_cld_frac_r;
auto qc_incld = temporaries.qc_incld;
auto qr_incld = temporaries.qr_incld;
auto qi_incld = temporaries.qi_incld;
auto qm_incld = temporaries.qm_incld;
auto nc_incld = temporaries.nc_incld;
auto nr_incld = temporaries.nr_incld;
auto ni_incld = temporaries.ni_incld;
auto bm_incld = temporaries.bm_incld;
auto inv_dz = temporaries.inv_dz;
auto inv_rho = temporaries.inv_rho;
auto ze_ice = temporaries.ze_ice;
auto ze_rain = temporaries.ze_rain;
auto prec = temporaries.prec;
auto rho = temporaries.rho;
auto rhofacr = temporaries.rhofacr;
auto rhofaci = temporaries.rhofaci;
auto acn = temporaries.acn;
auto qv_sat_l = temporaries.qv_sat_l;
auto qv_sat_i = temporaries.qv_sat_i;
auto sup = temporaries.sup;
auto qv_supersat_i = temporaries.qv_supersat_i;
auto tmparr2 = temporaries.tmparr2;
auto exner = temporaries.exner;
auto diag_equiv_reflectivity = temporaries.diag_equiv_reflectivity;
auto diag_vm_qi = temporaries.diag_vm_qi;
auto diag_diam_qi = temporaries.diag_diam_qi;
auto pratot = temporaries.pratot;
auto prctot = temporaries.prctot;
auto qtend_ignore = temporaries.qtend_ignore;
auto ntend_ignore = temporaries.ntend_ignore;
auto mu_c = temporaries.mu_c;
auto lamc = temporaries.lamc;
auto qr_evap_tend = temporaries.qr_evap_tend;
auto v_qc = temporaries.v_qc;
auto v_nc = temporaries.v_nc;
auto flux_qx = temporaries.flux_qx;
auto flux_nx = temporaries.flux_nx;
auto v_qit = temporaries.v_qit;
auto v_nit = temporaries.v_nit;
auto flux_nit = temporaries.flux_nit;
auto flux_bir = temporaries.flux_bir;
auto flux_qir = temporaries.flux_qir;
auto flux_qit = temporaries.flux_qit;
auto v_qr = temporaries.v_qr;
auto v_nr = temporaries.v_nr;

// we do not want to measure init stuff
auto start = std::chrono::steady_clock::now();
Expand All @@ -231,7 +240,7 @@ ::p3_main_internal_disp(
p3_main_part1_disp(
nj, nk, infrastructure.predictNc, infrastructure.prescribedCCN, infrastructure.dt,
pres, dpres, dz, nc_nuceat_tend, nccn_prescribed, inv_exner, exner, inv_cld_frac_l, inv_cld_frac_i,
inv_cld_frac_r, latent_heat_vapor, latent_heat_sublim, latent_heat_fusion,
inv_cld_frac_r,
T_atm, rho, inv_rho, qv_sat_l, qv_sat_i, qv_supersat_i, rhofacr,
rhofaci, acn, qv, th, qc, nc, qr, nr, qi, ni, qm,
bm, qc_incld, qr_incld, qi_incld, qm_incld, nc_incld, nr_incld,
Expand All @@ -242,12 +251,11 @@ ::p3_main_internal_disp(

p3_main_part2_disp(
nj, nk, runtime_options.max_total_ni, infrastructure.predictNc, infrastructure.prescribedCCN, infrastructure.dt, inv_dt,
lookup_tables.dnu_table_vals, lookup_tables.ice_table_vals, lookup_tables.collect_table_vals,
lookup_tables.dnu_table_vals, lookup_tables.ice_table_vals, lookup_tables.collect_table_vals,
lookup_tables.revap_table_vals, pres, dpres, dz, nc_nuceat_tend, inv_exner,
exner, inv_cld_frac_l, inv_cld_frac_i, inv_cld_frac_r, ni_activated, inv_qc_relvar, cld_frac_i,
cld_frac_l, cld_frac_r, qv_prev, t_prev, T_atm, rho, inv_rho, qv_sat_l, qv_sat_i, qv_supersat_i, rhofacr, rhofaci, acn,
qv, th, qc, nc, qr, nr, qi, ni, qm, bm, latent_heat_vapor,
latent_heat_sublim, latent_heat_fusion, qc_incld, qr_incld, qi_incld, qm_incld, nc_incld,
qv, th, qc, nc, qr, nr, qi, ni, qm, bm, qc_incld, qr_incld, qi_incld, qm_incld, nc_incld,
nr_incld, ni_incld, bm_incld, mu_c, nu, lamc, cdist, cdist1, cdistr,
mu_r, lamr, logn0r, qv2qi_depos_tend, precip_total_tend, nevapr, qr_evap_tend,
vap_liq_exchange, vap_ice_exchange, liq_ice_exchange,
Expand Down Expand Up @@ -288,7 +296,7 @@ ::p3_main_internal_disp(

// homogeneous freezing f cloud and rain
homogeneous_freezing_disp(
T_atm, inv_exner, latent_heat_fusion, nj, nk, ktop, kbot, kdir, qc, nc, qr, nr, qi,
T_atm, inv_exner, nj, nk, ktop, kbot, kdir, qc, nc, qr, nr, qi,
ni, qm, bm, th, nucleationPossible, hydrometeorsPresent);

//
Expand All @@ -298,7 +306,7 @@ ::p3_main_internal_disp(
p3_main_part3_disp(
nj, nk_pack, runtime_options.max_total_ni, lookup_tables.dnu_table_vals, lookup_tables.ice_table_vals, inv_exner, cld_frac_l, cld_frac_r, cld_frac_i,
rho, inv_rho, rhofaci, qv, th, qc, nc, qr, nr, qi, ni,
qm, bm, latent_heat_vapor, latent_heat_sublim, mu_c, nu, lamc, mu_r, lamr,
qm, bm, mu_c, nu, lamc, mu_r, lamr,
vap_liq_exchange, ze_rain, ze_ice, diag_vm_qi, diag_eff_radius_qi, diag_diam_qi,
rho_qi, diag_equiv_reflectivity, diag_eff_radius_qc, diag_eff_radius_qr, nucleationPossible, hydrometeorsPresent,
p3constants);
Expand Down
11 changes: 4 additions & 7 deletions components/eamxx/src/physics/p3/disp/p3_main_impl_part1_disp.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -31,9 +31,6 @@ ::p3_main_part1_disp(
const uview_2d<const Spack>& inv_cld_frac_l,
const uview_2d<const Spack>& inv_cld_frac_i,
const uview_2d<const Spack>& inv_cld_frac_r,
const uview_2d<const Spack>& latent_heat_vapor,
const uview_2d<const Spack>& latent_heat_sublim,
const uview_2d<const Spack>& latent_heat_fusion,
const uview_2d<Spack>& T_atm,
const uview_2d<Spack>& rho,
const uview_2d<Spack>& inv_rho,
Expand Down Expand Up @@ -73,17 +70,17 @@ ::p3_main_part1_disp(
policy, KOKKOS_LAMBDA(const MemberType& team) {

const Int i = team.league_rank();

p3_main_part1(
team, nk, predictNc, prescribedCCN, dt,
ekat::subview(pres, i), ekat::subview(dpres, i), ekat::subview(dz, i), ekat::subview(nc_nuceat_tend, i),
ekat::subview(nccn_prescribed, i), ekat::subview(inv_exner, i), ekat::subview(exner, i), ekat::subview(inv_cld_frac_l, i),
ekat::subview(inv_cld_frac_i, i), ekat::subview(inv_cld_frac_r, i), ekat::subview(latent_heat_vapor, i), ekat::subview(latent_heat_sublim, i),
ekat::subview(latent_heat_fusion, i), ekat::subview(T_atm, i), ekat::subview(rho, i), ekat::subview(inv_rho, i), ekat::subview(qv_sat_l, i),
ekat::subview(inv_cld_frac_i, i), ekat::subview(inv_cld_frac_r, i),
ekat::subview(T_atm, i), ekat::subview(rho, i), ekat::subview(inv_rho, i), ekat::subview(qv_sat_l, i),
ekat::subview(qv_sat_i, i), ekat::subview(qv_supersat_i, i), ekat::subview(rhofacr, i), ekat::subview(rhofaci, i), ekat::subview(acn, i),
ekat::subview(qv, i), ekat::subview(th_atm, i), ekat::subview(qc, i), ekat::subview(nc, i), ekat::subview(qr, i), ekat::subview(nr, i), ekat::subview(qi, i),
ekat::subview(ni, i), ekat::subview(qm, i), ekat::subview(bm, i), ekat::subview(qc_incld, i), ekat::subview(qr_incld, i), ekat::subview(qi_incld, i),
ekat::subview(qm_incld, i), ekat::subview(nc_incld, i), ekat::subview(nr_incld, i), ekat::subview(ni_incld, i), ekat::subview(bm_incld, i),
ekat::subview(qm_incld, i), ekat::subview(nc_incld, i), ekat::subview(nr_incld, i), ekat::subview(ni_incld, i), ekat::subview(bm_incld, i),
nucleationPossible(i), hydrometeorsPresent(i), p3constants);

});
Expand Down
Loading

0 comments on commit e65666f

Please sign in to comment.