Replace arrays of ColumnView with a 2D view. Use template parameters instead of raw pointers. Replace the single loop with two nested parallel loops.

Author: odiazib

src/mam4xx/gas_chem.hpp

@@ -77,15 +77,15 @@ void imp_slv_inti(Real epsilon[clscnt4]) {
     epsilon[i] = rel_err;
   }
 }
-
+template <typename VectorType>
 KOKKOS_INLINE_FUNCTION
 void newton_raphson_iter(const Real dti, const Real lin_jac[nzcnt],
                          const Real lrxt[rxntot],
                          const Real lhet[gas_pcnst],         // in
                          const Real iter_invariant[clscnt4], // in
                          const bool factor[itermax],
                          const int permute_4[gas_pcnst],
-                         const int clsmap_4[gas_pcnst], Real lsol[gas_pcnst],
+                         const int clsmap_4[gas_pcnst], VectorType& lsol,
                          Real solution[clscnt4],                     // inout
                          bool converged[clscnt4], bool &convergence, // out
                          Real prod[clscnt4], Real loss[clscnt4],
@@ -239,9 +239,9 @@ void newton_raphson_iter(const Real dti, const Real lin_jac[nzcnt],
     } // end if (nr_iter > 0)
   }   // end nr_iter loop
 } // newton_raphson_iter() function
-
+template <typename VectorType>
 KOKKOS_INLINE_FUNCTION
