From d6cb4c004a686664dc9d5885837f69c46be2f611 Mon Sep 17 00:00:00 2001
From: Oscar Diaz-Ibarra <odiazib@sandia.gov>
Date: Tue, 12 Nov 2024 17:42:32 -0700
Subject: [PATCH] MAM4xx: Implement find_season_index using RangePolicy.

---
 src/mam4xx/mo_drydep.hpp                      | 20 ++++++++++++++++++-
 .../mo_drydep/find_season_index.cpp           | 12 +++++++----
 2 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/src/mam4xx/mo_drydep.hpp b/src/mam4xx/mo_drydep.hpp
index 141b1077..d7c18891 100644
--- a/src/mam4xx/mo_drydep.hpp
+++ b/src/mam4xx/mo_drydep.hpp
@@ -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>
@@ -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) {
@@ -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],
diff --git a/src/validation/mo_drydep/find_season_index.cpp b/src/validation/mo_drydep/find_season_index.cpp
index b189f89d..286fdb1e 100644
--- a/src/validation/mo_drydep/find_season_index.cpp
+++ b/src/validation/mo_drydep/find_season_index.cpp
@@ -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;