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.

Sign up for GitHub

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

Jump to bottom

Xpetra&MueLu: Fix for diagonal extraction #12738

Merged
merged 2 commits into from
Feb 13, 2024
Merged

Xpetra&MueLu: Fix for diagonal extraction #12738

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
561af35
Xpetra: Expose getLocalDiagOffsets and getLocalDiagCopy with Kokkos v…
cgcgcg Feb 12, 2024
4ce5684
MueLu: Fix #12736 and use Kokkos views in GetMatrixDiagonal
cgcgcg Feb 12, 2024
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
31 changes: 13 additions & 18 deletions packages/muelu/src/Utils/MueLu_UtilitiesBase_def.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -197,7 +197,16 @@ UtilitiesBase<Scalar, LocalOrdinal, GlobalOrdinal, Node>::
const auto rowMap = A.getRowMap();
auto diag = Xpetra::VectorFactory<Scalar, LocalOrdinal, GlobalOrdinal, Node>::Build(rowMap, true);

A.getLocalDiagCopy(*diag);
const CrsMatrixWrap* crsOp = dynamic_cast<const CrsMatrixWrap*>(&A);
if ((crsOp != NULL) && (rowMap->lib() == Xpetra::UseTpetra)) {
using local_vector_type = typename Vector::dual_view_type::t_dev_um;
using execution_space = typename local_vector_type::execution_space;
Kokkos::View<size_t*, execution_space> offsets("offsets", rowMap->getLocalNumElements());
crsOp->getCrsGraph()->getLocalDiagOffsets(offsets);
crsOp->getCrsMatrix()->getLocalDiagCopy(*diag, offsets);
} else {
A.getLocalDiagCopy(*diag);
}

return diag;
}
Expand Down Expand Up @@ -621,24 +630,10 @@ template <class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
RCP<Xpetra::Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node>>
UtilitiesBase<Scalar, LocalOrdinal, GlobalOrdinal, Node>::
GetMatrixOverlappedDiagonal(const Matrix& A) {
// FIXME_KOKKOS
RCP<const Map> rowMap = A.getRowMap(), colMap = A.getColMap();
RCP<Vector> localDiag = VectorFactory::Build(rowMap);

const CrsMatrixWrap* crsOp = dynamic_cast<const CrsMatrixWrap*>(&A);
if ((crsOp != NULL) && (rowMap->lib() == Xpetra::UseTpetra)) {
Teuchos::ArrayRCP<size_t> offsets;
crsOp->getLocalDiagOffsets(offsets);
crsOp->getLocalDiagCopy(*localDiag, offsets());
} else {
auto localDiagVals = localDiag->getDeviceLocalView(Xpetra::Access::ReadWrite);
const auto diagVals = GetMatrixDiagonal(A)->getDeviceLocalView(Xpetra::Access::ReadOnly);
Kokkos::deep_copy(localDiagVals, diagVals);
}

RCP<Vector> diagonal = VectorFactory::Build(colMap);
RCP<const Import> importer;
importer = A.getCrsGraph()->getImporter();
RCP<Vector> localDiag = GetMatrixDiagonal(A);
RCP<Vector> diagonal = VectorFactory::Build(colMap);
RCP<const Import> importer = A.getCrsGraph()->getImporter();
if (importer == Teuchos::null) {
importer = ImportFactory::Build(rowMap, colMap);
}
Expand Down
3 changes: 3 additions & 0 deletions packages/xpetra/src/CrsGraph/Xpetra_CrsGraph.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -242,6 +242,9 @@ class CrsGraph
virtual typename local_graph_type::HostMirror getLocalGraphHost() const = 0;
virtual local_graph_type getLocalGraphDevice() const = 0;

//! Get offsets of the diagonal entries in the matrix.
virtual void getLocalDiagOffsets(const Kokkos::View<size_t *, device_type, Kokkos::MemoryUnmanaged> &offsets) const = 0;

#else
#ifdef __GNUC__
#warning "Xpetra Kokkos interface for CrsMatrix is enabled (HAVE_XPETRA_KOKKOS_REFACTOR) but Tpetra is disabled. The Kokkos interface needs Tpetra to be enabled, too."
Expand Down
14 changes: 14 additions & 0 deletions packages/xpetra/src/CrsGraph/Xpetra_EpetraCrsGraph.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -309,6 +309,11 @@ class EpetraCrsGraphT
"Xpetra::EpetraCrsGraph only available for GO=int or GO=long long with EpetraNode (Serial or OpenMP depending on configuration)");
TEUCHOS_UNREACHABLE_RETURN((local_graph_type()));
}

