MueLu_CoalesceDropFactory_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_COALESCEDROPFACTORY_KOKKOS_DEF_HPP
#define MUELU_COALESCEDROPFACTORY_KOKKOS_DEF_HPP

#include <Kokkos_Core.hpp>
#include <KokkosSparse_CrsMatrix.hpp>

#include "Xpetra_Matrix.hpp"

#include "MueLu_CoalesceDropFactory_kokkos_decl.hpp"

#include "MueLu_AmalgamationInfo_kokkos.hpp"
#include "MueLu_Exceptions.hpp"
#include "MueLu_Level.hpp"
#include "MueLu_LWGraph_kokkos.hpp"
#include "MueLu_MasterList.hpp"
#include "MueLu_Monitor.hpp"
#include "MueLu_Utilities_kokkos.hpp"

namespace MueLu {


  namespace CoalesceDrop_Kokkos_Details { // anonymous

    template<class LO, class RowType>
    class ScanFunctor {
    public:
      ScanFunctor(RowType rows_) : rows(rows_) { }

      KOKKOS_INLINE_FUNCTION
      void operator()(const LO i, LO& upd, const bool& final) const {
        upd += rows(i);
        if (final)
          rows(i) = upd;
      }

    private:
      RowType rows;
    };

    template<class LO, class GhostedViewType>
    class ClassicalDropFunctor {
    private:
      typedef typename GhostedViewType::value_type      SC;
      typedef          Kokkos::ArithTraits<SC>          ATS;
      typedef typename ATS::magnitudeType               magnitudeType;

      GhostedViewType   diag;   // corresponds to overlapped diagonal multivector (2D View)
      magnitudeType     eps;

    public:
      ClassicalDropFunctor(GhostedViewType ghostedDiag, magnitudeType threshold) :
          diag(ghostedDiag),
          eps(threshold)
      { }

      // Return true if we drop, false if not
      KOKKOS_FORCEINLINE_FUNCTION
      bool operator()(LO row, LO col, SC val) const {
        // We avoid square root by using squared values
        auto aiiajj = ATS::magnitude(diag(row, 0)) * ATS::magnitude(diag(col, 0));   // |a_ii|*|a_jj|
        auto aij2   = ATS::magnitude(val)          * ATS::magnitude(val);            // |a_ij|^2

        return (aij2 <= eps*eps * aiiajj);
      }
    };

    template<class LO, class CoordsType>
    class DistanceFunctor {
    private:
      typedef typename CoordsType::value_type           SC;
      typedef          Kokkos::ArithTraits<SC>          ATS;
      typedef typename ATS::magnitudeType               magnitudeType;

    public:
      typedef SC value_type;

    public:
      DistanceFunctor(CoordsType coords_) : coords(coords_) { }

      KOKKOS_INLINE_FUNCTION
      magnitudeType distance2(LO row, LO col) const {
        SC d = ATS::zero(), s;
        for (size_t j = 0; j < coords.extent(1); j++) {
          s = coords(row,j) - coords(col,j);
          d += s*s;
        }
        return ATS::magnitude(d);
      }
    private:
      CoordsType coords;
    };

    template<class LO, class GhostedViewType, class DistanceFunctor>
    class DistanceLaplacianDropFunctor {
    private:
      typedef typename GhostedViewType::value_type      SC;
      typedef          Kokkos::ArithTraits<SC>          ATS;
      typedef typename ATS::magnitudeType               magnitudeType;

    public:
      DistanceLaplacianDropFunctor(GhostedViewType ghostedLaplDiag, DistanceFunctor distFunctor_, magnitudeType threshold) :
          diag(ghostedLaplDiag),
          distFunctor(distFunctor_),
          eps(threshold)
      { }

      // Return true if we drop, false if not
      KOKKOS_INLINE_FUNCTION
      bool operator()(LO row, LO col, SC /* val */) const {
        // We avoid square root by using squared values

        // We ignore incoming value of val as we operate on an auxiliary
        // distance Laplacian matrix
        typedef typename DistanceFunctor::value_type      dSC;
        typedef          Kokkos::ArithTraits<dSC>         dATS;
        auto fval = dATS::one() / distFunctor.distance2(row, col);

        auto aiiajj = ATS::magnitude(diag(row, 0)) * ATS::magnitude(diag(col, 0));   // |a_ii|*|a_jj|
        auto aij2   = ATS::magnitude(fval)         * ATS::magnitude(fval);           // |a_ij|^2

        return (aij2 <= eps*eps * aiiajj);
      }

    private:
      GhostedViewType   diag;   // corresponds to overlapped diagonal multivector (2D View)
      DistanceFunctor   distFunctor;
      magnitudeType     eps;
    };

    template<class SC, class LO, class MatrixType, class BndViewType, class DropFunctorType>
    class ScalarFunctor {
    private:
      typedef typename MatrixType::StaticCrsGraphType   graph_type;
      typedef typename graph_type::row_map_type         rows_type;
      typedef typename graph_type::entries_type         cols_type;
      typedef typename MatrixType::values_type          vals_type;
      typedef          Kokkos::ArithTraits<SC>          ATS;
      typedef typename ATS::val_type                    impl_Scalar;
      typedef Kokkos::ArithTraits<impl_Scalar>          impl_ATS;
      typedef typename ATS::magnitudeType               magnitudeType;

