From d143e58169e930c59c571d82d102ace9571954b4 Mon Sep 17 00:00:00 2001
From: Christian Glusa <caglusa@sandia.gov>
Date: Wed, 20 Apr 2022 14:38:49 -0600
Subject: [PATCH] Xpetra MatrixUtils: Fix checkLocalRowMapMatchesColMap

---
 .../xpetra/sup/Utils/Xpetra_MatrixUtils.hpp   | 23 ++++++++++++-------
 1 file changed, 15 insertions(+), 8 deletions(-)

diff --git a/packages/xpetra/sup/Utils/Xpetra_MatrixUtils.hpp b/packages/xpetra/sup/Utils/Xpetra_MatrixUtils.hpp
index 3ce7e2f9d3f6..9a16770d2edd 100644
--- a/packages/xpetra/sup/Utils/Xpetra_MatrixUtils.hpp
+++ b/packages/xpetra/sup/Utils/Xpetra_MatrixUtils.hpp
@@ -740,15 +740,22 @@ class MatrixUtils {
   static void checkLocalRowMapMatchesColMap(const Xpetra::Matrix<Scalar, LocalOrdinal, GlobalOrdinal, Node> & A) {
     RCP<const Map> rowMap = A.getRowMap();
     RCP<const Map> colMap = A.getColMap();
-    RCP<const Teuchos::Comm<int> > comm = rowMap->getComm();
-    LO numRows = Teuchos::as<LocalOrdinal>(rowMap->getLocalNumElements());
     bool fail = false;
-    for (LO rowLID = 0; rowLID < numRows; rowLID++) {
-      GO rowGID = rowMap->getGlobalElement(rowLID);
-      LO colLID = colMap->getLocalElement(rowGID);
-      if (rowLID != colLID) {
-        fail = true;
-        std::cerr << "On rank " << comm->getRank() << ", GID " << rowGID << " is LID " << rowLID << "in the rowmap, but LID " << colLID << " in the column map.\n";
+    if (rowMap->lib() == Xpetra::UseTpetra) {
+      auto tpRowMap = Teuchos::rcp_dynamic_cast<const TpetraMap>(rowMap, true)->getTpetra_Map();
+      auto tpColMap = Teuchos::rcp_dynamic_cast<const TpetraMap>(colMap, true)->getTpetra_Map();
+      fail = ! tpColMap->isLocallyFitted(*tpRowMap);
+    } else {
+      RCP<const Teuchos::Comm<int> > comm = rowMap->getComm();
+      LO numRows = Teuchos::as<LocalOrdinal>(rowMap->getLocalNumElements());
+
+      for (LO rowLID = 0; rowLID < numRows; rowLID++) {
+        GO rowGID = rowMap->getGlobalElement(rowLID);
+        LO colGID = colMap->getGlobalElement(rowLID);
+        if (rowGID != colGID) {
+          fail = true;
+          std::cerr << "On rank " << comm->getRank() << ", LID " << rowLID << " is GID " << rowGID << " in the rowmap, but GID " << colGID << " in the column map.\n";
+        }
       }
     }
     TEUCHOS_TEST_FOR_EXCEPTION(fail, Exceptions::RuntimeError,