Global interpolation matrix --> MPI-distributed interpolation matrices

src/atlas/interpolation/AssembleGlobalMatrix.cc

@@ -11,15 +11,18 @@
 
 #include "AssembleGlobalMatrix.h"
 
+#include <unordered_map>
+
 #include "eckit/linalg/types.h"
 
 #include "atlas/array.h"
 #include "atlas/linalg/sparse/SparseMatrixToTriplets.h"
 #include "atlas/functionspace/StructuredColumns.h"
+#include "atlas/interpolation/Cache.h"
 
 namespace atlas::interpolation {
 
-atlas::linalg::SparseMatrixStorage assemble_global_matrix(const Interpolation& interpolation, int mpi_root) {
+linalg::SparseMatrixStorage assemble_global_matrix(const Interpolation& interpolation, int mpi_root) {
 
     auto src_fs = interpolation.source();
     auto tgt_fs = interpolation.target();
@@ -50,10 +53,10 @@ atlas::linalg::SparseMatrixStorage assemble_global_matrix(const Interpolation& i
                 return src_remote_index(idx) + one_if_structuredcolumns;
             };
 
-            const auto src_global_index = array::make_view<gidx_t,1>(src_fs.global_index());
-            const auto tgt_global_index = array::make_view<gidx_t,1>(tgt_fs.global_index());
-            const auto src_part         = array::make_view<int,1>(src_fs.partition());
-            const auto tgt_ghost        = array::make_view<int,1>(tgt_fs.ghost());
+            const auto src_global_index = array::make_view<gidx_t, 1>(src_fs.global_index());
+            const auto tgt_global_index = array::make_view<gidx_t, 1>(tgt_fs.global_index());
+            const auto src_part         = array::make_view<int, 1>(src_fs.partition());
+            const auto tgt_ghost        = array::make_view<int, 1>(tgt_fs.ghost());
 
             eckit::mpi::Buffer<gidx_t> recv_global_idx_buf(mpi_size);
             {
@@ -149,7 +152,7 @@ atlas::linalg::SparseMatrixStorage assemble_global_matrix(const Interpolation& i
     gidx_t tgt_max_gidx = compute_max_global_index(tgt_fs);
     gidx_t src_max_gidx = compute_max_global_index(src_fs);
 
-    atlas::linalg::SparseMatrixStorage global_matrix;
+    linalg::SparseMatrixStorage global_matrix;
     if (mpi_rank == mpi_root) {
         size_t nrows = tgt_max_gidx;
         size_t ncols = src_max_gidx;
@@ -160,6 +163,139 @@ atlas::linalg::SparseMatrixStorage assemble_global_matrix(const Interpolation& i
     return global_matrix;
 }
 
+template <typename ViewValue, typename ViewIndex, typename Value, typename Index> 
+void distribute_global_matrix(const linalg::SparseMatrixView<ViewValue,ViewIndex>& global_matrix,
+    const array::Array& partition, std::vector<Index>& rows, std::vector<Index>& cols, std::vector<Value>& vals, int mpi_root) {
+    ATLAS_TRACE("distribute_global_matrix_lowlevel");
+
+    const auto tgt_part_glb = array::make_view<int,1>(partition);
+
+    auto& mpi_comm = mpi::comm();
+    auto mpi_size  = mpi_comm.size();
+    auto mpi_rank  = mpi_comm.rank();
+
+    // compute how many nnz-entries each task gets
+    size_t nnz_loc = 0;
+    int mpi_tag = 0;
+    std::vector<std::size_t> nnz_per_task(mpi_size);
+    const auto outer  = global_matrix.outer();
+    const auto inner  = global_matrix.inner();
+    const auto value  = global_matrix.value();
+    if (mpi_rank == mpi_root) {        
+
+        for(std::size_t r = 0; r < global_matrix.rows(); ++r) {
+            nnz_per_task[tgt_part_glb(r)] += outer[r+1] - outer[r];
+        }
+        for (int jproc = 0; jproc < mpi::comm().size(); ++jproc) {
+            if (jproc != mpi_root) {
+                mpi::comm().send(nnz_per_task.data() + jproc, 1, jproc, mpi_tag);
+            }
+        }
+        nnz_loc = nnz_per_task[mpi_root];
+    }
+    else {
+        mpi_comm.receive(&nnz_loc, 1, mpi_root, mpi_tag);
+    }
+
+    rows.resize(nnz_loc);
+    cols.resize(nnz_loc);
+    vals.resize(nnz_loc);
+
+    if (mpi_rank == mpi_root) {
+        std::vector<std::vector<Index>> send_rows(mpi_size);
+        std::vector<std::vector<Index>> send_cols(mpi_size);
+        std::vector<std::vector<Value>> send_vals(mpi_size);
+        for(std::size_t jproc=0; jproc < mpi_size; ++jproc) {
+            send_rows[jproc].reserve(nnz_per_task[jproc]);
+            send_cols[jproc].reserve(nnz_per_task[jproc]);
+            send_vals[jproc].reserve(nnz_per_task[jproc]);
+        }
+        for(std::size_t r = 0; r < global_matrix.rows(); ++r) {
+            int jproc = tgt_part_glb(r);
+            for (auto c = outer[r]; c < outer[r + 1]; ++c) {
+                auto col = inner[c];
+                send_rows[jproc].emplace_back(r);
+                send_cols[jproc].emplace_back(col);
+                send_vals[jproc].emplace_back(value[c]);
+            }
+        }
+        for(std::size_t jproc = 0; jproc < mpi_size; ++jproc) {
+            if (jproc != mpi_root) {
+                mpi_comm.send(send_rows[jproc].data(), send_rows[jproc].size(), jproc, mpi_tag);
+                mpi_comm.send(send_cols[jproc].data(), send_cols[jproc].size(), jproc, mpi_tag);
+                mpi_comm.send(send_vals[jproc].data(), send_vals[jproc].size(), jproc, mpi_tag);
+            }
+            else {
+                rows = send_rows[jproc];
+                cols = send_cols[jproc];
+                vals = send_vals[jproc];
+            }
+        }
+    }
+    else {
+        mpi_comm.receive(rows.data(), nnz_loc, mpi_root, mpi_tag);
+        mpi_comm.receive(cols.data(), nnz_loc, mpi_root, mpi_tag);
+        mpi_comm.receive(vals.data(), nnz_loc, mpi_root, mpi_tag);
+    }
+}
+
+linalg::SparseMatrixStorage distribute_global_matrix(const FunctionSpace& src_fs, const FunctionSpace& tgt_fs, const linalg::SparseMatrixStorage& gmatrix, int mpi_root) {
+    ATLAS_TRACE("distribute_global_matrix");
+    Field field_tgt_part_glb = tgt_fs.createField(tgt_fs.partition(), option::global(mpi_root));
+    ATLAS_TRACE_SCOPE("gather partition") {
+        tgt_fs.gather(tgt_fs.partition(), field_tgt_part_glb);
+    }
+
+    using Index = eckit::linalg::Index;
+    using Value = eckit::linalg::Scalar;
+    std::vector<Index> rows, cols;
+    std::vector<Value> vals;
+    distribute_global_matrix(atlas::linalg::make_host_view<Value, Index>(gmatrix), field_tgt_part_glb, rows, cols, vals, mpi_root);
+
+    // map global index to local index
+    std::unordered_map<gidx_t, idx_t> to_local_rows;
+    std::unordered_map<gidx_t, idx_t> to_local_cols;
+
+    ATLAS_TRACE_SCOPE("convert to local indexing") {
+        auto tgt_gidx_exchanged = tgt_fs.createField(tgt_fs.global_index());
+        tgt_gidx_exchanged.array().copy(tgt_fs.global_index());
+        tgt_fs.haloExchange(tgt_gidx_exchanged);
+        const auto tgt_global_index = array::make_view<gidx_t, 1>(tgt_gidx_exchanged);
+        const auto tgt_ghost = array::make_view<int,1>(tgt_fs.ghost());
+
+        auto src_gidx_exchanged = src_fs.createField(src_fs.global_index());
+        src_gidx_exchanged.array().copy(src_fs.global_index());
+        src_fs.haloExchange(src_gidx_exchanged);
+        const auto src_global_index = array::make_view<gidx_t, 1>(src_gidx_exchanged);
+        const auto src_ghost = array::make_view<int,1>(src_fs.ghost());
+
+        for (idx_t r = 0; r < tgt_global_index.size(); ++r) {
+            auto gr = tgt_global_index(r);
+            if (tgt_ghost(r) && to_local_rows.find(gr) != to_local_rows.end()) {
+                continue;
+            }
+            to_local_rows[gr] = r;
+        }
+        for (idx_t c = 0; c < src_global_index.size(); ++c) {
+            auto gc = src_global_index(c);
+            if (src_ghost(c) && to_local_cols.find(gc) != to_local_cols.end()) {
+                continue;
+            }
+            to_local_cols[gc] = c;
+        }
+        for (int r = 0; r < rows.size(); ++r) {
+            rows[r] = to_local_rows[rows[r] + 1];
+            cols[r] = to_local_cols[cols[r] + 1];
+        }
+    }
+
+    linalg::SparseMatrixStorage matrix;
+    constexpr int index_base = 0;
+    constexpr bool is_sorted = false;
+    matrix = linalg::make_sparse_matrix_storage_from_rows_columns_values(tgt_fs.size(), src_fs.size(), rows, cols, vals, index_base, is_sorted);
+
+    return matrix;
+}
 
 
 } //end namespace

src/atlas/interpolation/AssembleGlobalMatrix.h

@@ -17,4 +17,6 @@ namespace atlas::interpolation {
 
     atlas::linalg::SparseMatrixStorage assemble_global_matrix(const Interpolation& interpolation, int mpi_root = 0);
 
+    atlas::linalg::SparseMatrixStorage distribute_global_matrix(const FunctionSpace& src_fs, const FunctionSpace& tgt_fs, const linalg::SparseMatrixStorage&, int mpi_root = 0);
+
 } // namespace atlas::interpolation

src/atlas/interpolation/Interpolation.cc

@@ -121,6 +121,16 @@ Interpolation::Interpolation(const Interpolation::Config& config, const Grid& so
         return impl;
     }()) {}
 
+Interpolation::Interpolation(const Interpolation::Config& config, const FunctionSpace& fs_in, const FunctionSpace& fs_out,
+                             const Interpolation::Cache& cache):
+    Handle([&]() -> Implementation* {
+        std::string type;
+        ATLAS_ASSERT(config.get("type", type));
+        Implementation* impl = interpolation::MethodFactory::build(type, config);
+        impl->setup(fs_in, fs_out, cache);
+        return impl;
+    }()) {}
+
 extern "C" {
 Interpolation::Implementation* atlas__Interpolation__new(const eckit::Parametrisation* config,
                                                          const functionspace::FunctionSpaceImpl* source,

src/atlas/interpolation/Interpolation.h

@@ -70,6 +70,8 @@ class Interpolation : DOXYGEN_HIDE(public util::ObjectHandle<interpolation::Meth
 
     Interpolation(const Config&, const Grid& source, const Grid& target, const Cache&) noexcept(false);
 
+    Interpolation(const Config&, const FunctionSpace& source, const FunctionSpace& target, const Cache&) noexcept(false);
+
     friend std::ostream& operator<<(std::ostream& out, const Interpolation& i) {
         i.print(out);
         return out;

src/atlas/interpolation/method/Method.cc

@@ -295,6 +295,11 @@ void Method::setup(const Grid& source, const Grid& target, const Cache& cache) {
     this->do_setup(source, target, cache);
 }
 
+void Method::setup(const FunctionSpace& source, const FunctionSpace& target, const Cache& cache) {
+    ATLAS_TRACE("atlas::interpolation::method::Method::setup(FunctionSpace, FunctionSpace, Cache)");
+    this->do_setup(source, target, cache);
+}
+
 Method::Metadata Method::execute(const FieldSet& source, FieldSet& target) const {
     ATLAS_TRACE("atlas::interpolation::method::Method::execute(FieldSet, FieldSet)");
     Metadata metadata;

src/atlas/interpolation/method/Method.h

@@ -57,6 +57,7 @@ class Method : public util::Object {
     void setup(const FunctionSpace& source, const Field& target);
     void setup(const FunctionSpace& source, const FieldSet& target);
     void setup(const Grid& source, const Grid& target, const Cache&);
+    void setup(const FunctionSpace& source, const FunctionSpace& target, const Cache&);
 
     Metadata execute(const FieldSet& source, FieldSet& target) const;
     Metadata execute(const Field& source, Field& target) const;
@@ -136,6 +137,7 @@ class Method : public util::Object {
 
     virtual void do_setup(const FunctionSpace& source, const FunctionSpace& target) = 0;
     virtual void do_setup(const Grid& source, const Grid& target, const Cache&)     = 0;
+    virtual void do_setup(const FunctionSpace& source, const FunctionSpace& target, const Cache&) = 0;
     virtual void do_setup(const FunctionSpace& source, const Field& target);
     virtual void do_setup(const FunctionSpace& source, const FieldSet& target);
 

src/atlas/interpolation/method/binning/Binning.cc

@@ -53,6 +53,12 @@ void Binning::do_setup(const Grid& source,
   ATLAS_NOTIMPLEMENTED;
 }
 
+void Binning::do_setup(const FunctionSpace& source,
+                       const FunctionSpace& target,
+                       const Cache&) {
+  ATLAS_NOTIMPLEMENTED;
+}
+
 
 void Binning::do_setup(const FunctionSpace& source,
                        const FunctionSpace& target) {

src/atlas/interpolation/method/binning/Binning.h

@@ -59,6 +59,7 @@ class Binning : public Method {
   void do_setup(const FunctionSpace& source,
                 const FunctionSpace& target) override;
   void do_setup(const Grid& source, const Grid& target, const Cache&) override;
+  void do_setup(const FunctionSpace& source, const FunctionSpace& target, const Cache&) override;
 
   std::vector<double> getAreaWeights(const FunctionSpace& source) const;
 

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

@@ -24,6 +24,10 @@ void CubedSphereBilinear::do_setup(const Grid& source, const Grid& target, const
     ATLAS_NOTIMPLEMENTED;
 }
 
+void CubedSphereBilinear::do_setup(const FunctionSpace& source, const FunctionSpace& target, const Cache&) {
+    ATLAS_NOTIMPLEMENTED;
+}
+
 void CubedSphereBilinear::do_setup(const FunctionSpace& source, const FunctionSpace& target) {
     source_ = source;
     target_ = target;

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

@@ -42,6 +42,7 @@ class CubedSphereBilinear : public Method {
     using Method::do_setup;
     void do_setup(const FunctionSpace& source, const FunctionSpace& target) override;
     void do_setup(const Grid& source, const Grid& target, const Cache&) override;
+     void do_setup(const FunctionSpace& source, const FunctionSpace& target, const Cache&) override;
 
     FunctionSpace source_;
     FunctionSpace target_;

src/atlas/interpolation/method/knn/GridBoxMethod.cc

@@ -181,6 +181,21 @@ void GridBoxMethod::do_setup(const Grid& source, const Grid& target, const Cache
     }
 }
 
+void GridBoxMethod::do_setup(const FunctionSpace& source, const FunctionSpace& target, const Cache& cache) {
+    ATLAS_TRACE("GridBoxMethod::setup()");
+
+    if (not matrixFree_ && interpolation::MatrixCache(cache)) {
+        setMatrix(cache);
+        source_ = source;
+        target_ = target;
+        ATLAS_ASSERT(matrix().rows() == target.size());
+        ATLAS_ASSERT(matrix().cols() == source.size());
+        return;
+    }
+
+    Log::warning() << "Can not create GridBoxMethod from (FunctionSpace, FunctionSpace, Cache). Use (Grid, Grid, Cache)";
+    ATLAS_NOTIMPLEMENTED;
+}
 
 void GridBoxMethod::giveUp(const std::forward_list<size_t>& failures) {
     Log::warning() << "Failed to intersect grid boxes: ";

src/atlas/interpolation/method/knn/GridBoxMethod.h

@@ -40,6 +40,7 @@ class GridBoxMethod : public KNearestNeighboursBase {
      */
     virtual void do_setup(const FunctionSpace& source, const FunctionSpace& target) override;
     virtual void do_setup(const Grid& source, const Grid& target, const Cache&) override;
+    virtual void do_setup(const FunctionSpace& source, const FunctionSpace& target, const Cache&) override;
 
     virtual const FunctionSpace& source() const override { return source_; }
     virtual const FunctionSpace& target() const override { return target_; }

src/atlas/interpolation/method/knn/KNearestNeighbours.cc

@@ -62,11 +62,19 @@ void KNearestNeighbours::do_setup(const Grid& source, const Grid& target, const
     do_setup(functionspace(source), functionspace(target));
 }
 
+void KNearestNeighbours::do_setup(const FunctionSpace& source, const FunctionSpace& target, const Cache& cache) {
+    if (interpolation::MatrixCache(cache)) {
+        setMatrix(cache);
+        source_ = source;
+        target_ = target;
+        buildPointSearchTree(source);
+    }
+}
+
 void KNearestNeighbours::do_setup(const FunctionSpace& source, const FunctionSpace& target) {
     source_                        = source;
     target_                        = target;
 
-    // build point-search tree
     buildPointSearchTree(source);
 
     array::ArrayView<double, 2> lonlat = array::make_view<double, 2>(target.lonlat());

src/atlas/interpolation/method/knn/KNearestNeighbours.h

@@ -39,6 +39,7 @@ class KNearestNeighbours : public KNearestNeighboursBase {
     using KNearestNeighboursBase::do_setup;
     virtual void do_setup(const FunctionSpace& source, const FunctionSpace& target) override;
     virtual void do_setup(const Grid& source, const Grid& target, const Cache&) override;
+    virtual void do_setup(const FunctionSpace& source, const FunctionSpace& target, const Cache&) override;
 
     FunctionSpace source_;
     FunctionSpace target_;

src/atlas/interpolation/method/knn/KNearestNeighboursBase.cc

@@ -31,7 +31,6 @@ void KNearestNeighboursBase::buildPointSearchTree(Mesh& meshSource, const mesh::
     ATLAS_TRACE();
     eckit::TraceTimer<Atlas> tim("KNearestNeighboursBase::buildPointSearchTree()");
 
-
     auto lonlat = array::make_view<double, 2>(meshSource.nodes().lonlat());
     auto halo   = array::make_view<int, 1>(meshSource.nodes().halo());
     int h       = _halo.size();
@@ -48,7 +47,6 @@ void KNearestNeighboursBase::buildPointSearchTree(Mesh& meshSource, const mesh::
     }
     pTree_.build();
 
-
 //    // generate 3D point coordinates
 //    mesh::actions::BuildXYZField("xyz")(meshSource);
 }

src/atlas/interpolation/method/knn/NearestNeighbour.cc

@@ -38,6 +38,15 @@ MethodBuilder<NearestNeighbour> __builder("nearest-neighbour");
 
 }  // namespace
 
+void NearestNeighbour::do_setup(const FunctionSpace& source, const FunctionSpace& target, const Cache& cache) {
+    if (interpolation::MatrixCache(cache)) {
+        setMatrix(cache);
+        source_ = source;
+        target_ = target;
+        buildPointSearchTree(source);
+    }
+}
+
 void NearestNeighbour::do_setup(const Grid& source, const Grid& target, const Cache&) {
     if (mpi::size() > 1) {
         ATLAS_NOTIMPLEMENTED;

src/atlas/interpolation/method/knn/NearestNeighbour.h

@@ -39,8 +39,8 @@ class NearestNeighbour : public KNearestNeighboursBase {
    * @param target functionspace containing target points
    */
     virtual void do_setup(const FunctionSpace& source, const FunctionSpace& target) override;
-
     virtual void do_setup(const Grid& source, const Grid& target, const Cache&) override;
+    virtual void do_setup(const FunctionSpace& source, const FunctionSpace& target, const Cache&) override;
 
     FunctionSpace source_;
     FunctionSpace target_;

src/atlas/interpolation/method/sphericalvector/SphericalVector.cc

@@ -53,6 +53,11 @@ void SphericalVector::do_setup(const Grid& source, const Grid& target,
   ATLAS_NOTIMPLEMENTED;
 }
 
+void SphericalVector::do_setup(const FunctionSpace& source, const FunctionSpace& target,
+                               const Cache&) {
+  ATLAS_NOTIMPLEMENTED;
+}
+
 void SphericalVector::do_setup(const FunctionSpace& source,
                                const FunctionSpace& target) {
   ATLAS_TRACE("interpolation::method::SphericalVector::do_setup");