Added OpenMP to the cubed-sphere interpolator and the matching cubed sphere partitioner. (#293)

* Added OpenMP to cubed sphere interpolator and matching mesh partitioner.

* Allow integer narrowing in triplet construction.

* Addressed race condition.

* Simplified weight triplet vector creation.

* Added const to lon-lat iterator declaration in partitioner.

* This might fail some CI...

* This should pass CI...

* Offending test case was an esoteric use case that wasn't that useful. Let's get rid of it.

* Removed redundant includes.

* Removed redundant helper function.

* Added test to demonstrate unusual lifetime features for const references.

* typo in case name

* Tidied up after review.

* Added check to deal with poorly performing CubedSphereIterator.

---------

Co-authored-by: Willem Deconinck <willem.deconinck@ecmwf.int>

Author: odlomax

src/atlas/grid/detail/partitioner/MatchingMeshPartitionerCubedSphere.cc

@@ -6,10 +6,17 @@
  */
 
 #include "atlas/grid/detail/partitioner/MatchingMeshPartitionerCubedSphere.h"
+
+#include <algorithm>
+#include <iterator>
+#include <vector>
+
 #include "atlas/grid/CubedSphereGrid.h"
 #include "atlas/grid/Iterator.h"
 #include "atlas/interpolation/method/cubedsphere/CellFinder.h"
 #include "atlas/parallel/mpi/mpi.h"
+#include "atlas/parallel/omp/omp.h"
+#include "atlas/util/Point.h"
 
 namespace atlas {
 namespace grid {
@@ -29,14 +36,24 @@ void MatchingMeshPartitionerCubedSphere::partition(const Grid& grid, int partiti
     const auto edgeEpsilon = epsilon;
     const size_t listSize  = 8;
 
-    // Loop over grid and set partioning[].
-    auto lonlatIt = grid.lonlat().begin();
-    for (gidx_t i = 0; i < grid.size(); ++i) {
-        // 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;
-        ++lonlatIt;
+
+    const auto setPartitioning = [&](const auto& lonLatIt) {
+        atlas_omp_parallel_for(gidx_t i = 0; i < grid.size(); ++i) {
+            const auto lonLat = *(lonLatIt + i);
+            // This is probably more expensive than it needs to be, as it performs
+            // a dry run of the cubedsphere interpolation method.
+            partitioning[i] = finder.getCell(lonLat, listSize, edgeEpsilon, epsilon).isect ? mpi_rank : -1;
+        }
+    };
+
+    // CubedSphereIterator::operator+=() is not implemented properly.
+    if (CubedSphereGrid(grid)) {
+        auto lonLats = std::vector<PointLonLat>{};
+        lonLats.reserve(grid.size());
+        std::copy(grid.lonlat().begin(), grid.lonlat().end(), std::back_inserter(lonLats));
+        setPartitioning(lonLats.begin());
+    } else {
+        setPartitioning(grid.lonlat().begin());
     }
 
     // AllReduce to get full partitioning array.

src/atlas/interpolation/method/cubedsphere/CubedSphereBilinear.cc

@@ -10,6 +10,7 @@
 #include "atlas/grid/CubedSphereGrid.h"
 #include "atlas/interpolation/method/MethodFactory.h"
 #include "atlas/interpolation/method/cubedsphere/CellFinder.h"
+#include "atlas/parallel/omp/omp.h"
 #include "atlas/util/CoordinateEnums.h"
 
 namespace atlas {
@@ -37,9 +38,6 @@ void CubedSphereBilinear::do_setup(const FunctionSpace& source, const FunctionSp
     ATLAS_ASSERT(ncSource);
     ATLAS_ASSERT(target_);
 
-    // Enable or disable halo exchange.
-    this->allow_halo_exchange_ = halo_exchange_;
-
 
     // return early if no output points on this partition reserve is called on
     // the triplets but also during the sparseMatrix constructor. This won't
@@ -54,16 +52,12 @@ void CubedSphereBilinear::do_setup(const FunctionSpace& source, const FunctionSp
     const auto N         = CubedSphereGrid(ncSource.mesh().grid()).N();
     const auto tolerance = 2. * std::numeric_limits<double>::epsilon() * N;
 
-    // Loop over target at calculate interpolation weights.
-    auto weights          = std::vector<Triplet>{};
+    // Weights vector is oversized. Empty triplets are ignored by Matrix constructor.
+    auto weights          = Triplets(4 * target_.size());
     const auto ghostView  = array::make_view<int, 1>(target_.ghost());
     const auto lonlatView = array::make_view<double, 2>(target_.lonlat());
-    const auto tijView    = array::make_view<idx_t, 2>(ncSource.mesh().cells().field("tij"));
-
-    // Make vector of tile indices for each target point (needed for vector field interpolation).
-    std::vector<idx_t> tileIndex{};
 
-    for (idx_t i = 0; i < target_.size(); ++i) {
+    atlas_omp_parallel_for(idx_t i = 0; i < target_.size(); ++i) {
         if (!ghostView(i)) {
             const auto cell =
                 finder.getCell(PointLonLat(lonlatView(i, LON), lonlatView(i, LAT)), listSize_, tolerance, tolerance);
@@ -75,24 +69,23 @@ void CubedSphereBilinear::do_setup(const FunctionSpace& source, const FunctionSp
                     std::to_string(i) + ".");
             }
 
-            tileIndex.push_back(tijView(cell.idx, 0));
             const auto& isect = cell.isect;
             const auto& j     = cell.nodes;
 
             switch (cell.nodes.size()) {
                 case (3): {
                     // Cell is a triangle.
-                    weights.emplace_back(i, j[0], 1. - isect.u - isect.v);
-                    weights.emplace_back(i, j[1], isect.u);
-                    weights.emplace_back(i, j[2], isect.v);
+                    weights[4 * i]     = Triplet(i, j[0], 1. - isect.u - isect.v);
+                    weights[4 * i + 1] = Triplet(i, j[1], isect.u);
+                    weights[4 * i + 2] = Triplet(i, j[2], isect.v);
                     break;
                 }
                 case (4): {
                     // Cell is quad.
-                    weights.emplace_back(i, j[0], (1. - isect.u) * (1. - isect.v));
-                    weights.emplace_back(i, j[1], isect.u * (1. - isect.v));
-                    weights.emplace_back(i, j[2], isect.u * isect.v);
-                    weights.emplace_back(i, j[3], (1. - isect.u) * isect.v);
+                    weights[4 * i]     = Triplet(i, j[0], (1. - isect.u) * (1. - isect.v));
+                    weights[4 * i + 1] = Triplet(i, j[1], isect.u * (1. - isect.v));
+                    weights[4 * i + 2] = Triplet(i, j[2], isect.u * isect.v);
+                    weights[4 * i + 3] = Triplet(i, j[3], (1. - isect.u) * isect.v);
                     break;
                 }
                 default: {
@@ -104,9 +97,6 @@ void CubedSphereBilinear::do_setup(const FunctionSpace& source, const FunctionSp
 
     // fill sparse matrix and return.
     setMatrix(target_.size(), source_.size(), weights);
-
-    // Add tile index metadata to target.
-    target_->metadata().set("tile index", tileIndex);
 }
 
 void CubedSphereBilinear::print(std::ostream&) const {

src/atlas/interpolation/method/cubedsphere/CubedSphereBilinear.h

@@ -30,7 +30,6 @@ class CubedSphereBilinear : public Method {
     CubedSphereBilinear(const Config& config): Method(config) {
         config.get("halo", halo_);
         config.get("list_size", listSize_);
-        config.get("halo_exchange", halo_exchange_);
     }
     virtual ~CubedSphereBilinear() override {}
 
@@ -49,7 +48,6 @@ class CubedSphereBilinear : public Method {
 
     int halo_{0};
     int listSize_{8};
-    bool halo_exchange_{true};
 };
 
 }  // namespace method

src/tests/interpolation/CMakeLists.txt

@@ -129,7 +129,8 @@ ecbuild_add_test( TARGET atlas_test_interpolation_structured2D_to_points
 ecbuild_add_test( TARGET atlas_test_interpolation_cubedsphere
   SOURCES  test_interpolation_cubedsphere.cc
   LIBS     atlas
-  MPI      6
+  MPI      3
+  OMP      3
   CONDITION eckit_HAVE_MPI AND MPI_SLOTS GREATER_EQUAL 6
   ENVIRONMENT ${ATLAS_TEST_ENVIRONMENT}
 )