Tpetra_RowMatrixTransposer_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_ROWMATRIXTRANSPOSER_DEF_HPP
#define TPETRA_ROWMATRIXTRANSPOSER_DEF_HPP

#include "Tpetra_CrsMatrix.hpp"
#include "Tpetra_BlockCrsMatrix.hpp"
#include "Tpetra_Export.hpp"
#include "Tpetra_Details_computeOffsets.hpp"
#include "Tpetra_Details_shortSort.hpp"
#include "Teuchos_ParameterList.hpp"
#include "Teuchos_TimeMonitor.hpp"
#include "KokkosSparse_Utils.hpp"
#include "KokkosSparse_SortCrs.hpp"

namespace Tpetra {

template<class Scalar,
         class LocalOrdinal,
         class GlobalOrdinal,
         class Node>
RowMatrixTransposer<Scalar, LocalOrdinal, GlobalOrdinal, Node>::
RowMatrixTransposer (const Teuchos::RCP<const crs_matrix_type>& origMatrix,
                     const std::string& label)
  : origMatrix_ (origMatrix), label_ (label)
{}

template<class Scalar,
         class LocalOrdinal,
         class GlobalOrdinal,
         class Node>
Teuchos::RCP<CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node> >
RowMatrixTransposer<Scalar, LocalOrdinal, GlobalOrdinal, Node>::
createTranspose (const Teuchos::RCP<Teuchos::ParameterList> &params)
{
  using Teuchos::RCP;
  // Do the local transpose
  RCP<crs_matrix_type> transMatrixWithSharedRows = createTransposeLocal (params);

#ifdef HAVE_TPETRA_MMM_TIMINGS
  const std::string prefix = std::string ("Tpetra ") + label_ + ": ";
  using Teuchos::TimeMonitor;
  TimeMonitor MM (*TimeMonitor::getNewTimer (prefix + "Transpose TAFC"));
#endif

  // If transMatrixWithSharedRows 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 =
    transMatrixWithSharedRows->getGraph ()->getExporter ();
  if (exporter.is_null ()) {
    return transMatrixWithSharedRows;
  }
  else {
    Teuchos::ParameterList labelList;
#ifdef HAVE_TPETRA_MMM_TIMINGS
    labelList.set("Timer Label", label_);
#endif
    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 matrix after communication, so
    //   no need to set "sorted = false" in parameters.
    return exportAndFillCompleteCrsMatrix<crs_matrix_type>
      (transMatrixWithSharedRows, *exporter, Teuchos::null,
       Teuchos::null, Teuchos::rcpFromRef (labelList));
  }
}


// mfh 03 Feb 2013: In a definition outside the class like this, the
// return value is considered outside the class scope (for things like
// resolving typedefs), but the arguments are considered inside the
// class scope.
template<class Scalar,
         class LocalOrdinal,
         class GlobalOrdinal,
         class Node>
Teuchos::RCP<CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node> >
RowMatrixTransposer<Scalar, LocalOrdinal, GlobalOrdinal, Node>::
createTransposeLocal (const Teuchos::RCP<Teuchos::ParameterList>& params)
{
  using Teuchos::Array;
  using Teuchos::ArrayRCP;
  using Teuchos::ArrayView;
  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>;

#ifdef HAVE_TPETRA_MMM_TIMINGS
  std::string prefix = std::string("Tpetra ") + label_ + ": ";
  using Teuchos::TimeMonitor;
  TimeMonitor MM (*TimeMonitor::getNewTimer (prefix + "Transpose Local"));
#endif

  const bool sort = [&] () {
    constexpr bool sortDefault = true; // see #4607 discussion
    const char sortParamName[] = "sort";
    return params.get () == nullptr ? sortDefault :
      params->get (sortParamName, sortDefault);
  } ();

  const LO lclNumRows (origMatrix_->getLocalNumRows ());

  RCP<const crs_matrix_type> crsMatrix =
    rcp_dynamic_cast<const crs_matrix_type> (origMatrix_);
  if (crsMatrix.is_null ()) {
    auto rowMap = origMatrix_->getRowMap ();
    if (rowMap->isOneToOne ()) {
      Teuchos::Array<size_t> numEntPerRow (lclNumRows);
      for (LO lclRow = 0; lclRow < lclNumRows; ++lclRow) {
        numEntPerRow[lclRow] = origMatrix_->getNumEntriesInLocalRow (lclRow);
      }
      auto colMap = origMatrix_->getColMap ();

      RCP<crs_matrix_type> crsMatrix_nc =
        rcp (new crs_matrix_type (rowMap, colMap, numEntPerRow ()));

      // When source & target Maps are same, Import just copies.
      import_type imp (rowMap, rowMap);
      crsMatrix_nc->doImport (*origMatrix_, imp, Tpetra::REPLACE);
      crsMatrix_nc->fillComplete (origMatrix_->getDomainMap (),
                                  origMatrix_->getRangeMap ());
      crsMatrix = crsMatrix_nc;
    }
    else {
      TEUCHOS_ASSERT( false ); // not implemented (it wasn't before)
    }
  }

  using local_matrix_device_type = typename crs_matrix_type::local_matrix_device_type;

  local_matrix_device_type lclMatrix = crsMatrix->getLocalMatrixDevice ();
  local_matrix_device_type lclTransposeMatrix = KokkosSparse::Impl::transpose_matrix(lclMatrix);
  if (sort)
    KokkosSparse::sort_crs_matrix(lclTransposeMatrix);

  // Prebuild the importers and exporters the no-communication way,
  // flipping the importers and exporters around.
  const auto origExport = origMatrix_->getGraph ()->getExporter ();
  RCP<const import_type> myImport = origExport.is_null () ?
    Teuchos::null : rcp (new import_type (*origExport));
  const auto origImport = origMatrix_->getGraph ()->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_matrix_type (lclTransposeMatrix,
                                   origMatrix_->getColMap (),
                                   origMatrix_->getRowMap (),
                                   origMatrix_->getRangeMap (),
                                   origMatrix_->getDomainMap (),
                                   myImport, myExport, graphParams));
}