    public:
      ScalarFunctor(MatrixType A_, BndViewType bndNodes_, DropFunctorType dropFunctor_,
                    typename rows_type::non_const_type rows_,
                    typename cols_type::non_const_type colsAux_,
                    typename vals_type::non_const_type valsAux_,
                    bool reuseGraph_, bool lumping_, SC /* threshold_ */,
                    bool aggregationMayCreateDirichlet_ ) :
          A(A_),
          bndNodes(bndNodes_),
          dropFunctor(dropFunctor_),
          rows(rows_),
          colsAux(colsAux_),
          valsAux(valsAux_),
          reuseGraph(reuseGraph_),
          lumping(lumping_),
          aggregationMayCreateDirichlet(aggregationMayCreateDirichlet_)
      {
        rowsA = A.graph.row_map;
        zero = impl_ATS::zero();
      }

      KOKKOS_INLINE_FUNCTION
      void operator()(const LO row, LO& nnz) const {
        auto rowView = A.rowConst(row);
        auto length  = rowView.length;
        auto offset  = rowsA(row);

        impl_Scalar diag = zero;
        LO rownnz = 0;
        LO diagID = -1;
        for (decltype(length) colID = 0; colID < length; colID++) {
          LO col = rowView.colidx(colID);
          impl_Scalar val = rowView.value (colID);

          if (!dropFunctor(row, col, rowView.value(colID)) || row == col) {
            colsAux(offset+rownnz) = col;

            LO valID = (reuseGraph ? colID : rownnz);
            valsAux(offset+valID) = val;
            if (row == col)
              diagID = valID;

            rownnz++;

          } else {
            // Rewrite with zeros (needed for reuseGraph)
            valsAux(offset+colID) = zero;
            diag += val;
          }
        }
        // How to assert on the device?
        // assert(diagIndex != -1);
        rows(row+1) = rownnz;
        // if (lumping && diagID != -1) {
        if (lumping) {
          // Add diag to the diagonal

          // NOTE_KOKKOS: valsAux was allocated with
          // ViewAllocateWithoutInitializing. This is not a problem here
          // because we explicitly set this value above.
          valsAux(offset+diagID) += diag;
        }

        // If the only element remaining after filtering is diagonal, mark node as boundary
        // FIXME: this should really be replaced by the following
        //    if (indices.size() == 1 && indices[0] == row)
        //        boundaryNodes[row] = true;
        // We do not do it this way now because there is no framework for distinguishing isolated
        // and boundary nodes in the aggregation algorithms
        bndNodes(row) = (rownnz == 1 && aggregationMayCreateDirichlet);

        nnz += rownnz;
      }

    private:
      MatrixType                            A;
      BndViewType                           bndNodes;
      DropFunctorType                       dropFunctor;

      rows_type                             rowsA;

      typename rows_type::non_const_type    rows;
      typename cols_type::non_const_type    colsAux;
      typename vals_type::non_const_type    valsAux;

      bool                                  reuseGraph;
      bool                                  lumping;
      bool                                  aggregationMayCreateDirichlet;
      impl_Scalar                           zero;
    };

    // collect number nonzeros of blkSize rows in nnz_(row+1)
    template<class MatrixType, class NnzType, class blkSizeType>
    class Stage1aVectorFunctor {
    private:
      typedef typename MatrixType::ordinal_type LO;

    public:
      Stage1aVectorFunctor(MatrixType kokkosMatrix_, NnzType nnz_, blkSizeType blkSize_) :
        kokkosMatrix(kokkosMatrix_),
        nnz(nnz_),
        blkSize(blkSize_) { }

      KOKKOS_INLINE_FUNCTION
      void operator()(const LO row, LO& totalnnz) const {

        // the following code is more or less what MergeRows is doing
        // count nonzero entries in all dof rows associated with node row
        LO nodeRowMaxNonZeros = 0;
        for (LO j = 0; j < blkSize; j++) {
          auto rowView = kokkosMatrix.row(row * blkSize + j);
          nodeRowMaxNonZeros += rowView.length;
        }
        nnz(row + 1) = nodeRowMaxNonZeros;
        totalnnz += nodeRowMaxNonZeros;
      }


    private:
      MatrixType kokkosMatrix; //< local matrix part
      NnzType nnz;             //< View containing number of nonzeros for current row
      blkSizeType blkSize;     //< block size (or partial block size in strided maps)
    };


    // build the dof-based column map containing the local dof ids belonging to blkSize rows in matrix
    // sort column ids
    // translate them into (unique) node ids
    // count the node column ids per node row
    template<class MatrixType, class NnzType, class blkSizeType, class ColDofType, class Dof2NodeTranslationType, class BdryNodeTypeConst, class BdryNodeType, class boolType>
    class Stage1bcVectorFunctor {
    private:
      typedef typename MatrixType::ordinal_type LO;

    private:
      MatrixType kokkosMatrix; //< local matrix part
      NnzType coldofnnz; //< view containing start and stop indices for subviews
      blkSizeType blkSize;     //< block size (or partial block size in strided maps)
      ColDofType coldofs;      //< view containing the local dof ids associated with columns for the blkSize rows (not sorted)
      Dof2NodeTranslationType dof2node; //< view containing the local node id associated with the local dof id
      NnzType colnodennz; //< view containing number of column nodes for each node row
      BdryNodeTypeConst  dirichletdof;  //< view containing with num dofs booleans. True if dof (not necessarily entire node) is dirichlet boundardy dof.
      BdryNodeType  bdrynode;  //< view containing with numNodes booleans. True if node is (full) dirichlet boundardy node.
      boolType usegreedydirichlet;    //< boolean for use of greedy Dirichlet (if any dof is Dirichlet, entire node is dirichlet) default false (need all dofs in node to be Dirichlet for node to be Dirichlet)

