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