EAMxx: avoid MPI tags in GridImportExport gather/scatter methods

components/eamxx/src/share/grid/abstract_grid.cpp

@@ -175,6 +175,8 @@ bool AbstractGrid::is_unique () const {
     return unique_gids==1;
   };
 
+
+  std::lock_guard<std::mutex> lock(m_mutex); // Lock the mutex
   if (not m_is_unique_computed) {
     m_is_unique = compute_is_unique();
     m_is_unique_computed = true;
@@ -218,6 +220,7 @@ is_valid_layout (const FieldLayout& layout) const
 auto AbstractGrid::
 get_global_min_dof_gid () const ->gid_type
 {
+  std::lock_guard<std::mutex> lock(m_mutex); // Lock the mutex
   // Lazy calculation
   if (m_global_min_dof_gid==std::numeric_limits<gid_type>::max()) {
     m_global_min_dof_gid = field_min<gid_type>(m_dofs_gids,&get_comm());
@@ -228,6 +231,7 @@ get_global_min_dof_gid () const ->gid_type
 auto AbstractGrid::
 get_global_max_dof_gid () const ->gid_type
 {
+  std::lock_guard<std::mutex> lock(m_mutex); // Lock the mutex
   // Lazy calculation
   if (m_global_max_dof_gid==-std::numeric_limits<gid_type>::max()) {
     m_global_max_dof_gid = field_max<gid_type>(m_dofs_gids,&get_comm());
@@ -238,6 +242,7 @@ get_global_max_dof_gid () const ->gid_type
 auto AbstractGrid::
 get_global_min_partitioned_dim_gid () const ->gid_type
 {
+  std::lock_guard<std::mutex> lock(m_mutex); // Lock the mutex
   // Lazy calculation
   if (m_global_min_partitioned_dim_gid==std::numeric_limits<gid_type>::max()) {
     m_global_min_partitioned_dim_gid = field_min<gid_type>(m_partitioned_dim_gids,&get_comm());
@@ -248,6 +253,7 @@ get_global_min_partitioned_dim_gid () const ->gid_type
 auto AbstractGrid::
 get_global_max_partitioned_dim_gid () const ->gid_type
 {
+  std::lock_guard<std::mutex> lock(m_mutex); // Lock the mutex
   // Lazy calculation
   if (m_global_max_partitioned_dim_gid==-std::numeric_limits<gid_type>::max()) {
     m_global_max_partitioned_dim_gid = field_max<gid_type>(m_partitioned_dim_gids,&get_comm());
@@ -550,14 +556,18 @@ void AbstractGrid::create_dof_fields (const int scalar2d_layout_rank)
 }
 
 auto AbstractGrid::get_gid2lid_map () const
- -> std::map<gid_type,int>
+ -> const std::map<gid_type,int>&
 {
-  std::map<gid_type,int> m;
-  auto gids_h = get_dofs_gids().get_view<const gid_type*,Host>();
-  for (int i=0; i<get_num_local_dofs(); ++i) {
-    m[gids_h[i]] = i;
+  std::lock_guard<std::mutex> lock(m_mutex); // Lock the mutex
+
+  int cur_sz = m_gid2lid.size();
+  if (cur_sz<get_num_local_dofs()) {
+    auto gids_h = get_dofs_gids().get_view<const gid_type*, Host>();
+    for (int i = 0; i < get_num_local_dofs(); ++i) {
+        m_gid2lid[gids_h[i]] = i; // Modify the mutable member
+    }
   }
-  return m;
+  return m_gid2lid; // Return the map
 }
 
 void AbstractGrid::copy_data (const AbstractGrid& src, const bool shallow)

components/eamxx/src/share/grid/abstract_grid.hpp

@@ -11,6 +11,7 @@
 #include <map>
 #include <list>
 #include <memory>
+#include <mutex>
 
 namespace scream
 {
@@ -205,7 +206,7 @@ class AbstractGrid : public ekat::enable_shared_from_this<AbstractGrid>
 
   int get_unique_grid_id () const { return m_unique_grid_id; }
 
-  std::map<gid_type,int> get_gid2lid_map () const;
+  const std::map<gid_type,int>& get_gid2lid_map () const;
 
 protected:
 
@@ -252,6 +253,12 @@ class AbstractGrid : public ekat::enable_shared_from_this<AbstractGrid>
 
   mutable std::map<std::string,Field>  m_geo_fields;
 
+  // Mutable, for lazy calculation
+  mutable std::map<gid_type,int> m_gid2lid;
+
+  // For thread safety in modifying mutable items (just in case someone ever runs this code in threaded regions)
+  mutable std::mutex m_mutex;
+
   // The MPI comm containing the ranks across which the global mesh is partitioned
   ekat::Comm            m_comm;
 };

components/eamxx/src/share/grid/grid_import_export.cpp

@@ -15,8 +15,6 @@ GridImportExport (const std::shared_ptr<const AbstractGrid>& unique,
   EKAT_REQUIRE_MSG (unique->is_unique(),
       "Error! GridImportExport unique grid is not unique.\n");
 
-  using gid_type = AbstractGrid::gid_type;
-
   m_unique = unique;
   m_overlapped = overlapped;
   m_comm = unique->get_comm();
@@ -153,22 +151,19 @@ GridImportExport (const std::shared_ptr<const AbstractGrid>& unique,
   Kokkos::deep_copy(m_export_lids,m_export_lids_h);
 
   // Compute counts per pid
-  {
-    m_num_exports_per_pid = view_1d<int>("",m_comm.size());
-    m_num_exports_per_pid_h = Kokkos::create_mirror_view(m_num_exports_per_pid);
-    for (size_t i=0; i<m_export_pids.size(); ++i) {
-      ++m_num_exports_per_pid_h[m_export_pids_h[i]];
-    }
-    Kokkos::deep_copy(m_num_exports_per_pid,m_num_exports_per_pid_h);
+  m_num_exports_per_pid = view_1d<int>("",m_comm.size());
+  m_num_exports_per_pid_h = Kokkos::create_mirror_view(m_num_exports_per_pid);
+  for (size_t i=0; i<m_export_pids.size(); ++i) {
+    ++m_num_exports_per_pid_h[m_export_pids_h[i]];
   }
-  {
-    m_num_imports_per_pid = view_1d<int>("",m_comm.size());
-    m_num_imports_per_pid_h = Kokkos::create_mirror_view(m_num_imports_per_pid);
-    for (size_t i=0; i<m_import_pids.size(); ++i) {
-      ++m_num_imports_per_pid_h[m_import_pids_h[i]];
-    }
-    Kokkos::deep_copy(m_num_imports_per_pid,m_num_imports_per_pid_h);
+  Kokkos::deep_copy(m_num_exports_per_pid,m_num_exports_per_pid_h);
+
+  m_num_imports_per_pid = view_1d<int>("",m_comm.size());
+  m_num_imports_per_pid_h = Kokkos::create_mirror_view(m_num_imports_per_pid);
+  for (size_t i=0; i<m_import_pids.size(); ++i) {
+    ++m_num_imports_per_pid_h[m_import_pids_h[i]];
   }
+  Kokkos::deep_copy(m_num_imports_per_pid,m_num_imports_per_pid_h);
 }
 
 } // namespace scream