Replace a single loop, and a few changes are needed for the small kernel PR

src/mam4xx/gas_chem.hpp

@@ -77,21 +77,18 @@ void imp_slv_inti(Real epsilon[clscnt4]) {
     epsilon[i] = rel_err;
   }
 }
-
-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],
-                         Real solution[clscnt4],                     // inout
-                         bool converged[clscnt4], bool &convergence, // out
-                         Real prod[clscnt4], Real loss[clscnt4],
-                         Real max_delta[clscnt4],
-                         // work array
-                         Real epsilon[clscnt4]) {
+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], VectorType &lsol,
+    Real solution[clscnt4],                     // inout
+    bool converged[clscnt4], bool &convergence, // out
+    Real prod[clscnt4], Real loss[clscnt4], Real max_delta[clscnt4],
+    // work array
+    Real epsilon[clscnt4]) {
 
   // dti := 1 / dt
   // lrxt := reaction rates in 1D array [1/cm^3/s]
@@ -239,14 +236,14 @@ 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
-
-KOKKOS_INLINE_FUNCTION
-void imp_sol(Real base_sol[gas_pcnst], // 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],
-             const bool factor[itermax], Real epsilon[clscnt4],
-             Real prod_out[clscnt4], Real loss_out[clscnt4]) {
+template <typename VectorType>
+KOKKOS_INLINE_FUNCTION 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],
+        const bool factor[itermax], Real epsilon[clscnt4],
+        Real prod_out[clscnt4], Real loss_out[clscnt4]) {
 
   // ---------------------------------------------------------------------------
   //  ... imp_sol advances the volumetric mixing ratio

src/mam4xx/gas_chem_mechanism.hpp

@@ -84,9 +84,10 @@ 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
-KOKKOS_INLINE_FUNCTION
-void imp_prod_loss(Real prod[clscnt4], Real loss[clscnt4], Real y[gas_pcnst],
-                   const Real rxt[rxntot], const Real het_rates[gas_pcnst]) {
+template <typename VectorType>
+KOKKOS_INLINE_FUNCTION 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]);
   prod[0] = zero;

src/mam4xx/gas_phase_chemistry.hpp

@@ -33,15 +33,15 @@ using mam4::gas_chemistry::indexm;
 
 // performs gas phase chemistry calculations on a single level of a single
 // atmospheric column
-KOKKOS_INLINE_FUNCTION
-void gas_phase_chemistry(
+template <typename VectorType>
+KOKKOS_INLINE_FUNCTION void gas_phase_chemistry(
     // in
     const Real temp, const Real dt,
     const Real photo_rates[mam4::mo_photo::phtcnt], const Real extfrc[extcnt],
     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,14 @@ void get_gcm_tend_diags_from_subareas(
 
 //--------------------------------------------------------------------------------
 //--------------------------------------------------------------------------------
-
-KOKKOS_INLINE_FUNCTION
-void modal_aero_amicphys_intr(
+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 +2129,9 @@ 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

@@ -33,8 +33,7 @@ KOKKOS_INLINE_FUNCTION
 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 int offset_aerosol, 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 +49,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 +203,8 @@ 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);
+  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,15 +173,14 @@ 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));
   });
 }
 

src/mam4xx/mo_sethet.hpp

