EAMxx: Using a regular range policy instead of MDRangePolicy, as it has been reported that the latter has performance issues. We can revert this commit once the Kokkos team addresses these concerns.

Author: odiazib

components/eamxx/src/physics/mam/eamxx_mam_microphysics_process_interface.cpp

@@ -463,7 +463,7 @@ void MAMMicrophysics::initialize_impl(const RunType run_type) {
   // ---------------------------------------------------------------
   populate_wet_atm(wet_atm_);
   populate_dry_atm(dry_atm_, buffer_);
-  
+
   // FIXME: we are using cldfrac_tot in other mam4xx process.
   dry_atm_.cldfrac = get_field_in("cldfrac_liq").get_view<const Real **>();
   // FIXME: phis is not populated by populate_wet_and_dry_atm.
@@ -954,9 +954,11 @@ void MAMMicrophysics::run_impl(const double dt) {
   // Transpose extfrc_ from internal layout [ncol][nlev][extcnt]
   // to output layout [ncol][extcnt][nlev]
   // This aligns with expected field storage in the EAMxx infrastructure.
-  Kokkos::parallel_for("transpose_extfrc", 
-    Kokkos::MDRangePolicy<Kokkos::Rank<3>>({0,0,0}, {ncol, extcnt, nlev}),
-    KOKKOS_LAMBDA(const int i, const int j, const int k) {
+  Kokkos::parallel_for("transpose_extfrc", ncol * extcnt * nlev,
+    KOKKOS_LAMBDA(const int linear_index) {
+      const int k = linear_index / (ncol * extcnt);
+      const int j = (linear_index / ncol) % extcnt;
+      const int i = linear_index % ncol;
       const int pcnst_idx = extfrc_pcnst_index[j];
       const Real molar_mass_g_per_mol = molar_mass_g_per_mol_tmp[pcnst_idx]; // g/mol
       // Modify units to MKS units: [molec/cm3/s] to [kg/m3/s]

components/eamxx/src/physics/mam/eamxx_mam_optics_process_interface.cpp

@@ -52,7 +52,7 @@ void MAMOptics::set_grids(
   FieldLayout scalar3d_int = grid_->get_3d_scalar_layout(false);
   add_tracers_wet_atm();
   add_fields_dry_atm();
-    
+
   // cloud liquid number mixing ratio [1/kg]
   add_tracer<Required>("nc", grid_, n_unit);
 
@@ -443,11 +443,14 @@ void MAMOptics::run_impl(const double dt) {
   Kokkos::fence();
 
   // nswbands loop is using rrtmg indexing.
+  const int ncol = ncol_;
+  const int nswbands = nswbands_;
   Kokkos::parallel_for(
-      "copying data from mam4xx to eamxx",
-      Kokkos::MDRangePolicy<Kokkos::Rank<3>>({0, 0, 0},
-                                             {ncol_, nswbands_, nlev_}),
-      KOKKOS_LAMBDA(const int icol, const int iswband, const int kk) {
+      "copying data from mam4xx to eamxx", ncol_*nswbands_*nlev_,
+      KOKKOS_LAMBDA(const int linear_index) {
+        const int kk = linear_index / (ncol * nswbands);
+        const int iswband = (linear_index / ncol) % nswbands;
+        const int icol = linear_index % ncol;
         // Extract single scattering albedo from the product-defined fields
         if(tau_sw(icol, iswband, kk + 1) > zero) {
           aero_ssa_sw_eamxx(icol, get_idx_rrtmgp_from_rrtmg_swbands(iswband),