diff --git a/packages/muelu/src/Utils/MueLu_UtilitiesBase_def.hpp b/packages/muelu/src/Utils/MueLu_UtilitiesBase_def.hpp index e101a7991e2a..e27fedd26abb 100644 --- a/packages/muelu/src/Utils/MueLu_UtilitiesBase_def.hpp +++ b/packages/muelu/src/Utils/MueLu_UtilitiesBase_def.hpp @@ -197,7 +197,16 @@ UtilitiesBase:: const auto rowMap = A.getRowMap(); auto diag = Xpetra::VectorFactory::Build(rowMap, true); - A.getLocalDiagCopy(*diag); + const CrsMatrixWrap* crsOp = dynamic_cast(&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 offsets("offsets", rowMap->getLocalNumElements()); + crsOp->getCrsGraph()->getLocalDiagOffsets(offsets); + crsOp->getCrsMatrix()->getLocalDiagCopy(*diag, offsets); + } else { + A.getLocalDiagCopy(*diag); + } return diag; } @@ -621,24 +630,10 @@ template RCP> UtilitiesBase:: GetMatrixOverlappedDiagonal(const Matrix& A) { - // FIXME_KOKKOS RCP rowMap = A.getRowMap(), colMap = A.getColMap(); - RCP localDiag = VectorFactory::Build(rowMap); - - const CrsMatrixWrap* crsOp = dynamic_cast(&A); - if ((crsOp != NULL) && (rowMap->lib() == Xpetra::UseTpetra)) { - Teuchos::ArrayRCP 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 diagonal = VectorFactory::Build(colMap); - RCP importer; - importer = A.getCrsGraph()->getImporter(); + RCP localDiag = GetMatrixDiagonal(A); + RCP diagonal = VectorFactory::Build(colMap); + RCP importer = A.getCrsGraph()->getImporter(); if (importer == Teuchos::null) { importer = ImportFactory::Build(rowMap, colMap); } diff --git a/packages/xpetra/src/CrsGraph/Xpetra_CrsGraph.hpp b/packages/xpetra/src/CrsGraph/Xpetra_CrsGraph.hpp index c03125feff35..d31186343baf 100644 --- a/packages/xpetra/src/CrsGraph/Xpetra_CrsGraph.hpp +++ b/packages/xpetra/src/CrsGraph/Xpetra_CrsGraph.hpp @@ -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 &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." diff --git a/packages/xpetra/src/CrsGraph/Xpetra_EpetraCrsGraph.hpp b/packages/xpetra/src/CrsGraph/Xpetra_EpetraCrsGraph.hpp index 99aefc2a4f16..ab9686b4399e 100644 --- a/packages/xpetra/src/CrsGraph/Xpetra_EpetraCrsGraph.hpp +++ b/packages/xpetra/src/CrsGraph/Xpetra_EpetraCrsGraph.hpp @@ -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 &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." @@ -838,6 +843,10 @@ class EpetraCrsGraphT return localGraph; } + void getLocalDiagOffsets(const Kokkos::View &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." @@ -1398,6 +1407,11 @@ class EpetraCrsGraphT return localGraph; } + + void getLocalDiagOffsets(const Kokkos::View &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." diff --git a/packages/xpetra/src/CrsGraph/Xpetra_TpetraCrsGraph_decl.hpp b/packages/xpetra/src/CrsGraph/Xpetra_TpetraCrsGraph_decl.hpp index bbc265ab37bc..926e79cd2f66 100644 --- a/packages/xpetra/src/CrsGraph/Xpetra_TpetraCrsGraph_decl.hpp +++ b/packages/xpetra/src/CrsGraph/Xpetra_TpetraCrsGraph_decl.hpp @@ -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& offsets) const; + //! Force the computation of global constants if we don't have them void computeGlobalConstants(); diff --git a/packages/xpetra/src/CrsGraph/Xpetra_TpetraCrsGraph_def.hpp b/packages/xpetra/src/CrsGraph/Xpetra_TpetraCrsGraph_def.hpp index 80a033373e5c..3a011a11f1b4 100644 --- a/packages/xpetra/src/CrsGraph/Xpetra_TpetraCrsGraph_def.hpp +++ b/packages/xpetra/src/CrsGraph/Xpetra_TpetraCrsGraph_def.hpp @@ -382,6 +382,11 @@ typename Xpetra::CrsGraph::local_graph_type T return getTpetra_CrsGraph()->getLocalGraphDevice(); } +template +void TpetraCrsGraph::getLocalDiagOffsets(const Kokkos::View &offsets) const { + getTpetra_CrsGraph()->getLocalDiagOffsets(offsets); +} + template void TpetraCrsGraph::computeGlobalConstants() { // mfh 07 May 2018: See GitHub Issue #2565. @@ -1097,6 +1102,11 @@ class TpetraCrsGraph TEUCHOS_UNREACHABLE_RETURN((local_graph_type())); } + void getLocalDiagOffsets(const Kokkos::View &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!"); diff --git a/packages/xpetra/src/CrsMatrix/Xpetra_CrsMatrix.hpp b/packages/xpetra/src/CrsMatrix/Xpetra_CrsMatrix.hpp index 5c8ff34bd9f9..2ec3c177a25d 100644 --- a/packages/xpetra/src/CrsMatrix/Xpetra_CrsMatrix.hpp +++ b/packages/xpetra/src/CrsMatrix/Xpetra_CrsMatrix.hpp @@ -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 &diag, const Teuchos::ArrayView &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 &diag, const Kokkos::View &offsets) const = 0; + //! Replace the diagonal entries of the matrix virtual void replaceDiag(const Vector &diag) = 0; diff --git a/packages/xpetra/src/CrsMatrix/Xpetra_EpetraCrsMatrix.hpp b/packages/xpetra/src/CrsMatrix/Xpetra_EpetraCrsMatrix.hpp index ea1634f6bbed..066f4d6edf98 100644 --- a/packages/xpetra/src/CrsMatrix/Xpetra_EpetraCrsMatrix.hpp +++ b/packages/xpetra/src/CrsMatrix/Xpetra_EpetraCrsMatrix.hpp @@ -199,6 +199,7 @@ class EpetraCrsMatrixT void getLocalDiagCopy(Vector &diag) const {} void getLocalDiagOffsets(Teuchos::ArrayRCP &offsets) const {} void getLocalDiagCopy(Vector &diag, const Teuchos::ArrayView &offsets) const {} + void getLocalDiagCopy(Vector &diag, const Kokkos::View &offsets) const {} void replaceDiag(const Vector &diag) {} void leftScale(const Vector &x){}; void rightScale(const Vector &x){}; @@ -941,6 +942,11 @@ class EpetraCrsMatrixT 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 & /* diag */, const Kokkos::View & /* 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 &diag) { mtx_->ReplaceDiagonalValues(toEpetra(diag)); @@ -2043,6 +2049,11 @@ class EpetraCrsMatrixT 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 & /* diag */, const Kokkos::View & /* 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 &diag) { mtx_->ReplaceDiagonalValues(toEpetra(diag)); diff --git a/packages/xpetra/src/CrsMatrix/Xpetra_TpetraBlockCrsMatrix_decl.hpp b/packages/xpetra/src/CrsMatrix/Xpetra_TpetraBlockCrsMatrix_decl.hpp index 43855021e187..f208e2432da5 100644 --- a/packages/xpetra/src/CrsMatrix/Xpetra_TpetraBlockCrsMatrix_decl.hpp +++ b/packages/xpetra/src/CrsMatrix/Xpetra_TpetraBlockCrsMatrix_decl.hpp @@ -323,6 +323,9 @@ class TpetraBlockCrsMatrix //! Get offsets of the diagonal entries in the matrix. void getLocalDiagOffsets(Teuchos::ArrayRCP &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 &offsets) const; + void replaceDiag(const Vector &diag); //! Left scale operator with given vector values diff --git a/packages/xpetra/src/CrsMatrix/Xpetra_TpetraBlockCrsMatrix_def.hpp b/packages/xpetra/src/CrsMatrix/Xpetra_TpetraBlockCrsMatrix_def.hpp index f0e38add2272..c131fe9435d3 100644 --- a/packages/xpetra/src/CrsMatrix/Xpetra_TpetraBlockCrsMatrix_def.hpp +++ b/packages/xpetra/src/CrsMatrix/Xpetra_TpetraBlockCrsMatrix_def.hpp @@ -581,6 +581,14 @@ void TpetraBlockCrsMatrix:: 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 +void TpetraBlockCrsMatrix:: + getLocalDiagCopy(Vector &diag, + const Kokkos::View &offsets) const { + throw std::runtime_error("Xpetra::TpetraBlockCrsMatrix function not implemented in " + std::string(__FILE__) + ":" + std::to_string(__LINE__)); +} + template void TpetraBlockCrsMatrix:: getLocalDiagOffsets(Teuchos::ArrayRCP &offsets) const { @@ -992,6 +1000,9 @@ class TpetraBlockCrsMatrix //! Get a copy of the diagonal entries owned by this node, with local row indices. void getLocalDiagCopy(Vector &diag, const Teuchos::ArrayView &offsets) const {} + //! Get a copy of the diagonal entries owned by this node, with local row indices. + void getLocalDiagCopy(Vector &diag, const Kokkos::View &offsets) const {} + void replaceDiag(const Vector &diag) {} void leftScale(const Vector &x) {} @@ -1289,6 +1300,9 @@ class TpetraBlockCrsMatrix //! Get a copy of the diagonal entries owned by this node, with local row indices. void getLocalDiagCopy(Vector &diag, const Teuchos::ArrayView &offsets) const {} + //! Get a copy of the diagonal entries owned by this node, with local row indices. + void getLocalDiagCopy(Vector &diag, const Kokkos::View &offsets) const {} + void replaceDiag(Vector &diag) const {} void leftScale(const Vector &x) {} diff --git a/packages/xpetra/src/CrsMatrix/Xpetra_TpetraCrsMatrix_decl.hpp b/packages/xpetra/src/CrsMatrix/Xpetra_TpetraCrsMatrix_decl.hpp index 514bc5dd3312..531b6f58726f 100644 --- a/packages/xpetra/src/CrsMatrix/Xpetra_TpetraCrsMatrix_decl.hpp +++ b/packages/xpetra/src/CrsMatrix/Xpetra_TpetraCrsMatrix_decl.hpp @@ -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 &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 &offsets) const; + //! Replace the diagonal entries of the matrix void replaceDiag(const Vector &diag); diff --git a/packages/xpetra/src/CrsMatrix/Xpetra_TpetraCrsMatrix_def.hpp b/packages/xpetra/src/CrsMatrix/Xpetra_TpetraCrsMatrix_def.hpp index dfad75406d47..3a424fe36899 100644 --- a/packages/xpetra/src/CrsMatrix/Xpetra_TpetraCrsMatrix_def.hpp +++ b/packages/xpetra/src/CrsMatrix/Xpetra_TpetraCrsMatrix_def.hpp @@ -556,6 +556,12 @@ void TpetraCrsMatrix::getLocalDiagCop mtx_->getLocalDiagCopy(*(toTpetra(diag)), offsets); } +template +void TpetraCrsMatrix::getLocalDiagCopy(Vector &diag, const Kokkos::View &offsets) const { + XPETRA_MONITOR("TpetraCrsMatrix::getLocalDiagCopy"); + mtx_->getLocalDiagCopy(*(toTpetra(diag)), offsets); +} + template void TpetraCrsMatrix::replaceDiag(const Vector &diag) { XPETRA_MONITOR("TpetraCrsMatrix::replaceDiag");