Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add GPU offloadable SparseMatrixStorage and read-only SparseMatrixView. #247

Merged
merged 2 commits into from
Dec 9, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions src/atlas/CMakeLists.txt
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion src/atlas/interpolation/Cache.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
4 changes: 2 additions & 2 deletions src/atlas/interpolation/Cache.h
Original file line number Diff line number Diff line change
Expand Up @@ -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"

Expand Down Expand Up @@ -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);
Expand Down
25 changes: 12 additions & 13 deletions src/atlas/interpolation/method/Method.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}

Expand Down Expand Up @@ -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());
}
}

Expand All @@ -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>
Expand All @@ -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>
Expand All @@ -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>
Expand All @@ -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 {
Expand Down Expand Up @@ -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
}
}
}
Expand Down
18 changes: 15 additions & 3 deletions src/atlas/interpolation/method/Method.h
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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);

Expand All @@ -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) {
Expand All @@ -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_; }
Expand Down
26 changes: 15 additions & 11 deletions src/atlas/interpolation/method/binning/Binning.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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;

Expand All @@ -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_);
Expand Down Expand Up @@ -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);
}


Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down
5 changes: 3 additions & 2 deletions src/atlas/interpolation/method/knn/GridBoxMaximum.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down Expand Up @@ -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();
Expand Down
3 changes: 1 addition & 2 deletions src/atlas/interpolation/method/knn/GridBoxMethod.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}

Expand Down
3 changes: 1 addition & 2 deletions src/atlas/interpolation/method/knn/KNearestNeighbours.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
3 changes: 1 addition & 2 deletions src/atlas/interpolation/method/knn/NearestNeighbour.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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"
Expand Down Expand Up @@ -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.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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);
}
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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: {
Expand Down Expand Up @@ -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;
Expand Down Expand Up @@ -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_;
Expand Down Expand Up @@ -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,
Expand Down
Loading
Loading