Tpetra_CrsGraphTransposer_def.hpp

// @HEADER
// *****************************************************************************
//          Tpetra: Templated Linear Algebra Services Package
//
// Copyright 2008 NTESS and the Tpetra contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER
 
#ifndef TPETRA_CRSGRAPHTRANSPOSER_DEF_HPP
#define TPETRA_CRSGRAPHTRANSPOSER_DEF_HPP
 
#include "Tpetra_CrsGraph.hpp"
#include "Tpetra_Export.hpp"
#include "Tpetra_Details_computeOffsets.hpp"
#include "Tpetra_Details_makeColMap.hpp"
#include "Tpetra_Details_shortSort.hpp"
#include "Teuchos_ParameterList.hpp"
#include "Teuchos_TimeMonitor.hpp"
#include "KokkosSparse_Utils.hpp"
#include "KokkosKernels_Handle.hpp"
#include "KokkosSparse_spadd.hpp"
 
namespace Tpetra {
 
template <typename GO,
          typename LocalIndicesType,
          typename GlobalIndicesType,
          typename ColMapType>
struct ConvertLocalToGlobalFunctor {
  ConvertLocalToGlobalFunctor(
      const LocalIndicesType& colindsOrig_,
      const GlobalIndicesType& colindsConverted_,
      const ColMapType& colmap_)
    : colindsOrig(colindsOrig_)
    , colindsConverted(colindsConverted_)
    , colmap(colmap_) {}
  KOKKOS_INLINE_FUNCTION void
  operator()(const GO i) const {
    colindsConverted(i) = colmap.getGlobalElement(colindsOrig(i));
  }
  LocalIndicesType colindsOrig;
  GlobalIndicesType colindsConverted;
  ColMapType colmap;
};
 
template <class LO, class GO, class LOView, class GOView, class LocalMap>
struct ConvertGlobalToLocalFunctor {
  ConvertGlobalToLocalFunctor(LOView& lids_, const GOView& gids_, const LocalMap localColMap_)
    : lids(lids_)
    , gids(gids_)
    , localColMap(localColMap_) {}
 
  KOKKOS_FUNCTION void operator()(const GO i) const {
    lids(i) = localColMap.getLocalElement(gids(i));
  }
 
  LOView lids;
  const GOView gids;
  const LocalMap localColMap;
};
 
template <typename size_type, typename ordinal_type,
          typename ArowptrsT, typename BrowptrsT, typename CrowptrsT,
          typename AcolindsT, typename BcolindsT, typename CcolindsT>
struct SortedNumericIndicesOnlyFunctor {
  SortedNumericIndicesOnlyFunctor(const ArowptrsT& Arowptrs_,
                                  const BrowptrsT& Browptrs_,
                                  const CrowptrsT& Crowptrs_,
                                  const AcolindsT& Acolinds_,
                                  const BcolindsT& Bcolinds_,
                                  const CcolindsT& Ccolinds_)
    : Arowptrs(Arowptrs_)
    , Browptrs(Browptrs_)
    , Crowptrs(Crowptrs_)
    , Acolinds(Acolinds_)
    , Bcolinds(Bcolinds_)
    , Ccolinds(Ccolinds_) {}
 
  KOKKOS_INLINE_FUNCTION void operator()(const ordinal_type i) const {
    const ordinal_type ORDINAL_MAX = KokkosKernels::ArithTraits<ordinal_type>::max();
 
    // count the union of nonzeros in Arow and Brow
    size_type ai        = 0;
    size_type bi        = 0;
    size_type Arowstart = Arowptrs(i);
    size_type Arowlen   = Arowptrs(i + 1) - Arowstart;
    size_type Browstart = Browptrs(i);
    size_type Browlen   = Browptrs(i + 1) - Browstart;
    ordinal_type Acol   = (Arowlen == 0) ? ORDINAL_MAX : Acolinds(Arowstart);
    ordinal_type Bcol   = (Browlen == 0) ? ORDINAL_MAX : Bcolinds(Browstart);
    size_type Coffset   = Crowptrs(i);
    while (Acol != ORDINAL_MAX || Bcol != ORDINAL_MAX) {
      ordinal_type Ccol = (Acol < Bcol) ? Acol : Bcol;
      while (Acol == Ccol) {
        ai++;
        if (ai == Arowlen)
          Acol = ORDINAL_MAX;
        else
          Acol = Acolinds(Arowstart + ai);
      }
      while (Bcol == Ccol) {
        bi++;
        if (bi == Browlen)
          Bcol = ORDINAL_MAX;
        else
          Bcol = Bcolinds(Browstart + bi);
      }
      Ccolinds(Coffset) = Ccol;
      Coffset++;
    }
  }
 