void getLocalDiagOffsets(const Kokkos::View<size_t *, typename Node::device_type, Kokkos::MemoryUnmanaged> &offsets) const {
TEUCHOS_TEST_FOR_EXCEPTION(true, Xpetra::Exceptions::RuntimeError,
"Epetra does not support getLocalDiagOffsets!");
}
#else
#ifdef __GNUC__
#warning "Xpetra Kokkos interface for CrsGraph is enabled (HAVE_XPETRA_KOKKOS_REFACTOR) but Tpetra is disabled. The Kokkos interface needs Tpetra to be enabled, too."
Expand Down Expand Up @@ -838,6 +843,10 @@ class EpetraCrsGraphT<int, EpetraNode>
return localGraph;
}

void getLocalDiagOffsets(const Kokkos::View<size_t *, typename Node::device_type, Kokkos::MemoryUnmanaged> &offsets) const {
TEUCHOS_TEST_FOR_EXCEPTION(true, Xpetra::Exceptions::RuntimeError,
"Epetra does not support getLocalDiagOffsets!");
}
#else
#ifdef __GNUC__
#warning "Xpetra Kokkos interface for CrsGraph is enabled (HAVE_XPETRA_KOKKOS_REFACTOR) but Tpetra is disabled. The Kokkos interface needs Tpetra to be enabled, too."
Expand Down Expand Up @@ -1398,6 +1407,11 @@ class EpetraCrsGraphT<long long, EpetraNode>

return localGraph;
}

void getLocalDiagOffsets(const Kokkos::View<size_t *, typename Node::device_type, Kokkos::MemoryUnmanaged> &offsets) const {
TEUCHOS_TEST_FOR_EXCEPTION(true, Xpetra::Exceptions::RuntimeError,
"Epetra does not support getLocalDiagOffsets!");
}
#else
#ifdef __GNUC__
#warning "Xpetra Kokkos interface for CrsGraph is enabled (HAVE_XPETRA_KOKKOS_REFACTOR) but Tpetra is disabled. The Kokkos interface needs Tpetra to be enabled, too."
Expand Down
3 changes: 3 additions & 0 deletions packages/xpetra/src/CrsGraph/Xpetra_TpetraCrsGraph_decl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -323,6 +323,9 @@ class TpetraCrsGraph
/// \brief Access the local KokkosSparse::StaticCrsGraph data for device use
local_graph_type getLocalGraphDevice() const;

//! Get offsets of the diagonal entries in the matrix.
void getLocalDiagOffsets(const Kokkos::View<size_t*, typename Node::device_type, Kokkos::MemoryUnmanaged>& offsets) const;

//! Force the computation of global constants if we don't have them
void computeGlobalConstants();

Expand Down
10 changes: 10 additions & 0 deletions packages/xpetra/src/CrsGraph/Xpetra_TpetraCrsGraph_def.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -382,6 +382,11 @@ typename Xpetra::CrsGraph<LocalOrdinal, GlobalOrdinal, Node>::local_graph_type T
return getTpetra_CrsGraph()->getLocalGraphDevice();
}

template <class LocalOrdinal, class GlobalOrdinal, class Node>
void TpetraCrsGraph<LocalOrdinal, GlobalOrdinal, Node>::getLocalDiagOffsets(const Kokkos::View<size_t *, typename Node::device_type, Kokkos::MemoryUnmanaged> &offsets) const {
getTpetra_CrsGraph()->getLocalDiagOffsets(offsets);
}

template <class LocalOrdinal, class GlobalOrdinal, class Node>
void TpetraCrsGraph<LocalOrdinal, GlobalOrdinal, Node>::computeGlobalConstants() {
// mfh 07 May 2018: See GitHub Issue #2565.
Expand Down Expand Up @@ -1097,6 +1102,11 @@ class TpetraCrsGraph<int, long long, EpetraNode>
TEUCHOS_UNREACHABLE_RETURN((local_graph_type()));
}

void getLocalDiagOffsets(const Kokkos::View<size_t *, typename Node::device_type, Kokkos::MemoryUnmanaged> &offsets) const {
TEUCHOS_TEST_FOR_EXCEPTION(true, Xpetra::Exceptions::RuntimeError,
"Epetra does not support getLocalDiagOffsets!");
}

