Use mpi_comm of partitioned in MatchingPartitioners

ecmwf · Aug 30, 2023 · a760ea9 · a760ea9
1 parent f1628a7
commit a760ea9
Show file tree

Hide file tree

Showing 6 changed files with 25 additions and 19 deletions.
diff --git a/src/atlas/grid/detail/partitioner/CubedSpherePartitioner.cc b/src/atlas/grid/detail/partitioner/CubedSpherePartitioner.cc
@@ -19,7 +19,6 @@
 #include <vector>
 
 #include "atlas/grid/CubedSphereGrid.h"
-#include "atlas/parallel/mpi/mpi.h"
 #include "atlas/runtime/Exception.h"
 #include "atlas/runtime/Log.h"
 

diff --git a/src/atlas/grid/detail/partitioner/MatchingFunctionSpacePartitionerLonLatPolygon.cc b/src/atlas/grid/detail/partitioner/MatchingFunctionSpacePartitionerLonLatPolygon.cc
@@ -40,14 +40,17 @@ void MatchingFunctionSpacePartitionerLonLatPolygon::partition(const Grid& grid,
     ATLAS_TRACE("MatchingFunctionSpacePartitionerLonLatPolygon");
     //atlas::vector<int> part( grid.size() );
 
-    if (mpi::size() == 1) {
+    const auto& comm   = mpi::comm(partitioned_.mpi_comm());
+    const int mpi_rank = int(comm.rank());
+    const int mpi_size = int(comm.size());
+
+    if (mpi_size == 1) {
         // shortcut
         omp::fill(part, part + grid.size(), 0);
     }
     else {
         const auto& p = partitioned_.polygon();
 
-        int rank = mpi::rank();
         util::PolygonXY poly{p};
         {
             ATLAS_TRACE("point-in-polygon check for entire grid (" + std::to_string(grid.size()) + " points)");
@@ -61,7 +64,7 @@ void MatchingFunctionSpacePartitionerLonLatPolygon::partition(const Grid& grid,
                 auto it            = grid.xy().begin() + chunk * grid.size() / chunks;
                 for (size_t n = begin; n < end; ++n) {
                     if (poly.contains(*it)) {
-                        part[n] = rank;
+                        part[n] = mpi_rank;
                     }
                     else {
                         part[n] = -1;
@@ -70,15 +73,15 @@ void MatchingFunctionSpacePartitionerLonLatPolygon::partition(const Grid& grid,
                 }
             }
         }
-        ATLAS_TRACE_MPI(ALLREDUCE) { mpi::comm().allReduceInPlace(part, grid.size(), eckit::mpi::max()); }
+        ATLAS_TRACE_MPI(ALLREDUCE) { comm.allReduceInPlace(part, grid.size(), eckit::mpi::max()); }
     }
 }
 
 
 #if 0
-    const eckit::mpi::Comm& comm = atlas::mpi::comm();
-    const int mpi_rank           = int( comm.rank() );
-    const int mpi_size           = int( comm.size() );
+    const auto& comm   = mpi::comm(partitioned_.mpi_comm());
+    const int mpi_rank = int(comm.rank());
+    const int mpi_size = int(comm.size());
 
     ATLAS_TRACE( "MatchingFunctionSpacePartitionerLonLatPolygon::partition" );
 

diff --git a/src/atlas/grid/detail/partitioner/MatchingMeshPartitionerBruteForce.cc b/src/atlas/grid/detail/partitioner/MatchingMeshPartitionerBruteForce.cc
@@ -75,9 +75,9 @@ bool point_in_quadrilateral(const PointLonLat& P, const PointLonLat& A, const Po
 void MatchingMeshPartitionerBruteForce::partition(const Grid& grid, int partitioning[]) const {
     ATLAS_TRACE("MatchingMeshPartitionerBruteForce::partition");
 
-
-    eckit::mpi::Comm& comm = eckit::mpi::comm();
-    const int mpi_rank     = int(comm.rank());
+    const auto& comm   = mpi::comm(prePartitionedMesh_.mpi_comm());
+    const int mpi_rank = int(comm.rank());
+    const int mpi_size = int(comm.size());
 
     // Point coordinates
     // - use a bounding box to quickly discard points,

diff --git a/src/atlas/grid/detail/partitioner/MatchingMeshPartitionerCubedSphere.cc b/src/atlas/grid/detail/partitioner/MatchingMeshPartitionerCubedSphere.cc
@@ -17,6 +17,10 @@ namespace detail {
 namespace partitioner {
 
 void MatchingMeshPartitionerCubedSphere::partition(const Grid& grid, int partitioning[]) const {
+    const auto& comm   = mpi::comm(prePartitionedMesh_.mpi_comm());
+    const int mpi_rank = int(comm.rank());
+    const int mpi_size = int(comm.size());
+
     // Make cell finder from owned mesh cells.
     const auto finder = interpolation::method::cubedsphere::CellFinder(prePartitionedMesh_);
 
@@ -32,12 +36,12 @@ void MatchingMeshPartitionerCubedSphere::partition(const Grid& grid, int partiti
         // This is probably more expensive than it needs to be, as it performs
         // a dry run of the cubedsphere interpolation method.
         const auto& lonlat = *lonlatIt;
-        partitioning[i]    = finder.getCell(lonlat, listSize, edgeEpsilon, epsilon).isect ? mpi::rank() : -1;
+        partitioning[i]    = finder.getCell(lonlat, listSize, edgeEpsilon, epsilon).isect ? mpi_rank : -1;
         ++lonlatIt;
     }
 
     // AllReduce to get full partitioning array.
-    mpi::comm().allReduceInPlace(partitioning, grid.size(), eckit::mpi::Operation::MAX);
+    comm.allReduceInPlace(partitioning, grid.size(), eckit::mpi::Operation::MAX);
     const auto misses = std::count_if(partitioning, partitioning + grid.size(), [](int elem) { return elem < 0; });
     if (misses > 0) {
         throw_Exception(

diff --git a/src/atlas/grid/detail/partitioner/MatchingMeshPartitionerLonLatPolygon.cc b/src/atlas/grid/detail/partitioner/MatchingMeshPartitionerLonLatPolygon.cc
@@ -38,9 +38,9 @@ PartitionerBuilder<MatchingMeshPartitionerLonLatPolygon> __builder("lonlat-polyg
 }
 
 void MatchingMeshPartitionerLonLatPolygon::partition(const Grid& grid, int partitioning[]) const {
-    const eckit::mpi::Comm& comm = atlas::mpi::comm();
-    const int mpi_rank           = int(comm.rank());
-    const int mpi_size           = int(comm.size());
+    const auto& comm   = mpi::comm(prePartitionedMesh_.mpi_comm());
+    const int mpi_rank = int(comm.rank());
+    const int mpi_size = int(comm.size());
 
     ATLAS_TRACE("MatchingMeshPartitionerLonLatPolygon::partition");
 

diff --git a/src/atlas/grid/detail/partitioner/MatchingMeshPartitionerSphericalPolygon.cc b/src/atlas/grid/detail/partitioner/MatchingMeshPartitionerSphericalPolygon.cc
@@ -34,9 +34,9 @@ PartitionerBuilder<MatchingMeshPartitionerSphericalPolygon> __builder("spherical
 }
 
 void MatchingMeshPartitionerSphericalPolygon::partition(const Grid& grid, int partitioning[]) const {
-    const eckit::mpi::Comm& comm = atlas::mpi::comm();
-    const int mpi_rank           = int(comm.rank());
-    const int mpi_size           = int(comm.size());
+    const auto& comm   = mpi::comm(prePartitionedMesh_.mpi_comm());
+    const int mpi_rank = int(comm.rank());
+    const int mpi_size = int(comm.size());
 
     ATLAS_TRACE("MatchingMeshPartitionerSphericalPolygon::partition");