Tpetra_replaceDiagonalCrsMatrix_def.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_REPLACEDIAGONALCRSMATRIX_DEF_HPP
#define TPETRA_REPLACEDIAGONALCRSMATRIX_DEF_HPP


#include "Tpetra_Details_Behavior.hpp"
#include "Tpetra_CrsMatrix.hpp"
#include "Tpetra_Vector.hpp"

namespace Tpetra {

template <class SC, class LO, class GO, class NT>
LO replaceDiagonalCrsMatrix(CrsMatrix<SC, LO, GO, NT>& matrix,
                            const Vector<SC, LO, GO, NT>& newDiag) {
  using map_type        = Map<LO, GO, NT>;
  using crs_matrix_type = CrsMatrix<SC, LO, GO, NT>;
  // using vec_type = ::Tpetra::Vector<SC, LO, GO, NT>;

  const LO oneLO = Teuchos::OrdinalTraits<LO>::one();

  // Count number of successfully replaced diagonal entries
  LO numReplacedDiagEntries = 0;

  // Extract some useful data
  auto rowMapPtr = matrix.getRowMap();
  if (rowMapPtr.is_null() || rowMapPtr->getComm().is_null()) {
    // Processes on which the row Map or its communicator is null
    // don't participate.  Users shouldn't even call this method on
    // those processes.
    return numReplacedDiagEntries;
  }
  auto colMapPtr = matrix.getColMap();

  TEUCHOS_TEST_FOR_EXCEPTION(colMapPtr.get() == nullptr, std::invalid_argument,
                             "replaceDiagonalCrsMatrix: "
                             "Input matrix must have a nonnull column Map.");

  const map_type& rowMap           = *rowMapPtr;
  const map_type& colMap           = *colMapPtr;
  const LO myNumRows               = static_cast<LO>(matrix.getLocalNumRows());
  const bool isFillCompleteOnInput = matrix.isFillComplete();

  if (Tpetra::Details::Behavior::debug()) {
    TEUCHOS_TEST_FOR_EXCEPTION(!rowMap.isSameAs(*newDiag.getMap()), std::runtime_error,
                               "Row map of matrix and map of input vector do not match.");
  }

  // KJ: This fence is necessary for UVM. Views used in the row map and colmap
  // can use UVM and they are accessed in the following routine. So, we need to
  // make sure that the values are available for touching in host.
  typename crs_matrix_type::execution_space().fence();

  if (isFillCompleteOnInput)
    matrix.resumeFill();

  Teuchos::ArrayRCP<const SC> newDiagData = newDiag.getVector(0)->getData();
  LO numReplacedEntriesPerRow             = 0;

  auto invalid = Teuchos::OrdinalTraits<LO>::invalid();

  // Loop over all local rows to replace the diagonal entry row by row
  for (LO lclRowInd = 0; lclRowInd < myNumRows; ++lclRowInd) {
    // Get row and column indices of the diagonal entry
    const GO gblInd    = rowMap.getGlobalElement(lclRowInd);
    const LO lclColInd = colMap.getLocalElement(gblInd);

    // If the row map is not one-to-one, the diagonal may not be on this proc.
    // Skip this row; some processor will have the diagonal for this row.
    if (lclColInd == invalid) continue;

    const SC vals[] = {static_cast<SC>(newDiagData[lclRowInd])};
    const LO cols[] = {lclColInd};

    // Do the actual replacement of the diagonal element, if on this proc
    numReplacedEntriesPerRow = matrix.replaceLocalValues(lclRowInd, oneLO,
                                                         vals, cols);

    // Check for success of replacement.
    // numReplacedEntriesPerRow is one if the diagonal was replaced.
    // numReplacedEntriesPerRow is zero if the diagonal is not on
    // this processor.  For example, in a 2D matrix distribution, gblInd may
    // be in both the row and column map, but the diagonal may not be on
    // this processor.
    if (numReplacedEntriesPerRow == oneLO) {
      ++numReplacedDiagEntries;
    }
  }

  if (isFillCompleteOnInput)
    matrix.fillComplete();

  return numReplacedDiagEntries;
}

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

#define TPETRA_REPLACEDIAGONALCRSMATRIX_INSTANT(SCALAR, LO, GO, NODE) \
                                                                      \
  template LO replaceDiagonalCrsMatrix(                               \
      CrsMatrix<SCALAR, LO, GO, NODE>& matrix,                        \
      const Vector<SCALAR, LO, GO, NODE>& newDiag);

#endif  // #ifndef TPETRA_REPLACEDIAGONALCRSMATRIX_DEF_HPP