-void imp_sol(Real base_sol[gas_pcnst], // inout - species mixing ratios [vmr]
+void imp_sol(VectorType& base_sol, // inout - species mixing ratios [vmr]
              const Real reaction_rates[rxntot], const Real het_rates[gas_pcnst],
              const Real extfrc[extcnt], const Real &delt,
              const int permute_4[gas_pcnst], const int clsmap_4[gas_pcnst],

src/mam4xx/gas_chem_mechanism.hpp

@@ -84,8 +84,9 @@ void adjrxt(Real rate[rxntot], const Real inv[nfs], const Real m) {
 // TODO: the lines of concern *kind of* bear resemblance to the similarly
 // concerning lines in linmat(), though it's difficult to tell if that results
 // in consistent units
+template <typename VectorType>
 KOKKOS_INLINE_FUNCTION
-void imp_prod_loss(Real prod[clscnt4], Real loss[clscnt4], Real y[gas_pcnst],
+void imp_prod_loss(Real prod[clscnt4], Real loss[clscnt4], VectorType& y,
                    const Real rxt[rxntot], const Real het_rates[gas_pcnst]) {
   const Real zero = 0;
   loss[0] = (+het_rates[1] + rxt[0] + rxt[2]) * (+y[1]);

src/mam4xx/gas_phase_chemistry.hpp

@@ -33,6 +33,7 @@ using mam4::gas_chemistry::indexm;
 
 // performs gas phase chemistry calculations on a single level of a single
 // atmospheric column
+template <typename VectorType>
 KOKKOS_INLINE_FUNCTION
 void gas_phase_chemistry(
     // in
@@ -41,7 +42,7 @@ void gas_phase_chemistry(
     const Real invariants[nfs], const int (&clsmap_4)[gas_pcnst],
     const int (&permute_4)[gas_pcnst], const Real het_rates[gas_pcnst],
     // out
-    Real (&qq)[gas_pcnst]) {
+    VectorType &qq) {
   //=====================================================================
   // ... set rates for "tabular" and user specified reactions
   //=====================================================================

src/mam4xx/mam4_amicphys.hpp

@@ -2113,15 +2113,15 @@ void get_gcm_tend_diags_from_subareas(
 
 //--------------------------------------------------------------------------------
 //--------------------------------------------------------------------------------
-
+template<typename VectorType, typename VectorTypeModes>
 KOKKOS_INLINE_FUNCTION
 void modal_aero_amicphys_intr(
     // in
     const AmicPhysConfig &config, const Real deltat, const Real temp,
     const Real pmid, const Real pdel, const Real zm, const Real pblh,
     const Real qv, const Real cld,
     // in/out
-    Real (&qq)[gas_pcnst], Real (&qqcw)[gas_pcnst],
+    VectorType &qq, VectorType &qqcw,
     // Diagnostics (out)
     const int kk, // level info needed for diagnistics output
     const View2D &gas_aero_exchange_condensation,
@@ -2130,11 +2130,12 @@ void modal_aero_amicphys_intr(
     const View2D &gas_aero_exchange_coagulation,
     const View2D &gas_aero_exchange_renaming_cloud_borne,
     // in
-    const Real (&q_pregaschem)[gas_pcnst],
-    const Real (&q_precldchem)[gas_pcnst],
-    const Real (&qqcw_precldchem)[gas_pcnst], const Real (&dgncur_a)[num_modes],
-    const Real (&dgncur_awet)[num_modes],
-    const Real (&wetdens_host)[num_modes]) {
+    const VectorType& q_pregaschem,
+    const VectorType& q_precldchem,
+    const VectorType& qqcw_precldchem,
+    const VectorTypeModes& dgncur_a,
+    const VectorTypeModes& dgncur_awet,
+    const VectorTypeModes& wetdens_host) {
   // deltat: time step
   // qq(ncol,pver,pcnst): current tracer mixing ratios (TMRs)
   //                           these values are updated (so out /= in)

src/mam4xx/mo_gas_phase_chemdr.hpp

@@ -34,7 +34,7 @@ void mmr2vmr_col(const ThreadTeam &team, const haero::Atmosphere &atm,
                  const mam4::Prognostics &progs,
                  const Real adv_mass_kg_per_moles[gas_pcnst],
                  const int offset_aerosol,
-                 const ColumnView vmr_col[gas_pcnst]) {
+                 const View2D vmr_col) {
   // Make a local copy of nlev to avoid the identifier "mam4::nlev" being
   // undefined in device code.
   constexpr int nlev_local = nlev;
@@ -50,7 +50,7 @@ void mmr2vmr_col(const ThreadTeam &team, const haero::Atmosphere &atm,
         // output (vmr)
         mam4::microphysics::mmr2vmr(qq, adv_mass_kg_per_moles, vmr);
         for (int i = 0; i < gas_pcnst; ++i) {
-          vmr_col[i](kk) = vmr[i];
+          vmr_col(kk,i) = vmr[i];
         }
       });
 }
@@ -204,11 +204,7 @@ void perform_atmospheric_chemistry_and_microphysics(
   work_set_het_ptr += nlev * gas_pcnst;
 
   // vmr0 stores mixing ratios before chemistry changes the mixing
-  ColumnView vmr_col[gas_pcnst];
-  for (int i = 0; i < gas_pcnst; ++i) {
-    vmr_col[i] = ColumnView(work_set_het_ptr, nlev);
-    work_set_het_ptr += nlev;
-  }
+  const auto vmr_col = View2D(work_set_het_ptr, nlev, gas_pcnst);
   const int sethet_work_len = mam4::mo_sethet::get_work_len_sethet();
   const auto work_sethet_call = View1D(work_set_het_ptr, sethet_work_len);
   work_set_het_ptr += sethet_work_len;

src/mam4xx/mo_photo.hpp

@@ -173,16 +173,17 @@ void set_ub_col(Real &o3_col_delta,
 KOKKOS_INLINE_FUNCTION
 void setcol(const ThreadTeam &team, const Real o3_col_deltas[mam4::nlev + 1],
             ColumnView &o3_col_dens) {
-  // we can probably accelerate this with a parallel_scan, but let's just do
-  // a simple loop for now
   constexpr int nlev = mam4::nlev;
-  Kokkos::single(Kokkos::PerTeam(team), [=]() {
-    o3_col_dens(0) = 0.5 * (o3_col_deltas[0] + o3_col_deltas[1]);
-    for (int k = 1; k < nlev; ++k) {
-      o3_col_dens(k) =
-          o3_col_dens(k - 1) + 0.5 * (o3_col_deltas[k] + o3_col_deltas[k + 1]);
-    }
-  });
+  Kokkos::parallel_for(
+    Kokkos::TeamThreadRange(team, nlev),
+    [&](int kk) {
+      Kokkos::parallel_reduce(
+          Kokkos::ThreadVectorRange(team, kk+1),
+          [&](int i, Real &lsum) {
+            lsum += 0.5 * (o3_col_deltas[i] + o3_col_deltas[i+1]);
+          },
+          o3_col_dens(kk));
+    });
 }
 
 KOKKOS_INLINE_FUNCTION

src/mam4xx/mo_sethet.hpp

@@ -215,7 +215,7 @@ void sethet_detail(
     const ConstColumnView &nevapr,   // evaporation [kg/kg/s] //in
     const Real delt,                 // time step [s] //in
     const View2D &invariants,        // total atms density [cm^-3] //in
-    const ColumnView qin[gas_pcnst], // xported species [vmr]  //in
+    const View2D qin, // xported species [vmr]  //in
     // working variables
     const ColumnView
         &t_factor, // temperature factor to calculate henry's law parameters
@@ -330,8 +330,8 @@ void sethet_detail(
     rain(kk) = mass_air * precip(kk) * invariants(kk, indexm) / mass_h2o;
     xliq(kk) =
         precip(kk) * delt * invariants(kk, indexm) / avo * mass_air * m3_2_cm3;
-    xh2o2(kk) = qin[spc_h2o2_ndx](kk) * invariants(kk, indexm);
-    xso2(kk) = qin[spc_so2_ndx](kk) * invariants(kk, indexm);
+    xh2o2(kk) = qin(kk,spc_h2o2_ndx) * invariants(kk, indexm);
+    xso2(kk) = qin(kk,spc_so2_ndx) * invariants(kk, indexm);
   });
   zsurf = m2km * phis * rga;
 
@@ -501,7 +501,7 @@ void sethet(
     const ConstColumnView &nevapr,   // evaporation [kg/kg/s] //in
     const Real dt,                   // time step [s] //in
     const View2D &invariants,        //
-    const ColumnView vmr[gas_pcnst], // xported species [vmr]  //in
+    const View2D& vmr, // xported species [vmr]  //in
     // working variables
     const View1D &work) {