Merge pull request #2976 from iulian787/iulian787/scream_with_moab_fields_rebase_var

EAMXX uses a different implementation of the atmosphere model, compared to EAM, and moab driver needs to be extended to support a coupled case

case tested is --compset F2010-SCREAMv1-MPASSI --res ne4pg2_oQU480 --driver moab
this case uses mpas sea ice, CICE is not yet supported by MOAB

the implementation follows similar strategy compared to the rest of the models
Atmosphere mesh in moab has 4 different instances, similar to EAM,

corresponding to coarse spectral mesh ( size 6nene),
fine quad mesh (corresponding to GLL quad mesh, size 6nene* (np-1)* (np-1) )
pg2 FV type mesh (size 6nene* pg2 * pg2)
point cloud corresponding to mct simplified vertex grid.

On the coupler side, we still use just the pg2 FV type mesh, used for intersection and coupling with other models

Author: rljacob

cime_config/machines/config_machines.xml

@@ -443,6 +443,9 @@
     <environment_variables compiler="nvidiagpu">
       <env name="MPICH_GPU_SUPPORT_ENABLED">1</env>
     </environment_variables>
+    <environment_variables compiler="gnugpu">
+      <env name="MOAB_ROOT">$SHELL{if [ -z "$MOAB_ROOT" ]; then echo /global/cfs/cdirs/e3sm/software/moab/gnugpu ; else echo "$MOAB_ROOT"; fi}</env>
+    </environment_variables>
     <environment_variables compiler="gnu.*" mpilib="mpich">
       <env name="ADIOS2_ROOT">$SHELL{if [ -z "$ADIOS2_ROOT" ]; then echo /global/cfs/cdirs/e3sm/3rdparty/adios2/2.9.1/cray-mpich-8.1.25/gcc-11.2.0; else echo "$ADIOS2_ROOT"; fi}</env>
     </environment_variables>

components/eam/src/dynamics/se/dyn_comp.F90

@@ -269,7 +269,7 @@ subroutine dyn_init1(fh, NLFileName, dyn_in, dyn_out)
     end if
 
 #ifdef HAVE_MOAB
-    call create_moab_meshes(par, elem)
+    call create_moab_meshes(par, elem, fv_nphys)
 #endif
     ! Define the CAM grids (this has to be after dycore spinup).
     ! Physics-grid will be defined later by phys_grid_init

components/eamxx/src/control/atmosphere_driver.cpp

@@ -300,6 +300,9 @@ void AtmosphereDriver::create_grids()
 void AtmosphereDriver::setup_surface_coupling_data_manager(SurfaceCouplingTransferType transfer_type,
                                                            const int num_cpl_fields, const int num_scream_fields,
                                                            const int field_size, Real* data_ptr,
+#ifdef HAVE_MOAB
+                                                           Real* data_ptr_moab,
+#endif
                                                            char* names_ptr, int* cpl_indices_ptr, int* vec_comps_ptr,
                                                            Real* constant_multiple_ptr, bool* do_transfer_during_init_ptr)
 {
@@ -318,6 +321,9 @@ void AtmosphereDriver::setup_surface_coupling_data_manager(SurfaceCouplingTransf
   } else EKAT_ERROR_MSG("Error! Unexpected SurfaceCouplingTransferType.");
 
   sc_data_mgr->setup_internals(num_cpl_fields, num_scream_fields, field_size, data_ptr,
+#ifdef HAVE_MOAB
+                               data_ptr_moab,
+#endif
                                names_ptr, cpl_indices_ptr, vec_comps_ptr,
                                constant_multiple_ptr, do_transfer_during_init_ptr);
 }

components/eamxx/src/control/atmosphere_driver.hpp

@@ -87,6 +87,9 @@ class AtmosphereDriver
   void setup_surface_coupling_data_manager(SurfaceCouplingTransferType transfer_type,
                                            const int num_cpl_fields, const int num_scream_fields,
                                            const int field_size, Real* data_ptr,
+#ifdef HAVE_MOAB
+                                           Real* data_ptr_moab,
+#endif
                                            char* names_ptr, int* cpl_indices_ptr, int* vec_comps_ptr,
                                            Real* constant_multiple_ptr, bool* do_transfer_during_init_ptr);
 