    public:
      Stage1bcVectorFunctor(MatrixType kokkosMatrix_,
                           NnzType coldofnnz_,
                           blkSizeType blkSize_,
                           ColDofType coldofs_,
                           Dof2NodeTranslationType dof2node_,
                           NnzType colnodennz_,
                           BdryNodeTypeConst dirichletdof_,
                           BdryNodeType bdrynode_,
                           boolType usegreedydirichlet_) :
        kokkosMatrix(kokkosMatrix_),
        coldofnnz(coldofnnz_),
        blkSize(blkSize_),
        coldofs(coldofs_),
        dof2node(dof2node_),
        colnodennz(colnodennz_),
        dirichletdof(dirichletdof_),
        bdrynode(bdrynode_),
        usegreedydirichlet(usegreedydirichlet_) {
      }

      KOKKOS_INLINE_FUNCTION
      void operator()(const LO rowNode, LO& nnz) const {

        LO pos = coldofnnz(rowNode);
        if( usegreedydirichlet ){
          bdrynode(rowNode) = false;
          for (LO j = 0; j < blkSize; j++) {
            auto rowView = kokkosMatrix.row(rowNode * blkSize + j);
            auto numIndices = rowView.length;

            // if any dof in the node is Dirichlet
            if( dirichletdof(rowNode * blkSize + j) )
              bdrynode(rowNode) = true;

            for (decltype(numIndices) k = 0; k < numIndices; k++) {
              auto dofID = rowView.colidx(k);
              coldofs(pos) = dofID;
              pos ++;
            }
          }
        }else{
          bdrynode(rowNode) = true;
          for (LO j = 0; j < blkSize; j++) {
            auto rowView = kokkosMatrix.row(rowNode * blkSize + j);
            auto numIndices = rowView.length;

            // if any dof in the node is not Dirichlet
            if( dirichletdof(rowNode * blkSize + j) == false )
              bdrynode(rowNode) = false;

            for (decltype(numIndices) k = 0; k < numIndices; k++) {
              auto dofID = rowView.colidx(k);
              coldofs(pos) = dofID;
              pos ++;
            }
          }
        }

        // sort coldofs
        LO begin = coldofnnz(rowNode);
        LO end   = coldofnnz(rowNode+1);
        LO n     = end - begin;
        for (LO i = 0; i < (n-1); i++) {
          for (LO j = 0; j < (n-i-1); j++) {
            if (coldofs(j+begin) > coldofs(j+begin+1)) {
              LO temp = coldofs(j+begin);
              coldofs(j+begin) = coldofs(j+begin+1);
              coldofs(j+begin+1) = temp;
            }
          }
        }
        size_t cnt = 0;
        LO lastNodeID = -1;
        for (LO i = 0; i < n; i++) {
          LO dofID  = coldofs(begin + i);
          LO nodeID = dof2node(dofID);
          if(nodeID != lastNodeID) {
            lastNodeID = nodeID;
            coldofs(begin+cnt) = nodeID;
            cnt++;
          }
        }
        colnodennz(rowNode+1) = cnt;
        nnz += cnt;
      }

    };

    // fill column node id view
    template<class MatrixType, class ColDofNnzType, class ColDofType, class ColNodeNnzType, class ColNodeType>
    class Stage1dVectorFunctor {
    private:
      typedef typename MatrixType::ordinal_type LO;
      typedef typename MatrixType::value_type SC;

    private:
      ColDofType     coldofs;    //< view containing mixed node and dof indices (only input)
      ColDofNnzType  coldofnnz;  //< view containing the start and stop indices for subviews (dofs)
      ColNodeType colnodes;      //< view containing the local node ids associated with columns
      ColNodeNnzType colnodennz; //< view containing start and stop indices for subviews

    public:
      Stage1dVectorFunctor(ColDofType coldofs_, ColDofNnzType coldofnnz_, ColNodeType colnodes_, ColNodeNnzType colnodennz_) :
        coldofs(coldofs_),
        coldofnnz(coldofnnz_),
        colnodes(colnodes_),
        colnodennz(colnodennz_) {
      }

      KOKKOS_INLINE_FUNCTION
      void operator()(const LO rowNode) const {
        auto dofbegin  = coldofnnz(rowNode);
        auto nodebegin = colnodennz(rowNode);
        auto nodeend   = colnodennz(rowNode+1);
        auto n = nodeend - nodebegin;

        for (decltype(nodebegin) i = 0; i < n; i++) {
          colnodes(nodebegin + i) = coldofs(dofbegin + i);
        }
      }
    };


  } // namespace