/*************************************************************************/

template<class Scalar,
         class LocalOrdinal,
         class GlobalOrdinal,
         class Node>
BlockCrsMatrixTransposer<Scalar, LocalOrdinal, GlobalOrdinal, Node>::
BlockCrsMatrixTransposer (const Teuchos::RCP<const bcrs_matrix_type>& origMatrix,
                     const std::string& label)
  : origMatrix_ (origMatrix), label_ (label)
{}

template<class Scalar,
         class LocalOrdinal,
         class GlobalOrdinal,
         class Node>
Teuchos::RCP<BlockCrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node> >
BlockCrsMatrixTransposer<Scalar, LocalOrdinal, GlobalOrdinal, Node>::
createTranspose (const Teuchos::RCP<Teuchos::ParameterList> &params)
{
  using Teuchos::RCP;
  // Do the local transpose
  RCP<bcrs_matrix_type> transMatrixWithSharedRows = createTransposeLocal (params);

#ifdef HAVE_TPETRA_MMM_TIMINGS
  const std::string prefix = std::string ("Tpetra ") + label_ + ": ";
  using Teuchos::TimeMonitor;
  TimeMonitor MM (*TimeMonitor::getNewTimer (prefix + "Transpose TAFC"));
#endif

  // If transMatrixWithSharedRows 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 =
    transMatrixWithSharedRows->getGraph ()->getExporter ();
  if (exporter.is_null ()) {
    return transMatrixWithSharedRows;
  }
  else {
    Teuchos::ParameterList labelList;
#ifdef HAVE_TPETRA_MMM_TIMINGS
    labelList.set("Timer Label", label_);
#endif
    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 matrix after communication, so
    //   no need to set "sorted = false" in parameters.
    return exportAndFillCompleteBlockCrsMatrix<bcrs_matrix_type>
      (transMatrixWithSharedRows, *exporter);
  }
}


// mfh 03 Feb 2013: In a definition outside the class like this, the
// return value is considered outside the class scope (for things like
// resolving typedefs), but the arguments are considered inside the
// class scope.
template<class Scalar,
         class LocalOrdinal,
         class GlobalOrdinal,
         class Node>
Teuchos::RCP<BlockCrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node> >
BlockCrsMatrixTransposer<Scalar, LocalOrdinal, GlobalOrdinal, Node>::
createTransposeLocal (const Teuchos::RCP<Teuchos::ParameterList>& params)
{
  using Teuchos::Array;
  using Teuchos::ArrayRCP;
  using Teuchos::ArrayView;
  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>;
  using crs_graph_type = typename bcrs_matrix_type::crs_graph_type;

#ifdef HAVE_TPETRA_MMM_TIMINGS
  std::string prefix = std::string("Tpetra ") + label_ + ": ";
  using Teuchos::TimeMonitor;
  TimeMonitor MM (*TimeMonitor::getNewTimer (prefix + "Transpose Local"));
#endif

  RCP<const bcrs_matrix_type> crsMatrix =
    rcp_dynamic_cast<const bcrs_matrix_type> (origMatrix_);

  if(crsMatrix.is_null()) 
    TEUCHOS_ASSERT( false ); // not implemented 

  using local_matrix_device_type = typename bcrs_matrix_type::local_matrix_device_type;

  typename local_matrix_device_type::values_type values ;
  RCP<const crs_graph_type> graph;
  {
    local_matrix_device_type lclMatrix = crsMatrix->getLocalMatrixDevice ();

    local_matrix_device_type lclTransposeMatrix = KokkosSparse::Impl::transpose_bsr_matrix(lclMatrix);    

    // BlockCrs requires that we sort stuff
    KokkosSparse::sort_crs_matrix(lclTransposeMatrix);
    values = lclTransposeMatrix.values;
    
    // Prebuild the importers and exporters the no-communication way,
    // flipping the importers and exporters around.
    const auto origExport = origMatrix_->getGraph ()->getExporter ();
    RCP<const import_type> myImport = origExport.is_null () ?
      Teuchos::null : rcp (new import_type (*origExport));
    const auto origImport = origMatrix_->getGraph ()->getImporter ();
    RCP<const export_type> myExport = origImport.is_null () ?
      Teuchos::null : rcp (new export_type (*origImport));

    RCP<Teuchos::ParameterList> graphParams = Teuchos::null;
    
    // Make the Transpose Graph
    graph = rcp(new crs_graph_type(lclTransposeMatrix.graph,                                   
                                   origMatrix_->getColMap (),
                                   origMatrix_->getRowMap (),
                                   origMatrix_->getGraph()->getRangeMap (),
                                   origMatrix_->getGraph()->getDomainMap (),
                                   myImport,
                                   myExport,
                                   graphParams));
  }
  // Now make the matrix
  return rcp (new bcrs_matrix_type (*graph,
                                    values,
                                    origMatrix_->getBlockSize()));
}
//


//
// Explicit instantiation macro
//
// Must be expanded from within the Tpetra namespace!
//

#define TPETRA_ROWMATRIXTRANSPOSER_INSTANT(SCALAR,LO,GO,NODE) \
  template class RowMatrixTransposer< SCALAR, LO , GO , NODE >;\
  template class BlockCrsMatrixTransposer< SCALAR, LO , GO , NODE >;

} // namespace Tpetra

#endif