components/eamxx/src/control/atmosphere_surface_coupling_exporter.cpp

@@ -137,6 +137,13 @@ void SurfaceCouplingExporter::setup_surface_coupling_data(const SCDataManager &s
                                                          m_num_cols, m_num_cpl_exports);
   m_cpl_exports_view_d = Kokkos::create_mirror_view(DefaultDevice(), m_cpl_exports_view_h);
 
+#ifdef HAVE_MOAB
+  // The export data is of size num_cpl_exports,ncols. All other data is of size num_scream_exports
+  m_moab_cpl_exports_view_h = decltype(m_moab_cpl_exports_view_h) (sc_data_manager.get_field_data_moab_ptr(),
+                                                         m_num_cpl_exports, m_num_cols);
+  m_moab_cpl_exports_view_d = Kokkos::create_mirror_view(DefaultDevice(), m_moab_cpl_exports_view_h);
+#endif
+
   m_export_field_names = new name_t[m_num_scream_exports];
   std::memcpy(m_export_field_names, sc_data_manager.get_field_name_ptr(), m_num_scream_exports*32*sizeof(char));
 
@@ -536,6 +543,12 @@ void SurfaceCouplingExporter::do_export_to_cpl(const bool called_during_initiali
   // Any field not exported by scream, or not exported
   // during initialization, is set to 0.0
   Kokkos::deep_copy(m_cpl_exports_view_d, 0.0);
+#ifdef HAVE_MOAB
+  // Any field not exported by scream, or not exported
+  // during initialization, is set to 0.0
+  Kokkos::deep_copy(m_moab_cpl_exports_view_d, 0.0);
+  const auto moab_cpl_exports_view_d = m_moab_cpl_exports_view_d;
+#endif
   const auto cpl_exports_view_d = m_cpl_exports_view_d;
   const int  num_exports        = m_num_scream_exports;
   const int  num_cols           = m_num_cols;
@@ -554,9 +567,27 @@ void SurfaceCouplingExporter::do_export_to_cpl(const bool called_during_initiali
       cpl_exports_view_d(icol,info.cpl_indx) = info.constant_multiple*info.data[offset];
     }
   });
+#ifdef HAVE_MOAB
+  Kokkos::parallel_for(export_policy, KOKKOS_LAMBDA(const int& i) {
+    const int ifield = i / num_cols;
+    const int icol   = i % num_cols;
+    const auto& info = col_info(ifield);
+    const auto offset = icol*info.col_stride + info.col_offset;
 
+    // if this is during initialization, check whether or not the field should be exported
+    bool do_export = (not called_during_initialization || info.transfer_during_initialization);
+    if (do_export) {
+      moab_cpl_exports_view_d(info.cpl_indx, icol) = info.constant_multiple*info.data[offset];
+    }
+  });
+#endif
   // Deep copy fields from device to cpl host array
   Kokkos::deep_copy(m_cpl_exports_view_h,m_cpl_exports_view_d);
+#ifdef HAVE_MOAB
+  // Deep copy fields from device to cpl host array
+  Kokkos::deep_copy(m_moab_cpl_exports_view_h,m_moab_cpl_exports_view_d);
+#endif
+
 }
 // =========================================================================================
 void SurfaceCouplingExporter::finalize_impl()

components/eamxx/src/control/atmosphere_surface_coupling_exporter.hpp

@@ -141,6 +141,13 @@ class SurfaceCouplingExporter : public AtmosphereProcess
   view_2d <DefaultDevice, Real> m_cpl_exports_view_d;
   uview_2d<HostDevice,    Real> m_cpl_exports_view_h;
 
