Final: mo_sethet (#373)

* convert fortran of sethet to c++

* use pcnst value from modal_aer_opt

* polish sethet

* compile

* add parameters

* fix parens

* remove extra const

* fix includes

* code compiles and test runs (and fails)

* change gas_pcnst and tolerance

* format

* add comments

* fix gas_washout, indexing issues in sethet

* format

* fix indexes and validation

* remove kokkos for from gas_washout

* format and cleanup

* remove changes in mo_chm_diags

* cleanup var name [skip-ci]

* cleanup

* add barrier to stop intermittent race condition(?)

* add more parallelism in sethet and calc_precip_rescale

* format

* cleanup

* fix compile errors

* change loops to vector ranges

* change tolerance

* format

* add additional parallel fors

* format

* add parallel reduces to calc_precip_rescale

* Starting from the commit where tests are passing on GPUs and CPUs, we first switch from ThreadVectorRange to TeamVectorRange; remove unneeded initialization. Remove team and parallelism from gas_washout.

* Fix compilation error in CUDA.

* Remove parallel_for to fix CUDA sethet test.

* Remove team_barriers and combine parallels.

* Clang format.

* address review comments and format

* address bugs from reviews

---------

Co-authored-by: Oscar Diaz-Ibarra <[email protected]>

src/mam4xx/mo_chm_diags.hpp

@@ -352,4 +352,4 @@ void chm_diags(
 } // chm_diags
 } // namespace mo_chm_diags
 } // namespace mam4
-#endif
+#endif

src/validation/mo_sethet/CMakeLists.txt

@@ -14,6 +14,7 @@ add_executable(sethet_driver
                calc_precip_rescale.cpp
                find_ktop.cpp
                gas_washout.cpp
+               sethet.cpp
                )
 
 target_link_libraries(sethet_driver skywalker;validation;${HAERO_LIBRARIES})
@@ -27,22 +28,24 @@ foreach(script
     COPYONLY
   )
 endforeach()
-# stand_calc_sum_wght_ts_355
+
 # Run the driver in several configurations to produce datasets.
 set(TEST_LIST
     calc_het_rates_ts_355
     calc_precip_rescale_ts_355
     find_ktop_ts_355
     gas_washout_ts_355
+    sethet_ts_355
     )
 # # matching the tests and errors, just for convenience
 
 set(DEFAULT_TOL 1e-13)
 set(ERROR_THRESHOLDS
-    9e-8 #this is due to using a more specific value for mass_h2o
-    ${DEFAULT_TOL}
-    ${DEFAULT_TOL} # output is an int
-    9e-3
+    9e-8 # calc_het_rates, this is due to using a more specific value for mass_h2o
+    ${DEFAULT_TOL} # calc_precip_rescale
+    ${DEFAULT_TOL} # find_ktop, output is an int
+    9e-8 # gas_washout
+    9e-8 # sethet
    )
 foreach(input tol IN ZIP_LISTS TEST_LIST ERROR_THRESHOLDS)
   # copy the baseline file into place.

src/validation/mo_sethet/calc_precip_rescale.cpp

@@ -42,9 +42,10 @@ void calc_precip_rescale(Ensemble *ensemble) {
     Kokkos::deep_copy(precip, precip_host);
     // std::vector<Real> precip(pver, zero);
     DeviceType::view_1d<Real> trp_out_val("Return from Device", 1);
+    auto team_policy = ThreadTeamPolicy(1u, Kokkos::AUTO);
     Kokkos::parallel_for(
-        "calc_precip_rescale", 1, KOKKOS_LAMBDA(int i) {
-          calc_precip_rescale(cmfdqr, nrain, nevapr, precip);
+        team_policy, KOKKOS_LAMBDA(const ThreadTeam &team) {
+          calc_precip_rescale(team, cmfdqr, nrain, nevapr, precip);
         });
 
     Kokkos::deep_copy(precip_host, precip);

src/validation/mo_sethet/gas_washout.cpp

@@ -42,21 +42,11 @@ void gas_washout(Ensemble *ensemble) {
     Kokkos::deep_copy(delz_i, delz_i_host);
     Kokkos::deep_copy(xgas, xgas_host);
 
-    // initialize internal veriables
-    ColumnView xeqca, xca;
-    std::vector<Real> vector0(pver, 0);
-    auto xeqca_host = View1DHost(vector0.data(), pver);
-    auto xca_host = View1DHost(vector0.data(), pver);
-    xeqca = haero::testing::create_column_view(pver);
-    xca = haero::testing::create_column_view(pver);
-    Kokkos::deep_copy(xeqca, xeqca_host);
-    Kokkos::deep_copy(xca, xca_host);
-
     auto team_policy = ThreadTeamPolicy(1u, Kokkos::AUTO);
     Kokkos::parallel_for(
         team_policy, KOKKOS_LAMBDA(const ThreadTeam &team) {
           gas_washout(team, plev - 1, xkgm, xliq_ik, xhen_i, tfld_i, delz_i,
-                      xeqca, xca, xgas);
+                      xgas);
         });
 
     Kokkos::deep_copy(xgas_host, xgas);
@@ -66,4 +56,4 @@ void gas_washout(Ensemble *ensemble) {
 
     output.set("xgas", xgas_out);
   });
-}
+}

