MueLu_AggregationStructuredAlgorithm_kokkos_def.hpp

Go to the documentation of this file.

// @HEADER
//
// ***********************************************************************
//
//        MueLu: A package for multigrid based preconditioning
//                  Copyright 2012 Sandia Corporation
//
// Under the terms of Contract DE-AC04-94AL85000 with Sandia Corporation,
// the U.S. Government retains certain rights in this software.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
// 1. Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// 3. Neither the name of the Corporation nor the names of the
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY SANDIA CORPORATION "AS IS" AND ANY
// EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
// PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL SANDIA CORPORATION OR THE
// CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
// EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
// PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
// PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
// LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
// NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// Questions? Contact
//                    Jonathan Hu       (jhu@sandia.gov)
//                    Andrey Prokopenko (aprokop@sandia.gov)
//                    Ray Tuminaro      (rstumin@sandia.gov)
//
// ***********************************************************************
//
// @HEADER
#ifndef MUELU_AGGREGATIONSTRUCTUREDALGORITHM_KOKKOS_DEF_HPP
#define MUELU_AGGREGATIONSTRUCTUREDALGORITHM_KOKKOS_DEF_HPP


#include <Teuchos_Comm.hpp>
#include <Teuchos_CommHelpers.hpp>

#include <Xpetra_MapFactory.hpp>
#include <Xpetra_Map.hpp>
#include <Xpetra_CrsGraphFactory.hpp>
#include <Xpetra_CrsGraph.hpp>

#include "MueLu_Exceptions.hpp"
#include "MueLu_Monitor.hpp"

#include "MueLu_LWGraph_kokkos.hpp"
#include "MueLu_Aggregates_kokkos.hpp"
#include "MueLu_IndexManager_kokkos.hpp"
#include "MueLu_AggregationStructuredAlgorithm_kokkos_decl.hpp"

namespace MueLu {

  template <class LocalOrdinal, class GlobalOrdinal, class Node>
  void AggregationStructuredAlgorithm_kokkos<LocalOrdinal, GlobalOrdinal, Node>::
  BuildAggregates(const Teuchos::ParameterList& /* params */, const LWGraph_kokkos& graph,
                  Aggregates_kokkos& aggregates,
                  Kokkos::View<unsigned*, memory_space>& aggStat,
                  LO& numNonAggregatedNodes) const {
    Monitor m(*this, "BuildAggregates");

    RCP<Teuchos::FancyOStream> out;
    if(const char* dbg = std::getenv("MUELU_STRUCTUREDALGORITHM_DEBUG")) {
      out = Teuchos::fancyOStream(Teuchos::rcpFromRef(std::cout));
      out->setShowAllFrontMatter(false).setShowProcRank(true);
    } else {
      out = Teuchos::getFancyOStream(rcp(new Teuchos::oblackholestream()));
    }

    RCP<IndexManager_kokkos> geoData = aggregates.GetIndexManager();
    const LO numLocalFineNodes= geoData->getNumLocalFineNodes();
    const LO numCoarseNodes   = geoData->getNumCoarseNodes();
    LOVectorView vertex2AggId = aggregates.GetVertex2AggId()->template getLocalView<memory_space>();
    LOVectorView procWinner   = aggregates.GetProcWinner()  ->template getLocalView<memory_space>();

    *out << "Loop over fine nodes and assign them to an aggregate and a rank" << std::endl;
    LO numAggregatedNodes;
    fillAggregatesFunctor fillAggregates(geoData,
                                         graph.GetComm()->getRank(),
                                         aggStat,
                                         vertex2AggId,
                                         procWinner);
    Kokkos::parallel_reduce("StructuredAggregation: fill aggregates data",
                            Kokkos::RangePolicy<execution_space>(0, numLocalFineNodes),
                            fillAggregates,
                            numAggregatedNodes);

    *out << "numCoarseNodes= " << numCoarseNodes
         << ", numAggregatedNodes= " << numAggregatedNodes << std::endl;
    numNonAggregatedNodes = numNonAggregatedNodes - numAggregatedNodes;

  } // BuildAggregates()


