Merge Pull Request #12738 from cgcgcg/Trilinos/mueluUtilsFix

Automatically Merged using Trilinos Pull Request AutoTester PR Title: b'Xpetra&MueLu: Fix for diagonal extraction' PR Author: cgcgcg
trilinos · Feb 13, 2024 · e4be19f · e4be19f
2 parents ffdcd00 + 4ce5684
commit e4be19f
Show file tree

Hide file tree

Showing 11 changed files with 83 additions and 18 deletions.
diff --git a/packages/muelu/src/Utils/MueLu_UtilitiesBase_def.hpp b/packages/muelu/src/Utils/MueLu_UtilitiesBase_def.hpp
@@ -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;
 }
@@ -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);
   }

diff --git 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<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."

diff --git 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<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."
@@ -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."
@@ -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."

diff --git 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<size_t*, typename Node::device_type, Kokkos::MemoryUnmanaged>& 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
@@ -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.
@@ -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!");

diff --git 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<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;
 

diff --git 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<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){};
@@ -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));
@@ -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));

diff --git 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<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

diff --git a/packages/xpetra/src/CrsMatrix/Xpetra_TpetraBlockCrsMatrix_def.hpp b/packages/xpetra/src/CrsMatrix/Xpetra_TpetraBlockCrsMatrix_def.hpp
@@ -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 {
@@ -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) {}
@@ -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) {}

diff --git 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<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);
 

diff --git a/packages/xpetra/src/CrsMatrix/Xpetra_TpetraCrsMatrix_def.hpp b/packages/xpetra/src/CrsMatrix/Xpetra_TpetraCrsMatrix_def.hpp
@@ -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");