From ccf7b5da024c4b195a4b1807482fa93c7f7be98a Mon Sep 17 00:00:00 2001
From: Chris Siefert <csiefer@sandia.gov>
Date: Thu, 31 Aug 2023 11:25:13 -0600
Subject: [PATCH] MueLu: Adding finer grained profiling regions for H2D/D2H
 diagnosis

---
 .../MueLu_RebalanceAcFactory_def.hpp          |  5 ++-
 .../MueLu_RepartitionFactory_def.hpp          | 39 ++++++++++++-------
 2 files changed, 28 insertions(+), 16 deletions(-)

diff --git a/packages/muelu/src/Rebalancing/MueLu_RebalanceAcFactory_def.hpp b/packages/muelu/src/Rebalancing/MueLu_RebalanceAcFactory_def.hpp
index b41f93cc2e49..81f8f5eb0130 100644
--- a/packages/muelu/src/Rebalancing/MueLu_RebalanceAcFactory_def.hpp
+++ b/packages/muelu/src/Rebalancing/MueLu_RebalanceAcFactory_def.hpp
@@ -126,7 +126,10 @@ namespace MueLu {
         }
         // NOTE: If the communicator is restricted away, Build returns Teuchos::null.
         XpetraList.set("Timer Label","MueLu::RebalanceAc-" + Teuchos::toString(coarseLevel.GetLevelID()));
-	rebalancedAc = MatrixFactory::Build(originalAc, *rebalanceImporter, *rebalanceImporter, targetMap, targetMap, rcp(&XpetraList,false));
+        {
+          SubFactoryMonitor subM2(*this, "Rebalancing existing Ac: MatrixFactory::Build", coarseLevel);
+          rebalancedAc = MatrixFactory::Build(originalAc, *rebalanceImporter, *rebalanceImporter, targetMap, targetMap, rcp(&XpetraList,false));
+        }
 
         if (!rebalancedAc.is_null()) {
           rebalancedAc->SetFixedBlockSize(originalAc->GetFixedBlockSize());
diff --git a/packages/muelu/src/Rebalancing/MueLu_RepartitionFactory_def.hpp b/packages/muelu/src/Rebalancing/MueLu_RepartitionFactory_def.hpp
index 71c08c37a84c..73b1b5291f1e 100644
--- a/packages/muelu/src/Rebalancing/MueLu_RepartitionFactory_def.hpp
+++ b/packages/muelu/src/Rebalancing/MueLu_RepartitionFactory_def.hpp
@@ -296,22 +296,26 @@ namespace MueLu {
 
     // Step 1: Find out how many processors send me data
     // partsIndexBase starts from zero, as the processors ids start from zero
-    GO partsIndexBase = 0;
-    RCP<Map>    partsIHave  = MapFactory   ::Build(lib, Teuchos::OrdinalTraits<Xpetra::global_size_t>::invalid(), myParts(), partsIndexBase, comm);
-    RCP<Map>    partsIOwn   = MapFactory   ::Build(lib,                                                 numProcs,  myPart(), partsIndexBase, comm);
-    RCP<Export> partsExport = ExportFactory::Build(partsIHave, partsIOwn);
-
-    RCP<GOVector> partsISend    = Xpetra::VectorFactory<GO, LO, GO, NO>::Build(partsIHave);
-    RCP<GOVector> numPartsIRecv = Xpetra::VectorFactory<GO, LO, GO, NO>::Build(partsIOwn);
-    if (numSend) {
-      ArrayRCP<GO> partsISendData = partsISend->getDataNonConst(0);
-      for (int i = 0; i < numSend; i++)
-        partsISendData[i] = 1;
+    {
+      SubFactoryMonitor m1(*this, "Mapping Step 1", currentLevel);
+      GO partsIndexBase = 0;
+      RCP<Map>    partsIHave  = MapFactory   ::Build(lib, Teuchos::OrdinalTraits<Xpetra::global_size_t>::invalid(), myParts(), partsIndexBase, comm);
+      RCP<Map>    partsIOwn   = MapFactory   ::Build(lib,                                                 numProcs,  myPart(), partsIndexBase, comm);
+      RCP<Export> partsExport = ExportFactory::Build(partsIHave, partsIOwn);
+      
+      RCP<GOVector> partsISend    = Xpetra::VectorFactory<GO, LO, GO, NO>::Build(partsIHave);
+      RCP<GOVector> numPartsIRecv = Xpetra::VectorFactory<GO, LO, GO, NO>::Build(partsIOwn);
+      if (numSend) {
+        ArrayRCP<GO> partsISendData = partsISend->getDataNonConst(0);
+        for (int i = 0; i < numSend; i++)
+          partsISendData[i] = 1;
+      }
+      (numPartsIRecv->getDataNonConst(0))[0] = 0;
+      
+      numPartsIRecv->doExport(*partsISend, *partsExport, Xpetra::ADD);
+      numRecv = (numPartsIRecv->getData(0))[0];
     }
-    (numPartsIRecv->getDataNonConst(0))[0] = 0;
 
-    numPartsIRecv->doExport(*partsISend, *partsExport, Xpetra::ADD);
-    numRecv = (numPartsIRecv->getData(0))[0];
 
     // Step 2: Get my GIDs from everybody else
     MPI_Datatype MpiType = Teuchos::Details::MpiTypeTraits<GO>::getType();
@@ -359,7 +363,12 @@ namespace MueLu {
     std::sort(myGIDs.begin(), myGIDs.end());
 
     // Step 3: Construct importer
-    RCP<Map>          newRowMap      = MapFactory   ::Build(lib, rowMap->getGlobalNumElements(), myGIDs(), indexBase, origComm);
+    RCP<Map>          newRowMap;
+    {
+      SubFactoryMonitor m1(*this, "Map construction", currentLevel);
+      newRowMap = MapFactory   ::Build(lib, rowMap->getGlobalNumElements(), myGIDs(), indexBase, origComm);
+    }
+
     RCP<const Import> rowMapImporter;
 
     RCP<const BlockedMap> blockedRowMap = Teuchos::rcp_dynamic_cast<const BlockedMap>(rowMap);