Skip to content

Commit

Permalink
MAM4xx: Implement find_season_index using RangePolicy.
Browse files Browse the repository at this point in the history
  • Loading branch information
odiazib committed Nov 13, 2024
1 parent df0e17c commit d6cb4c0
Show file tree
Hide file tree
Showing 2 changed files with 27 additions and 5 deletions.
20 changes: 19 additions & 1 deletion src/mam4xx/mo_drydep.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
#ifndef MAM4XX_MO_DRYDEP_HPP
#define MAM4XX_MO_DRYDEP_HPP

#include <ekat/kokkos/ekat_subview_utils.hpp>
#include <haero/math.hpp>
#include <mam4xx/aero_config.hpp>
#include <mam4xx/gas_chem_mechanism.hpp>
Expand Down Expand Up @@ -43,10 +44,14 @@ constexpr int nddvels = mam4::seq_drydep::n_drydep;

// find_season_index only needs to be executed one time and is small.
// Thus, execute it only on the host.

using View1DHost = DeviceType::view_1d<Real>::HostMirror;
using ConstView1DHost = DeviceType::view_1d<const Real>::HostMirror;
using View1DIntHost = DeviceType::view_1d<int>::HostMirror;
using View2DIntHost = DeviceType::view_2d<int>::HostMirror;
using View3DIntHost = DeviceType::view_3d<int>::HostMirror;

using KTH = ekat::KokkosTypes<ekat::HostDevice>;

inline void find_season_index(const Real clat_j, const View1DHost &lat_lai,
const int nlat_lai, const View3DIntHost &wk_lai,
const View1DIntHost &index_season_lai) {
Expand Down Expand Up @@ -100,6 +105,19 @@ inline void find_season_index(const Real clat_j, const View1DHost &lat_lai,
} // m
} // findSeasonIndex

inline void find_season_index(const ConstView1DHost clat,
const View1DHost &lat_lai, const int nlat_lai,
const View3DIntHost &wk_lai,
const View2DIntHost &index_season_lai) {
const int plon = clat.extent(0);
auto policy = KTH::RangePolicy(0, plon);
Kokkos::parallel_for(policy, [&](const int &j) {
const auto index_season_lai_at_j = ekat::subview(index_season_lai, j);
mo_drydep::find_season_index(clat(j), lat_lai, nlat_lai, wk_lai,
index_season_lai_at_j);
});
}

KOKKOS_INLINE_FUNCTION
void calculate_uustar(
const seq_drydep::Data &drydep_data, const int index_season[n_land_type],
Expand Down
12 changes: 8 additions & 4 deletions src/validation/mo_drydep/find_season_index.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -48,13 +48,17 @@ void find_season_index(Ensemble *ensemble) {

View2DIntHost index_season_lai("index_season_lai", plon, 12);

// convert to radians
for (int j = 0; j < plon; ++j) {
clat(j) *= r2d;
}

mo_drydep::find_season_index(clat, lat_lai, nlat_lai, wk_lai,
index_season_lai);

auto policy = KTH::RangePolicy(0, plon);
Kokkos::parallel_for(policy, [&](const int &j) {
const auto index_season_lai_at_j = ekat::subview(index_season_lai, j);
// convert to radians
const Real clat_rads = clat(j) * r2d;
mo_drydep::find_season_index(clat_rads, lat_lai, nlat_lai, wk_lai,
index_season_lai_at_j);
// c++ to Fortran; only for validation.
for (int i = 0; i < 12; ++i) {
index_season_lai_at_j(i) = index_season_lai_at_j(i) + 1;
Expand Down

0 comments on commit d6cb4c0

Please sign in to comment.