+#ifdef HAVE_MOAB
+  // Views storing a 2d array with dims (num_fields, num_cols) for moab cpl export data.
+  // The field cols strides faster, since that's what moab does (so we can "view" the
+  // pointer to the whole a2x_am(:,:) array from Fortran)
+  view_2d <DefaultDevice, Real> m_moab_cpl_exports_view_d;
+  uview_2d<HostDevice,    Real> m_moab_cpl_exports_view_h;
+#endif
   // Array storing the field names for exports
   name_t*                   m_export_field_names;
   std::vector<std::string>  m_export_field_names_vector;

components/eamxx/src/control/atmosphere_surface_coupling_importer.cpp

@@ -74,6 +74,13 @@ void SurfaceCouplingImporter::set_grids(const std::shared_ptr<const GridsManager
                                                          m_num_cols, m_num_cpl_imports);
   m_cpl_imports_view_d = Kokkos::create_mirror_view_and_copy(DefaultDevice(),
                                                              m_cpl_imports_view_h);
+#ifdef HAVE_MOAB
+  // The import data is of size num_cpl_imports, ncol. All other data is of size num_scream_imports
+  m_moab_cpl_imports_view_h = decltype(m_moab_cpl_imports_view_h) (sc_data_manager.get_field_data_moab_ptr(),
+                                                         m_num_cpl_imports, m_num_cols);
+  m_moab_cpl_imports_view_d = Kokkos::create_mirror_view_and_copy(DefaultDevice(),
+                                                             m_moab_cpl_imports_view_h);
+#endif
   m_import_field_names = new name_t[m_num_scream_imports];
   std::memcpy(m_import_field_names, sc_data_manager.get_field_name_ptr(), m_num_scream_imports*32*sizeof(char));
 
@@ -153,6 +160,12 @@ void SurfaceCouplingImporter::do_import(const bool called_during_initialization)
 
   // Deep copy cpl host array to devic
   Kokkos::deep_copy(m_cpl_imports_view_d,m_cpl_imports_view_h);
+#ifdef HAVE_MOAB
+  // Deep copy cpl host array to device
+  const auto moab_cpl_imports_view_d = m_moab_cpl_imports_view_d;
+  Kokkos::deep_copy(m_moab_cpl_imports_view_d,m_moab_cpl_imports_view_h);
+#endif
+
 
   // Unpack the fields
   auto unpack_policy = policy_type(0,num_imports*num_cols);
@@ -171,6 +184,24 @@ void SurfaceCouplingImporter::do_import(const bool called_during_initialization)
     }
   });
 
+#ifdef HAVE_MOAB
+  Kokkos::parallel_for(unpack_policy, KOKKOS_LAMBDA(const int& i) {
+
+    const int icol   = i / num_imports;
+    const int ifield = i % num_imports;
+
+    const auto& info = col_info(ifield);
+
+    auto offset = icol*info.col_stride + info.col_offset;
+
+    // if this is during initialization, check whether or not the field should be imported
+    bool do_import = (not called_during_initialization || info.transfer_during_initialization);
+    if (do_import) {
+      info.data[offset] = moab_cpl_imports_view_d(info.cpl_indx, icol)*info.constant_multiple;
+    }
+  });
+#endif
+
   if (m_iop) {
     if (m_iop->get_params().get<bool>("iop_srf_prop")) {
       // Overwrite imports with data from IOP file
@@ -232,6 +263,9 @@ void SurfaceCouplingImporter::overwrite_iop_imports (const bool called_during_in
       const auto& info_d = col_info_d(ifield);
       const auto offset = icol*info_d.col_stride + info_d.col_offset;
       info_d.data[offset] = col_val;
+#ifdef HAVE_MOAB
+   //  TODO
+#endif
     });
   }
 }

components/eamxx/src/control/atmosphere_surface_coupling_importer.hpp

@@ -83,6 +83,14 @@ class SurfaceCouplingImporter : public AtmosphereProcess
   view_2d <DefaultDevice, Real> m_cpl_imports_view_d;
   uview_2d<HostDevice,    Real> m_cpl_imports_view_h;
 
