Tpetra_Details_getNumDiags.hpp

Go to the documentation of this file.

// @HEADER
// *****************************************************************************
//          Tpetra: Templated Linear Algebra Services Package
//
// Copyright 2008 NTESS and the Tpetra contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER

#ifndef TPETRA_DETAILS_GETNUMDIAGS_HPP
#define TPETRA_DETAILS_GETNUMDIAGS_HPP


#include "Tpetra_CrsGraph.hpp"
#include "Teuchos_CommHelpers.hpp"
#include "Tpetra_Details_OrdinalTraits.hpp"

namespace Tpetra {
namespace Details {

namespace Impl {
template <class LocalGraphType, class LocalMapType>
class CountLocalNumDiags {
 public:
  CountLocalNumDiags(const LocalGraphType& G,
                     const LocalMapType& rowMap,
                     const LocalMapType& colMap)
    : G_(G)
    , rowMap_(rowMap)
    , colMap_(colMap) {}

  // Result can't be more than the number of local rows, so
  // local_ordinal_type is appropriate.
  using result_type = typename LocalMapType::local_ordinal_type;

  KOKKOS_INLINE_FUNCTION void
  operator()(const typename LocalMapType::local_ordinal_type lclRow,
             result_type& diagCount) const {
    using LO  = typename LocalMapType::local_ordinal_type;
    using LOT = typename ::Tpetra::Details::OrdinalTraits<LO>;

    auto G_row      = G_.rowConst(lclRow);
    const LO numEnt = G_row.length;
    if (numEnt != 0) {
      // Use global row and column indices to find the diagonal
      // entry.  Caller promises that local row index is in the row
      // Map on the calling process.
      const LO lclDiagCol = colMap_.getLocalElement(rowMap_.getGlobalElement(lclRow));
      // If it's not in the column Map, then there's no diagonal entry.
      if (lclDiagCol != LOT::invalid()) {
        // TODO (mfh 25 Apr 2018) Use findRelOffset to optimize for
        // the sorted case, but note that it requires operator[].
        for (LO k = 0; k < numEnt; ++k) {
          if (lclDiagCol == G_row(k)) {
            ++diagCount;
            break;  // don't count duplicates
          }
        }
      }
    }
  }

