Teko_TpetraThyraConverter.cpp

// @HEADER
// *****************************************************************************
//      Teko: A package for block and physics based preconditioning
//
// Copyright 2010 NTESS and the Teko contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER

#include "Teko_TpetraThyraConverter.hpp"
#include "Tpetra_Core.hpp"

// Teuchos includes
#include "Teuchos_Array.hpp"
#include "Teuchos_ArrayRCP.hpp"

// Thyra includes
#include "Thyra_DefaultProductVectorSpace.hpp"
#include "Thyra_DefaultProductMultiVector.hpp"
#include "Thyra_SpmdMultiVectorBase.hpp"
#include "Thyra_SpmdVectorSpaceBase.hpp"
#include "Thyra_MultiVectorStdOps.hpp"

#include "Thyra_TpetraMultiVector.hpp"

#include <iostream>
#include <vector>

using Teuchos::null;
using Teuchos::Ptr;
using Teuchos::ptr_dynamic_cast;
using Teuchos::RCP;
using Teuchos::rcp;
using Teuchos::rcp_dynamic_cast;
using Teuchos::rcpFromRef;

namespace Teko {
namespace TpetraHelpers {

// const Teuchos::RCP<const Thyra::MultiVectorBase<double> >
// blockTpetraToThyra(int numVectors,const double * tpetraData,int leadingDim,const
// Teuchos::RCP<const Thyra::VectorSpaceBase<double> > & vs,int & localDim)

void blockTpetraToThyra(int numVectors, Teuchos::ArrayRCP<const ST> tpetraData, int leadingDim,
                        const Teuchos::Ptr<Thyra::MultiVectorBase<ST>>& mv, int& localDim) {
  localDim = 0;

  // check the base case
  const Ptr<Thyra::ProductMultiVectorBase<ST>> prodMV =
      ptr_dynamic_cast<Thyra::ProductMultiVectorBase<ST>>(mv);
  if (prodMV == Teuchos::null) {
    // VS object must be a SpmdMultiVector object
    const Ptr<Thyra::SpmdMultiVectorBase<ST>> spmdX =
        ptr_dynamic_cast<Thyra::SpmdMultiVectorBase<ST>>(mv, true);
    const RCP<const Thyra::SpmdVectorSpaceBase<ST>> spmdVS = spmdX->spmdSpace();

    int localSubDim = spmdVS->localSubDim();

    Thyra::Ordinal thyraLeadingDim = 0;

    Teuchos::ArrayRCP<ST> thyraData_arcp;
    Teuchos::ArrayView<ST> thyraData;
    spmdX->getNonconstLocalData(Teuchos::outArg(thyraData_arcp), Teuchos::outArg(thyraLeadingDim));
    thyraData = thyraData_arcp();  // build array view

    for (int i = 0; i < localSubDim; i++) {
      // copy each vector
      for (int v = 0; v < numVectors; v++)
        thyraData[i + thyraLeadingDim * v] = tpetraData[i + leadingDim * v];
    }

    // set the local dimension
    localDim = localSubDim;

    return;
  }

  // this keeps track of current location in the tpetraData vector
  Teuchos::ArrayRCP<const ST> localData = tpetraData;

  // loop over all the blocks in the vector space
  for (int blkIndex = 0; blkIndex < prodMV->productSpace()->numBlocks(); blkIndex++) {
    int subDim                                     = 0;
    const RCP<Thyra::MultiVectorBase<ST>> blockVec = prodMV->getNonconstMultiVectorBlock(blkIndex);

    // perform the recusive copy
    blockTpetraToThyra(numVectors, localData, leadingDim, blockVec.ptr(), subDim);

    // shift to the next block
    localData += subDim;

    // account for the size of this subblock
    localDim += subDim;
  }
}

void blockTpetraToThyraTpetraVec(const Tpetra::MultiVector<ST, LO, GO, NT>& tpetraX,
                                 const Teuchos::Ptr<Thyra::ProductMultiVectorBase<ST>>& prodMV) {
  auto view = tpetraX.getLocalViewDevice(Tpetra::Access::ReadOnly);

  // loop over all the blocks in the vector space
  size_t offset        = 0;
  const auto numBlocks = prodMV->productSpace()->numBlocks();
  for (int blkIndex = 0; blkIndex < numBlocks; blkIndex++) {
    const auto blockVec = prodMV->getNonconstMultiVectorBlock(blkIndex);
    const auto blockMV  = rcp_dynamic_cast<Thyra::TpetraMultiVector<ST, LO, GO, NT>>(blockVec, true)
                             ->getTpetraMultiVector();

    auto blockView = blockMV->getLocalViewDevice(Tpetra::Access::OverwriteAll);
    auto subView =
        Kokkos::subview(view, Kokkos::pair(offset, blockView.extent(0) + offset), Kokkos::ALL);
    Kokkos::deep_copy(decltype(blockView)::execution_space{}, blockView, subView);

    offset += blockView.extent(0);
  }
}

// Convert a Tpetra_MultiVector with assumed block structure dictated by the
// vector space into a Thyra::MultiVectorBase object.
// const Teuchos::RCP<const Thyra::MultiVectorBase<double> > blockTpetraToThyra(const
// Tpetra_MultiVector & e,const Teuchos::RCP<const Thyra::VectorSpaceBase<double> > & vs)
void blockTpetraToThyra(const Tpetra::MultiVector<ST, LO, GO, NT>& tpetraX,
                        const Teuchos::Ptr<Thyra::MultiVectorBase<ST>>& thyraX) {
  TEUCHOS_ASSERT((Tpetra::global_size_t)thyraX->range()->dim() == tpetraX.getGlobalLength());

  const auto prodMV = ptr_dynamic_cast<Thyra::ProductMultiVectorBase<ST>>(thyraX);
  bool allTpetra    = [&]() {
    if (prodMV == Teuchos::null) return false;

    for (int blkIndex = 0; blkIndex < prodMV->productSpace()->numBlocks(); blkIndex++) {
      const auto blockVec = prodMV->getNonconstMultiVectorBlock(blkIndex);
      const auto blockMV =
          Teuchos::rcp_dynamic_cast<Thyra::TpetraMultiVector<ST, LO, GO, NT>>(blockVec);
      if (blockMV == Teuchos::null) return false;
    }

    return true;
  }();

  if (allTpetra) {
    blockTpetraToThyraTpetraVec(tpetraX, prodMV);
    return;
  }

  // extract local information from the Tpetra_MultiVector
  LO leadingDim = 0, localDim = 0;
  leadingDim                             = tpetraX.getStride();
  Teuchos::ArrayRCP<const ST> tpetraData = tpetraX.get1dView();

  int numVectors = tpetraX.getNumVectors();

  blockTpetraToThyra(numVectors, tpetraData, leadingDim, thyraX.ptr(), localDim);

  TEUCHOS_ASSERT((size_t)localDim == tpetraX.getLocalLength());
}

void blockThyraToTpetra(LO numVectors, Teuchos::ArrayRCP<ST> tpetraData, LO leadingDim,
                        const Teuchos::RCP<const Thyra::MultiVectorBase<ST>>& tX, LO& localDim) {
  localDim = 0;

  // check the base case
  const RCP<const Thyra::ProductMultiVectorBase<ST>> prodX =
      rcp_dynamic_cast<const Thyra::ProductMultiVectorBase<ST>>(tX);
  if (prodX == Teuchos::null) {
    // the base case

    // VS object must be a SpmdMultiVector object
    RCP<const Thyra::SpmdMultiVectorBase<ST>> spmdX =
        rcp_dynamic_cast<const Thyra::SpmdMultiVectorBase<ST>>(tX, true);
    RCP<const Thyra::SpmdVectorSpaceBase<ST>> spmdVS = spmdX->spmdSpace();

    Thyra::Ordinal thyraLeadingDim = 0;
    Teuchos::ArrayView<const ST> thyraData;
    Teuchos::ArrayRCP<const ST> thyraData_arcp;
    spmdX->getLocalData(Teuchos::outArg(thyraData_arcp), Teuchos::outArg(thyraLeadingDim));
    thyraData = thyraData_arcp();  // grab the array view

    LO localSubDim = spmdVS->localSubDim();
    for (LO i = 0; i < localSubDim; i++) {
      // copy each vector
      for (LO v = 0; v < numVectors; v++) {
        tpetraData[i + leadingDim * v] = thyraData[i + thyraLeadingDim * v];
      }
    }

    // set the local dimension
    localDim = localSubDim;

    return;
  }

  const RCP<const Thyra::ProductVectorSpaceBase<ST>> prodVS = prodX->productSpace();

  // this keeps track of current location in the tpetraData vector
  Teuchos::ArrayRCP<ST> localData = tpetraData;

  // loop over all the blocks in the vector space
  for (int blkIndex = 0; blkIndex < prodVS->numBlocks(); blkIndex++) {
    int subDim = 0;

    // construct the block vector
    blockThyraToTpetra(numVectors, localData, leadingDim, prodX->getMultiVectorBlock(blkIndex),
                       subDim);

    // shift to the next block
    localData += subDim;

    // account for the size of this subblock
    localDim += subDim;
  }

  return;
}

void blockThyraToTpetraTpetraVec(
    const Teuchos::RCP<const Thyra::ProductMultiVectorBase<ST>>& prodMV,
    Tpetra::MultiVector<ST, LO, GO, NT>& tpetraX) {
  auto view = tpetraX.getLocalViewDevice(Tpetra::Access::OverwriteAll);

  // loop over all the blocks in the vector space
  size_t offset        = 0;
  const auto numBlocks = prodMV->productSpace()->numBlocks();
  for (int blkIndex = 0; blkIndex < numBlocks; blkIndex++) {
    const auto blockVec = prodMV->getMultiVectorBlock(blkIndex);
    const auto blockMV =
        rcp_dynamic_cast<const Thyra::TpetraMultiVector<ST, LO, GO, NT>>(blockVec, true)
            ->getConstTpetraMultiVector();

    auto blockView = blockMV->getLocalViewDevice(Tpetra::Access::ReadOnly);
    auto subView =
        Kokkos::subview(view, Kokkos::pair(offset, blockView.extent(0) + offset), Kokkos::ALL);
    Kokkos::deep_copy(decltype(blockView)::execution_space{}, subView, blockView);

    offset += blockView.extent(0);
  }
}

// Convert a Thyra::MultiVectorBase object to a Tpetra_MultiVector object with
// the map defined by the Tpetra_Map.
// const Teuchos::RCP<const Tpetra_MultiVector>
// blockThyraToTpetra(const Teuchos::RCP<const Thyra::MultiVectorBase<double> > & tX,const RCP<const
// Tpetra_Map> & map)
void blockThyraToTpetra(const Teuchos::RCP<const Thyra::MultiVectorBase<ST>>& thyraX,
                        Tpetra::MultiVector<ST, LO, GO, NT>& tpetraX) {
  const auto prodMV = rcp_dynamic_cast<const Thyra::ProductMultiVectorBase<ST>>(thyraX);
  bool allTpetra    = [&]() {
    if (prodMV == Teuchos::null) return false;

    for (int blkIndex = 0; blkIndex < prodMV->productSpace()->numBlocks(); blkIndex++) {
      const auto blockVec = prodMV->getMultiVectorBlock(blkIndex);
      const auto blockMV =
          Teuchos::rcp_dynamic_cast<const Thyra::TpetraMultiVector<ST, LO, GO, NT>>(blockVec);
      if (blockMV == Teuchos::null) return false;
    }

    return true;
  }();

  if (allTpetra) {
    blockThyraToTpetraTpetraVec(prodMV, tpetraX);
    return;
  }

  // build an Tpetra_MultiVector object
  LO numVectors = thyraX->domain()->dim();

  // make sure the number of vectors are the same
  TEUCHOS_ASSERT((size_t)numVectors == tpetraX.getNumVectors());
  TEUCHOS_ASSERT((Tpetra::global_size_t)thyraX->range()->dim() == tpetraX.getGlobalLength());

  // extract local information from the Tpetra_MultiVector
  LO leadingDim = 0, localDim = 0;
  leadingDim                       = tpetraX.getStride();
  Teuchos::ArrayRCP<ST> tpetraData = tpetraX.get1dViewNonConst();

  // perform recursive copy
  blockThyraToTpetra(numVectors, tpetraData, leadingDim, thyraX, localDim);

  // sanity check
  TEUCHOS_ASSERT((size_t)localDim == tpetraX.getLocalLength());
}

void thyraVSToTpetraMap(std::vector<GO>& myIndicies, int blockOffset,
                        const Thyra::VectorSpaceBase<ST>& vs, int& localDim) {
  // zero out set local dimension
  localDim = 0;

  const RCP<const Thyra::ProductVectorSpaceBase<ST>> prodVS =
      rcp_dynamic_cast<const Thyra::ProductVectorSpaceBase<ST>>(rcpFromRef(vs));

  // is more recursion needed?
  if (prodVS == Teuchos::null) {
    // base case

    // try to cast to an SPMD capable vector space
    const RCP<const Thyra::SpmdVectorSpaceBase<ST>> spmdVS =
        rcp_dynamic_cast<const Thyra::SpmdVectorSpaceBase<ST>>(rcpFromRef(vs));
    TEUCHOS_TEST_FOR_EXCEPTION(spmdVS == Teuchos::null, std::runtime_error,
                               "thyraVSToTpetraMap requires all subblocks to be SPMD");

    // get local data storage information
    int localOffset = spmdVS->localOffset();
    int localSubDim = spmdVS->localSubDim();

    // add indicies to matrix
    for (int i = 0; i < localSubDim; i++) myIndicies.push_back(blockOffset + localOffset + i);

    localDim += localSubDim;

    return;
  }

  // loop over all the blocks in the vector space
  for (int blkIndex = 0; blkIndex < prodVS->numBlocks(); blkIndex++) {
    int subDim = 0;

    // construct the block vector
    thyraVSToTpetraMap(myIndicies, blockOffset, *prodVS->getBlock(blkIndex), subDim);

    blockOffset += prodVS->getBlock(blkIndex)->dim();

    // account for the size of this subblock
    localDim += subDim;
  }
}

// From a Thyra vector space create a compatable Tpetra_Map
const RCP<Tpetra::Map<LO, GO, NT>> thyraVSToTpetraMap(
    const Thyra::VectorSpaceBase<ST>& vs, const RCP<const Teuchos::Comm<Thyra::Ordinal>>& comm) {
  int localDim = 0;
  std::vector<GO> myGIDs;

  // call recursive routine that constructs the mapping
  thyraVSToTpetraMap(myGIDs, 0, vs, localDim);

  TEUCHOS_ASSERT(myGIDs.size() == (size_t)localDim);

  auto tpetraComm = Thyra::convertThyraToTpetraComm(comm);

  // create the map
  return rcp(
      new Tpetra::Map<LO, GO, NT>(vs.dim(), Teuchos::ArrayView<const GO>(myGIDs), 0, tpetraComm));
}

}  // namespace TpetraHelpers
}  // end namespace Teko