MueLu_RebalanceTransferFactory_def.hpp

Go to the documentation of this file.

// @HEADER
// *****************************************************************************
//        MueLu: A package for multigrid based preconditioning
//
// Copyright 2012 NTESS and the MueLu contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER
 
#ifndef MUELU_REBALANCETRANSFERFACTORY_DEF_HPP
#define MUELU_REBALANCETRANSFERFACTORY_DEF_HPP
 
#include <sstream>
#include <Teuchos_Tuple.hpp>
 
#include "Xpetra_MultiVector.hpp"
#include "Xpetra_MultiVectorFactory.hpp"
#include "Xpetra_Vector.hpp"
#include "Xpetra_VectorFactory.hpp"
#include <Xpetra_Matrix.hpp>
#include <Xpetra_MapFactory.hpp>
#include <Xpetra_MatrixFactory.hpp>
#include <Xpetra_Import.hpp>
#include <Xpetra_ImportFactory.hpp>
#include <Xpetra_IO.hpp>
 
#include "MueLu_RebalanceTransferFactory_decl.hpp"
 
#include "MueLu_AmalgamationFactory.hpp"
#include "MueLu_Level.hpp"
#include "MueLu_MasterList.hpp"
#include "MueLu_Monitor.hpp"
#include "MueLu_PerfUtils.hpp"
 