typename local_graph_type::HostMirror getLocalGraphHost() const {
TEUCHOS_TEST_FOR_EXCEPTION(true, Xpetra::Exceptions::RuntimeError,
"Epetra does not support Kokkos::StaticCrsGraph!");
Expand Down
3 changes: 3 additions & 0 deletions packages/xpetra/src/CrsMatrix/Xpetra_CrsMatrix.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -227,6 +227,9 @@ class CrsMatrix
//! Get a copy of the diagonal entries owned by this node, with local row indices, using row offsets.
virtual void getLocalDiagCopy(Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &diag, const Teuchos::ArrayView<const size_t> &offsets) const = 0;

//! Get a copy of the diagonal entries owned by this node, with local row indices, using row offsets.
virtual void getLocalDiagCopy(Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &diag, const Kokkos::View<const size_t *, typename Node::device_type, Kokkos::MemoryUnmanaged> &offsets) const = 0;

//! Replace the diagonal entries of the matrix
virtual void replaceDiag(const Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &diag) = 0;

Expand Down
11 changes: 11 additions & 0 deletions packages/xpetra/src/CrsMatrix/Xpetra_EpetraCrsMatrix.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -199,6 +199,7 @@ class EpetraCrsMatrixT
void getLocalDiagCopy(Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &diag) const {}
void getLocalDiagOffsets(Teuchos::ArrayRCP<size_t> &offsets) const {}
void getLocalDiagCopy(Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &diag, const Teuchos::ArrayView<const size_t> &offsets) const {}
void getLocalDiagCopy(Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &diag, const Kokkos::View<const size_t *, typename Node::device_type, Kokkos::MemoryUnmanaged> &offsets) const {}
void replaceDiag(const Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &diag) {}
void leftScale(const Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &x){};
void rightScale(const Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &x){};
Expand Down Expand Up @@ -941,6 +942,11 @@ class EpetraCrsMatrixT<int, EpetraNode>
TEUCHOS_TEST_FOR_EXCEPTION(true, Xpetra::Exceptions::NotImplemented, "Xpetra::EpetraCrsMatrixT.getLocalDiagCopy using offsets is not implemented or supported.");
}

//! Get a copy of the diagonal entries owned by this node, with local row indices, using row offsets.
void getLocalDiagCopy(Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> & /* diag */, const Kokkos::View<const size_t *, typename Node::device_type, Kokkos::MemoryUnmanaged> & /* offsets */) const {
TEUCHOS_TEST_FOR_EXCEPTION(true, Xpetra::Exceptions::NotImplemented, "Xpetra::EpetraCrsMatrixT.getLocalDiagCopy using offsets is not implemented or supported.");
}

//! Replace the diagonal entries of the matrix
void replaceDiag(const Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &diag) {
mtx_->ReplaceDiagonalValues(toEpetra<GlobalOrdinal, Node>(diag));
Expand Down Expand Up @@ -2043,6 +2049,11 @@ class EpetraCrsMatrixT<long long, EpetraNode>
TEUCHOS_TEST_FOR_EXCEPTION(true, Xpetra::Exceptions::NotImplemented, "Xpetra::EpetraCrsMatrixT.getLocalDiagCopy using offsets is not implemented or supported.");
}

//! Get a copy of the diagonal entries owned by this node, with local row indices, using row offsets.
void getLocalDiagCopy(Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> & /* diag */, const Kokkos::View<const size_t *, typename Node::device_type, Kokkos::MemoryUnmanaged> & /* offsets */) const {
TEUCHOS_TEST_FOR_EXCEPTION(true, Xpetra::Exceptions::NotImplemented, "Xpetra::EpetraCrsMatrixT.getLocalDiagCopy using offsets is not implemented or supported.");
}

