Ifpack2_Container_def.hpp

/*@HEADER
// ***********************************************************************
//
//       Ifpack2: Templated Object-Oriented Algebraic Preconditioner Package
//                 Copyright (2009) Sandia Corporation
//
// Under terms of Contract DE-AC04-94AL85000, there is a non-exclusive
// license for use of this work by or on behalf of the U.S. Government.
//
// 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 Michael A. Heroux (maherou@sandia.gov)
//
// ***********************************************************************
//@HEADER
*/

#ifndef IFPACK2_CONTAINER_DEF_HPP
#define IFPACK2_CONTAINER_DEF_HPP

#include <Ifpack2_Details_MultiVectorLocalGatherScatter.hpp>
#include <Teuchos_Time.hpp>

namespace Ifpack2 {

//Implementation of Ifpack2::Container

template<class MatrixType>
Container<MatrixType>::Container(
    const Teuchos::RCP<const row_matrix_type>& matrix,
    const Teuchos::Array<Teuchos::Array<LO> >& partitions,
    bool pointIndexed) :
  inputMatrix_ (matrix),
  inputCrsMatrix_ (Teuchos::rcp_dynamic_cast<const crs_matrix_type>(inputMatrix_)),
  inputBlockMatrix_ (Teuchos::rcp_dynamic_cast<const block_crs_matrix_type>(inputMatrix_)),
  pointIndexed_(pointIndexed),
  IsInitialized_(false),
  IsComputed_(false)
{
  using Teuchos::Ptr;
  using Teuchos::RCP;
  using Teuchos::Array;
  using Teuchos::ArrayView;
  using Teuchos::Comm;
  NumLocalRows_ = inputMatrix_->getNodeNumRows();
  NumGlobalRows_ = inputMatrix_->getGlobalNumRows();
  NumGlobalNonzeros_ = inputMatrix_->getGlobalNumEntries();
  IsParallel_ = inputMatrix_->getRowMap()->getComm()->getSize() != 1;
  hasBlockCrs_ = !inputBlockMatrix_.is_null();
  if(hasBlockCrs_)
    bcrsBlockSize_ = inputBlockMatrix_->getBlockSize();
  else
    bcrsBlockSize_ = 1;
  if(hasBlockCrs_ && !pointIndexed_)
    scalarsPerRow_ = bcrsBlockSize_;
  else
    scalarsPerRow_ = 1;
  setBlockSizes(partitions);
  //Sanity check the partitions
  #ifdef HAVE_IFPACK2_DEBUG
  // Check whether the input set of local row indices is correct.
  const map_type& rowMap = *inputMatrix_->getRowMap();
  for(int i = 0; i < numBlocks_; i++)
  {
    Teuchos::ArrayView<const LO> blockRows = getBlockRows(i);
    for(LO j = 0; j < blockSizes_[i]; j++)
    {
      LO row = blockRows[j];
      if(pointIndexed)
      {
        //convert the point row to the corresponding block row
        row /= bcrsBlockSize_;
      }
      TEUCHOS_TEST_FOR_EXCEPTION(
          !rowMap.isNodeLocalElement(row),
          std::invalid_argument, "Ifpack2::Container: "
          "On process " << rowMap.getComm()->getRank() << " of "
          << rowMap.getComm()->getSize() << ", in the given set of local row "
          "indices blockRows = " << Teuchos::toString(blockRows) << ", the following "
          "entries is not valid local row index on the calling process: "
          << row << ".");
    }
  }
  #endif
}

template<class MatrixType>
Container<MatrixType>::
~Container() {}

template<class MatrixType>
Teuchos::ArrayView<const typename MatrixType::local_ordinal_type>
Container<MatrixType>::getBlockRows(int blockIndex) const
{
  return Teuchos::ArrayView<const LO>
    (&blockRows_[blockOffsets_[blockIndex]], blockSizes_[blockIndex]);
}

template<class MatrixType>
void Container<MatrixType>::setBlockSizes(const Teuchos::Array<Teuchos::Array<LO> >& partitions)
{
  //First, create a grand total of all the rows in all the blocks
  //Note: If partitioner allowed overlap, this could be greater than the # of local rows
  LO totalBlockRows = 0;
  numBlocks_ = partitions.size();
  blockSizes_.resize(numBlocks_);
  blockOffsets_.resize(numBlocks_);
  maxBlockSize_ = 0;
  for(int i = 0; i < numBlocks_; i++)
  {
    LO rowsInBlock = partitions[i].size();
    blockSizes_[i] = rowsInBlock;
    blockOffsets_[i] = totalBlockRows;
    totalBlockRows += rowsInBlock;
    maxBlockSize_ = std::max(maxBlockSize_, rowsInBlock * scalarsPerRow_);
  }
  blockRows_.resize(totalBlockRows);
  //set blockRows_: each entry is the partition/block of the row
  LO iter = 0;
  for(int i = 0; i < numBlocks_; i++)
  {
    for(int j = 0; j < blockSizes_[i]; j++)
    {
      blockRows_[iter++] = partitions[i][j];
    }
  }
}

template<class MatrixType>
void Container<MatrixType>::getMatDiag() const
{
  if(Diag_.is_null())
  {
    Diag_ = rcp(new vector_type(inputMatrix_->getDomainMap()));
    inputMatrix_->getLocalDiagCopy(*Diag_);
  }
}

template<class MatrixType>
bool Container<MatrixType>::isInitialized() const {
  return IsInitialized_;
}

template<class MatrixType>
bool Container<MatrixType>::isComputed () const {
  return IsComputed_;
}

template<class MatrixType>
void Container<MatrixType>::
applyMV(mv_type& X, mv_type& Y) const
{
  TEUCHOS_TEST_FOR_EXCEPT_MSG(true, "Not implemented.");
}

template<class MatrixType>
void Container<MatrixType>::
weightedApplyMV(mv_type& X, mv_type& Y, vector_type& W) const
{
  TEUCHOS_TEST_FOR_EXCEPT_MSG(true, "Not implemented.");
}

template<class MatrixType>
std::string Container<MatrixType>::
getName()
{
  return "Generic";
}

template<class MatrixType>
void Container<MatrixType>::DoGSBlock(HostView X, HostView Y, HostView Y2, HostView Resid,
    SC dampingFactor, LO i) const
{
  TEUCHOS_TEST_FOR_EXCEPT_MSG(true, "Not implemented.");
}

template <class MatrixType>
void Container<MatrixType>::DoJacobi(HostView X, HostView Y, SC dampingFactor) const
{
  using STS = Teuchos::ScalarTraits<ISC>;
  const ISC one = STS::one();
  // use blockRows_ and blockSizes_
  size_t numVecs = X.extent(1);
  // Non-overlapping Jacobi
  for (LO i = 0; i < numBlocks_; i++)
  {
    // may happen that a partition is empty
    if(blockSizes_[i] != 1 || hasBlockCrs_)
    {
      if(blockSizes_[i] == 0 )
        continue;
      apply(X, Y, i, Teuchos::NO_TRANS, dampingFactor, one);
    }
    else    // singleton, can't access Containers_[i] as it was never filled and may be null.
    {
      LO LRID = blockRows_[blockOffsets_[i]];
      getMatDiag();
      HostView diagView = Diag_->getLocalViewHost();
      ISC d = one / diagView(LRID, 0);
      for(size_t nv = 0; nv < numVecs; nv++)
      {
        ISC x = X(LRID, nv);
        Y(LRID, nv) = x * d;
      }
    }
  }
}

template <class MatrixType>
void Container<MatrixType>::DoOverlappingJacobi(HostView X, HostView Y, HostView W, SC dampingFactor) const
{
  using STS = Teuchos::ScalarTraits<SC>;
  // Overlapping Jacobi
  for(LO i = 0; i < numBlocks_; i++)
  {
    // may happen that a partition is empty
    if(blockSizes_[i] == 0)
      continue;
    if(blockSizes_[i] != 1)
      weightedApply(X, Y, W, i, Teuchos::NO_TRANS, dampingFactor, STS::one());
  }
}

//Do Gauss-Seidel with just block i
//This is used 3 times: once in DoGaussSeidel and twice in DoSGS
template<class MatrixType, typename LocalScalarType>
void ContainerImpl<MatrixType, LocalScalarType>::DoGSBlock(
    HostView X, HostView Y, HostView Y2, HostView Resid,
    SC dampingFactor, LO i) const
{
  using Teuchos::ArrayView;
  using STS = Teuchos::ScalarTraits<ISC>;
  size_t numVecs = X.extent(1);
  const ISC one = STS::one();
  if(this->blockSizes_[i] == 0)
    return; // Skip empty partitions
  if(this->hasBlockCrs_ && !this->pointIndexed_)
  {
    //Use efficient blocked version
    ArrayView<const LO> blockRows = this->getBlockRows(i);
    const size_t localNumRows = this->blockSizes_[i];
    for(size_t j = 0; j < localNumRows; j++)
    {
      LO row = blockRows[j]; // Containers_[i]->ID (j);
      LO numEntries;
      SC* values;
      const LO* colinds;
      this->inputBlockMatrix_->getLocalRowView(row, colinds, values, numEntries);
      for(size_t m = 0; m < numVecs; m++)
      {
        for (int localR = 0; localR < this->bcrsBlockSize_; localR++)
          Resid(row * this->bcrsBlockSize_ + localR, m) = X(row * this->bcrsBlockSize_ + localR, m);
        for (LO k = 0; k < numEntries; ++k)
        {
          const LO col = colinds[k];
          for(int localR = 0; localR < this->bcrsBlockSize_; localR++)
          {
            for(int localC = 0; localC < this->bcrsBlockSize_; localC++)
            {
              Resid(row * this->bcrsBlockSize_ + localR, m) -=
                values[k * this->bcrsBlockSize_ * this->bcrsBlockSize_ + localR + localC * this->bcrsBlockSize_]
                * Y2(col * this->bcrsBlockSize_ + localC, m); }
          }
        }
      }
    }
    // solve with this block
    //
    // Note: I'm abusing the ordering information, knowing that X/Y
    // and Y2 have the same ordering for on-proc unknowns.
    //
    // Note: Add flop counts for inverse apply
    apply(Resid, Y2, i, Teuchos::NO_TRANS, dampingFactor, one);
  }
  else if(!this->hasBlockCrs_ && this->blockSizes_[i] == 1)
  {
    // singleton, can't access Containers_[i] as it was never filled and may be null.
    // a singleton calculation (just using matrix diagonal) is exact, all residuals should be zero.
    LO LRID = this->blockOffsets_[i];  // by definition, a singleton 1 row in block.
    HostView diagView = this->Diag_->getLocalViewHost();
    ISC d = one / diagView(LRID, 0);
    for(size_t m = 0; m < numVecs; m++)
    {
      ISC x = X(LRID, m);
      ISC newy = x * d;
      Y2(LRID, m) = newy;
    }
  }
  else if(!this->inputCrsMatrix_.is_null() &&
      std::is_same<typename crs_matrix_type::device_type::memory_space, Kokkos::HostSpace>::value)
  {
    //Use the KokkosSparse internal matrix for low-overhead values/indices access
    //But, can only do this if the matrix is accessible directly from host, since it's not a DualView
    using container_exec_space = typename ContainerImpl<MatrixType, LocalScalarType>::crs_matrix_type::execution_space;
    container_exec_space().fence();
    auto localA = this->inputCrsMatrix_->getLocalMatrix();
    using size_type = typename crs_matrix_type::local_matrix_type::size_type;
    const auto& rowmap = localA.graph.row_map;
    const auto& entries = localA.graph.entries;
    const auto& values = localA.values;
    ArrayView<const LO> blockRows = this->getBlockRows(i);
    for(size_t j = 0; j < size_t(blockRows.size()); j++)
    {
      const LO row = blockRows[j];
      for(size_t m = 0; m < numVecs; m++)
      {
        ISC r = X(row, m);
        for(size_type k = rowmap(row); k < rowmap(row + 1); k++)
        {
          const LO col = entries(k);
          r -= values(k) * Y2(col, m);
        }
        Resid(row, m) = r;
      }
    }
    // solve with this block
    //
    // Note: I'm abusing the ordering information, knowing that X/Y
    // and Y2 have the same ordering for on-proc unknowns.
    //
    // Note: Add flop counts for inverse apply
    apply(Resid, Y2, i, Teuchos::NO_TRANS, dampingFactor, one);
  }
  else
  {
    //Either a point-indexed block matrix, or a normal row matrix
    //that doesn't support getLocalMatrix
    ArrayView<const LO> blockRows = this->getBlockRows(i);
    for(size_t j = 0; j < size_t(blockRows.size()); j++)
    {
      const LO row = blockRows[j];
      auto rowView = getInputRowView(row);
      for(size_t m = 0; m < numVecs; m++)
      {
        Resid(row, m) = X(row, m);
        for (size_t k = 0; k < rowView.size(); ++k)
        {
          const LO col = rowView.ind(k);
          Resid(row, m) -= rowView.val(k) * Y2(col, m);
        }
      }
    }
    // solve with this block
    //
    // Note: I'm abusing the ordering information, knowing that X/Y
    // and Y2 have the same ordering for on-proc unknowns.
    //
    // Note: Add flop counts for inverse apply
    apply(Resid, Y2, i, Teuchos::NO_TRANS, dampingFactor, one);
  }
}

template<class MatrixType>
void Container<MatrixType>::
DoGaussSeidel(HostView X, HostView Y, HostView Y2, SC dampingFactor) const
{
  using Teuchos::Array;
  using Teuchos::ArrayRCP;
  using Teuchos::ArrayView;
  using Teuchos::Ptr;
  using Teuchos::RCP;
  using Teuchos::rcp;
  using Teuchos::rcpFromRef;
  //This function just extracts the diagonal if it hasn't already.
  getMatDiag();
  auto numVecs = X.extent(1);
  // X = RHS, Y = initial guess
  HostView Resid("", X.extent(0), X.extent(1));
  for(LO i = 0; i < numBlocks_; i++)
  {
    DoGSBlock(X, Y, Y2, Resid, dampingFactor, i);
  }
  if(IsParallel_)
  {
    auto numMyRows = inputMatrix_->getNodeNumRows();
    for (size_t m = 0; m < numVecs; ++m)
    {
      for (size_t i = 0; i < numMyRows * bcrsBlockSize_; ++i)
      {
        Y(i, m) = Y2(i, m);
      }
    }
  }
}

template<class MatrixType>
void Container<MatrixType>::
DoSGS(HostView X, HostView Y, HostView Y2, SC dampingFactor) const
{
  // X = RHS, Y = initial guess
  using Teuchos::Array;
  using Teuchos::ArrayRCP;
  using Teuchos::ArrayView;
  using Teuchos::Ptr;
  using Teuchos::ptr;
  using Teuchos::RCP;
  using Teuchos::rcp;
  using Teuchos::rcpFromRef;
  auto numVecs = X.extent(1);
  HostView Resid("", X.extent(0), X.extent(1));
  // Forward Sweep
  for(LO i = 0; i < numBlocks_; i++)
  {
    DoGSBlock(X, Y, Y2, Resid, dampingFactor, i);
  }
  static_assert(std::is_signed<LO>::value,
      "Local ordinal must be signed (unsigned breaks reverse iteration to 0)");
  // Reverse Sweep
  for(LO i = numBlocks_ - 1; i >= 0; --i)
  {
    DoGSBlock(X, Y, Y2, Resid, dampingFactor, i);
  }
  if(IsParallel_)
  {
    auto numMyRows = inputMatrix_->getNodeNumRows();
    for (size_t m = 0; m < numVecs; ++m)
    {
      for (size_t i = 0; i < numMyRows * bcrsBlockSize_; ++i)
      {
        Y(i, m) = Y2(i, m);
      }
    }
  }
}

template<class MatrixType>
void Container<MatrixType>::
clearBlocks()
{
  numBlocks_ = 0;
  blockRows_.clear();
  blockSizes_.clear();
  blockOffsets_.clear();
  Diag_ = Teuchos::null;      //Diag_ will be recreated if needed
}

//Implementation of Ifpack2::ContainerImpl

template<class MatrixType, class LocalScalarType>
ContainerImpl<MatrixType, LocalScalarType>::
ContainerImpl(
      const Teuchos::RCP<const row_matrix_type>& matrix,
      const Teuchos::Array<Teuchos::Array<LO> >& partitions,
      bool pointIndexed)
  : Container<MatrixType>(matrix, partitions, pointIndexed) {}

template<class MatrixType, class LocalScalarType>
ContainerImpl<MatrixType, LocalScalarType>::
~ContainerImpl() {}

template<class MatrixType, class LocalScalarType>
void ContainerImpl<MatrixType, LocalScalarType>::
setParameters (const Teuchos::ParameterList& List) {}

template<class MatrixType, class LocalScalarType>
void ContainerImpl<MatrixType, LocalScalarType>::
applyInverseJacobi (const mv_type& /* X */, mv_type& /* Y */,
                         SC dampingFactor,
                         bool /* zeroStartingSolution = false */,
                         int /* numSweeps = 1 */) const
{
  TEUCHOS_TEST_FOR_EXCEPT_MSG(true, "Not implemented.");
}

template<class MatrixType, class LocalScalarType>
void ContainerImpl<MatrixType, LocalScalarType>::
applyMV (mv_type& X, mv_type& Y) const
{
  HostView XView = X.getLocalViewHost();
  HostView YView = Y.getLocalViewHost();
  this->apply (XView, YView, 0);
}

template<class MatrixType, class LocalScalarType>
void ContainerImpl<MatrixType, LocalScalarType>::
weightedApplyMV (mv_type& X,
                 mv_type& Y,
                 vector_type& W) const
{
  HostView XView = X.getLocalViewHost();
  HostView YView = Y.getLocalViewHost();
  HostView WView = W.getLocalViewHost();
  weightedApply (XView, YView, WView, 0);
}

template<class MatrixType, class LocalScalarType>
std::string ContainerImpl<MatrixType, LocalScalarType>::
getName()
{
  return "Generic";
}

template<class MatrixType, class LocalScalarType>
void ContainerImpl<MatrixType, LocalScalarType>::
solveBlock(HostSubviewLocal X,
           HostSubviewLocal Y,
           int blockIndex,
           Teuchos::ETransp mode,
           const LSC alpha,
           const LSC beta) const
{
  TEUCHOS_TEST_FOR_EXCEPT_MSG(true, "Not implemented.");
}

template<class MatrixType, class LocalScalarType>
typename ContainerImpl<MatrixType, LocalScalarType>::LO
ContainerImpl<MatrixType, LocalScalarType>::
translateRowToCol(LO row)
{
  LO LO_INVALID = Teuchos::OrdinalTraits<LO>::invalid();
  GO GO_INVALID = Teuchos::OrdinalTraits<GO>::invalid();
  const map_type& globalRowMap = *(this->inputMatrix_->getRowMap());
  const map_type& globalColMap = *(this->inputMatrix_->getColMap());
  LO rowLID = row;
  LO dofOffset = 0;
  if(this->pointIndexed_)
  {
    rowLID = row / this->bcrsBlockSize_;
    dofOffset = row % this->bcrsBlockSize_;
  }
  GO diagGID = globalRowMap.getGlobalElement(rowLID);
  TEUCHOS_TEST_FOR_EXCEPTION(
    diagGID == GO_INVALID,
    std::runtime_error, "Ifpack2::Container::translateRowToCol: "
    "On process " << this->inputMatrix_->getRowMap()->getComm()->getRank() <<
    ", at least one row index in the set of local "
    "row indices given to the constructor is not a valid local row index in "
    "the input matrix's row Map on this process.  This should be impossible "
    "because the constructor checks for this case.  Here is the complete set "
    "of invalid local row indices: " << rowLID << ".  "
    "Please report this bug to the Ifpack2 developers.");
  //now, can translate diagGID (both a global row AND global col ID) to local column
  LO colLID = globalColMap.getLocalElement(diagGID);
  TEUCHOS_TEST_FOR_EXCEPTION(
    colLID == LO_INVALID,
    std::runtime_error, "Ifpack2::Container::translateRowToCol: "
    "On process " << this->inputMatrix_->getRowMap()->getComm()->getRank() << ", "
    "at least one row index in the set of row indices given to the constructor "
    "does not have a corresponding column index in the input matrix's column "
    "Map.  This probably means that the column(s) in question is/are empty on "
    "this process, which would make the submatrix to extract structurally "
    "singular. The invalid global column index is " << diagGID << ".");
  //colLID could identify a block column - translate to split column if needed
  if(this->pointIndexed_)
    return colLID * this->bcrsBlockSize_ + dofOffset;
  return colLID;
}

template<class MatrixType, class LocalScalarType>
void ContainerImpl<MatrixType, LocalScalarType>::
apply (HostView X,
       HostView Y,
       int blockIndex,
       Teuchos::ETransp mode,
       SC alpha,
       SC beta) const
{
  using Teuchos::ArrayView;
  using Teuchos::as;
  using Teuchos::RCP;
  using Teuchos::rcp;

  // The local operator might have a different Scalar type than
  // MatrixType.  This means that we might have to convert X and Y to
  // the Tpetra::MultiVector specialization that the local operator
  // wants.  This class' X_ and Y_ internal fields are of the right
  // type for the local operator, so we can use those as targets.

  Details::MultiVectorLocalGatherScatter<mv_type, local_mv_type> mvgs;

  TEUCHOS_TEST_FOR_EXCEPTION(
    ! this->IsComputed_, std::runtime_error, "Ifpack2::Container::apply: "
    "You must have called the compute() method before you may call apply().  "
    "You may call the apply() method as many times as you want after calling "
    "compute() once, but you must have called compute() at least once.");

  const size_t numVecs = X.extent(1);

  if(numVecs == 0) {
    return; // done! nothing to do
  }

  // The local operator works on a permuted subset of the local parts
  // of X and Y.  The subset and permutation are defined by the index
  // array returned by getBlockRows().  If the permutation is trivial
  // and the subset is exactly equal to the local indices, then we
  // could use the local parts of X and Y exactly, without needing to
  // permute.  Otherwise, we have to use temporary storage to permute
  // X and Y.  For now, we always use temporary storage.
  //
  // Create temporary permuted versions of the input and output.
  // (Re)allocate X_ and/or Y_ only if necessary.  We'll use them to
  // store the permuted versions of X resp. Y.  Note that X_local has
  // the domain Map of the operator, which may be a permuted subset of
  // the local Map corresponding to X.getMap().  Similarly, Y_local
  // has the range Map of the operator, which may be a permuted subset
  // of the local Map corresponding to Y.getMap().  numRows_ here
  // gives the number of rows in the row Map of the local Inverse_
  // operator.
  //
  // FIXME (mfh 20 Aug 2013) There might be an implicit assumption
  // here that the row Map and the range Map of that operator are
  // the same.
  //
  // FIXME (mfh 20 Aug 2013) This "local permutation" functionality
  // really belongs in Tpetra.

  if(X_localBlocks_.size() == 0 || X.extent(1) != X_local_.extent(1))
  {
    //need to resize (or create for the first time) the three scratch arrays
    X_localBlocks_.clear();
    Y_localBlocks_.clear();
    X_localBlocks_.reserve(this->numBlocks_);
    Y_localBlocks_.reserve(this->numBlocks_);

    X_local_ = HostViewLocal("X_local", this->blockRows_.size() * this->scalarsPerRow_, numVecs);
    Y_local_ = HostViewLocal("Y_local", this->blockRows_.size() * this->scalarsPerRow_, numVecs);

    //create all X_local and Y_local managed Views at once, are
    //reused in subsequent apply() calls
    for(int i = 0; i < this->numBlocks_; i++)
    {
      auto blockBounds = std::make_pair(this->blockOffsets_[i] * this->scalarsPerRow_,
          (this->blockOffsets_[i] + this->blockSizes_[i]) * this->scalarsPerRow_);
      X_localBlocks_.emplace_back(X_local_, blockBounds, Kokkos::ALL());
      Y_localBlocks_.emplace_back(Y_local_, blockBounds, Kokkos::ALL());
    }
  }

  const ArrayView<const LO> blockRows = this->getBlockRows(blockIndex);

  if(this->scalarsPerRow_ == 1)
    mvgs.gatherViewToView (X_localBlocks_[blockIndex], X, blockRows);
  else
    mvgs.gatherViewToViewBlock (X_localBlocks_[blockIndex], X, blockRows, this->scalarsPerRow_);

  // We must gather the contents of the output multivector Y even on
  // input to solveBlock(), since the inverse operator might use it as
  // an initial guess for a linear solve.  We have no way of knowing
  // whether it does or does not.

  if(this->scalarsPerRow_ == 1)
    mvgs.gatherViewToView (Y_localBlocks_[blockIndex], Y, blockRows);
  else
    mvgs.gatherViewToViewBlock (Y_localBlocks_[blockIndex], Y, blockRows, this->scalarsPerRow_);

  // Apply the local operator:
  // Y_local := beta*Y_local + alpha*M^{-1}*X_local
  this->solveBlock (X_localBlocks_[blockIndex], Y_localBlocks_[blockIndex], blockIndex, mode,
                   LSC(alpha), LSC(beta));

  // Scatter the permuted subset output vector Y_local back into the
  // original output multivector Y.
  if(this->scalarsPerRow_ == 1)
    mvgs.scatterViewToView (Y, Y_localBlocks_[blockIndex], blockRows);
  else
    mvgs.scatterViewToViewBlock (Y, Y_localBlocks_[blockIndex], blockRows, this->scalarsPerRow_);
}

template<class MatrixType, class LocalScalarType>
void ContainerImpl<MatrixType, LocalScalarType>::
weightedApply(HostView X,
              HostView Y,
              HostView D,
              int blockIndex,
              Teuchos::ETransp mode,
              SC alpha,
              SC beta) const
{
  using Teuchos::ArrayRCP;
  using Teuchos::ArrayView;
  using Teuchos::Range1D;
  using Teuchos::Ptr;
  using Teuchos::ptr;
  using Teuchos::RCP;
  using Teuchos::rcp;
  using Teuchos::rcp_const_cast;
  using std::endl;
  using STS = Teuchos::ScalarTraits<SC>;

  // The local operator template parameter might have a different
  // Scalar type than MatrixType.  This means that we might have to
  // convert X and Y to the Tpetra::MultiVector specialization that
  // the local operator wants.  This class' X_ and Y_ internal fields
  // are of the right type for the local operator, so we can use those
  // as targets.

  const char prefix[] = "Ifpack2::Container::weightedApply: ";
  TEUCHOS_TEST_FOR_EXCEPTION(
    ! this->IsComputed_, std::runtime_error, prefix << "You must have called the "
    "compute() method before you may call this method.  You may call "
    "weightedApply() as many times as you want after calling compute() once, "
    "but you must have called compute() at least once first.");
  
  //bmk 7-2019: BlockRelaxation already checked this, but if that changes...
  TEUCHOS_TEST_FOR_EXCEPTION(
    this->scalarsPerRow_ > 1, std::logic_error, prefix << "Use of block rows isn't allowed "
    "in overlapping Jacobi (the only method that uses weightedApply");

  const size_t numVecs = X.extent(1);

  TEUCHOS_TEST_FOR_EXCEPTION(
    X.extent(1) != Y.extent(1), std::runtime_error,
    prefix << "X and Y have different numbers of vectors (columns).  X has "
    << X.extent(1) << ", but Y has " << Y.extent(1) << ".");

  if(numVecs == 0) {
    return; // done! nothing to do
  }

  const size_t numRows = this->blockSizes_[blockIndex];

  // The local operator works on a permuted subset of the local parts
  // of X and Y.  The subset and permutation are defined by the index
  // array returned by getBlockRows().  If the permutation is trivial
  // and the subset is exactly equal to the local indices, then we
  // could use the local parts of X and Y exactly, without needing to
  // permute.  Otherwise, we have to use temporary storage to permute
  // X and Y.  For now, we always use temporary storage.
  //
  // Create temporary permuted versions of the input and output.
  // (Re)allocate X_ and/or Y_ only if necessary.  We'll use them to
  // store the permuted versions of X resp. Y.  Note that X_local has
  // the domain Map of the operator, which may be a permuted subset of
  // the local Map corresponding to X.getMap().  Similarly, Y_local
  // has the range Map of the operator, which may be a permuted subset
  // of the local Map corresponding to Y.getMap().  numRows_ here
  // gives the number of rows in the row Map of the local operator.
  //
  // FIXME (mfh 20 Aug 2013) There might be an implicit assumption
  // here that the row Map and the range Map of that operator are
  // the same.
  //
  // FIXME (mfh 20 Aug 2013) This "local permutation" functionality
  // really belongs in Tpetra.
  if(X_localBlocks_.size() == 0 || X.extent(1) != X_local_.extent(1))
  {
    //need to resize (or create for the first time) the three scratch arrays
    X_localBlocks_.clear();
    Y_localBlocks_.clear();
    X_localBlocks_.reserve(this->numBlocks_);
    Y_localBlocks_.reserve(this->numBlocks_);

    X_local_ = HostViewLocal("X_local", this->blockRows_.size() * this->scalarsPerRow_, numVecs);
    Y_local_ = HostViewLocal("Y_local", this->blockRows_.size() * this->scalarsPerRow_, numVecs);

    //create all X_local and Y_local managed Views at once, are
    //reused in subsequent apply() calls
    for(int i = 0; i < this->numBlocks_; i++)
    {
      auto blockBounds = std::make_pair(this->blockOffsets_[i] * this->scalarsPerRow_,
          (this->blockOffsets_[i] + this->blockSizes_[i]) * this->scalarsPerRow_);
      X_localBlocks_.emplace_back(X_local_, blockBounds, Kokkos::ALL());
      Y_localBlocks_.emplace_back(Y_local_, blockBounds, Kokkos::ALL());
    }
  }
  if(int(weightedApplyScratch_.extent(0)) != 3 * this->maxBlockSize_ ||
      weightedApplyScratch_.extent(1) != numVecs)
  {
    weightedApplyScratch_ = HostViewLocal("weightedApply scratch", 3 * this->maxBlockSize_, numVecs);
  }

  ArrayView<const LO> blockRows = this->getBlockRows(blockIndex);

  Details::MultiVectorLocalGatherScatter<mv_type, local_mv_type> mvgs;

  //note: BlockCrs w/ weighted Jacobi isn't allowed, so no need to use block gather/scatter
  mvgs.gatherViewToView (X_localBlocks_[blockIndex], X, blockRows);
  // We must gather the output multivector Y even on input to
  // solveBlock(), since the local operator might use it as an initial
  // guess for a linear solve.  We have no way of knowing whether it
  // does or does not.

  mvgs.gatherViewToView (Y_localBlocks_[blockIndex], Y, blockRows);

  // Apply the diagonal scaling D to the input X.  It's our choice
  // whether the result has the original input Map of X, or the
  // permuted subset Map of X_local.  If the latter, we also need to
  // gather D into the permuted subset Map.  We choose the latter, to
  // save memory and computation.  Thus, we do the following:
  //
  // 1. Gather D into a temporary vector D_local.
  // 2. Create a temporary X_scaled to hold diag(D_local) * X_local.
  // 3. Compute X_scaled := diag(D_loca) * X_local.
  auto maxBS = this->maxBlockSize_;
  auto bs = this->blockSizes_[blockIndex] * this->scalarsPerRow_;

  HostSubviewLocal D_local(weightedApplyScratch_, std::make_pair(0, bs), std::make_pair(0, 1));
  mvgs.gatherViewToView (D_local, D, blockRows);
  HostSubviewLocal X_scaled(weightedApplyScratch_, std::make_pair(maxBS, maxBS + bs), Kokkos::ALL());
  for(size_t j = 0; j < numVecs; j++)
    for(size_t i = 0; i < numRows; i++)
      X_scaled(i, j) = X_localBlocks_[blockIndex](i, j) * D_local(i, 0);

  HostSubviewLocal Y_temp(weightedApplyScratch_, std::make_pair(maxBS * 2, maxBS * 2 + bs), Kokkos::ALL());
  // Apply the local operator: Y_temp := M^{-1} * X_scaled
  this->solveBlock (X_scaled, Y_temp, blockIndex, mode, STS::one(), STS::zero());
  // Y_local := beta * Y_local + alpha * diag(D_local) * Y_temp.
  //
  // Note that we still use the permuted subset scaling D_local here,
  // because Y_temp has the same permuted subset Map.  That's good, in
  // fact, because it's a subset: less data to read and multiply.
  LISC a = alpha;
  LISC b = beta;
  for(size_t j = 0; j < numVecs; j++)
    for(size_t i = 0; i < numRows; i++)
      Y_localBlocks_[blockIndex](i, j) = b * Y_localBlocks_[blockIndex](i, j) + a * Y_temp(i, j) * D_local(i, 0);

  // Copy the permuted subset output vector Y_local into the original
  // output multivector Y.
  mvgs.scatterViewToView (Y, Y_localBlocks_[blockIndex], blockRows);
}

template<class MatrixType, class LocalScalarType>
Details::StridedRowView<
  typename ContainerImpl<MatrixType, LocalScalarType>::SC,
  typename ContainerImpl<MatrixType, LocalScalarType>::LO,
  typename ContainerImpl<MatrixType, LocalScalarType>::GO,
  typename ContainerImpl<MatrixType, LocalScalarType>::NO>
ContainerImpl<MatrixType, LocalScalarType>::
getInputRowView(LO row) const
{
  if(this->hasBlockCrs_)
  {
    const LO* colinds;
    SC* values;
    LO numEntries;
    this->inputBlockMatrix_->getLocalRowView(row / this->bcrsBlockSize_, colinds, values, numEntries);
    return StridedRowView(values + row % this->bcrsBlockSize_, colinds, this->bcrsBlockSize_, numEntries * this->bcrsBlockSize_);
  }
  else if(!this->inputMatrix_->supportsRowViews())
  {
    size_t maxEntries = this->inputMatrix_->getNodeMaxNumRowEntries();
    Teuchos::Array<LO> indsCopy(maxEntries);
    Teuchos::Array<SC> valsCopy(maxEntries);
    size_t numEntries;
    this->inputMatrix_->getLocalRowCopy(row, indsCopy, valsCopy, numEntries);
    indsCopy.resize(numEntries);
    valsCopy.resize(numEntries);
    return StridedRowView(valsCopy, indsCopy);
  }
  else
  {
    const LO* colinds;
    const SC* values;
    LO numEntries;
    this->inputMatrix_->getLocalRowViewRaw(row, numEntries, colinds, values);
    return StridedRowView(values, colinds, 1, numEntries);
  }
}

template<class MatrixType, class LocalScalarType>
void ContainerImpl<MatrixType, LocalScalarType>::
clearBlocks()
{
  X_localBlocks_.clear();
  Y_localBlocks_.clear();
  X_local_ = HostViewLocal();
  Y_local_ = HostViewLocal();
  Container<MatrixType>::clearBlocks();
}

namespace Details {

//Implementation of Ifpack2::Details::StridedRowView
template<class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
StridedRowView<Scalar, LocalOrdinal, GlobalOrdinal, Node>::
StridedRowView(const SC* vals_, const LO* inds_, int blockSize_, size_t nnz_)
  : vals(vals_), inds(inds_), blockSize(blockSize_), nnz(nnz_)
{}

template<class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
StridedRowView<Scalar, LocalOrdinal, GlobalOrdinal, Node>::
StridedRowView(Teuchos::Array<SC>& vals_, Teuchos::Array<LO>& inds_)
  : vals(nullptr), inds(nullptr), blockSize(1), nnz(vals_.size())
{
  valsCopy.swap(vals_);
  indsCopy.swap(inds_);
}

template<class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
Scalar StridedRowView<Scalar, LocalOrdinal, GlobalOrdinal, Node>::
val(size_t i) const
{
  #ifdef HAVE_IFPACK2_DEBUG
  TEUCHOS_TEST_FOR_EXCEPTION(i >= nnz, std::runtime_error,
      "Out-of-bounds access into Ifpack2::Container::StridedRowView");
  #endif
  if(vals)
  {
    if(blockSize == 1)
      return vals[i];
    else
      return vals[i * blockSize];
  }
  else
    return valsCopy[i];
}

template<class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
LocalOrdinal StridedRowView<Scalar, LocalOrdinal, GlobalOrdinal, Node>::
ind(size_t i) const
{
  #ifdef HAVE_IFPACK2_DEBUG
    TEUCHOS_TEST_FOR_EXCEPTION(i >= nnz, std::runtime_error,
        "Out-of-bounds access into Ifpack2::Container::StridedRowView");
  #endif
  //inds is smaller than vals by a factor of the block size (dofs/node)
  if(inds)
  {
    if(blockSize == 1)
      return inds[i];
    else
      return inds[i / blockSize] * blockSize + i % blockSize;
  }
  else
    return indsCopy[i];
}

template<class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
size_t StridedRowView<Scalar, LocalOrdinal, GlobalOrdinal, Node>::
size() const
{
  return nnz;
}
}

}

template <class MatrixType>
std::ostream& operator<<(std::ostream& os, const Ifpack2::Container<MatrixType>& obj)
{
  return obj.print(os);
}

#define IFPACK2_CONTAINER_INSTANT(S,LO,GO,N) \
  template class Ifpack2::Container<Tpetra::RowMatrix<S, LO, GO, N>>; \
  template class Ifpack2::ContainerImpl<Tpetra::RowMatrix<S, LO, GO, N>, S>; \
  template class Ifpack2::Details::StridedRowView<S, LO, GO, N>; \
  template std::ostream& operator<< <Tpetra::RowMatrix<S, LO, GO, N>>( \
      std::ostream& os, const Ifpack2::Container<Tpetra::RowMatrix<S, LO, GO, N>>& obj);

#endif