  template <class Scalar, class LocalOrdinal, class GlobalOrdinal, class DeviceType>
  RCP<const ParameterList> CoalesceDropFactory_kokkos<Scalar,LocalOrdinal,GlobalOrdinal,Kokkos::Compat::KokkosDeviceWrapperNode<DeviceType>>::GetValidParameterList() const {
    RCP<ParameterList> validParamList = rcp(new ParameterList());

#define SET_VALID_ENTRY(name) validParamList->setEntry(name, MasterList::getEntry(name))
    SET_VALID_ENTRY("aggregation: drop tol");
    SET_VALID_ENTRY("aggregation: Dirichlet threshold");
    SET_VALID_ENTRY("aggregation: drop scheme");
    SET_VALID_ENTRY("aggregation: dropping may create Dirichlet");
    SET_VALID_ENTRY("aggregation: greedy Dirichlet");
    SET_VALID_ENTRY("filtered matrix: use lumping");
    SET_VALID_ENTRY("filtered matrix: reuse graph");
    SET_VALID_ENTRY("filtered matrix: reuse eigenvalue");
    {
      typedef Teuchos::StringToIntegralParameterEntryValidator<int> validatorType;
      validParamList->getEntry("aggregation: drop scheme").setValidator(
        rcp(new validatorType(Teuchos::tuple<std::string>("classical", "distance laplacian"), "aggregation: drop scheme")));
    }
#undef  SET_VALID_ENTRY
    validParamList->set< RCP<const FactoryBase> >("A",                  Teuchos::null, "Generating factory of the matrix A");
    validParamList->set< RCP<const FactoryBase> >("UnAmalgamationInfo", Teuchos::null, "Generating factory for UnAmalgamationInfo");
    validParamList->set< RCP<const FactoryBase> >("Coordinates",        Teuchos::null, "Generating factory for Coordinates");

    return validParamList;
  }

  template <class Scalar,class LocalOrdinal, class GlobalOrdinal, class DeviceType>
  void CoalesceDropFactory_kokkos<Scalar,LocalOrdinal,GlobalOrdinal,Kokkos::Compat::KokkosDeviceWrapperNode<DeviceType>>::DeclareInput(Level &currentLevel) const {
    Input(currentLevel, "A");
    Input(currentLevel, "UnAmalgamationInfo");

    const ParameterList& pL = GetParameterList();
    if (pL.get<std::string>("aggregation: drop scheme") == "distance laplacian")
      Input(currentLevel, "Coordinates");
  }