//! Replace the diagonal entries of the matrix
void replaceDiag(const Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &diag) {
mtx_->ReplaceDiagonalValues(toEpetra<GlobalOrdinal, Node>(diag));
Expand Down
3 changes: 3 additions & 0 deletions packages/xpetra/src/CrsMatrix/Xpetra_TpetraBlockCrsMatrix_decl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -323,6 +323,9 @@ class TpetraBlockCrsMatrix
//! Get offsets of the diagonal entries in the matrix.
void getLocalDiagOffsets(Teuchos::ArrayRCP<size_t> &offsets) const;

//! Get a copy of the diagonal entries owned by this node, with local row indices, using row offsets.
void getLocalDiagCopy(Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &diag, const Kokkos::View<const size_t *, typename Node::device_type, Kokkos::MemoryUnmanaged> &offsets) const;

void replaceDiag(const Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &diag);

//! Left scale operator with given vector values
Expand Down
14 changes: 14 additions & 0 deletions packages/xpetra/src/CrsMatrix/Xpetra_TpetraBlockCrsMatrix_def.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -581,6 +581,14 @@ void TpetraBlockCrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>::
throw std::runtime_error("Xpetra::TpetraBlockCrsMatrix function not implemented in " + std::string(__FILE__) + ":" + std::to_string(__LINE__));
}

//! Get a copy of the diagonal entries owned by this node, with local row indices.
template <class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
void TpetraBlockCrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>::
getLocalDiagCopy(Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &diag,
const Kokkos::View<const size_t *, typename Node::device_type, Kokkos::MemoryUnmanaged> &offsets) const {
throw std::runtime_error("Xpetra::TpetraBlockCrsMatrix function not implemented in " + std::string(__FILE__) + ":" + std::to_string(__LINE__));
}

template <class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
void TpetraBlockCrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>::
getLocalDiagOffsets(Teuchos::ArrayRCP<size_t> &offsets) const {
Expand Down Expand Up @@ -992,6 +1000,9 @@ class TpetraBlockCrsMatrix<Scalar, int, int, EpetraNode>
//! Get a copy of the diagonal entries owned by this node, with local row indices.
void getLocalDiagCopy(Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &diag, const Teuchos::ArrayView<const size_t> &offsets) const {}

//! Get a copy of the diagonal entries owned by this node, with local row indices.
void getLocalDiagCopy(Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &diag, const Kokkos::View<const size_t *, typename Node::device_type, Kokkos::MemoryUnmanaged> &offsets) const {}

void replaceDiag(const Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &diag) {}

void leftScale(const Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &x) {}
Expand Down Expand Up @@ -1289,6 +1300,9 @@ class TpetraBlockCrsMatrix<Scalar, int, long long, EpetraNode>
//! Get a copy of the diagonal entries owned by this node, with local row indices.
void getLocalDiagCopy(Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &diag, const Teuchos::ArrayView<const size_t> &offsets) const {}

//! Get a copy of the diagonal entries owned by this node, with local row indices.
void getLocalDiagCopy(Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &diag, const Kokkos::View<const size_t *, typename Node::device_type, Kokkos::MemoryUnmanaged> &offsets) const {}

void replaceDiag(Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &diag) const {}

void leftScale(const Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> &x) {}
Expand Down
3 changes: 3 additions & 0 deletions packages/xpetra/src/CrsMatrix/Xpetra_TpetraCrsMatrix_decl.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -377,6 +377,9 @@ class TpetraCrsMatrix
//! Get a copy of the diagonal entries owned by this node, with local row indices.
void getLocalDiagCopy(Vector &diag, const Teuchos::ArrayView<const size_t> &offsets) const;

//! Get a copy of the diagonal entries owned by this node, with local row indices, using row offsets.
void getLocalDiagCopy(Vector &diag, const Kokkos::View<const size_t *, typename Node::device_type, Kokkos::MemoryUnmanaged> &offsets) const;

//! Replace the diagonal entries of the matrix
void replaceDiag(const Vector &diag);

Expand Down
6 changes: 6 additions & 0 deletions packages/xpetra/src/CrsMatrix/Xpetra_TpetraCrsMatrix_def.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -556,6 +556,12 @@ void TpetraCrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>::getLocalDiagCop
mtx_->getLocalDiagCopy(*(toTpetra(diag)), offsets);
}

template <class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
void TpetraCrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>::getLocalDiagCopy(Vector &diag, const Kokkos::View<const size_t *, typename Node::device_type, Kokkos::MemoryUnmanaged> &offsets) const {
XPETRA_MONITOR("TpetraCrsMatrix::getLocalDiagCopy");
mtx_->getLocalDiagCopy(*(toTpetra(diag)), offsets);
}

template <class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
void TpetraCrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>::replaceDiag(const Vector &diag) {
XPETRA_MONITOR("TpetraCrsMatrix::replaceDiag");
Expand Down
Loading