  template <class LocalOrdinal, class GlobalOrdinal, class Node>
  void AggregationStructuredAlgorithm_kokkos<LocalOrdinal, GlobalOrdinal, Node>::
  BuildGraph(const LWGraph_kokkos& graph, RCP<IndexManager_kokkos>& geoData, const LO dofsPerNode,
             RCP<CrsGraph>& myGraph) const {
    Monitor m(*this, "BuildGraphP");

    RCP<Teuchos::FancyOStream> out;
    if(const char* dbg = std::getenv("MUELU_STRUCTUREDALGORITHM_DEBUG")) {
      out = Teuchos::fancyOStream(Teuchos::rcpFromRef(std::cout));
      out->setShowAllFrontMatter(false).setShowProcRank(true);
    } else {
      out = Teuchos::getFancyOStream(rcp(new Teuchos::oblackholestream()));
    }

    typedef typename CrsGraph::local_graph_type local_graph_type;

    // Compute the number of coarse points needed to interpolate quantities to a fine point
    int numInterpolationPoints = 0;
    if(geoData->getInterpolationOrder() == 0) {
      numInterpolationPoints = 1;
    } else if(geoData->getInterpolationOrder() == 1) {
      // Compute 2^numDimensions using bit logic to avoid round-off errors from std::pow()
      numInterpolationPoints = 1 << geoData->getNumDimensions();
    }
    *out << "numInterpolationPoints=" << numInterpolationPoints << std::endl;

    const LO numLocalFineNodes = geoData->getNumLocalFineNodes();
    const LO numCoarseNodes    = geoData->getNumCoarseNodes();
    const LO numNnzEntries     = dofsPerNode*(numCoarseNodes + numInterpolationPoints
                                              *(numLocalFineNodes - numCoarseNodes));

    non_const_row_map_type rowPtr("Prolongator graph, rowPtr", dofsPerNode*(numLocalFineNodes + 1));
    entries_type colIndex("Prolongator graph, colIndices", numNnzEntries);

    *out << "Compute prolongatorGraph data" << std::endl;
    if(geoData->getInterpolationOrder() == 0) {
      computeGraphDataConstantFunctor computeGraphData(geoData,
                                                       numCoarseNodes,
                                                       dofsPerNode,
                                                       geoData->getCoarseningRates(),
                                                       geoData->getCoarseningEndRates(),
                                                       geoData->getLocalFineNodesPerDir(),
                                                       rowPtr,
                                                       colIndex);
      Kokkos::parallel_for("Structured Aggregation: compute loca graph data",
                           Kokkos::RangePolicy<execution_space>(0, numLocalFineNodes),
                           computeGraphData);
    } else if(geoData->getInterpolationOrder() == 1) {
      // Note, lbv 2018-11-08: in the piece-wise linear case I am computing the rowPtr
      // using a parallel scan, it might be possible to do something faster than that
      // by including this calculation in computeGraphDataLinearFunctor but at the moment
      // all the ideas I have include a bunch of if statements which I would like to avoid.
      computeGraphRowPtrFunctor computeGraphRowPtr(geoData,
                                                   dofsPerNode,
                                                   numInterpolationPoints,
                                                   numLocalFineNodes,
                                                   geoData->getCoarseningRates(),
                                                   geoData->getLocalFineNodesPerDir(),
                                                   rowPtr);
      Kokkos::parallel_scan("Structured Aggregation: compute rowPtr for prolongator graph",
                            Kokkos::RangePolicy<execution_space>(0, numLocalFineNodes + 1),
                            computeGraphRowPtr);

      computeGraphDataLinearFunctor computeGraphData(geoData,
                                                     geoData->getNumDimensions(),
                                                     numCoarseNodes,
                                                     dofsPerNode,
                                                     numInterpolationPoints,
                                                     geoData->getCoarseningRates(),
                                                     geoData->getCoarseningEndRates(),
                                                     geoData->getLocalFineNodesPerDir(),
                                                     geoData->getCoarseNodesPerDir(),
                                                     rowPtr,
                                                     colIndex);
      Kokkos::parallel_for("Structured Aggregation: compute loca graph data",
                           Kokkos::RangePolicy<execution_space>(0, numLocalFineNodes),
                           computeGraphData);
    }

    local_graph_type myLocalGraph(colIndex, rowPtr);

    // Compute graph's colMap and domainMap
    RCP<Map> colMap, domainMap;
    *out << "Compute domain and column maps of the CrsGraph" << std::endl;
    colMap = MapFactory::Build(graph.GetDomainMap()->lib(),
                               Teuchos::OrdinalTraits<GO>::invalid(),
                               numCoarseNodes,
                               graph.GetDomainMap()->getIndexBase(),
                               graph.GetDomainMap()->getComm());
    domainMap = colMap;

    myGraph = CrsGraphFactory::Build(myLocalGraph, graph.GetDomainMap(), colMap,
                                     colMap, graph.GetDomainMap());

  } // BuildGraph()