src/validation/mo_sethet/sethet.cpp

@@ -0,0 +1,156 @@
+// mam4xx: Copyright (c) 2022,
+// Battelle Memorial Institute and
+// National Technology & Engineering Solutions of Sandia, LLC (NTESS)
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include <mam4xx/mam4.hpp>
+
+#include <mam4xx/aero_config.hpp>
+#include <skywalker.hpp>
+#include <validation.hpp>
+
+using namespace skywalker;
+using namespace mam4;
+using namespace haero;
+using namespace mo_sethet;
+
+void sethet(Ensemble *ensemble) {
+  ensemble->process([=](const Input &input, Output &output) {
+    using View1DHost = typename HostType::view_1d<Real>;
+    using ColumnView = haero::ColumnView;
+    constexpr int pver = mam4::nlev;
+    constexpr int gas_pcnst = mam4::gas_chemistry::gas_pcnst;
+
+    // non-ColumnView input values
+    const Real rlat = input.get_array("rlat")[0]; // need
+    const Real phis = input.get_array("phis")[0];
+    const Real delt = input.get_array("delt")[0];
+
+    const int spc_h2o2_ndx = input.get_array("spc_h2o2_ndx")[0] - 1;
+    const int spc_so2_ndx = input.get_array("spc_so2_ndx")[0] - 1;
+    const int h2o2_ndx = input.get_array("h2o2_ndx")[0] - 1;
+    const int so2_ndx = input.get_array("so2_ndx")[0] - 1;
+    const int h2so4_ndx = input.get_array("h2so4_ndx")[0] - 1;
+    const int gas_wetdep_cnt = input.get_array("gas_wetdep_cnt")[0];
+
+    int wetdep_map[3];
+    const auto wetdep_map_in = input.get_array("wetdep_map");
+    wetdep_map[0] = wetdep_map_in[0] - 1;
+    wetdep_map[1] = wetdep_map_in[1] - 1;
+    wetdep_map[2] = wetdep_map_in[2] - 1;
+
+    const auto press_in = input.get_array("press");
+    const auto zmid_in = input.get_array("zmid");
+    const auto tfld_in = input.get_array("tfld");
+    const auto cmfdqr_in = input.get_array("cmfdqr");
+    const auto nrain_in = input.get_array("nrain");
+    const auto nevapr_in = input.get_array("nevapr");
+    const auto xhnm_in = input.get_array("xhnm");
+    const auto qin_in = input.get_array("qin");
+
+    // ColumnView input values
+    ColumnView press, zmid, tfld, cmfdqr, nrain, nevapr, xhnm;
+
+    auto press_host = View1DHost((Real *)press_in.data(), pver);
+    auto zmid_host = View1DHost((Real *)zmid_in.data(), pver);
+    auto tfld_host = View1DHost((Real *)tfld_in.data(), pver);
+    auto cmfdqr_host = View1DHost((Real *)cmfdqr_in.data(), pver);
+    auto nrain_host = View1DHost((Real *)nrain_in.data(), pver);
+    auto nevapr_host = View1DHost((Real *)nevapr_in.data(), pver);
+    auto xhnm_host = View1DHost((Real *)xhnm_in.data(), pver);
+
+    press = haero::testing::create_column_view(pver);
+    zmid = haero::testing::create_column_view(pver);
+    tfld = haero::testing::create_column_view(pver);
+    cmfdqr = haero::testing::create_column_view(pver);
+    nrain = haero::testing::create_column_view(pver);
+    nevapr = haero::testing::create_column_view(pver);
+    xhnm = haero::testing::create_column_view(pver);
+
+    Kokkos::deep_copy(press, press_host);
+    Kokkos::deep_copy(zmid, zmid_host);
+    Kokkos::deep_copy(tfld, tfld_host);
+    Kokkos::deep_copy(cmfdqr, cmfdqr_host);
+    Kokkos::deep_copy(nrain, nrain_host);
+    Kokkos::deep_copy(nevapr, nevapr_host);
+    Kokkos::deep_copy(xhnm, xhnm_host);
+
+    // working var inputs
+    ColumnView xgas2, xgas3, delz, xh2o2, xso2, xliq, rain, precip, xhen_h2o2,
+        xhen_hno3, xhen_so2, t_factor, xk0_hno3, xk0_so2, so2_diss;
+
+    xgas2 = haero::testing::create_column_view(pver);
+    xgas3 = haero::testing::create_column_view(pver);
+    delz = haero::testing::create_column_view(pver);
+    xh2o2 = haero::testing::create_column_view(pver);
+    xso2 = haero::testing::create_column_view(pver);
+    xliq = haero::testing::create_column_view(pver);
+    rain = haero::testing::create_column_view(pver);
+    precip = haero::testing::create_column_view(pver);
+    xhen_h2o2 = haero::testing::create_column_view(pver);
+    xhen_hno3 = haero::testing::create_column_view(pver);
+    xhen_so2 = haero::testing::create_column_view(pver);
+    t_factor = haero::testing::create_column_view(pver);
+    xk0_hno3 = haero::testing::create_column_view(pver);
+    xk0_so2 = haero::testing::create_column_view(pver);
+    so2_diss = haero::testing::create_column_view(pver);
+
+    ColumnView het_rates[gas_pcnst];
+    ColumnView tmp_hetrates[gas_pcnst];
+    ColumnView qin[gas_pcnst];
+    View1DHost het_rates_host[gas_pcnst];
+    View1DHost tmp_hetrates_host[gas_pcnst];
+    View1DHost qin_host[gas_pcnst];
+
+    for (int mm = 0; mm < gas_pcnst; ++mm) {
+      het_rates[mm] = haero::testing::create_column_view(pver);
+      tmp_hetrates[mm] = haero::testing::create_column_view(pver);
+      qin[mm] = haero::testing::create_column_view(pver);
+
+      het_rates_host[mm] = View1DHost("het_rates_host", 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];
+        count++;
+      }
+    }
+
+    // transfer data to GPU.
+    for (int mm = 0; mm < gas_pcnst; ++mm) {
+      Kokkos::deep_copy(het_rates[mm], 0.0);
+      Kokkos::deep_copy(tmp_hetrates[mm], 0.0);
+      Kokkos::deep_copy(qin[mm], qin_host[mm]);
+    }
+
+    auto team_policy = ThreadTeamPolicy(1u, Kokkos::AUTO);
+    Kokkos::parallel_for(
+        team_policy, KOKKOS_LAMBDA(const ThreadTeam &team) {
+          mo_sethet::sethet(
+              team, het_rates, rlat, press, zmid, phis, tfld, cmfdqr, nrain,
+              nevapr, delt, xhnm, qin, t_factor, xk0_hno3, xk0_so2, so2_diss,
+              xgas2, xgas3, delz, xh2o2, xso2, xliq, rain, precip, xhen_h2o2,
+              xhen_hno3, xhen_so2, tmp_hetrates, spc_h2o2_ndx, spc_so2_ndx,
+              h2o2_ndx, so2_ndx, h2so4_ndx, gas_wetdep_cnt, wetdep_map);
+        });
+
+    // transfer data to GPU.
+    for (int mm = 0; mm < gas_pcnst; ++mm) {
+      Kokkos::deep_copy(het_rates_host[mm], het_rates[mm]);
+    }
+    std::vector<Real> het_rates_out(pver * gas_pcnst);
+    count = 0;
+    for (int mm = 0; mm < gas_pcnst; ++mm) {
+      for (int kk = 0; kk < pver; ++kk) {
+        het_rates_out[count] = het_rates_host[mm](kk);
+        count++;
+      }
+    }
+
+    output.set("het_rates", het_rates_out);
+  });
+}

src/validation/mo_sethet/sethet_driver.cpp

@@ -26,6 +26,7 @@ void calc_het_rates(Ensemble *ensemble);
 void calc_precip_rescale(Ensemble *ensemble);
 void find_ktop(Ensemble *ensemble);
 void gas_washout(Ensemble *ensemble);
+void sethet(Ensemble *ensemble);
 
 int main(int argc, char **argv) {
   if (argc == 1) {
@@ -57,6 +58,8 @@ int main(int argc, char **argv) {
       find_ktop(ensemble);
     } else if (func_name == "gas_washout") {
       gas_washout(ensemble);
+    } else if (func_name == "sethet") {
+      sethet(ensemble);
     } else {
       std::cerr << "Error: Function name '" << func_name
                 << "' does not have an implemented test!" << std::endl;