  const ArowptrsT Arowptrs;
  const BrowptrsT Browptrs;
  const CrowptrsT Crowptrs;
  const AcolindsT Acolinds;
  const BcolindsT Bcolinds;
  CcolindsT Ccolinds;
};
 
template <typename size_type, typename ordinal_type,
          typename ArowptrsT, typename BrowptrsT, typename CrowptrsT,
          typename AcolindsT, typename BcolindsT, typename CcolindsT>
struct UnsortedNumericIndicesOnlyFunctor {
  UnsortedNumericIndicesOnlyFunctor(
      const ArowptrsT Arowptrs_, const BrowptrsT Browptrs_, const CrowptrsT Crowptrs_,
      const AcolindsT Acolinds_, const BcolindsT Bcolinds_, CcolindsT Ccolinds_,
      const CcolindsT Apos_, const CcolindsT Bpos_)
    : Arowptrs(Arowptrs_)
    , Browptrs(Browptrs_)
    , Crowptrs(Crowptrs_)
    , Acolinds(Acolinds_)
    , Bcolinds(Bcolinds_)
    , Ccolinds(Ccolinds_)
    , Apos(Apos_)
    , Bpos(Bpos_) {}
 
  KOKKOS_INLINE_FUNCTION void operator()(const ordinal_type i) const {
    size_type CrowStart = Crowptrs(i);
    size_type ArowStart = Arowptrs(i);
    size_type ArowEnd   = Arowptrs(i + 1);
    size_type BrowStart = Browptrs(i);
    size_type BrowEnd   = Browptrs(i + 1);
    // add in A entries, while setting C colinds
    for (size_type j = ArowStart; j < ArowEnd; j++) {
      Ccolinds(CrowStart + Apos(j)) = Acolinds(j);
    }
    // add in B entries, while setting C colinds
    for (size_type j = BrowStart; j < BrowEnd; j++) {
      Ccolinds(CrowStart + Bpos(j)) = Bcolinds(j);
    }
  }
  const ArowptrsT Arowptrs;
  const BrowptrsT Browptrs;
  const CrowptrsT Crowptrs;
  const AcolindsT Acolinds;
  const BcolindsT Bcolinds;
  CcolindsT Ccolinds;
  const CcolindsT Apos;
  const CcolindsT Bpos;
};
 
template <class LocalOrdinal,
          class GlobalOrdinal,
          class Node>
CrsGraphTransposer<LocalOrdinal, GlobalOrdinal, Node>::
    CrsGraphTransposer(const Teuchos::RCP<const crs_graph_type>& origGraph)
  : origGraph_(origGraph) {}
 
template <class LocalOrdinal,
          class GlobalOrdinal,
          class Node>
Teuchos::RCP<CrsGraph<LocalOrdinal, GlobalOrdinal, Node> >
CrsGraphTransposer<LocalOrdinal, GlobalOrdinal, Node>::
    symmetrize(const Teuchos::RCP<Teuchos::ParameterList>& params) {
  using Teuchos::RCP;
  using device_type             = typename Node::device_type;
  using execution_space         = typename device_type::execution_space;
  using range_type              = Kokkos::RangePolicy<execution_space, size_t>;
  using local_graph_device_type = typename crs_graph_type::local_graph_device_type;
  using impl_scalar_type        = ::Tpetra::Details::DefaultTypes::scalar_type;
  using row_ptrs_array          = typename local_graph_device_type::row_map_type::non_const_type;
  using col_inds_array          = typename local_graph_device_type::entries_type::non_const_type;
  using local_map_type          = typename map_type::local_map_type;
  using global_col_inds_array   = typename Kokkos::View<GlobalOrdinal*, device_type>;
 
  auto graph      = origGraph_;
  auto domain_map = graph->getDomainMap();
  auto range_map  = graph->getRangeMap();
  auto row_map    = graph->getRowMap();
  auto col_map    = graph->getColMap();
  RCP<const map_type> col_map_sym;
  RCP<const Import<LocalOrdinal, GlobalOrdinal, Node> > importer;
 
  TEUCHOS_ASSERT(domain_map->isSameAs(*range_map));
  TEUCHOS_ASSERT(domain_map->isSameAs(*row_map));
 
  // Do the transpose
  RCP<crs_graph_type> graphT = createTranspose(params);
 
  auto col_map_T = graphT->getColMap();
  TEUCHOS_ASSERT(!col_map_T.is_null());
  TEUCHOS_ASSERT(domain_map->isSameAs(*graphT->getDomainMap()));
 
  bool graphSorted     = graph->isSorted();
  bool graphTSorted    = graphT->isSorted();
  bool sorted          = graphSorted && graphTSorted;
  bool matchingColMaps = col_map->isSameAs(*col_map_T);
 
  auto lclGraph  = graph->getLocalGraphDevice();
  auto lclGraphT = graphT->getLocalGraphDevice();
 
  using KKH_LO = KokkosKernels::Experimental::KokkosKernelsHandle<size_t, LocalOrdinal, impl_scalar_type,
                                                                  typename Node::execution_space, typename Node::memory_space, typename Node::memory_space>;
  using KKH_GO = KokkosKernels::Experimental::KokkosKernelsHandle<size_t, GlobalOrdinal, impl_scalar_type,
                                                                  typename Node::execution_space, typename Node::memory_space, typename Node::memory_space>;
 
  auto rowptrs  = lclGraph.row_map;
  auto rowptrsT = lclGraphT.row_map;
  auto colinds  = lclGraph.entries;
  auto colindsT = lclGraphT.entries;
 
  auto nrows      = rowptrs.extent(0) - 1;
  auto rowptrsSym = row_ptrs_array(Kokkos::ViewAllocateWithoutInitializing("row ptrs sym"), nrows + 1);
 
  col_inds_array colindsSym;
 
  if (!matchingColMaps) {
    // convert indices of local graph to GlobalOrdinal
    auto lclColmap = col_map->getLocalMap();
    global_col_inds_array colindsConverted(Kokkos::ViewAllocateWithoutInitializing("colinds (converted)"), colinds.extent(0));
    ConvertLocalToGlobalFunctor<GlobalOrdinal, col_inds_array, global_col_inds_array, local_map_type> convert(colinds, colindsConverted, lclColmap);
    Kokkos::parallel_for("colInds (converted)", range_type(0, colinds.extent(0)), convert);
 
    // convert indices of local graphT to GlobalOrdinal
    auto lclColmapT = col_map_T->getLocalMap();
    global_col_inds_array colindsTConverted(Kokkos::ViewAllocateWithoutInitializing("colindsT (converted)"), colindsT.extent(0));
    ConvertLocalToGlobalFunctor<GlobalOrdinal, col_inds_array, global_col_inds_array, local_map_type> convertT(colindsT, colindsTConverted, lclColmapT);
    Kokkos::parallel_for("colIndsT (converted)", range_type(0, colindsT.extent(0)), convertT);
 
    // sum graph and graphT in GlobalOrdinal
    KKH_GO handle;
    handle.create_spadd_handle(false);
    auto addHandle = handle.get_spadd_handle();
 
    global_col_inds_array globalColindsSym;
 
    KokkosSparse::spadd_symbolic(&handle,
                                 nrows, graph->getGlobalNumCols(),
                                 rowptrs, colindsConverted, rowptrsT, colindsTConverted, rowptrsSym);
    globalColindsSym = global_col_inds_array(Kokkos::ViewAllocateWithoutInitializing("global colinds sym"), addHandle->get_c_nnz());
 
    UnsortedNumericIndicesOnlyFunctor<
        size_t, GlobalOrdinal,
        typename row_ptrs_array::const_type, typename row_ptrs_array::const_type, row_ptrs_array,
        typename global_col_inds_array::const_type, typename global_col_inds_array::const_type, global_col_inds_array>
        unsortedNumeric(rowptrs, rowptrsT, rowptrsSym,
                        colindsConverted, colindsTConverted, globalColindsSym,
                        addHandle->get_a_pos(), addHandle->get_b_pos());
    Kokkos::parallel_for("KokkosSparse::SpAdd:Numeric::InputNotSorted",
                         range_type(0, nrows), unsortedNumeric);
 
    // build column map for graphSym
    Tpetra::Details::makeColMap<LocalOrdinal, GlobalOrdinal, Node>(col_map_sym, domain_map, globalColindsSym);
 
    // convert indices of local graphSym to LocalOrdinal
    auto lclColmapSym = col_map_sym->getLocalMap();
    colindsSym        = col_inds_array("colindsSym", globalColindsSym.extent(0));
    ConvertGlobalToLocalFunctor<LocalOrdinal, GlobalOrdinal, col_inds_array, global_col_inds_array, typename map_type::local_map_type> convertSym(colindsSym, globalColindsSym, lclColmapSym);
    Kokkos::parallel_for(range_type(0, globalColindsSym.extent(0)), convertSym);
 
  } else {
    // sum graph and graphT in LocalOrdinal
    KKH_LO handle;
    handle.create_spadd_handle(sorted);
    auto addHandle = handle.get_spadd_handle();
 
    KokkosSparse::spadd_symbolic(&handle,
                                 nrows, graph->getGlobalNumCols(),
                                 rowptrs, colinds, rowptrsT, colindsT, rowptrsSym);
    colindsSym = col_inds_array(Kokkos::ViewAllocateWithoutInitializing("C colinds"), addHandle->get_c_nnz());
 
    if (sorted) {
      SortedNumericIndicesOnlyFunctor<
          size_t, LocalOrdinal,
          typename row_ptrs_array::const_type, typename row_ptrs_array::const_type, row_ptrs_array,
          typename col_inds_array::const_type, typename col_inds_array::const_type, col_inds_array>
          sortedNumeric(rowptrs, rowptrsT, rowptrsSym,
                        colinds, colindsT, colindsSym);
      Kokkos::parallel_for("KokkosSparse::SpAdd:Numeric::InputSorted",
                           range_type(0, nrows), sortedNumeric);
 
    } else {
      UnsortedNumericIndicesOnlyFunctor<
          size_t, LocalOrdinal,
          typename row_ptrs_array::const_type, typename row_ptrs_array::const_type, row_ptrs_array,
          typename col_inds_array::const_type, typename col_inds_array::const_type, col_inds_array>
          unsortedNumeric(rowptrs, rowptrsT, rowptrsSym,
                          colinds, colindsT, colindsSym,
                          addHandle->get_a_pos(), addHandle->get_b_pos());
      Kokkos::parallel_for("KokkosSparse::SpAdd:Numeric::InputNotSorted",
                           range_type(0, nrows), unsortedNumeric);
    }
 
    // column map for graphSym is graph's column map
    col_map_sym = col_map;
    importer    = graph->getImporter();
  }
 
  bool sort = true;
  if (sort)
    KokkosSparse::sort_crs_graph<execution_space, row_ptrs_array, col_inds_array>(rowptrsSym, colindsSym);
 
  local_graph_device_type lclGraphSym = local_graph_device_type(colindsSym, rowptrsSym);
 
  RCP<Teuchos::ParameterList> graphParams = Teuchos::null;
  if (!sort) {
    graphParams = rcp(new Teuchos::ParameterList);
    graphParams->set("sorted", false);
  }
 
  return rcp(new crs_graph_type(lclGraphSym,
                                row_map,
                                col_map_sym,
                                domain_map,
                                range_map,
                                importer,
                                Teuchos::null,
                                graphParams));
}
 
template <class LocalOrdinal,
          class GlobalOrdinal,
          class Node>
Teuchos::RCP<CrsGraph<LocalOrdinal, GlobalOrdinal, Node> >
CrsGraphTransposer<LocalOrdinal, GlobalOrdinal, Node>::
    createTranspose(const Teuchos::RCP<Teuchos::ParameterList>& params) {
  using Teuchos::RCP;
  // Do the local transpose
  RCP<crs_graph_type> transGraphWithSharedRows = createTransposeLocal(params);
 
  Tpetra::Details::ProfilingRegion MM("Tpetra: Transpose TAFC");
 
  // If transGraphWithSharedRows has an exporter, that's what we
  // want.  If it doesn't, the rows aren't actually shared, and we're
  // done!
  using export_type = Export<LocalOrdinal, GlobalOrdinal, Node>;
  RCP<const export_type> exporter =
      transGraphWithSharedRows->getExporter();
  if (exporter.is_null()) {
    return transGraphWithSharedRows;
  } else {
    Teuchos::ParameterList labelList;
    if (!params.is_null()) {
      const char paramName[] = "compute global constants";
      labelList.set(paramName, params->get(paramName, true));
    }
    // Use the Export object to do a fused Export and fillComplete.
    // This always sorts the local graph after communication, so
    //   no need to set "sorted = false" in parameters.
    return exportAndFillCompleteCrsGraph<crs_graph_type>(transGraphWithSharedRows, *exporter, Teuchos::null,
                                                         Teuchos::null, Teuchos::rcpFromRef(labelList));
  }
}
 
template <class LocalOrdinal,
          class GlobalOrdinal,
          class Node>
Teuchos::RCP<CrsGraph<LocalOrdinal, GlobalOrdinal, Node> >
CrsGraphTransposer<LocalOrdinal, GlobalOrdinal, Node>::
    createTransposeLocal(const Teuchos::RCP<Teuchos::ParameterList>& params) {
  using Teuchos::RCP;
  using Teuchos::rcp;
  using Teuchos::rcp_dynamic_cast;
  using LO          = LocalOrdinal;
  using GO          = GlobalOrdinal;
  using import_type = Tpetra::Import<LO, GO, Node>;
  using export_type = Tpetra::Export<LO, GO, Node>;
 
  Tpetra::Details::ProfilingRegion MM("Tpetra: Transpose Local");
 
  const bool sort = [&]() {
    constexpr bool sortDefault = true;  // see #4607 discussion
    const char sortParamName[] = "sort";
    return params.get() == nullptr ? sortDefault : params->get(sortParamName, sortDefault);
  }();
 
  using local_graph_device_type    = typename crs_graph_type::local_graph_device_type;
  local_graph_device_type lclGraph = origGraph_->getLocalGraphDevice();
 
  // Allocate views and call the other version of transpose_graph
  using c_rowmap_t     = typename local_graph_device_type::row_map_type;
  using c_entries_t    = typename local_graph_device_type::entries_type;
  using rowmap_t       = typename local_graph_device_type::row_map_type::non_const_type;
  using entries_t      = typename local_graph_device_type::entries_type::non_const_type;
  LocalOrdinal numCols = origGraph_->getColMap()->getLocalNumElements();
  rowmap_t lclGraphT_rowmap("Transpose rowmap", numCols + 1);
  entries_t lclGraphT_entries(
      Kokkos::ViewAllocateWithoutInitializing("Transpose entries"), lclGraph.entries.extent(0));
  KokkosSparse::Impl::transpose_graph<
      c_rowmap_t, c_entries_t,
      rowmap_t, entries_t,
      rowmap_t, typename local_graph_device_type::execution_space>(
      lclGraph.numRows(), numCols,
      lclGraph.row_map, lclGraph.entries,
      lclGraphT_rowmap, lclGraphT_entries);
 
  if (sort)
    KokkosSparse::sort_crs_graph<
        typename local_graph_device_type::execution_space,
        rowmap_t, entries_t>(
        lclGraphT_rowmap,
        lclGraphT_entries);
 
  // And construct the transpose local_graph_device_type
  local_graph_device_type lclGraphT = local_graph_device_type(lclGraphT_entries, lclGraphT_rowmap);
 
  // Prebuild the importers and exporters the no-communication way,
  // flipping the importers and exporters around.
  const auto origExport           = origGraph_->getExporter();
  RCP<const import_type> myImport = origExport.is_null() ? Teuchos::null : rcp(new import_type(*origExport));
  const auto origImport           = origGraph_->getImporter();
  RCP<const export_type> myExport = origImport.is_null() ? Teuchos::null : rcp(new export_type(*origImport));
 
  RCP<Teuchos::ParameterList> graphParams = Teuchos::null;
  if (!sort) {
    graphParams = rcp(new Teuchos::ParameterList);
    graphParams->set("sorted", false);
  }
 
  return rcp(new crs_graph_type(lclGraphT,
                                origGraph_->getColMap(),
                                origGraph_->getRowMap(),
                                origGraph_->getRangeMap(),
                                origGraph_->getDomainMap(),
                                myImport, myExport, graphParams));
}
 
//
// Explicit instantiation macro
//
// Must be expanded from within the Tpetra namespace!
//
 
#define TPETRA_CRSGRAPHTRANSPOSER_INSTANT(LO, GO, NODE) \
  template class CrsGraphTransposer<LO, GO, NODE>;
 
}  // namespace Tpetra
 
#endif