 private:
  LocalGraphType G_;
  LocalMapType rowMap_;
  LocalMapType colMap_;
};

template <class LO, class GO, class NT>
typename ::Tpetra::CrsGraph<LO, GO, NT>::local_ordinal_type
countLocalNumDiagsInFillCompleteGraph(const ::Tpetra::CrsGraph<LO, GO, NT>& G) {
  using crs_graph_type          = ::Tpetra::CrsGraph<LO, GO, NT>;
  using local_map_type          = typename crs_graph_type::map_type::local_map_type;
  using local_graph_device_type = typename crs_graph_type::local_graph_device_type;
  using functor_type            = CountLocalNumDiags<local_graph_device_type, local_map_type>;
  using execution_space         = typename crs_graph_type::device_type::execution_space;
  using policy_type             = Kokkos::RangePolicy<execution_space, LO>;

  const auto rowMap = G.getRowMap();
  const auto colMap = G.getColMap();
  if (rowMap.get() == nullptr || colMap.get() == nullptr) {
    return 0;  // this process does not participate
  } else {
    LO lclNumDiags{0};
    functor_type f(G.getLocalGraphDevice(), rowMap->getLocalMap(), colMap->getLocalMap());
    Kokkos::parallel_reduce(policy_type(0, G.getLocalNumRows()), f, lclNumDiags);
    return lclNumDiags;
  }
}

template <class MapType>
typename MapType::local_ordinal_type
getLocalDiagonalColumnIndex(const typename MapType::local_ordinal_type lclRow,
                            const MapType& rowMap,
                            const MapType& colMap) {
  return colMap.getLocalElement(rowMap.getGlobalElement(lclRow));
}

template <class LO, class GO, class NT>
typename ::Tpetra::RowGraph<LO, GO, NT>::local_ordinal_type
countLocalNumDiagsInNonFillCompleteLocallyIndexedGraphWithRowViews(const ::Tpetra::RowGraph<LO, GO, NT>& G) {
  using LOT = typename ::Tpetra::Details::OrdinalTraits<LO>;

  const auto rowMap = G.getRowMap();
  const auto colMap = G.getColMap();
  if (rowMap.get() == nullptr || colMap.get() == nullptr) {
    return 0;  // this process does not participate
  } else {
    TEUCHOS_TEST_FOR_EXCEPTION(!G.supportsRowViews(), std::logic_error, "Not implemented!");

    typename ::Tpetra::RowGraph<LO, GO, NT>::local_inds_host_view_type
        lclColInds;
    const LO lclNumRows = static_cast<LO>(G.getLocalNumRows());

    LO diagCount = 0;
    for (LO lclRow = 0; lclRow < lclNumRows; ++lclRow) {
      G.getLocalRowView(lclRow, lclColInds);
      const LO numEnt = static_cast<LO>(lclColInds.size());
      if (numEnt != 0) {
        const LO lclDiagCol = colMap->getLocalElement(rowMap->getGlobalElement(lclRow));
        // If it's not in the column Map, then there's no diagonal entry.
        if (lclDiagCol != LOT::invalid()) {
          // TODO (mfh 25 Apr 2018) Use findRelOffset to optimize
          // for the sorted case.
          for (LO k = 0; k < numEnt; ++k) {
            if (lclDiagCol == lclColInds[k]) {
              ++diagCount;
              break;  // don't count duplicate entries
            }
          }  // for each columm index in lclRow
        }    // if lclDiagCol is valid
      }      // numEnt != 0
    }        // for each lclRow

    return diagCount;
  }  // if-else
}

template <class LO, class GO, class NT>
typename ::Tpetra::RowGraph<LO, GO, NT>::local_ordinal_type
countLocalNumDiagsInNonFillCompleteLocallyIndexedGraphWithoutRowViews(const ::Tpetra::RowGraph<LO, GO, NT>& G) {
  using LOT = typename ::Tpetra::Details::OrdinalTraits<LO>;

  const auto rowMap = G.getRowMap();
  const auto colMap = G.getColMap();
  if (rowMap.get() == nullptr || colMap.get() == nullptr) {
    return 0;  // this process does not participate
  } else {
    using inds_type = typename ::Tpetra::RowGraph<LO, GO, NT>::nonconst_local_inds_host_view_type;
    inds_type lclColIndsBuf("lclColIndsBuf", G.getLocalMaxNumRowEntries());
    const LO lclNumRows = static_cast<LO>(G.getLocalNumRows());

    LO diagCount = 0;
    for (LO lclRow = 0; lclRow < lclNumRows; ++lclRow) {
      size_t numEntSizeT = G.getNumEntriesInLocalRow(lclRow);
      const LO numEnt    = static_cast<LO>(numEntSizeT);

      inds_type lclColInds = Kokkos::subview(lclColIndsBuf, std::make_pair(0, numEnt));
      G.getLocalRowCopy(lclRow, lclColInds, numEntSizeT);

      if (numEnt != 0) {
        const LO lclDiagCol =
            colMap->getLocalElement(rowMap->getGlobalElement(lclRow));
        // If it's not in the column Map, then there's no diagonal entry.
        if (lclDiagCol != LOT::invalid()) {
          // TODO (mfh 25 Apr 2018) Use findRelOffset to optimize
          // for the sorted case.
          for (LO k = 0; k < numEnt; ++k) {
            if (lclDiagCol == lclColInds[k]) {
              ++diagCount;
              break;  // don't count duplicate entries
            }
          }  // for each columm index in lclRow
        }    // if lclDiagCol is valid
      }      // numEnt != 0
    }        // for each lclRow

    return diagCount;
  }  // if-else
}

template <class LO, class GO, class NT>
typename ::Tpetra::RowGraph<LO, GO, NT>::local_ordinal_type
countLocalNumDiagsInNonFillCompleteGloballyIndexedGraphWithRowViews(const ::Tpetra::RowGraph<LO, GO, NT>& G) {
  const auto rowMap = G.getRowMap();
  if (rowMap.get() == nullptr) {
    return 0;  // this process does not participate
  } else {
    typename ::Tpetra::RowGraph<LO, GO, NT>::global_inds_host_view_type
        gblColInds;
    const LO lclNumRows = static_cast<LO>(G.getLocalNumRows());

    LO diagCount = 0;
    for (LO lclRow = 0; lclRow < lclNumRows; ++lclRow) {
      const GO gblRow = rowMap->getGlobalElement(lclRow);
      G.getGlobalRowView(gblRow, gblColInds);
      const LO numEnt = static_cast<LO>(gblColInds.size());
      if (numEnt != 0) {
        // TODO (mfh 25 Apr 2018) Use findRelOffset to optimize for
        // the sorted case.
        for (LO k = 0; k < numEnt; ++k) {
          if (gblRow == gblColInds[k]) {
            ++diagCount;
            break;  // don't count duplicate entries
          }
        }  // for each column index in lclRow
      }    // if numEnt != 0
    }      // for each lclRow

    return diagCount;
  }  // if-else
}

template <class LO, class GO, class NT>
typename ::Tpetra::RowGraph<LO, GO, NT>::local_ordinal_type
countLocalNumDiagsInNonFillCompleteGloballyIndexedGraphWithoutRowViews(const ::Tpetra::RowGraph<LO, GO, NT>& G) {
  using gids_type   = typename ::Tpetra::RowGraph<LO, GO, NT>::nonconst_global_inds_host_view_type;
  const auto rowMap = G.getRowMap();
  if (rowMap.get() == nullptr) {
    return 0;  // this process does not participate
  } else {
    gids_type gblColIndsBuf;
    const LO lclNumRows = static_cast<LO>(G.getLocalNumRows());

    LO diagCount = 0;
    for (LO lclRow = 0; lclRow < lclNumRows; ++lclRow) {
      size_t numEntSizeT = G.getNumEntriesInLocalRow(lclRow);
      const LO numEnt    = static_cast<LO>(numEntSizeT);
      if (static_cast<LO>(gblColIndsBuf.size()) < numEnt) {
        Kokkos::resize(gblColIndsBuf, numEnt);
      }

      gids_type gblColInds = Kokkos::subview(gblColIndsBuf, std::make_pair((LO)0, numEnt));
      const GO gblRow      = rowMap->getGlobalElement(lclRow);
      G.getGlobalRowCopy(gblRow, gblColInds, numEntSizeT);

      if (numEnt != 0) {
        // TODO (mfh 25 Apr 2018) Use findRelOffset to optimize for
        // the sorted case.
        for (LO k = 0; k < numEnt; ++k) {
          if (gblRow == gblColInds[k]) {
            ++diagCount;
            break;  // don't count duplicate entries
          }
        }  // for each column index in lclRow
      }    // if numEnt != 0
    }      // for each lclRow

    return diagCount;
  }  // if-else
}

template <class MatrixType>
struct GetLocalNumDiags {
  static typename MatrixType::local_ordinal_type
  getLocalNumDiags(const MatrixType& A) {
    using LO             = typename MatrixType::local_ordinal_type;
    using GO             = typename MatrixType::global_ordinal_type;
    using NT             = typename MatrixType::node_type;
    using row_graph_type = ::Tpetra::RowGraph<LO, GO, NT>;

    auto G = A.getGraph();
    if (G.get() == nullptr) {
      return 0;
    } else {
      return GetLocalNumDiags<row_graph_type>::getLocalNumDiags(*G);
    }
  }
};

template <class LO, class GO, class NT>
struct GetLocalNumDiags< ::Tpetra::RowGraph<LO, GO, NT> > {
  static LO
  getLocalNumDiags(const ::Tpetra::RowGraph<LO, GO, NT>& G) {
    using crs_graph_type = ::Tpetra::CrsGraph<LO, GO, NT>;

    const crs_graph_type* G_crs = dynamic_cast<const crs_graph_type*>(&G);
    if (G_crs != nullptr && G_crs->isFillComplete()) {
      return countLocalNumDiagsInFillCompleteGraph(*G_crs);
    } else {
      if (G.isLocallyIndexed()) {
        if (G.supportsRowViews()) {
          return countLocalNumDiagsInNonFillCompleteLocallyIndexedGraphWithRowViews(G);
        } else {
          return countLocalNumDiagsInNonFillCompleteLocallyIndexedGraphWithoutRowViews(G);
        }
      } else if (G.isGloballyIndexed()) {
        if (G.supportsRowViews()) {
          return countLocalNumDiagsInNonFillCompleteGloballyIndexedGraphWithRowViews(G);
        } else {
          return countLocalNumDiagsInNonFillCompleteGloballyIndexedGraphWithoutRowViews(G);
        }
      } else {  // G is empty
        return 0;
      }
    }
  }
};

template <class LO, class GO, class NT>
struct GetLocalNumDiags< ::Tpetra::CrsGraph<LO, GO, NT> > {
  static LO
  getLocalNumDiags(const ::Tpetra::CrsGraph<LO, GO, NT>& G) {
    using row_graph_type = ::Tpetra::RowGraph<LO, GO, NT>;
    return GetLocalNumDiags<row_graph_type>::getLocalNumDiags(G);
  }
};
}  // namespace Impl

template <class CrsGraphType>
typename CrsGraphType::local_ordinal_type
getLocalNumDiags(const CrsGraphType& G) {
  return Impl::GetLocalNumDiags<CrsGraphType>::getLocalNumDiags(G);
}

template <class CrsGraphType>
typename CrsGraphType::global_ordinal_type
getGlobalNumDiags(const CrsGraphType& G) {
  using GO = typename CrsGraphType::global_ordinal_type;

  const auto map = G.getRowMap();
  if (map.get() == nullptr) {
    return GO(0);  // this process should not participate
  } else {
    const auto comm = map->getComm();
    if (comm.get() == nullptr) {
      return GO(0);  // this process should not participate
    } else {
      const GO lclNumDiags = static_cast<GO>(getLocalNumDiags(G));

      using Teuchos::outArg;
      using Teuchos::REDUCE_SUM;
      using Teuchos::reduceAll;

      GO gblNumDiags{0};
      reduceAll<int, GO>(*comm, REDUCE_SUM, lclNumDiags, outArg(gblNumDiags));
      return gblNumDiags;
    }
  }
}

}  // namespace Details
}  // namespace Tpetra

#endif  // TPETRA_DETAILS_GETNUMDIAGS_HPP