Add GPU offload capable sparse matrix

hic/tests/test_hicsparse.cc

@@ -283,6 +283,7 @@ int test_hicsparseSpMM() {
     }
 
     // Clean up
+    HIC_CALL(hicFree(buffer));
     HIC_CALL(hicFree(dC));
     HIC_CALL(hicFree(dB));
     HIC_CALL(hicFree(dcolumn_indices));

src/atlas/CMakeLists.txt

@@ -703,6 +703,8 @@ linalg/Indexing.h
 linalg/Introspection.h
 linalg/View.h
 linalg/sparse.h
+linalg/SparseMatrix.h
+linalg/SparseMatrix.cc
 linalg/sparse/Backend.h
 linalg/sparse/Backend.cc
 linalg/sparse/SparseMatrixMultiply.h

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/SparseMatrix.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::SparseMatrix;
     ~MatrixCacheEntry() override;
     MatrixCacheEntry(const Matrix* matrix, const std::string& uid = ""): matrix_{matrix}, uid_(uid) {
         ATLAS_ASSERT(matrix_ != nullptr);

src/atlas/interpolation/method/Method.h

@@ -16,10 +16,10 @@
 
 #include "atlas/interpolation/Cache.h"
 #include "atlas/interpolation/NonLinear.h"
+#include "atlas/linalg/SparseMatrix.h"
 #include "atlas/util/Metadata.h"
 #include "atlas/util/Object.h"
 #include "eckit/config/Configuration.h"
-#include "eckit/linalg/SparseMatrix.h"
 
 namespace atlas {
 class Field;
@@ -87,7 +87,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::SparseMatrix;
 
     static void normalise(Triplets& triplets);
 

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

@@ -12,12 +12,12 @@
 #include "atlas/interpolation/Interpolation.h"
 #include "atlas/interpolation/method/binning/Binning.h"
 #include "atlas/interpolation/method/MethodFactory.h"
+#include "atlas/linalg/SparseMatrix.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"
 
@@ -58,7 +58,7 @@ void Binning::do_setup(const FunctionSpace& source,
 
   using Index = eckit::linalg::Index;
   using Triplet = eckit::linalg::Triplet;
-  using SMatrix = eckit::linalg::SparseMatrix;
+  using SMatrix = atlas::linalg::SparseMatrix;
 
   source_ = source;
   target_ = target;

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

@@ -61,7 +61,7 @@ void GridBoxMaximum::do_execute(const Field& source, Field& target, Metadata&) c
 
     if (!matrixFree_) {
         const Matrix& m = matrix();
-        Matrix::const_iterator k(m);
+        Matrix::const_iterator 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/unstructured/ConservativeSphericalPolygonInterpolation.cc

@@ -1248,7 +1248,7 @@ void ConservativeSphericalPolygonInterpolation::intersect_polygons(const CSPolyg
     }
 }
 
-eckit::linalg::SparseMatrix ConservativeSphericalPolygonInterpolation::compute_1st_order_matrix() {
+atlas::linalg::SparseMatrix ConservativeSphericalPolygonInterpolation::compute_1st_order_matrix() {
     ATLAS_TRACE("ConservativeMethod::setup: build cons-1 interpolant matrix");
     ATLAS_ASSERT(not matrix_free_);
     Triplets triplets;
@@ -1358,7 +1358,7 @@ eckit::linalg::SparseMatrix ConservativeSphericalPolygonInterpolation::compute_1
     return Matrix(n_tpoints_, n_spoints_, triplets);
 }
 
-eckit::linalg::SparseMatrix ConservativeSphericalPolygonInterpolation::compute_2nd_order_matrix() {
+atlas::linalg::SparseMatrix ConservativeSphericalPolygonInterpolation::compute_2nd_order_matrix() {
     ATLAS_TRACE("ConservativeMethod::setup: build cons-2 interpolant matrix");
     ATLAS_ASSERT(not matrix_free_);
     const auto& src_points_ = data_->src_points_;

src/atlas/interpolation/nonlinear/Missing.cc

@@ -81,7 +81,7 @@ bool MissingIfAllMissing::executeT(NonLinear::Matrix& W, const Field& field) con
     bool zeros = false;
 
     Size i = 0;
-    Matrix::iterator it(W);
+    Matrix::iterator it = W.begin();
     for (Size r = 0; r < W.rows(); ++r) {
         const Matrix::iterator end = W.end(r);
 
@@ -128,6 +128,9 @@ bool MissingIfAllMissing::executeT(NonLinear::Matrix& W, const Field& field) con
             modif = true;
         }
     }
+    if (modif) {
+        W.setDeviceNeedsUpdate(true);
+    }
 
     if (zeros && missingValue.isnan()) {
         W.prune(0.);
@@ -161,7 +164,7 @@ bool MissingIfAnyMissing::executeT(NonLinear::Matrix& W, const Field& field) con
     bool zeros = false;
 
     Size i = 0;
-    Matrix::iterator it(W);
+    Matrix::iterator it = W.begin();
     for (Size r = 0; r < W.rows(); ++r) {
         const Matrix::iterator end = W.end(r);
 
@@ -195,6 +198,9 @@ bool MissingIfAnyMissing::executeT(NonLinear::Matrix& W, const Field& field) con
             modif = true;
         }
     }
+    if (modif) {
+        W.setDeviceNeedsUpdate(true);
+    }
 
     if (zeros && missingValue.isnan()) {
         W.prune(0.);
@@ -228,7 +234,7 @@ bool MissingIfHeaviestMissing::executeT(NonLinear::Matrix& W, const Field& field
     bool zeros = false;
 
     Size i = 0;
-    Matrix::iterator it(W);
+    Matrix::iterator it = W.begin();
     for (Size r = 0; r < W.rows(); ++r) {
         const Matrix::iterator end = W.end(r);
 
@@ -282,6 +288,9 @@ bool MissingIfHeaviestMissing::executeT(NonLinear::Matrix& W, const Field& field
             modif = true;
         }
     }
+    if (modif) {
+        W.setDeviceNeedsUpdate(true);
+    }
 
     if (zeros && missingValue.isnan()) {
         W.prune(0.);

src/atlas/interpolation/nonlinear/NonLinear.h

@@ -16,10 +16,10 @@
 #include <type_traits>
 
 #include "eckit/config/Parametrisation.h"
-#include "eckit/linalg/SparseMatrix.h"
 
 #include "atlas/array.h"
 #include "atlas/field/Field.h"
+#include "atlas/linalg/SparseMatrix.h"
 #include "atlas/runtime/Exception.h"
 #include "atlas/util/Factory.h"
 #include "atlas/util/ObjectHandle.h"
@@ -37,9 +37,9 @@ namespace nonlinear {
 class NonLinear : public util::Object {
 public:
     using Config = eckit::Parametrisation;
-    using Matrix = eckit::linalg::SparseMatrix;
-    using Scalar = eckit::linalg::Scalar;
-    using Size   = eckit::linalg::Size;
+    using Matrix = atlas::linalg::SparseMatrix;
+    using Scalar = atlas::linalg::SparseMatrix::Scalar;
+    using Size   = atlas::linalg::SparseMatrix::Size;
 
     /**
      * @brief ctor

src/atlas/linalg/SparseMatrix.cc

@@ -0,0 +1,84 @@
+
+#include "atlas/linalg/SparseMatrix.h"
+
+#include <algorithm>
+#include <cstring>
+#include <limits>
+
+#include "atlas/array/helpers/ArrayCopier.h"
+
+#include "eckit/exception/Exceptions.h"
+
+namespace atlas {
+namespace linalg {
+
+//----------------------------------------------------------------------------------------------------------------------
+
+SparseMatrix::SparseMatrix() : host_matrix_() {}
+
+
+SparseMatrix::SparseMatrix(Size nrows, Size ncols, const std::vector<eckit::linalg::Triplet>& triplets) :
+    host_matrix_(nrows, ncols, triplets) {
+    resetDeviceStorage();
+}
+
+
+SparseMatrix::SparseMatrix(const SparseMatrix& other) : host_matrix_(other.host_matrix_) {
+    if (!other.empty()) {  // in case we copy an other that was constructed empty
+        resetDeviceStorage();
+    }
+}
+
+
+SparseMatrix& SparseMatrix::operator=(const SparseMatrix& other) {
+    SparseMatrix copy(other);
+    swap(copy);
+    return *this;
+}
+
+
+void SparseMatrix::swap(SparseMatrix& other) {
+    host_matrix_.swap(other.host_matrix_);
+    outer_.swap(other.outer_);
+    inner_.swap(other.inner_);
+    data_.swap(other.data_);
+}
+
+
+size_t SparseMatrix::footprint() const {
+    return host_matrix_.footprint() +
+           outer_->footprint() +
+           inner_->footprint() +
+           data_->footprint();
+}
+
+
+SparseMatrix& SparseMatrix::prune(Scalar val) {
+    host_matrix_.prune(val);
+    resetDeviceStorage();
+    return *this;
+}
+
+
+SparseMatrix& SparseMatrix::transpose() {
+    host_matrix_.transpose();
+    resetDeviceStorage();
+    return *this;
+}
+
+SparseMatrix& SparseMatrix::setIdentity(Size nrows, Size ncols) {
+    host_matrix_.setIdentity(nrows, ncols);
+    resetDeviceStorage();
+    return *this;
+}
+
+void SparseMatrix::resetDeviceStorage() {
+    outer_.reset(atlas::array::Array::wrap<Index>(const_cast<Index*>(outer()), atlas::array::make_shape(rows() + 1)));
+    inner_.reset(atlas::array::Array::wrap<Index>(const_cast<Index*>(inner()), atlas::array::make_shape(nonZeros())));
+    data_.reset(atlas::array::Array::wrap<Scalar>(const_cast<Scalar*>(data()), atlas::array::make_shape(nonZeros())));
+    setDeviceNeedsUpdate(true);
+}
+
+//----------------------------------------------------------------------------------------------------------------------
+}  // namespace linalg
+}  // namespace atlas