namespace MueLu {
 
template <class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
RebalanceTransferFactory<Scalar, LocalOrdinal, GlobalOrdinal, Node>::RebalanceTransferFactory() = default;
 
template <class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
RebalanceTransferFactory<Scalar, LocalOrdinal, GlobalOrdinal, Node>::~RebalanceTransferFactory() = default;
 
template <class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
RCP<const ParameterList> RebalanceTransferFactory<Scalar, LocalOrdinal, GlobalOrdinal, Node>::GetValidParameterList() const {
  RCP<ParameterList> validParamList = rcp(new ParameterList());
 
#define SET_VALID_ENTRY(name) validParamList->setEntry(name, MasterList::getEntry(name))
  SET_VALID_ENTRY("repartition: rebalance P and R");
  SET_VALID_ENTRY("repartition: explicit via new copy rebalance P and R");
  SET_VALID_ENTRY("repartition: rebalance Nullspace");
  SET_VALID_ENTRY("transpose: use implicit");
  SET_VALID_ENTRY("repartition: use subcommunicators");
  SET_VALID_ENTRY("repartition: send type");
#undef SET_VALID_ENTRY
 
  {
    typedef Teuchos::StringValidator validatorType;
    RCP<validatorType> typeValidator = rcp(new validatorType(Teuchos::tuple<std::string>("Interpolation", "Restriction")));
    validParamList->set("type", "Interpolation", "Type of the transfer operator that need to be rebalanced (Interpolation or Restriction)", typeValidator);
  }
 
  validParamList->set<RCP<const FactoryBase> >("P", null, "Factory of the prolongation operator that need to be rebalanced (only used if type=Interpolation)");
  validParamList->set<RCP<const FactoryBase> >("R", null, "Factory of the restriction operator that need to be rebalanced (only used if type=Restriction)");
  validParamList->set<RCP<const FactoryBase> >("Nullspace", null, "Factory of the nullspace that need to be rebalanced (only used if type=Interpolation)");
  validParamList->set<RCP<const FactoryBase> >("Coordinates", null, "Factory of the coordinates that need to be rebalanced (only used if type=Interpolation)");
  validParamList->set<RCP<const FactoryBase> >("Material", null, "Factory of the material that need to be rebalanced (only used if type=Interpolation)");
  validParamList->set<RCP<const FactoryBase> >("BlockNumber", null, "Factory of the block ids that need to be rebalanced (only used if type=Interpolation)");
  validParamList->set<RCP<const FactoryBase> >("Importer", null, "Factory of the importer object used for the rebalancing");
  validParamList->set<int>("write start", -1, "First level at which coordinates should be written to file");
  validParamList->set<int>("write end", -1, "Last level at which coordinates should be written to file");
 
  // TODO validation: "P" parameter valid only for type="Interpolation" and "R" valid only for type="Restriction". Like so:
  // if (paramList.isEntry("type") && paramList.get("type) == "Interpolation) {
  //     validParamList->set< RCP<const FactoryBase> >("P",              Teuchos::null, "Factory of the prolongation operator that need to be rebalanced (only used if type=Interpolation)");
 
  return validParamList;
}
 
template <class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
void RebalanceTransferFactory<Scalar, LocalOrdinal, GlobalOrdinal, Node>::DeclareInput(Level& /* fineLevel */, Level& coarseLevel) const {
  const ParameterList& pL = GetParameterList();
 
  if (pL.get<std::string>("type") == "Interpolation") {
    Input(coarseLevel, "P");
    if (pL.get<bool>("repartition: rebalance Nullspace"))
      Input(coarseLevel, "Nullspace");
    if (pL.get<RCP<const FactoryBase> >("Coordinates") != Teuchos::null)
      Input(coarseLevel, "Coordinates");
    if (pL.get<RCP<const FactoryBase> >("Material") != Teuchos::null)
      Input(coarseLevel, "Material");
    if (pL.get<RCP<const FactoryBase> >("BlockNumber") != Teuchos::null)
      Input(coarseLevel, "BlockNumber");
 
  } else {
    if (pL.get<bool>("transpose: use implicit") == false)
      Input(coarseLevel, "R");
  }
 
  Input(coarseLevel, "Importer");
}
 
template <class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
void RebalanceTransferFactory<Scalar, LocalOrdinal, GlobalOrdinal, Node>::Build(Level& fineLevel, Level& coarseLevel) const {
  FactoryMonitor m(*this, "Build", coarseLevel);
  typedef Xpetra::MultiVector<typename Teuchos::ScalarTraits<Scalar>::magnitudeType, LO, GO, NO> xdMV;
 
  const ParameterList& pL = GetParameterList();
 
  RCP<Matrix> originalP = Get<RCP<Matrix> >(coarseLevel, "P");
  // If we don't have a valid P (e.g., # global aggregates is 0), skip this rebalancing. This level will
  // ultimately be removed in MueLu_Hierarchy_defs.h via a resize()
  if (originalP == Teuchos::null) {
    Set(coarseLevel, "P", originalP);
    return;
  }
  int implicit       = !pL.get<bool>("repartition: rebalance P and R");
  int reallyExplicit = pL.get<bool>("repartition: explicit via new copy rebalance P and R");
  int writeStart     = pL.get<int>("write start");
  int writeEnd       = pL.get<int>("write end");
 
  if (writeStart == 0 && fineLevel.GetLevelID() == 0 && writeStart <= writeEnd && IsAvailable(fineLevel, "Coordinates")) {
    std::string fileName = "coordinates_level_0.m";
    RCP<xdMV> fineCoords = fineLevel.Get<RCP<xdMV> >("Coordinates");
    if (fineCoords != Teuchos::null)
      Xpetra::IO<typename Teuchos::ScalarTraits<Scalar>::magnitudeType, LO, GO, NO>::Write(fileName, *fineCoords);
  }
 
  if (writeStart == 0 && fineLevel.GetLevelID() == 0 && writeStart <= writeEnd && IsAvailable(fineLevel, "BlockNumber")) {
    std::string fileName                    = "BlockNumber_level_0.m";
    RCP<LocalOrdinalVector> fineBlockNumber = fineLevel.Get<RCP<LocalOrdinalVector> >("BlockNumber");
    if (fineBlockNumber != Teuchos::null)
      Xpetra::IO<SC, LO, GO, NO>::WriteLOMV(fileName, *fineBlockNumber);
  }
 
  RCP<const Import> importer = Get<RCP<const Import> >(coarseLevel, "Importer");
  if (implicit) {
    // Save the importer, we'll need it for solve
    coarseLevel.Set("Importer", importer, NoFactory::get());
  }
 
  RCP<ParameterList> params = rcp(new ParameterList());
  if (IsPrint(Statistics2)) {
    params->set("printLoadBalancingInfo", true);
    params->set("printCommInfo", true);
  }
 
  std::string transferType = pL.get<std::string>("type");
  if (transferType == "Interpolation") {
    originalP = Get<RCP<Matrix> >(coarseLevel, "P");
 
    {
      // This line must be after the Get call
      if (implicit || importer.is_null()) {
        GetOStream(Runtime0) << "Using original prolongator" << std::endl;
        Set(coarseLevel, "P", originalP);
 
      } else {
        SubFactoryMonitor m1(*this, "Rebalancing prolongator", coarseLevel);
        // There are two version of an explicit rebalanced P and R.
        // The !reallyExplicit way, is sufficient for all MueLu purposes
        // with the exception of the CombinePFactory that needs true domain
        // and column maps.
        //   !reallyExplicit:
        //     Rather than calling fillComplete (which would entail creating a new
        //     column map), it's sufficient to replace the domain map and importer.
        //     Note that this potentially violates the assumption that in the
        //     column map, local IDs appear before any off-rank IDs.
        //
        //  reallyExplicit:
        //     P transfers from coarse grid to the fine grid. Here, we change
        //     the domain map (coarse) of Paccording to the new partition. The
        //     range map (fine) is kept unchanged.
        //
        //     The domain map of P must match the range map of R. To change the
        //     domain map of P, P needs to be fillCompleted again with the new
        //     domain map.  To achieve this, P is copied into a new matrix that
        //     is not fill-completed.  The doImport() operation is just used
        //     here to make a copy of P: the importer is trivial and there is
        //     no data movement involved.  The reordering actually happens during
        //     fillComplete() with domainMap == importer->getTargetMap().
 
        RCP<Matrix> rebalancedP;
        if (reallyExplicit) {
          size_t totalMaxPerRow = 0;
          ArrayRCP<size_t> nnzPerRow(originalP->getRowMap()->getLocalNumElements(), 0);
          for (size_t i = 0; i < originalP->getRowMap()->getLocalNumElements(); ++i) {
            nnzPerRow[i] = originalP->getNumEntriesInLocalRow(i);
            if (nnzPerRow[i] > totalMaxPerRow) totalMaxPerRow = nnzPerRow[i];
          }
 
          rebalancedP = MatrixFactory::Build(originalP->getRowMap(), totalMaxPerRow);
 
          {
            RCP<Import> trivialImporter = ImportFactory::Build(originalP->getRowMap(), originalP->getRowMap());
 
            std::string monitorLabel{"Rebalancing prolongator -- import only"};
            if (auto sendType = pL.get<std::string>("repartition: send type"); sendType != "") {
              auto sendTypeList = Teuchos::rcp(new Teuchos::ParameterList());
              sendTypeList->set("Send type", sendType);
              trivialImporter->setDistributorParameters(sendTypeList);
              monitorLabel += " [" + sendType + "]";
            }
 
            SubFactoryMonitor m2(*this, monitorLabel.c_str(), coarseLevel);
            rebalancedP->doImport(*originalP, *trivialImporter, Xpetra::INSERT);
          }
          rebalancedP->fillComplete(importer->getTargetMap(), originalP->getRangeMap());
 
        } else {
          rebalancedP                    = originalP;
          RCP<const CrsMatrixWrap> crsOp = rcp_dynamic_cast<const CrsMatrixWrap>(originalP);
          TEUCHOS_TEST_FOR_EXCEPTION(crsOp == Teuchos::null, Exceptions::BadCast, "Cast from Xpetra::Matrix to Xpetra::CrsMatrixWrap failed");
 
          RCP<CrsMatrix> rebalancedP2 = crsOp->getCrsMatrix();
          TEUCHOS_TEST_FOR_EXCEPTION(rebalancedP2 == Teuchos::null, std::runtime_error, "Xpetra::CrsMatrixWrap doesn't have a CrsMatrix");
 
          {
            SubFactoryMonitor subM(*this, "Rebalancing prolongator -- fast map replacement", coarseLevel);
 
            RCP<const Import> newImporter;
            {
              SubFactoryMonitor subM2(*this, "Import construction", coarseLevel);
              newImporter = ImportFactory::Build(importer->getTargetMap(), rebalancedP->getColMap());
            }
            rebalancedP2->replaceDomainMapAndImporter(importer->getTargetMap(), newImporter);
          }
        }
        // TODO FIXME somehow we have to transfer the striding information of the permuted domain/range maps.
        // That is probably something for an external permutation factory
        //   if (originalP->IsView("stridedMaps"))
        //     rebalancedP->CreateView("stridedMaps", originalP);
        if (!rebalancedP.is_null()) {
          std::ostringstream oss;
          oss << "P_" << coarseLevel.GetLevelID();
          rebalancedP->setObjectLabel(oss.str());
        }
        Set(coarseLevel, "P", rebalancedP);
 
        if (IsPrint(Statistics2))
          GetOStream(Statistics2) << PerfUtils::PrintMatrixInfo(*rebalancedP, "P (rebalanced)", params);
      }
    }
 
    if (importer.is_null()) {
      if (IsAvailable(coarseLevel, "Nullspace"))
        Set(coarseLevel, "Nullspace", Get<RCP<MultiVector> >(coarseLevel, "Nullspace"));
 
      if (pL.isParameter("Coordinates") && pL.get<RCP<const FactoryBase> >("Coordinates") != Teuchos::null)
        if (IsAvailable(coarseLevel, "Coordinates"))
          Set(coarseLevel, "Coordinates", Get<RCP<xdMV> >(coarseLevel, "Coordinates"));
 
      if (pL.isParameter("Material") && pL.get<RCP<const FactoryBase> >("Material") != Teuchos::null)
        if (IsAvailable(coarseLevel, "Material"))
          Set(coarseLevel, "Material", Get<RCP<MultiVector> >(coarseLevel, "Material"));
 
      if (pL.isParameter("BlockNumber") && pL.get<RCP<const FactoryBase> >("BlockNumber") != Teuchos::null)
        if (IsAvailable(coarseLevel, "BlockNumber"))
          Set(coarseLevel, "BlockNumber", Get<RCP<LocalOrdinalVector> >(coarseLevel, "BlockNumber"));
 
      return;
    }
 
    RCP<const Map> subCommMap;
    bool subCommMapConstructed = false;
 
    if (pL.isParameter("Coordinates") &&
        pL.get<RCP<const FactoryBase> >("Coordinates") != Teuchos::null &&
        IsAvailable(coarseLevel, "Coordinates")) {
      RCP<xdMV> coords = Get<RCP<xdMV> >(coarseLevel, "Coordinates");
 
      // This line must be after the Get call
      SubFactoryMonitor subM(*this, "Rebalancing coordinates", coarseLevel);
 
      GO numElts = coords->getMap()->getGlobalNumElements();
      LO blkSize = Teuchos::as<LO>(importer->getSourceMap()->getGlobalNumElements() / numElts);
 
      RCP<const Import> coordImporter;
 
      if (blkSize == 1) {
        coordImporter = importer;
      } else {
        RCP<const Map> origMap = coords->getMap();
        std::vector<size_t> stridingInfo{Teuchos::as<size_t>(blkSize)};
        RCP<const Map> targetMap = StridedMapFactory::Build(origMap->lib(), Teuchos::OrdinalTraits<Xpetra::global_size_t>::invalid(),
                                                            importer->getTargetMap()->getLocalElementList(), origMap->getIndexBase(), stridingInfo, origMap->getComm());
        RCP<const Map> targetVectorMap;
 
        AmalgamationFactory<SC, LO, GO, NO>::AmalgamateMap(rcp_dynamic_cast<const StridedMap>(targetMap), targetVectorMap);
        coordImporter = ImportFactory::Build(origMap, targetVectorMap);
      }
 
      RCP<xdMV> permutedCoords = Xpetra::MultiVectorFactory<typename Teuchos::ScalarTraits<Scalar>::magnitudeType, LO, GO, NO>::Build(coordImporter->getTargetMap(), coords->getNumVectors(), false);
      permutedCoords->doImport(*coords, *coordImporter, Xpetra::INSERT);
 
      if (pL.isParameter("repartition: use subcommunicators") && pL.get<bool>("repartition: use subcommunicators")) {
        if (blkSize == 1) {
          if (!subCommMapConstructed) {
            subCommMap            = permutedCoords->getMap()->removeEmptyProcesses();
            subCommMapConstructed = true;
          }
          permutedCoords->replaceMap(subCommMap);
        } else {
          permutedCoords->replaceMap(permutedCoords->getMap()->removeEmptyProcesses());
        }
      }
 
      if (permutedCoords->getMap() == Teuchos::null)
        permutedCoords = Teuchos::null;
 
      Set(coarseLevel, "Coordinates", permutedCoords);
 
      std::string fileName = "rebalanced_coordinates_level_" + toString(coarseLevel.GetLevelID()) + ".m";
      if (writeStart <= coarseLevel.GetLevelID() && coarseLevel.GetLevelID() <= writeEnd && permutedCoords->getMap() != Teuchos::null)
        Xpetra::IO<typename Teuchos::ScalarTraits<Scalar>::magnitudeType, LO, GO, NO>::Write(fileName, *permutedCoords);
    }
 
    if (IsAvailable(coarseLevel, "Material")) {
      RCP<MultiVector> material = Get<RCP<MultiVector> >(coarseLevel, "Material");
 
      // This line must be after the Get call
      SubFactoryMonitor subM(*this, "Rebalancing material", coarseLevel);
 
      GO numElts = material->getMap()->getGlobalNumElements();
      LO blkSize = Teuchos::as<LO>(importer->getSourceMap()->getGlobalNumElements() / numElts);
 
      RCP<const Import> materialImporter;
 
      if (blkSize == 1) {
        materialImporter = importer;
      } else {
        RCP<const Map> origMap = material->getMap();
        std::vector<size_t> stridingInfo{Teuchos::as<size_t>(blkSize)};
        RCP<const Map> targetMap = StridedMapFactory::Build(origMap->lib(), Teuchos::OrdinalTraits<Xpetra::global_size_t>::invalid(),
                                                            importer->getTargetMap()->getLocalElementList(), origMap->getIndexBase(), stridingInfo, origMap->getComm());
        RCP<const Map> targetVectorMap;
 
        AmalgamationFactory<SC, LO, GO, NO>::AmalgamateMap(rcp_dynamic_cast<const StridedMap>(targetMap), targetVectorMap);
        materialImporter = ImportFactory::Build(origMap, targetVectorMap);
      }
 
      RCP<MultiVector> permutedMaterial = MultiVectorFactory::Build(materialImporter->getTargetMap(), material->getNumVectors(), false);
      permutedMaterial->doImport(*material, *materialImporter, Xpetra::INSERT);
 
      if (pL.get<bool>("repartition: use subcommunicators")) {
        if (blkSize == 1) {
          if (!subCommMapConstructed) {
            subCommMap            = permutedMaterial->getMap()->removeEmptyProcesses();
            subCommMapConstructed = true;
          }
          permutedMaterial->replaceMap(subCommMap);
        } else {
          permutedMaterial->replaceMap(permutedMaterial->getMap()->removeEmptyProcesses());
        }
      }
 
      if (permutedMaterial->getMap() == Teuchos::null)
        permutedMaterial = Teuchos::null;
 
      Set(coarseLevel, "Material", permutedMaterial);
    }
 
    if (pL.isParameter("BlockNumber") &&
        pL.get<RCP<const FactoryBase> >("BlockNumber") != Teuchos::null &&
        IsAvailable(coarseLevel, "BlockNumber")) {
      RCP<LocalOrdinalVector> BlockNumber = Get<RCP<LocalOrdinalVector> >(coarseLevel, "BlockNumber");
 
      // This line must be after the Get call
      SubFactoryMonitor subM(*this, "Rebalancing BlockNumber", coarseLevel);
 
      RCP<LocalOrdinalVector> permutedBlockNumber = LocalOrdinalVectorFactory::Build(importer->getTargetMap(), false);
      permutedBlockNumber->doImport(*BlockNumber, *importer, Xpetra::INSERT);
 
      if (pL.isParameter("repartition: use subcommunicators") && pL.get<bool>("repartition: use subcommunicators")) {
        if (!subCommMapConstructed) {
          subCommMap            = permutedBlockNumber->getMap()->removeEmptyProcesses();
          subCommMapConstructed = true;
        }
        permutedBlockNumber->replaceMap(subCommMap);
      }
 
      if (permutedBlockNumber->getMap() == Teuchos::null)
        permutedBlockNumber = Teuchos::null;
 
      Set(coarseLevel, "BlockNumber", permutedBlockNumber);
 
      std::string fileName = "rebalanced_BlockNumber_level_" + toString(coarseLevel.GetLevelID()) + ".m";
      if (writeStart <= coarseLevel.GetLevelID() && coarseLevel.GetLevelID() <= writeEnd && permutedBlockNumber->getMap() != Teuchos::null)
        Xpetra::IO<SC, LO, GO, NO>::WriteLOMV(fileName, *permutedBlockNumber);
    }
 
    if (IsAvailable(coarseLevel, "Nullspace")) {
      RCP<MultiVector> nullspace = Get<RCP<MultiVector> >(coarseLevel, "Nullspace");
 
      // This line must be after the Get call
      SubFactoryMonitor subM(*this, "Rebalancing nullspace", coarseLevel);
 
      RCP<MultiVector> permutedNullspace = MultiVectorFactory::Build(importer->getTargetMap(), nullspace->getNumVectors(), false);
      permutedNullspace->doImport(*nullspace, *importer, Xpetra::INSERT);
 
      if (pL.get<bool>("repartition: use subcommunicators")) {
        if (!subCommMapConstructed) {
          subCommMap            = permutedNullspace->getMap()->removeEmptyProcesses();
          subCommMapConstructed = true;
        }
        permutedNullspace->replaceMap(subCommMap);
      }
 
      if (permutedNullspace->getMap() == Teuchos::null)
        permutedNullspace = Teuchos::null;
 
      Set(coarseLevel, "Nullspace", permutedNullspace);
    }
 
  } else {
    if (!pL.get<bool>("transpose: use implicit")) {
      RCP<Matrix> originalR = Get<RCP<Matrix> >(coarseLevel, "R");
 
      if (implicit || importer.is_null()) {
        GetOStream(Runtime0) << "Using original restrictor" << std::endl;
        Set(coarseLevel, "R", originalR);
 
      } else {
        SubFactoryMonitor m2(*this, "Rebalancing restrictor", coarseLevel);
 
        RCP<Matrix> rebalancedR;
        {
          SubFactoryMonitor subM(*this, "Rebalancing restriction -- fusedImport", coarseLevel);
 
          RCP<Map> dummy;  // meaning: use originalR's domain map.
          Teuchos::ParameterList listLabel;
          listLabel.set("Timer Label", "MueLu::RebalanceR-" + Teuchos::toString(coarseLevel.GetLevelID()));
          rebalancedR = MatrixFactory::Build(originalR, *importer, dummy, importer->getTargetMap(), Teuchos::rcp(&listLabel, false));
        }
        if (!rebalancedR.is_null()) {
          std::ostringstream oss;
          oss << "R_" << coarseLevel.GetLevelID();
          rebalancedR->setObjectLabel(oss.str());
        }
        Set(coarseLevel, "R", rebalancedR);
 
        // TODO FIXME somehow we have to transfer the striding information of the permuted domain/range maps.
        // That is probably something for an external permutation factory
        // if (originalR->IsView("stridedMaps"))
        //   rebalancedR->CreateView("stridedMaps", originalR);
 
        if (IsPrint(Statistics2))
          GetOStream(Statistics2) << PerfUtils::PrintMatrixInfo(*rebalancedR, "R (rebalanced)", params);
      }
    }
  }
}
 
}  // namespace MueLu
 
#endif  // MUELU_REBALANCETRANSFERFACTORY_DEF_HPP