Merge pull request #159 from ecmwf/feature/mpi_comm

Register and use non-default MPI communicators in Atlas objects
ecmwf · Sep 21, 2023 · d02f399 · d02f399
2 parents 1633483 + 1d98b4b
commit d02f399
Show file tree

Hide file tree

Showing 97 changed files with 1,768 additions and 388 deletions.
diff --git a/src/atlas/functionspace/CellColumns.cc b/src/atlas/functionspace/CellColumns.cc
@@ -95,7 +95,8 @@ class CellColumnsHaloExchangeCache : public util::Cache<std::string, parallel::H
 
     static value_type* create(const Mesh& mesh) {
         value_type* value = new value_type();
-        value->setup(array::make_view<int, 1>(mesh.cells().partition()).data(),
+        value->setup(mesh.mpi_comm(),
+                     array::make_view<int, 1>(mesh.cells().partition()).data(),
                      array::make_view<idx_t, 1>(mesh.cells().remote_index()).data(), REMOTE_IDX_BASE,
                      mesh.cells().size());
         return value;
@@ -129,7 +130,8 @@ class CellColumnsGatherScatterCache : public util::Cache<std::string, parallel::
 
     static value_type* create(const Mesh& mesh) {
         value_type* value = new value_type();
-        value->setup(array::make_view<int, 1>(mesh.cells().partition()).data(),
+        value->setup(mesh.mpi_comm(),
+                     array::make_view<int, 1>(mesh.cells().partition()).data(),
                      array::make_view<idx_t, 1>(mesh.cells().remote_index()).data(), REMOTE_IDX_BASE,
                      array::make_view<gidx_t, 1>(mesh.cells().global_index()).data(), mesh.cells().size());
         return value;
@@ -193,14 +195,14 @@ void CellColumns::set_field_metadata(const eckit::Configuration& config, Field&
 }
 
 idx_t CellColumns::config_size(const eckit::Configuration& config) const {
-    const idx_t rank = mpi::rank();
     idx_t size       = nb_cells();
     bool global(false);
     if (config.get("global", global)) {
         if (global) {
             idx_t owner(0);
             config.get("owner", owner);
             idx_t _nb_cells_global(nb_cells_global());
+            const idx_t rank = mpi::comm(mpi_comm()).rank();
             size = (rank == owner ? _nb_cells_global : 0);
         }
     }
@@ -269,7 +271,7 @@ CellColumns::CellColumns(const Mesh& mesh, const eckit::Configuration& config):
         return nb_cells;
     };
 
-    mesh::actions::build_nodes_parallel_fields(mesh_.nodes());
+    mesh::actions::build_nodes_parallel_fields(mesh_);
     mesh::actions::build_cells_parallel_fields(mesh_);
     mesh::actions::build_periodic_boundaries(mesh_);
 

diff --git a/src/atlas/functionspace/CellColumns.h b/src/atlas/functionspace/CellColumns.h
@@ -99,6 +99,8 @@ class CellColumns : public functionspace::FunctionSpaceImpl {
 
     Field partition() const override;
 
+    std::string mpi_comm() const override { return mesh_.mpi_comm(); }
+
 private:  // methods
     idx_t config_size(const eckit::Configuration& config) const;
     array::DataType config_datatype(const eckit::Configuration&) const;

diff --git a/src/atlas/functionspace/EdgeColumns.cc b/src/atlas/functionspace/EdgeColumns.cc
@@ -96,7 +96,8 @@ class EdgeColumnsHaloExchangeCache : public util::Cache<std::string, parallel::H
 
     static value_type* create(const Mesh& mesh) {
         value_type* value = new value_type();
-        value->setup(array::make_view<int, 1>(mesh.edges().partition()).data(),
+        value->setup(mesh.mpi_comm(),
+                     array::make_view<int, 1>(mesh.edges().partition()).data(),
                      array::make_view<idx_t, 1>(mesh.edges().remote_index()).data(), REMOTE_IDX_BASE,
                      mesh.edges().size());
         return value;
@@ -130,7 +131,8 @@ class EdgeColumnsGatherScatterCache : public util::Cache<std::string, parallel::
 
     static value_type* create(const Mesh& mesh) {
         value_type* value = new value_type();
-        value->setup(array::make_view<int, 1>(mesh.edges().partition()).data(),
+        value->setup(mesh.mpi_comm(),
+                     array::make_view<int, 1>(mesh.edges().partition()).data(),
                      array::make_view<idx_t, 1>(mesh.edges().remote_index()).data(), REMOTE_IDX_BASE,
                      array::make_view<gidx_t, 1>(mesh.edges().global_index()).data(), mesh.edges().size());
         return value;
@@ -194,14 +196,14 @@ void EdgeColumns::set_field_metadata(const eckit::Configuration& config, Field&
 }
 
 idx_t EdgeColumns::config_size(const eckit::Configuration& config) const {
-    const idx_t rank = mpi::rank();
     idx_t size       = nb_edges();
     bool global(false);
     if (config.get("global", global)) {
         if (global) {
             idx_t owner(0);
             config.get("owner", owner);
             idx_t _nb_edges_global(nb_edges_global());
+            const idx_t rank = mpi::comm(mpi_comm()).rank();
             size = (rank == owner ? _nb_edges_global : 0);
         }
     }
@@ -266,7 +268,7 @@ EdgeColumns::EdgeColumns(const Mesh& mesh, const eckit::Configuration& config):
         return _nb_edges;
     };
 
-    mesh::actions::build_nodes_parallel_fields(mesh_.nodes());
+    mesh::actions::build_nodes_parallel_fields(mesh_);
     mesh::actions::build_periodic_boundaries(mesh_);
 
     if (halo_.size() > 0) {

diff --git a/src/atlas/functionspace/EdgeColumns.h b/src/atlas/functionspace/EdgeColumns.h
@@ -93,6 +93,8 @@ class EdgeColumns : public functionspace::FunctionSpaceImpl {
 
     Field partition() const override;
 
+    std::string mpi_comm() const override { return mesh_.mpi_comm(); }
+
 private:  // methods
     idx_t config_size(const eckit::Configuration& config) const;
     array::DataType config_datatype(const eckit::Configuration&) const;

diff --git a/src/atlas/functionspace/FunctionSpace.cc b/src/atlas/functionspace/FunctionSpace.cc
@@ -58,8 +58,12 @@ idx_t FunctionSpace::size() const {
     return get()->size();
 }
 
-idx_t FunctionSpace::nb_partitions() const {
-    return get()->nb_partitions();
+idx_t FunctionSpace::part() const {
+    return get()->part();
+}
+
+idx_t FunctionSpace::nb_parts() const {
+    return get()->nb_parts();
 }
 
 Field FunctionSpace::lonlat() const {
@@ -114,6 +118,10 @@ const parallel::GatherScatter& FunctionSpace::scatter() const {
     return get()->scatter();
 }
 
+std::string FunctionSpace::mpi_comm() const {
+    return get()->mpi_comm();
+}
+
 const util::PartitionPolygon& FunctionSpace::polygon(idx_t halo) const {
     return get()->polygon(halo);
 }

diff --git a/src/atlas/functionspace/FunctionSpace.h b/src/atlas/functionspace/FunctionSpace.h
@@ -79,7 +79,9 @@ class FunctionSpace : DOXYGEN_HIDE(public util::ObjectHandle<functionspace::Func
 
     const Projection& projection() const;
 
-    idx_t nb_partitions() const;
+    idx_t part() const;
+
+    idx_t nb_parts() const;
 
     idx_t size() const;
 
@@ -95,6 +97,8 @@ class FunctionSpace : DOXYGEN_HIDE(public util::ObjectHandle<functionspace::Func
 
     const parallel::GatherScatter& gather() const;
     const parallel::GatherScatter& scatter() const;
+
+    std::string mpi_comm() const;
 };
 
 //------------------------------------------------------------------------------------------------------

diff --git a/src/atlas/functionspace/NodeColumns.cc b/src/atlas/functionspace/NodeColumns.cc
@@ -105,7 +105,7 @@ class NodeColumnsHaloExchangeCache : public util::Cache<std::string, parallel::H
         idx_t nb_nodes(mesh.nodes().size());
         mesh.metadata().get(ss.str(), nb_nodes);
 
-        value->setup(array::make_view<int, 1>(mesh.nodes().partition()).data(),
+        value->setup(mesh.mpi_comm(), array::make_view<int, 1>(mesh.nodes().partition()).data(),
                      array::make_view<idx_t, 1>(mesh.nodes().remote_index()).data(), REMOTE_IDX_BASE, nb_nodes);
 
         return value;
@@ -154,7 +154,8 @@ class NodeColumnsGatherScatterCache : public util::Cache<std::string, parallel::
             //}
         }
 
-        value->setup(array::make_view<int, 1>(mesh.nodes().partition()).data(),
+        value->setup(mesh.mpi_comm(),
+                     array::make_view<int, 1>(mesh.nodes().partition()).data(),
                      array::make_view<idx_t, 1>(mesh.nodes().remote_index()).data(), REMOTE_IDX_BASE,
                      array::make_view<gidx_t, 1>(mesh.nodes().global_index()).data(), mask.data(), mesh.nodes().size());
         return value;
@@ -206,7 +207,7 @@ NodeColumns::NodeColumns(Mesh mesh, const eckit::Configuration& config):
     else {
         halo_ = mesh::Halo(mesh);
     }
-    mesh::actions::build_nodes_parallel_fields(mesh_.nodes());
+    mesh::actions::build_nodes_parallel_fields(mesh_);
     mesh::actions::build_periodic_boundaries(mesh_);
 
     if (halo_.size() > 0) {
@@ -255,7 +256,8 @@ idx_t NodeColumns::config_nb_nodes(const eckit::Configuration& config) const {
             idx_t owner(0);
             config.get("owner", owner);
             idx_t _nb_nodes_global = nb_nodes_global();
-            size                   = (mpi::rank() == owner ? _nb_nodes_global : 0);
+            idx_t rank = mpi::comm(mpi_comm()).rank();
+            size = (rank == owner ? _nb_nodes_global : 0);
         }
     }
     return size;

diff --git a/src/atlas/functionspace/NodeColumns.h b/src/atlas/functionspace/NodeColumns.h
@@ -258,7 +258,9 @@ class NodeColumns : public functionspace::FunctionSpaceImpl {
 
     virtual idx_t size() const override { return nb_nodes_; }
 
-    idx_t nb_partitions() const override { return mesh_.nb_partitions(); }
+    idx_t part() const override { return mesh_.part(); }
+
+    idx_t nb_parts() const override { return mesh_.nb_parts(); }
 
     Field lonlat() const override { return nodes_.lonlat(); }
 
@@ -276,6 +278,8 @@ class NodeColumns : public functionspace::FunctionSpaceImpl {
 
     const Projection& projection() const override { return mesh_.projection(); }
 
+    std::string mpi_comm() const override { return mesh_.mpi_comm(); }
+
 private:  // methods
     void constructor();