  template <class LocalOrdinal, class GlobalOrdinal, class Node>
  AggregationStructuredAlgorithm_kokkos<LocalOrdinal, GlobalOrdinal, Node>::
  fillAggregatesFunctor::fillAggregatesFunctor(RCP<IndexManager_kokkos> geoData,
                                               const int myRank,
                                               Kokkos::View<unsigned*, memory_space> aggStat,
                                               LOVectorView vertex2AggID,
                                               LOVectorView procWinner) :
    geoData_(*geoData), myRank_(myRank), aggStat_(aggStat),
    vertex2AggID_(vertex2AggID), procWinner_(procWinner) {}

  template <class LocalOrdinal, class GlobalOrdinal, class Node>
  KOKKOS_INLINE_FUNCTION
  void AggregationStructuredAlgorithm_kokkos<LocalOrdinal, GlobalOrdinal, Node>::
  fillAggregatesFunctor::operator() (const LO nodeIdx, LO& lNumAggregatedNodes) const {
    // Compute coarse ID associated with fine LID
    LO rem, rate;
    LO coarseNodeCoarseLID;
    LO nodeFineTuple[3], coarseIdx[3];
    auto coarseRate          = geoData_.getCoarseningRates();
    auto endRate             = geoData_.getCoarseningEndRates();
    auto lFineNodesPerDir    = geoData_.getLocalFineNodesPerDir();
    // Compute coarse ID associated with fine LID
    geoData_.getFineLID2FineTuple(nodeIdx, nodeFineTuple);

    for(int dim = 0; dim < 3; ++dim) {
      coarseIdx[dim] = nodeFineTuple[dim] / coarseRate(dim);
      rem  = nodeFineTuple[dim] % coarseRate(dim);
      rate = (nodeFineTuple[dim] < lFineNodesPerDir(dim) - endRate(dim)) ? coarseRate(dim) : endRate(dim);
      if(rem > (rate / 2)) {++coarseIdx[dim];}
    }

    geoData_.getCoarseTuple2CoarseLID(coarseIdx[0], coarseIdx[1], coarseIdx[2],
                                     coarseNodeCoarseLID);

    vertex2AggID_(nodeIdx, 0) = coarseNodeCoarseLID;
    procWinner_(nodeIdx, 0)   = myRank_;
    aggStat_(nodeIdx)         = AGGREGATED;
    ++lNumAggregatedNodes;

  } // fillAggregatesFunctor::operator()

  template <class LocalOrdinal, class GlobalOrdinal, class Node>
  AggregationStructuredAlgorithm_kokkos<LocalOrdinal, GlobalOrdinal, Node>::
  computeGraphDataConstantFunctor::
  computeGraphDataConstantFunctor(RCP<IndexManager_kokkos> geoData,
                                  const LO NumGhostedNodes,
                                  const LO dofsPerNode,
                                  constIntTupleView coarseRate,
                                  constIntTupleView endRate,
                                  constLOTupleView  lFineNodesPerDir,
                                  non_const_row_map_type rowPtr,
                                  entries_type colIndex) : geoData_(*geoData),
    numGhostedNodes_(NumGhostedNodes), dofsPerNode_(dofsPerNode),
    coarseRate_(coarseRate), endRate_(endRate), lFineNodesPerDir_(lFineNodesPerDir),
    rowPtr_(rowPtr), colIndex_(colIndex) {

  } // computeGraphDataConstantFunctor()

  template <class LocalOrdinal, class GlobalOrdinal, class Node>
  KOKKOS_INLINE_FUNCTION
  void AggregationStructuredAlgorithm_kokkos<LocalOrdinal, GlobalOrdinal, Node>::
  computeGraphDataConstantFunctor::operator() (const LO nodeIdx) const {
    LO nodeFineTuple[3]   = {0, 0, 0};
    LO nodeCoarseTuple[3] = {0, 0, 0};

    // Compute ghosted tuple associated with fine LID
    geoData_.getFineLID2FineTuple(nodeIdx, nodeFineTuple);

    // Compute coarse tuple associated with fine point
    // then overwrite it with tuple associated with aggregate
    LO rem, rate, coarseNodeCoarseLID;
    for(int dim = 0; dim < 3; ++dim) {
      nodeCoarseTuple[dim] = nodeFineTuple[dim] / coarseRate_(dim);
      rem            = nodeFineTuple[dim] % coarseRate_(dim);
      if( nodeFineTuple[dim] < (lFineNodesPerDir_(dim) - endRate_(dim)) ) {
        rate = coarseRate_(dim);
      } else {
        rate = endRate_(dim);
      }
      if(rem > (rate / 2)) {++nodeCoarseTuple[dim];}
    }

    // get LID associted with aggregate
    geoData_.getCoarseTuple2CoarseLID(nodeCoarseTuple[0], nodeCoarseTuple[1], nodeCoarseTuple[2],
                                      coarseNodeCoarseLID);

    // store data into CrsGraph taking care of multiple dofs case
    for(LO dof = 0; dof < dofsPerNode_; ++dof) {
      rowPtr_(nodeIdx*dofsPerNode_ + dof + 1) = nodeIdx*dofsPerNode_ + dof + 1;
      colIndex_(nodeIdx*dofsPerNode_ + dof)   = coarseNodeCoarseLID*dofsPerNode_ + dof;
    }

  } // computeGraphDataConstantFunctor::operator()