@@ -204,18 +204,18 @@ int get_total_work_len_sethet() {
 KOKKOS_INLINE_FUNCTION
 void sethet_detail(
     const ThreadTeam &team,
-    const View2D &het_rates,         // rainout rates [1/s] //out
-    const Real rlat,                 // latitude in radians for columns
-    const ConstColumnView &press,    // pressure [pascals] //in
-    const ConstColumnView &zmid,     // midpoint geopot [km]  //in
-    const Real phis,                 // surf geopotential //in
-    const ConstColumnView &tfld,     // temperature [K]  //in
-    const ColumnView &cmfdqr,        // dq/dt for convection [kg/kg/s] //in
-    const ConstColumnView &nrain,    // stratoform precip [kg/kg/s] //in
-    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 &het_rates,       // rainout rates [1/s] //out
+    const Real rlat,               // latitude in radians for columns
+    const ConstColumnView &press,  // pressure [pascals] //in
+    const ConstColumnView &zmid,   // midpoint geopot [km]  //in
+    const Real phis,               // surf geopotential //in
+    const ConstColumnView &tfld,   // temperature [K]  //in
+    const ColumnView &cmfdqr,      // dq/dt for convection [kg/kg/s] //in
+    const ConstColumnView &nrain,  // stratoform precip [kg/kg/s] //in
+    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 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;
 
@@ -498,10 +498,10 @@ void sethet(
     const Real phis,              // surf geopotential //in
     const ColumnView &cmfdqr,     // dq/dt for convection [kg/kg/s] //in
     const ConstColumnView &prain, // stratoform precip [kg/kg/s] //in
-    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 ConstColumnView &nevapr, // evaporation [kg/kg/s] //in
+    const Real dt,                 // time step [s] //in
+    const View2D &invariants,      //
+    const View2D &vmr,             // xported species [vmr]  //in
     // working variables
     const View1D &work) {
 

src/validation/gas_chem/imp_prod_loss.cpp

@@ -25,8 +25,7 @@ void imp_prod_loss(Ensemble *ensemble) {
     const auto het_rates = input.get_array("het_rates");
     auto y = input.get_array("y");
 
-    imp_prod_loss(prod.data(), loss.data(), y.data(), rxt.data(),
-                  het_rates.data());
+    imp_prod_loss(prod.data(), loss.data(), y, rxt.data(), het_rates.data());
     output.set("prod", prod);
     output.set("loss", loss);
   });

src/validation/gas_chem/imp_sol.cpp

@@ -35,7 +35,7 @@ void imp_sol(Ensemble *ensemble) {
       factor[i] = true;
     }
 
-    imp_sol(base_sol.data(), //    ! species mixing ratios [vmr] & !
+    imp_sol(base_sol, //    ! species mixing ratios [vmr] & !
             reaction_rates.data(), het_rates.data(), extfrc.data(), delt,
             permute_4, clsmap_4, factor, epsilon, prod_out.data(),
             loss_out.data());

src/validation/gas_chem/newton_raphson_iter.cpp

@@ -40,7 +40,7 @@ void newton_raphson_iter(Ensemble *ensemble) {
     newton_raphson_iter(dti, lin_jac.data(), lrxt.data(),
                         lhet.data(),           // & ! in
                         iter_invariant.data(), //              & ! in
-                        factor, permute_4, clsmap_4, lsol.data(),
+                        factor, permute_4, clsmap_4, lsol,
                         solution.data(),        //              & ! inout
                         converged, convergence, //         & ! out
                         prod.data(), loss.data(), max_delta.data(),

src/validation/mo_sethet/sethet.cpp

@@ -17,6 +17,7 @@ using namespace mo_sethet;
 void sethet(Ensemble *ensemble) {
   ensemble->process([=](const Input &input, Output &output) {
     using View1DHost = typename HostType::view_1d<Real>;
+    using View2DHost = typename HostType::view_2d<Real>;
     using ColumnView = haero::ColumnView;
     constexpr int pver = mam4::nlev;
     constexpr int gas_pcnst = mam4::gas_chemistry::gas_pcnst;
@@ -101,34 +102,32 @@ void sethet(Ensemble *ensemble) {
     so2_diss = haero::testing::create_column_view(pver);
 
     ColumnView tmp_hetrates[gas_pcnst];
-    ColumnView qin[gas_pcnst];
     View1DHost tmp_hetrates_host[gas_pcnst];
-    View1DHost qin_host[gas_pcnst];
+    View2DHost qin_host("qin_host", pver, gas_pcnst);
 
     View2D het_rates("het_rates", pver, gas_pcnst);
+    View2D qin("qin", pver, gas_pcnst);
     auto het_rates_host = Kokkos::create_mirror_view(het_rates);
 
     for (int mm = 0; mm < gas_pcnst; ++mm) {
 
       tmp_hetrates[mm] = haero::testing::create_column_view(pver);
-      qin[mm] = haero::testing::create_column_view(pver);
       tmp_hetrates_host[mm] = View1DHost("tmp_hetrates_host", pver);
-      qin_host[mm] = View1DHost("qin_host", pver);
     }
 
     int count = 0;
     for (int mm = 0; mm < gas_pcnst; ++mm) {
       for (int kk = 0; kk < pver; ++kk) {
-        qin_host[mm](kk) = qin_in[count];
+        qin_host(kk, mm) = qin_in[count];
         count++;
       }
     }
 
     // transfer data to GPU.
     for (int mm = 0; mm < gas_pcnst; ++mm) {
       Kokkos::deep_copy(tmp_hetrates[mm], 0.0);
-      Kokkos::deep_copy(qin[mm], qin_host[mm]);
     }
+    Kokkos::deep_copy(qin, qin_host);
 
     auto team_policy = ThreadTeamPolicy(1u, Kokkos::AUTO);
     Kokkos::parallel_for(