Refactor sparse_matrix_multiply using SparseMatrixView. Adapt interpo…

…lation to hold SparseMatrixStorage and create views

Co-authored-by: Liam Adams <[email protected]>

src/atlas/CMakeLists.txt

@@ -705,12 +705,18 @@ linalg/View.h
 linalg/sparse.h
 linalg/sparse/Backend.h
 linalg/sparse/Backend.cc
+linalg/sparse/MakeEckitSparseMatrix.h
+linalg/sparse/MakeSparseMatrixStorageEckit.h
+linalg/sparse/MakeSparseMatrixStorageEigen.h
 linalg/sparse/SparseMatrixMultiply.h
 linalg/sparse/SparseMatrixMultiply.tcc
 linalg/sparse/SparseMatrixMultiply_EckitLinalg.h
 linalg/sparse/SparseMatrixMultiply_EckitLinalg.cc
 linalg/sparse/SparseMatrixMultiply_OpenMP.h
 linalg/sparse/SparseMatrixMultiply_OpenMP.cc
+linalg/sparse/SparseMatrixStorage.cc
+linalg/sparse/SparseMatrixStorage.h
+linalg/sparse/SparseMatrixView.h
 linalg/dense.h
 linalg/dense/Backend.h
 linalg/dense/Backend.cc

src/atlas/interpolation/Cache.cc