  template <class LocalOrdinal, class GlobalOrdinal, class Node>
  AggregationStructuredAlgorithm_kokkos<LocalOrdinal, GlobalOrdinal, Node>::
  computeGraphRowPtrFunctor::computeGraphRowPtrFunctor(RCP<IndexManager_kokkos> geoData,
                                                       const LO dofsPerNode,
                                                       const int numInterpolationPoints,
                                                       const LO  numLocalRows,
                                                       constIntTupleView coarseRate,
                                                       constLOTupleView  lFineNodesPerDir,
                                                       non_const_row_map_type rowPtr) :
  geoData_(*geoData), dofsPerNode_(dofsPerNode),
    numInterpolationPoints_(numInterpolationPoints), numLocalRows_(numLocalRows),
    coarseRate_(coarseRate), lFineNodesPerDir_(lFineNodesPerDir), rowPtr_(rowPtr) {}

  template <class LocalOrdinal, class GlobalOrdinal, class Node>
  KOKKOS_INLINE_FUNCTION
  void AggregationStructuredAlgorithm_kokkos<LocalOrdinal, GlobalOrdinal, Node>::
  computeGraphRowPtrFunctor::operator() (const LO rowIdx, GO& update, const bool final) const {
    if (final) {
      // Kokkos uses a multipass algorithm to implement scan.
      // Only update the array on the final pass.  Updating the
      // array before changing 'update' means that we do an
      // exclusive scan.  Update the array after for an inclusive
      // scan.
      rowPtr_(rowIdx) = update;
    }
    if (rowIdx < numLocalRows_) {
      LO nodeIdx = rowIdx / dofsPerNode_;
      bool allCoarse = true;
      LO nodeFineTuple[3] = {0, 0, 0};
      geoData_.getFineLID2FineTuple(nodeIdx, nodeFineTuple);
      for(int dim = 0; dim < 3; ++dim) {
        const LO rem = nodeFineTuple[dim] % coarseRate_(dim);

        // Check if Fine node lies on Coarse Node
        allCoarse = (allCoarse && ((rem == 0) || (nodeFineTuple[dim] == lFineNodesPerDir_(dim) - 1)));
      }
      update += (allCoarse ? 1 : numInterpolationPoints_);
    }
  } // computeGraphRowPtrFunctor::operator()

  template <class LocalOrdinal, class GlobalOrdinal, class Node>
  AggregationStructuredAlgorithm_kokkos<LocalOrdinal, GlobalOrdinal, Node>::
  computeGraphDataLinearFunctor::computeGraphDataLinearFunctor(RCP<IndexManager_kokkos> geoData,
                                                               const int numDimensions,
                                                               const LO numGhostedNodes,
                                                               const LO dofsPerNode,
                                                               const int numInterpolationPoints,
                                                               constIntTupleView coarseRate,
                                                               constIntTupleView endRate,
                                                               constLOTupleView  lFineNodesPerDir,
                                                               constLOTupleView  ghostedNodesPerDir,
                                                               non_const_row_map_type rowPtr,
                                                               entries_type colIndex) :
  geoData_(*geoData), numDimensions_(numDimensions),
    numGhostedNodes_(numGhostedNodes),
    dofsPerNode_(dofsPerNode), numInterpolationPoints_(numInterpolationPoints),
    coarseRate_(coarseRate), endRate_(endRate), lFineNodesPerDir_(lFineNodesPerDir),
    ghostedNodesPerDir_(ghostedNodesPerDir), rowPtr_(rowPtr), colIndex_(colIndex) {

  } // computeGraphDataLinearFunctor()