+#ifdef HAVE_MOAB
+  // Views storing a 2d array with dims (num_fields,num_cols) for import data.
+  // The colums index strides faster, since that's what moab does (so we can "view" the
+  // pointer to the whole x2a_am(:,:) array from Fortran)
+  view_2d <DefaultDevice, Real> m_moab_cpl_imports_view_d;
+  uview_2d<HostDevice,    Real> m_moab_cpl_imports_view_h;
+#endif
+
   // Array storing the field names for imports
   name_t* m_import_field_names;
 

components/eamxx/src/dynamics/homme/interface/phys_grid_mod.F90

@@ -549,6 +549,14 @@ subroutine phys_grid_init (pgN)
     use gllfvremap_mod,    only: gfr_init
     use homme_context_mod, only: elem, par
     use dimensions_mod,    only: nelem, nelemd
+#ifdef HAVE_MOAB
+    use seq_comm_mct,      only: MHID, MHFID  ! id of homme moab coarse and fine applications
+    use seq_comm_mct,      only: ATMID
+    use seq_comm_mct,      only: mhpgid       ! id of pgx moab application
+    use semoab_mod,        only: create_moab_meshes
+    use iMOAB, only : iMOAB_RegisterApplication
+    use iso_c_binding
+#endif
     !
     ! Input(s)
     !
@@ -561,6 +569,10 @@ subroutine phys_grid_init (pgN)
     character(2) :: str
     type(pg_specs_t), pointer :: pg
 
+#ifdef HAVE_MOAB
+    integer :: ATM_ID1
+    character*32  appname
+#endif
     pg => pg_specs(pgN)
 
     if (pg%inited) then
@@ -612,6 +624,48 @@ subroutine phys_grid_init (pgN)
     call compute_global_dofs (pg)
     call compute_global_coords (pg)
     call compute_global_area (pg)
+#ifdef HAVE_MOAB
+    if (pgN > 0) then
+       appname="HM_COARSE"//C_NULL_CHAR
+       ATM_ID1 = 120 !
+       ierr = iMOAB_RegisterApplication(appname, par%comm, ATM_ID1, MHID)
+       if (ierr > 0 )  &
+       call abortmp('Error: cannot register moab app')
+       if(par%masterproc) then
+           write(iulog,*) " "
+           write(iulog,*) "register MOAB app:", trim(appname), "  MHID=", MHID
+           write(iulog,*) " "
+       endif
+       appname="HM_FINE"//C_NULL_CHAR
+       ATM_ID1 = 119 ! this number should not conflict with other components IDs; how do we know?
+       ierr = iMOAB_RegisterApplication(appname, par%comm, ATM_ID1, MHFID)
+       if (ierr > 0 )  &
+       call abortmp('Error: cannot register moab app for fine mesh')
+       if(par%masterproc) then
+           write(iulog,*) " "
+           write(iulog,*) "register MOAB app:", trim(appname), "  MHFID=", MHFID
+           write(iulog,*) " "
+       endif
+         appname="HM_PGX"//C_NULL_CHAR
+         ATM_ID1 =  ATMID(1) ! this number should not conflict with other components IDs; how do we know?
+         !
+         ! in this case, we reuse the main atm id, mhid will not be used for intersection anymore
+         ! still, need to be careful
+         ierr = iMOAB_RegisterApplication(appname, par%comm, ATM_ID1, mhpgid)
+         if (ierr > 0 )  &
+             call abortmp('Error: cannot register moab app for fine mesh')
+         if(par%masterproc) then
+             write(iulog,*) " "
+             write(iulog,*) "register MOAB app:", trim(appname), "  MHPGID=", mhpgid
+             write(iulog,*) " "
+         endif
+! instance distributed moab meshes from elem structures
+!    1 ) spectral coarse mesh
+!    2 ) GLL fine quad mesh (used mostly for visualization)
+!    3 ) pgN FV type mesh, (most of the time pg2 mesh), used for coupling with other components;
+       call create_moab_meshes(par, elem, pgN)
+     endif
+#endif
   end subroutine phys_grid_init