  template <class Scalar,class LocalOrdinal, class GlobalOrdinal, class DeviceType>
  void CoalesceDropFactory_kokkos<Scalar,LocalOrdinal,GlobalOrdinal,Kokkos::Compat::KokkosDeviceWrapperNode<DeviceType>>::
  Build(Level& currentLevel) const {
    FactoryMonitor m(*this, "Build", currentLevel);

    typedef Teuchos::ScalarTraits<SC> STS;
    typedef typename STS::magnitudeType MT;
    const MT zero = Teuchos::ScalarTraits<MT>::zero();

    auto A         = Get< RCP<Matrix> >(currentLevel, "A");


    /* NOTE: storageblocksize (from GetStorageBlockSize()) is the size of a block in the chosen storage scheme.
       blkSize is the number of storage blocks that must kept together during the amalgamation process.

       Both of these quantities may be different than numPDEs (from GetFixedBlockSize()), but the following must always hold:

       numPDEs = blkSize * storageblocksize.
       
       If numPDEs==1
         Matrix is point storage (classical CRS storage).  storageblocksize=1 and  blkSize=1
         No other values makes sense.

       If numPDEs>1
         If matrix uses point storage, then storageblocksize=1  and blkSize=numPDEs.
         If matrix uses block storage, with block size of n, then storageblocksize=n, and blkSize=numPDEs/n.
         Thus far, only storageblocksize=numPDEs and blkSize=1 has been tested.
      */
 
    TEUCHOS_TEST_FOR_EXCEPTION(A->GetFixedBlockSize() % A->GetStorageBlockSize() != 0,Exceptions::RuntimeError,"A->GetFixedBlockSize() needs to be a multiple of A->GetStorageBlockSize()");
    LO   blkSize   = A->GetFixedBlockSize() / A->GetStorageBlockSize();

    auto amalInfo = Get< RCP<AmalgamationInfo_kokkos> >(currentLevel, "UnAmalgamationInfo");

    const ParameterList& pL = GetParameterList();

    std::string algo = pL.get<std::string>("aggregation: drop scheme");

    double threshold = pL.get<double>("aggregation: drop tol");
    GetOStream(Runtime0) << "algorithm = \"" << algo << "\": threshold = " << threshold
                         << ", blocksize = " << A->GetFixedBlockSize() << std::endl;

    const typename STS::magnitudeType dirichletThreshold =
        STS::magnitude(as<SC>(pL.get<double>("aggregation: Dirichlet threshold")));

    GO numDropped = 0, numTotal = 0;

    RCP<LWGraph_kokkos> graph;
    LO                  dofsPerNode = -1;

    typedef typename LWGraph_kokkos::boundary_nodes_type boundary_nodes_type;
    boundary_nodes_type boundaryNodes;

    RCP<Matrix> filteredA;
    if (blkSize == 1 && threshold == zero) {
      // Scalar problem without dropping

      // Detect and record rows that correspond to Dirichlet boundary conditions
      boundaryNodes = Utilities_kokkos::DetectDirichletRows(*A, dirichletThreshold);

      // Trivial LWGraph construction
      graph = rcp(new LWGraph_kokkos(A->getCrsGraph()->getLocalGraphDevice(), A->getRowMap(), A->getColMap(), "graph of A"));
      graph->getLocalLWGraph().SetBoundaryNodeMap(boundaryNodes);

      numTotal = A->getLocalNumEntries();
      dofsPerNode = 1;

      filteredA = A;

    } else if (blkSize == 1 && threshold != zero) {
      // Scalar problem with dropping

      typedef typename Matrix::local_matrix_type          local_matrix_type;
      typedef typename LWGraph_kokkos::local_graph_type   kokkos_graph_type;
      typedef typename kokkos_graph_type::row_map_type::non_const_type    rows_type;
      typedef typename kokkos_graph_type::entries_type::non_const_type    cols_type;
      typedef typename local_matrix_type::values_type::non_const_type     vals_type;

      LO   numRows      = A->getLocalNumRows();
      local_matrix_type kokkosMatrix = A->getLocalMatrixDevice();
      auto nnzA  = kokkosMatrix.nnz();
      auto rowsA = kokkosMatrix.graph.row_map;


      typedef Kokkos::ArithTraits<SC>          ATS;
      typedef typename ATS::val_type           impl_Scalar;
      typedef Kokkos::ArithTraits<impl_Scalar> impl_ATS;

      bool reuseGraph = pL.get<bool>("filtered matrix: reuse graph");
      bool lumping    = pL.get<bool>("filtered matrix: use lumping");
      if (lumping)
        GetOStream(Runtime0) << "Lumping dropped entries" << std::endl;

      const bool aggregationMayCreateDirichlet = pL.get<bool>("aggregation: dropping may create Dirichlet");

      // FIXME_KOKKOS: replace by ViewAllocateWithoutInitializing + setting a single value
      rows_type rows   ("FA_rows",     numRows+1);
      cols_type colsAux(Kokkos::ViewAllocateWithoutInitializing("FA_aux_cols"), nnzA);
      vals_type valsAux;
      if (reuseGraph) {
        SubFactoryMonitor m2(*this, "CopyMatrix", currentLevel);

        // Share graph with the original matrix
        filteredA = MatrixFactory::Build(A->getCrsGraph());

        // Do a no-op fill-complete
        RCP<ParameterList> fillCompleteParams(new ParameterList);
        fillCompleteParams->set("No Nonlocal Changes", true);
        filteredA->fillComplete(fillCompleteParams);

        // No need to reuseFill, just modify in place
        valsAux = filteredA->getLocalMatrixDevice().values;

      } else {
        // Need an extra array to compress
        valsAux = vals_type(Kokkos::ViewAllocateWithoutInitializing("FA_aux_vals"), nnzA);
      }

      typename boundary_nodes_type::non_const_type bndNodes(Kokkos::ViewAllocateWithoutInitializing("boundaryNodes"), numRows);

      LO nnzFA = 0;
      {
        if (algo == "classical") {
          // Construct overlapped matrix diagonal
          RCP<Vector> ghostedDiag;
          {
            kokkosMatrix = local_matrix_type();
            SubFactoryMonitor m2(*this, "Ghosted diag construction", currentLevel);
            ghostedDiag     = Utilities_kokkos::GetMatrixOverlappedDiagonal(*A);
            kokkosMatrix=A->getLocalMatrixDevice();
          }

          // Filter out entries
          {
            SubFactoryMonitor m2(*this, "MainLoop", currentLevel);

            auto ghostedDiagView = ghostedDiag->getDeviceLocalView(Xpetra::Access::ReadWrite);

            CoalesceDrop_Kokkos_Details::ClassicalDropFunctor<LO, decltype(ghostedDiagView)> dropFunctor(ghostedDiagView, threshold);
            CoalesceDrop_Kokkos_Details::ScalarFunctor<typename ATS::val_type, LO, local_matrix_type, decltype(bndNodes), decltype(dropFunctor)>
              scalarFunctor(kokkosMatrix, bndNodes, dropFunctor, rows, colsAux, valsAux, reuseGraph, lumping, threshold, aggregationMayCreateDirichlet);

            Kokkos::parallel_reduce("MueLu:CoalesceDropF:Build:scalar_filter:main_loop", range_type(0,numRows),
                                    scalarFunctor, nnzFA);
          }

        } else if (algo == "distance laplacian") {
          typedef Xpetra::MultiVector<typename Teuchos::ScalarTraits<Scalar>::magnitudeType,LO,GO,NO> doubleMultiVector;
          auto coords = Get<RCP<doubleMultiVector> >(currentLevel, "Coordinates");

          auto uniqueMap    = A->getRowMap();
          auto nonUniqueMap = A->getColMap();

          // Construct ghosted coordinates
          RCP<const Import> importer;
          {
            SubFactoryMonitor m2(*this, "Coords Import construction", currentLevel);
            importer = ImportFactory::Build(uniqueMap, nonUniqueMap);
          }
          RCP<doubleMultiVector> ghostedCoords;
          {
            SubFactoryMonitor m2(*this, "Ghosted coords construction", currentLevel);
            ghostedCoords = Xpetra::MultiVectorFactory<typename Teuchos::ScalarTraits<Scalar>::magnitudeType,LO,GO,NO>::Build(nonUniqueMap, coords->getNumVectors());
            ghostedCoords->doImport(*coords, *importer, Xpetra::INSERT);
          }

          auto ghostedCoordsView = ghostedCoords->getDeviceLocalView(Xpetra::Access::ReadWrite);
          CoalesceDrop_Kokkos_Details::DistanceFunctor<LO, decltype(ghostedCoordsView)> distFunctor(ghostedCoordsView);

          // Construct Laplacian diagonal
          RCP<Vector> localLaplDiag;
          {
            SubFactoryMonitor m2(*this, "Local Laplacian diag construction", currentLevel);

            localLaplDiag = VectorFactory::Build(uniqueMap);

            auto localLaplDiagView = localLaplDiag->getDeviceLocalView(Xpetra::Access::OverwriteAll);
            auto kokkosGraph = kokkosMatrix.graph;

            Kokkos::parallel_for("MueLu:CoalesceDropF:Build:scalar_filter:laplacian_diag", range_type(0,numRows),
              KOKKOS_LAMBDA(const LO row) {
                auto rowView = kokkosGraph.rowConst(row);
                auto length  = rowView.length;

                impl_Scalar d = impl_ATS::zero();
                for (decltype(length) colID = 0; colID < length; colID++) {
                  auto col = rowView(colID);
                  if (row != col)
                    d += impl_ATS::one()/distFunctor.distance2(row, col);
                }
                localLaplDiagView(row,0) = d;
              });
          }

          // Construct ghosted Laplacian diagonal
          RCP<Vector> ghostedLaplDiag;
          {
            SubFactoryMonitor m2(*this, "Ghosted Laplacian diag construction", currentLevel);
            ghostedLaplDiag = VectorFactory::Build(nonUniqueMap);
            ghostedLaplDiag->doImport(*localLaplDiag, *importer, Xpetra::INSERT);
          }

          // Filter out entries
          {
            SubFactoryMonitor m2(*this, "MainLoop", currentLevel);

            auto ghostedLaplDiagView = ghostedLaplDiag->getDeviceLocalView(Xpetra::Access::ReadWrite);

            CoalesceDrop_Kokkos_Details::DistanceLaplacianDropFunctor<LO, decltype(ghostedLaplDiagView), decltype(distFunctor)>
                dropFunctor(ghostedLaplDiagView, distFunctor, threshold);
            CoalesceDrop_Kokkos_Details::ScalarFunctor<SC, LO, local_matrix_type, decltype(bndNodes), decltype(dropFunctor)>
                scalarFunctor(kokkosMatrix, bndNodes, dropFunctor, rows, colsAux, valsAux, reuseGraph, lumping, threshold, true);

            Kokkos::parallel_reduce("MueLu:CoalesceDropF:Build:scalar_filter:main_loop", range_type(0,numRows),
                                    scalarFunctor, nnzFA);
          }
        }

      }
      numDropped = nnzA - nnzFA;

      boundaryNodes = bndNodes;

      {
        SubFactoryMonitor m2(*this, "CompressRows", currentLevel);

        // parallel_scan (exclusive)
        Kokkos::parallel_scan("MueLu:CoalesceDropF:Build:scalar_filter:compress_rows", range_type(0,numRows+1),
          KOKKOS_LAMBDA(const LO i, LO& update, const bool& final_pass) {
            update += rows(i);
            if (final_pass)
              rows(i) = update;
          });
      }

      // Compress cols (and optionally vals)
      // We use a trick here: we moved all remaining elements to the beginning
      // of the original row in the main loop, so we don't need to check for
      // INVALID here, and just stop when achieving the new number of elements
      // per row.
      cols_type cols(Kokkos::ViewAllocateWithoutInitializing("FA_cols"), nnzFA);
      vals_type vals;
      if (reuseGraph) {
        GetOStream(Runtime1) << "reuse matrix graph for filtering (compress matrix columns only)" << std::endl;
        // Only compress cols
        SubFactoryMonitor m2(*this, "CompressColsAndVals", currentLevel);

        Kokkos::parallel_for("MueLu:TentativePF:Build:compress_cols", range_type(0,numRows),
          KOKKOS_LAMBDA(const LO i) {
            // Is there Kokkos memcpy?
            LO rowStart  = rows(i);
            LO rowAStart = rowsA(i);
            size_t rownnz = rows(i+1) - rows(i);
            for (size_t j = 0; j < rownnz; j++)
              cols(rowStart+j) = colsAux(rowAStart+j);
          });
      } else {
        // Compress cols and vals
        GetOStream(Runtime1) << "new matrix graph for filtering (compress matrix columns and values)" << std::endl;
        SubFactoryMonitor m2(*this, "CompressColsAndVals", currentLevel);

        vals = vals_type(Kokkos::ViewAllocateWithoutInitializing("FA_vals"), nnzFA);

        Kokkos::parallel_for("MueLu:TentativePF:Build:compress_cols", range_type(0,numRows),
          KOKKOS_LAMBDA(const LO i) {
            LO rowStart  = rows(i);
            LO rowAStart = rowsA(i);
            size_t rownnz = rows(i+1) - rows(i);
            for (size_t j = 0; j < rownnz; j++) {
              cols(rowStart+j) = colsAux(rowAStart+j);
              vals(rowStart+j) = valsAux(rowAStart+j);
            }
          });
      }

      kokkos_graph_type kokkosGraph(cols, rows);

      {
        SubFactoryMonitor m2(*this, "LWGraph construction", currentLevel);

        graph = rcp(new LWGraph_kokkos(kokkosGraph, A->getRowMap(), A->getColMap(), "filtered graph of A"));
        graph->getLocalLWGraph().SetBoundaryNodeMap(boundaryNodes);
      }

      numTotal = A->getLocalNumEntries();

      dofsPerNode = 1;

      if (!reuseGraph) {
        SubFactoryMonitor m2(*this, "LocalMatrix+FillComplete", currentLevel);

        local_matrix_type localFA = local_matrix_type("A", numRows, A->getLocalMatrixDevice().numCols(), nnzFA, vals, rows, cols);
        auto filteredACrs = CrsMatrixFactory::Build(localFA, A->getRowMap(), A->getColMap(), A->getDomainMap(), A->getRangeMap(),
                                                    A->getCrsGraph()->getImporter(), A->getCrsGraph()->getExporter());
        filteredA = rcp(new CrsMatrixWrap(filteredACrs));
      }

      filteredA->SetFixedBlockSize(A->GetFixedBlockSize());

      if (pL.get<bool>("filtered matrix: reuse eigenvalue")) {
        // Reuse max eigenvalue from A
        // It is unclear what eigenvalue is the best for the smoothing, but we already may have
        // the D^{-1}A estimate in A, may as well use it.
        // NOTE: ML does that too
        filteredA->SetMaxEigenvalueEstimate(A->GetMaxEigenvalueEstimate());
      } else {
        filteredA->SetMaxEigenvalueEstimate(-Teuchos::ScalarTraits<SC>::one());
      }

    } else if (blkSize > 1 && threshold == zero) {
      // Case 3:  block problem without filtering
      //
      // FIXME_KOKKOS: this code is completely unoptimized. It really should do
      // a very simple thing: merge rows and produce nodal graph. But the code
      // seems very complicated. Can we do better?

      TEUCHOS_TEST_FOR_EXCEPTION(A->getRowMap()->getLocalNumElements() % blkSize != 0, MueLu::Exceptions::RuntimeError, "MueLu::CoalesceDropFactory: Number of local elements is " << A->getRowMap()->getLocalNumElements() << " but should be a multiply of " << blkSize);

      const RCP<const Map> rowMap = A->getRowMap();
      const RCP<const Map> colMap = A->getColMap();

      // build a node row map (uniqueMap = non-overlapping) and a node column map
      // (nonUniqueMap = overlapping). The arrays rowTranslation and colTranslation
      // stored in the AmalgamationInfo class container contain the local node id
      // given a local dof id. The data is calculated in the AmalgamationFactory and
      // stored in the variable "UnAmalgamationInfo" (which is of type AmalagamationInfo)
      const RCP<const Map> uniqueMap = amalInfo->getNodeRowMap();
      const RCP<const Map> nonUniqueMap = amalInfo->getNodeColMap();
      Array<LO> rowTranslationArray = *(amalInfo->getRowTranslation()); // TAW should be transform that into a View?
      Array<LO> colTranslationArray = *(amalInfo->getColTranslation());

      Kokkos::View<LO*, Kokkos::MemoryUnmanaged>
        rowTranslationView(rowTranslationArray.getRawPtr(),rowTranslationArray.size() );
      Kokkos::View<LO*, Kokkos::MemoryUnmanaged>
        colTranslationView(colTranslationArray.getRawPtr(),colTranslationArray.size() );

      // get number of local nodes
      LO numNodes = Teuchos::as<LocalOrdinal>(uniqueMap->getLocalNumElements());
      typedef typename Kokkos::View<LocalOrdinal*, DeviceType> id_translation_type;
      id_translation_type rowTranslation("dofId2nodeId",rowTranslationArray.size());
      id_translation_type colTranslation("ov_dofId2nodeId",colTranslationArray.size());
      Kokkos::deep_copy(rowTranslation, rowTranslationView);
      Kokkos::deep_copy(colTranslation, colTranslationView);

      // extract striding information
      blkSize = A->GetFixedBlockSize();  //< the full block size (number of dofs per node in strided map)
      LocalOrdinal blkId   = -1;         //< the block id within a strided map or -1 if it is a full block map
      LocalOrdinal blkPartSize = A->GetFixedBlockSize(); //< stores block size of part blkId (or the full block size)
      if(A->IsView("stridedMaps") == true) {
        const RCP<const Map> myMap = A->getRowMap("stridedMaps");
        const RCP<const StridedMap> strMap = Teuchos::rcp_dynamic_cast<const StridedMap>(myMap);
        TEUCHOS_TEST_FOR_EXCEPTION(strMap.is_null() == true, Exceptions::RuntimeError, "Map is not of type stridedMap");
        blkSize = Teuchos::as<const LocalOrdinal>(strMap->getFixedBlockSize());
        blkId   = strMap->getStridedBlockId();
        if (blkId > -1)
          blkPartSize = Teuchos::as<LocalOrdinal>(strMap->getStridingData()[blkId]);
      }
      auto kokkosMatrix = A->getLocalMatrixDevice(); // access underlying kokkos data

      //
      typedef typename LWGraph_kokkos::local_graph_type   kokkos_graph_type;
      typedef typename kokkos_graph_type::row_map_type    row_map_type;
      //typedef typename row_map_type::HostMirror           row_map_type_h;
      typedef typename kokkos_graph_type::entries_type    entries_type;

      // Stage 1c: get number of dof-nonzeros per blkSize node rows
      typename row_map_type::non_const_type dofNnz("nnz_map", numNodes + 1);
      LO numDofCols = 0;
      CoalesceDrop_Kokkos_Details::Stage1aVectorFunctor<decltype(kokkosMatrix), decltype(dofNnz), decltype(blkPartSize)> stage1aFunctor(kokkosMatrix, dofNnz, blkPartSize);
      Kokkos::parallel_reduce("MueLu:CoalesceDropF:Build:scalar_filter:stage1a", range_type(0,numNodes), stage1aFunctor, numDofCols);
      // parallel_scan (exclusive)
      CoalesceDrop_Kokkos_Details::ScanFunctor<LO,decltype(dofNnz)> scanFunctor(dofNnz);
      Kokkos::parallel_scan("MueLu:CoalesceDropF:Build:scalar_filter:stage1_scan", range_type(0,numNodes+1), scanFunctor);

      // Detect and record dof rows that correspond to Dirichlet boundary conditions
      boundary_nodes_type singleEntryRows = Utilities_kokkos::DetectDirichletRows(*A, dirichletThreshold);

      typename entries_type::non_const_type dofcols("dofcols", numDofCols/*dofNnz(numNodes)*/); // why does dofNnz(numNodes) work? should be a parallel reduce, i guess

      // we have dofcols and dofids from Stage1dVectorFunctor
      LO numNodeCols = 0;
      typename row_map_type::non_const_type rows("nnz_nodemap", numNodes + 1);
      typename boundary_nodes_type::non_const_type bndNodes("boundaryNodes", numNodes);

      CoalesceDrop_Kokkos_Details::Stage1bcVectorFunctor <decltype(kokkosMatrix), decltype(dofNnz), decltype(blkPartSize), decltype(dofcols), decltype(colTranslation), decltype(singleEntryRows), decltype(bndNodes), bool> stage1bcFunctor(kokkosMatrix, dofNnz, blkPartSize, dofcols, colTranslation, rows, singleEntryRows, bndNodes, pL.get<bool>("aggregation: greedy Dirichlet"));
      Kokkos::parallel_reduce("MueLu:CoalesceDropF:Build:scalar_filter:stage1c", range_type(0,numNodes), stage1bcFunctor,numNodeCols);

      // parallel_scan (exclusive)
      CoalesceDrop_Kokkos_Details::ScanFunctor<LO,decltype(rows)> scanNodeFunctor(rows);
      Kokkos::parallel_scan("MueLu:CoalesceDropF:Build:scalar_filter:stage1_scan", range_type(0,numNodes+1), scanNodeFunctor);

      // create column node view
      typename entries_type::non_const_type cols("nodecols", numNodeCols);


      CoalesceDrop_Kokkos_Details::Stage1dVectorFunctor <decltype(kokkosMatrix), decltype(dofNnz), decltype(dofcols), decltype(rows), decltype(cols)> stage1dFunctor(dofcols, dofNnz, cols, rows);
      Kokkos::parallel_for("MueLu:CoalesceDropF:Build:scalar_filter:stage1c", range_type(0,numNodes), stage1dFunctor);
      kokkos_graph_type kokkosGraph(cols, rows);

      // create LW graph
      graph = rcp(new LWGraph_kokkos(kokkosGraph, uniqueMap, nonUniqueMap, "amalgamated graph of A"));

      boundaryNodes = bndNodes;
      graph->getLocalLWGraph().SetBoundaryNodeMap(boundaryNodes);
      numTotal = A->getLocalNumEntries();

      dofsPerNode = blkSize;

      filteredA = A;

    } else {
      TEUCHOS_TEST_FOR_EXCEPTION(true, Exceptions::RuntimeError, "MueLu: CoalesceDropFactory_kokkos: Block filtering is not implemented");
    }

    if (GetVerbLevel() & Statistics1) {
      GO numLocalBoundaryNodes  = 0;
      GO numGlobalBoundaryNodes = 0;

      Kokkos::parallel_reduce("MueLu:CoalesceDropF:Build:bnd", range_type(0, boundaryNodes.extent(0)),
        KOKKOS_LAMBDA(const LO i, GO& n) {
          if (boundaryNodes(i))
            n++;
        }, numLocalBoundaryNodes);

      auto comm = A->getRowMap()->getComm();
      MueLu_sumAll(comm, numLocalBoundaryNodes, numGlobalBoundaryNodes);
      GetOStream(Statistics1) << "Detected " << numGlobalBoundaryNodes << " Dirichlet nodes" << std::endl;
    }

    if ((GetVerbLevel() & Statistics1) && threshold != zero) {
      auto comm = A->getRowMap()->getComm();

      GO numGlobalTotal, numGlobalDropped;
      MueLu_sumAll(comm, numTotal,   numGlobalTotal);
      MueLu_sumAll(comm, numDropped, numGlobalDropped);

      if (numGlobalTotal != 0) {
        GetOStream(Statistics1) << "Number of dropped entries: "
            << numGlobalDropped << "/" << numGlobalTotal
            << " (" << 100*Teuchos::as<double>(numGlobalDropped)/Teuchos::as<double>(numGlobalTotal) << "%)" << std::endl;
      }
    }

    Set(currentLevel, "DofsPerNode",  dofsPerNode);
    Set(currentLevel, "Graph",        graph);
    Set(currentLevel, "A",            filteredA);
  }
}
#endif // MUELU_COALESCEDROPFACTORY_KOKKOS_DEF_HPP