  template <class LocalOrdinal, class GlobalOrdinal, class Node>
  KOKKOS_INLINE_FUNCTION
  void AggregationStructuredAlgorithm_kokkos<LocalOrdinal, GlobalOrdinal, Node>::
  computeGraphDataLinearFunctor::operator() (const LO nodeIdx) const {
    LO nodeFineTuple[3]   = {0, 0, 0};
    LO nodeCoarseTuple[3] = {0, 0, 0};

    // Compute coarse ID associated with fine LID
    geoData_.getFineLID2FineTuple(nodeIdx, nodeFineTuple);

    LO coarseNodeCoarseLID;
    bool allCoarse = false;
    for(int dim = 0; dim < 3; ++dim) {
      nodeCoarseTuple[dim] = nodeFineTuple[dim] / coarseRate_(dim);
    }
    if(rowPtr_(nodeIdx + 1) == rowPtr_(nodeIdx) + 1) {allCoarse = true;}

    geoData_.getCoarseTuple2CoarseLID(nodeCoarseTuple[0], nodeCoarseTuple[1], nodeCoarseTuple[2],
                                      coarseNodeCoarseLID);

    if(allCoarse) {
      // Fine node lies on Coarse node, easy case, we only need the LID of the coarse node.
      for(LO dof = 0; dof < dofsPerNode_; ++dof) {
        colIndex_(rowPtr_(nodeIdx*dofsPerNode_ + dof)) = coarseNodeCoarseLID*dofsPerNode_ + dof;
      }
    } else {

      for(int dim = 0; dim < numDimensions_; ++dim) {
        if(nodeCoarseTuple[dim] == ghostedNodesPerDir_(dim) - 1) { --nodeCoarseTuple[dim]; }
      }
      // Compute Coarse Node LID
      // Note lbv 10-06-2018: it is likely benefitial to remove the two if statments and somehow
      // find out the number of dimensions before calling the opertor() of the functor.
      for(LO dof = 0; dof < dofsPerNode_; ++dof) {
        geoData_.getCoarseTuple2CoarseLID(    nodeCoarseTuple[0],   nodeCoarseTuple[1],   nodeCoarseTuple[2],   colIndex_(rowPtr_(nodeIdx*dofsPerNode_ + dof)+0));
        geoData_.getCoarseTuple2CoarseLID(    nodeCoarseTuple[0]+1, nodeCoarseTuple[1],   nodeCoarseTuple[2],   colIndex_(rowPtr_(nodeIdx*dofsPerNode_ + dof)+1));
        if(numDimensions_ > 1) {
          geoData_.getCoarseTuple2CoarseLID(  nodeCoarseTuple[0],   nodeCoarseTuple[1]+1, nodeCoarseTuple[2],   colIndex_(rowPtr_(nodeIdx*dofsPerNode_ + dof)+2));
          geoData_.getCoarseTuple2CoarseLID(  nodeCoarseTuple[0]+1, nodeCoarseTuple[1]+1, nodeCoarseTuple[2],   colIndex_(rowPtr_(nodeIdx*dofsPerNode_ + dof)+3));
          if(numDimensions_ > 2) {
            geoData_.getCoarseTuple2CoarseLID(nodeCoarseTuple[0],   nodeCoarseTuple[1],   nodeCoarseTuple[2]+1, colIndex_(rowPtr_(nodeIdx*dofsPerNode_ + dof)+4));
            geoData_.getCoarseTuple2CoarseLID(nodeCoarseTuple[0]+1, nodeCoarseTuple[1],   nodeCoarseTuple[2]+1, colIndex_(rowPtr_(nodeIdx*dofsPerNode_ + dof)+5));
            geoData_.getCoarseTuple2CoarseLID(nodeCoarseTuple[0],   nodeCoarseTuple[1]+1, nodeCoarseTuple[2]+1, colIndex_(rowPtr_(nodeIdx*dofsPerNode_ + dof)+6));
            geoData_.getCoarseTuple2CoarseLID(nodeCoarseTuple[0]+1, nodeCoarseTuple[1]+1, nodeCoarseTuple[2]+1, colIndex_(rowPtr_(nodeIdx*dofsPerNode_ + dof)+7));
          }
        }
      }
    }
  } // computeGraphDataLinearFunctor::operator()

} // end namespace


#endif /* MUELU_AGGREGATIONSTRUCTUREDALGORITHM_DEF_HPP_ */