@@ -54,7 +54,7 @@ MatrixCacheEntry::~MatrixCacheEntry() = default;
 class MatrixCacheEntryOwned : public MatrixCacheEntry {
 public:
     MatrixCacheEntryOwned(Matrix&& matrix): MatrixCacheEntry(&matrix_) {
-        const_cast<Matrix&>(matrix_).swap(reinterpret_cast<Matrix&>(matrix));
+        const_cast<Matrix&>(matrix_) = std::move(matrix);
     }
 
 private:

src/atlas/interpolation/Cache.h

@@ -17,8 +17,8 @@
 #include "eckit/filesystem/PathName.h"
 #include "eckit/io/Buffer.h"
 
-#include "eckit/linalg/SparseMatrix.h"
 
+#include "atlas/linalg/sparse.h"
 #include "atlas/runtime/Exception.h"
 #include "atlas/util/KDTree.h"
 
@@ -82,7 +82,7 @@ class Cache {
 
 class MatrixCacheEntry : public InterpolationCacheEntry {
 public:
-    using Matrix = eckit::linalg::SparseMatrix;
+    using Matrix = atlas::linalg::SparseMatrixStorage;
     ~MatrixCacheEntry() override;
     MatrixCacheEntry(const Matrix* matrix, const std::string& uid = ""): matrix_{matrix}, uid_(uid) {
         ATLAS_ASSERT(matrix_ != nullptr);

src/atlas/interpolation/method/Method.cc

@@ -105,12 +105,12 @@ void Method::interpolate_field_rank1(const Field& src, Field& tgt, const Matrix&
     auto tgt_v   = array::make_view<Value, 1>(tgt);
 
     if (nonLinear_(src)) {
-        Matrix W_nl(W);  // copy (a big penalty -- copy-on-write would definitely be better)
+        eckit::linalg::SparseMatrix W_nl = atlas::linalg::make_eckit_sparse_matrix(W);
         nonLinear_->execute(W_nl, src);
         sparse_matrix_multiply(W_nl, src_v, tgt_v, backend);
     }
     else {
-        sparse_matrix_multiply(W, src_v, tgt_v, backend);
+        sparse_matrix_multiply(make_host_view<eckit::linalg::Scalar,eckit::linalg::Index>(W), src_v, tgt_v, backend);
     }
 }
 
@@ -150,7 +150,7 @@ void Method::interpolate_field_rank2(const Field& src, Field& tgt, const Matrix&
         }
     }
     else {
-        sparse_matrix_multiply(W, src_v, tgt_v, sparse::backend::openmp());
+        sparse_matrix_multiply(make_host_view<eckit::linalg::Scalar,eckit::linalg::Index>(W), src_v, tgt_v, sparse::backend::openmp());
     }
 }
 
@@ -163,7 +163,7 @@ void Method::interpolate_field_rank3(const Field& src, Field& tgt, const Matrix&
     if (not W.empty() && nonLinear_(src)) {
         ATLAS_ASSERT(false, "nonLinear interpolation not supported for rank-3 fields.");
     }
-    sparse_matrix_multiply(W, src_v, tgt_v, sparse::backend::openmp());
+    sparse_matrix_multiply(make_host_view<eckit::linalg::Scalar,eckit::linalg::Index>(W), src_v, tgt_v, sparse::backend::openmp());
 }
 
 template <typename Value>
@@ -173,7 +173,7 @@ void Method::adjoint_interpolate_field_rank1(Field& src, const Field& tgt, const
     auto src_v = array::make_view<Value, 1>(src);
     auto tgt_v = array::make_view<Value, 1>(tgt);
 
-    sparse_matrix_multiply_add(W, tgt_v, src_v, backend);
+    sparse_matrix_multiply_add(make_host_view<eckit::linalg::Scalar,eckit::linalg::Index>(W), tgt_v, src_v, backend);
 }
 
 template <typename Value>
@@ -182,7 +182,7 @@ void Method::adjoint_interpolate_field_rank2(Field& src, const Field& tgt, const
     auto src_v = array::make_view<Value, 2>(src);
     auto tgt_v = array::make_view<Value, 2>(tgt);
 
-    sparse_matrix_multiply_add(W, tgt_v, src_v, sparse::backend::openmp());
+    sparse_matrix_multiply_add(make_host_view<eckit::linalg::Scalar,eckit::linalg::Index>(W), tgt_v, src_v, sparse::backend::openmp());
 }
 
 template <typename Value>
@@ -191,7 +191,7 @@ void Method::adjoint_interpolate_field_rank3(Field& src, const Field& tgt, const
     auto src_v = array::make_view<Value, 3>(src);
     auto tgt_v = array::make_view<Value, 3>(tgt);
 
-    sparse_matrix_multiply_add(W, tgt_v, src_v, sparse::backend::openmp());
+    sparse_matrix_multiply_add(make_host_view<eckit::linalg::Scalar,eckit::linalg::Index>(W), tgt_v, src_v, sparse::backend::openmp());
 }
 
 void Method::check_compatibility(const Field& src, const Field& tgt, const Matrix& W) const {
@@ -266,12 +266,11 @@ void Method::setup(const FunctionSpace& source, const FunctionSpace& target) {
     ATLAS_TRACE("atlas::interpolation::method::Method::setup(FunctionSpace, FunctionSpace)");
     this->do_setup(source, target);
 
-    if (adjoint_ && target.size() > 0) {
-        Matrix tmp(*matrix_);
-
-        // if interpolation is matrix free then matrix->nonZeros() will be zero.
-        if (tmp.nonZeros() > 0) {
-            matrix_transpose_ = tmp.transpose();
+    if (adjoint_ && target.size() > 0 && matrixAllocated()) {
+        if (not matrix_->empty()) {
+            eckit::linalg::SparseMatrix matrix_copy = make_eckit_sparse_matrix(*matrix_); // Makes a copy!
+            matrix_copy.transpose(); // transpose the copy in place
+            matrix_transpose_ = linalg::make_sparse_matrix_storage(std::move(matrix_copy)); // Move the copy into storage
         }
     }
 }

src/atlas/interpolation/method/Method.h

@@ -16,10 +16,12 @@
 
 #include "atlas/interpolation/Cache.h"
 #include "atlas/interpolation/NonLinear.h"
+#include "atlas/linalg/sparse/SparseMatrixStorage.h"
 #include "atlas/util/Metadata.h"
 #include "atlas/util/Object.h"
 #include "eckit/config/Configuration.h"
 #include "eckit/linalg/SparseMatrix.h"
+#include "atlas/linalg/sparse/MakeSparseMatrixStorageEckit.h"
 
 namespace atlas {
 class Field;
@@ -87,7 +89,7 @@ class Method : public util::Object {
 
     using Triplet  = eckit::linalg::Triplet;
     using Triplets = std::vector<Triplet>;
-    using Matrix   = eckit::linalg::SparseMatrix;
+    using Matrix   = atlas::linalg::SparseMatrixStorage;
 
     static void normalise(Triplets& triplets);
 
@@ -102,13 +104,15 @@ class Method : public util::Object {
     friend class interpolation::MatrixCache;
 
 protected:
-    void setMatrix(Matrix& m, const std::string& uid = "") {
+
+    void setMatrix(Matrix&& m, const std::string& uid = "") {
         if (not matrix_shared_) {
             matrix_shared_ = std::make_shared<Matrix>();
         }
-        matrix_shared_->swap(m);
+        *matrix_shared_ = std::move(m);
         matrix_cache_ = interpolation::MatrixCache(matrix_shared_, uid);
         matrix_       = &matrix_cache_.matrix();
+
     }
 
     void setMatrix(interpolation::MatrixCache matrix_cache) {
@@ -118,6 +122,14 @@ class Method : public util::Object {
         matrix_shared_.reset();
     }
 
+    void setMatrix(eckit::linalg::SparseMatrix&& m, const std::string& uid = "") {
+        setMatrix( linalg::make_sparse_matrix_storage(std::move(m)), uid );
+    }
+
+    void setMatrix(std::size_t rows, std::size_t cols, const Triplets& triplets, const std::string& uid = "") {
+        setMatrix( eckit::linalg::SparseMatrix{rows, cols, triplets}, uid);
+    }
+
     bool matrixAllocated() const { return matrix_shared_.use_count(); }
 
     const Matrix& matrix() const { return *matrix_; }

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

@@ -5,21 +5,23 @@
  * which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
  */
 
+#include "eckit/config/LocalConfiguration.h"
+#include "eckit/linalg/Triplet.h"
+#include "eckit/linalg/SparseMatrix.h"
+#include "eckit/mpi/Comm.h"
 
 #include "atlas/functionspace/NodeColumns.h"
 #include "atlas/functionspace/StructuredColumns.h"
 #include "atlas/grid.h"
 #include "atlas/interpolation/Interpolation.h"
 #include "atlas/interpolation/method/binning/Binning.h"
 #include "atlas/interpolation/method/MethodFactory.h"
+#include "atlas/linalg/sparse/SparseMatrixStorage.h"
+#include "atlas/linalg/sparse/MakeEckitSparseMatrix.h"
 #include "atlas/mesh.h"
 #include "atlas/mesh/actions/GetCubedSphereNodalArea.h"
 #include "atlas/runtime/Trace.h"
 
-#include "eckit/config/LocalConfiguration.h"
-#include "eckit/linalg/SparseMatrix.h"
-#include "eckit/linalg/Triplet.h"
-#include "eckit/mpi/Comm.h"
 
 
 namespace atlas {
@@ -56,7 +58,8 @@ void Binning::do_setup(const FunctionSpace& source,
                        const FunctionSpace& target) {
   ATLAS_TRACE("atlas::interpolation::method::Binning::do_setup()");
 
-  using Index = eckit::linalg::Index;
+  using Scalar  = eckit::linalg::Scalar;
+  using Index   = eckit::linalg::Index;
   using Triplet = eckit::linalg::Triplet;
   using SMatrix = eckit::linalg::SparseMatrix;
 
@@ -80,14 +83,16 @@ void Binning::do_setup(const FunctionSpace& source,
 
   auto smx_interp = smx_interp_cache.matrix();
 
-  auto smx_interp_tr = smx_interp.transpose();
+  auto eckit_smx_interp = make_non_owning_eckit_sparse_matrix(smx_interp);
+  SMatrix smx_interp_tr(eckit_smx_interp); // copy
+  smx_interp_tr.transpose(); // transpose the copy in-place
 
   const auto rows_tamx = smx_interp_tr.rows();
   const auto cols_tamx = smx_interp_tr.cols();
 
-  const double* ptr_tamx_data = smx_interp_tr.data();
-  const Index* ptr_tamx_idxs_col = smx_interp_tr.inner();
-  const Index* ptr_tamx_o = smx_interp_tr.outer();
+  const Scalar* ptr_tamx_data = smx_interp_tr.data();
+  const Index*  ptr_tamx_idxs_col = smx_interp_tr.inner();
+  const Index*  ptr_tamx_o = smx_interp_tr.outer();
 
   // diagonal of 'area weights matrix', W
   auto ds_aweights = getAreaWeights(source_);
@@ -123,8 +128,7 @@ void Binning::do_setup(const FunctionSpace& source,
   }
 
   // 'binning matrix' (sparse matrix), B = N A^T W
-  SMatrix smx_binning{rows_tamx, cols_tamx, smx_binning_els};
-  setMatrix(smx_binning);
+  setMatrix(rows_tamx, cols_tamx, smx_binning_els);
 }
 
 

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

@@ -99,8 +99,7 @@ void CubedSphereBilinear::do_setup(const FunctionSpace& source, const FunctionSp
     }
 
     // fill sparse matrix and return.
-    Matrix A(target_.size(), source_.size(), weights);
-    setMatrix(A);
+    setMatrix(target_.size(), source_.size(), weights);
 
     // Add tile index metadata to target.
     target_->metadata().set("tile index", tileIndex);

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

@@ -22,6 +22,7 @@
 #include "atlas/functionspace/PointCloud.h"
 #include "atlas/interpolation/method/MethodFactory.h"
 #include "atlas/runtime/Exception.h"
+#include "atlas/linalg/sparse/MakeEckitSparseMatrix.h"
 
 
 namespace atlas {
@@ -60,8 +61,8 @@ void GridBoxMaximum::do_execute(const Field& source, Field& target, Metadata&) c
 
 
     if (!matrixFree_) {
-        const Matrix& m = matrix();
-        Matrix::const_iterator k(m);
+        const auto m = atlas::linalg::make_non_owning_eckit_sparse_matrix(matrix());
+        auto k = m.begin();
 
         for (decltype(m.rows()) i = 0, j = 0; i < m.rows(); ++i) {
             double max = std::numeric_limits<double>::lowest();

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

@@ -177,8 +177,7 @@ void GridBoxMethod::do_setup(const Grid& source, const Grid& target, const Cache
 
     {
         ATLAS_TRACE("GridBoxMethod::setup: build interpolant matrix");
-        Matrix A(targetBoxes_.size(), sourceBoxes_.size(), allTriplets);
-        setMatrix(A);
+        setMatrix(targetBoxes_.size(), sourceBoxes_.size(), allTriplets);
     }
 }
 

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

@@ -130,8 +130,7 @@ void KNearestNeighbours::do_setup(const FunctionSpace& source, const FunctionSpa
     }
 
     // fill sparse matrix and return
-    Matrix A(out_npts, inp_npts, weights_triplets);
-    setMatrix(A);
+    setMatrix(out_npts, inp_npts, weights_triplets);
 }
 
 }  // namespace method

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

@@ -102,8 +102,7 @@ void NearestNeighbour::do_setup(const FunctionSpace& source, const FunctionSpace
     }
 
     // fill sparse matrix and return
-    Matrix A(out_npts, inp_npts, weights_triplets);
-    setMatrix(A);
+    setMatrix( out_npts, inp_npts, weights_triplets );
 }
 
 }  // namespace method

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

@@ -18,6 +18,7 @@
 #include "atlas/interpolation/method/MethodFactory.h"
 #include "atlas/interpolation/method/sphericalvector/ComplexMatrixMultiply.h"
 #include "atlas/interpolation/method/sphericalvector/Types.h"
+#include "atlas/linalg/sparse/MakeEckitSparseMatrix.h"
 #include "atlas/option/Options.h"
 #include "atlas/parallel/omp/omp.h"
 #include "atlas/runtime/Exception.h"
@@ -65,12 +66,13 @@ void SphericalVector::do_setup(const FunctionSpace& source,
   setMatrix(Interpolation(interpolationScheme_, source_, target_));
 
   // Get matrix data.
-  const auto nRows = static_cast<Index>(matrix().rows());
-  const auto nCols = static_cast<Index>(matrix().cols());
-  const auto nNonZeros = static_cast<std::size_t>(matrix().nonZeros());
-  const auto* outerIndices = matrix().outer();
-  const auto* innerIndices = matrix().inner();
-  const auto* baseWeights = matrix().data();
+  const auto m = atlas::linalg::make_non_owning_eckit_sparse_matrix(matrix());
+  const auto nRows = static_cast<Index>(m.rows());
+  const auto nCols = static_cast<Index>(m.cols());
+  const auto nNonZeros = static_cast<std::size_t>(m.nonZeros());
+  const auto* outerIndices = m.outer();
+  const auto* innerIndices = m.inner();
+  const auto* baseWeights  = m.data();
 
   // Note: need to store copy of weights as Eigen3 sorts compressed rows by j
   // whereas eckit does not.

src/atlas/interpolation/method/structured/StructuredInterpolation2D.tcc

@@ -428,8 +428,7 @@ void StructuredInterpolation2D<Kernel>::setup( const FunctionSpace& source ) {
         // fill sparse matrix
         if( failed_points.empty() && out_npts_) {
             idx_t inp_npts = source.size();
-            Matrix A( out_npts_, inp_npts, triplets );
-            setMatrix(A);
+            setMatrix(out_npts_, inp_npts, triplets);
         }
     }
 }

src/atlas/interpolation/method/unstructured/ConservativeSphericalPolygonInterpolation.cc

@@ -870,13 +870,11 @@ void ConservativeSphericalPolygonInterpolation::do_setup(const FunctionSpace& sr
         stopwatch.start();
         switch (order_) {
             case 1: {
-                auto M = compute_1st_order_matrix();
-                setMatrix(M, "1");
+                setMatrix(n_tpoints_, n_spoints_, compute_1st_order_triplets(), "1");
                 break;
             }
             case 2: {
-                auto M = compute_2nd_order_matrix();
-                setMatrix(M, "2");
+                setMatrix(n_tpoints_, n_spoints_, compute_2nd_order_triplets(), "2");
                 break;
             }
             default: {
@@ -1248,7 +1246,7 @@ void ConservativeSphericalPolygonInterpolation::intersect_polygons(const CSPolyg
     }
 }
 
-eckit::linalg::SparseMatrix ConservativeSphericalPolygonInterpolation::compute_1st_order_matrix() {
+ConservativeSphericalPolygonInterpolation::Triplets ConservativeSphericalPolygonInterpolation::compute_1st_order_triplets() {
     ATLAS_TRACE("ConservativeMethod::setup: build cons-1 interpolant matrix");
     ATLAS_ASSERT(not matrix_free_);
     Triplets triplets;
@@ -1355,10 +1353,10 @@ eckit::linalg::SparseMatrix ConservativeSphericalPolygonInterpolation::compute_1
 //             }
 //         }
 //     }
-    return Matrix(n_tpoints_, n_spoints_, triplets);
+    return triplets;
 }
 
-eckit::linalg::SparseMatrix ConservativeSphericalPolygonInterpolation::compute_2nd_order_matrix() {
+ConservativeSphericalPolygonInterpolation::Triplets ConservativeSphericalPolygonInterpolation::compute_2nd_order_triplets() {
     ATLAS_TRACE("ConservativeMethod::setup: build cons-2 interpolant matrix");
     ATLAS_ASSERT(not matrix_free_);
     const auto& src_points_ = data_->src_points_;
@@ -1562,7 +1560,7 @@ eckit::linalg::SparseMatrix ConservativeSphericalPolygonInterpolation::compute_2
         }
     }
     sort_and_accumulate_triplets(triplets);  // Very expensive!!! (90% of this routine). We need to avoid it
-    return Matrix(n_tpoints_, n_spoints_, triplets);
+    return triplets;
 }
 
 void ConservativeSphericalPolygonInterpolation::do_execute(const FieldSet& src_fields, FieldSet& tgt_fields,