Skip to content

Commit

Permalink
Merge pull request trilinos#9510 from NexGenAnalytics/JD_GetMatrixDia…
Browse files Browse the repository at this point in the history 
…gonal_Kokkos

Update MueLu's GetMatrixDiagonal function (Kokkos version), to return RCP<Vector>
  • Loading branch information
@cgcgcg
cgcgcg authored Aug 2, 2021
2 parents 286bab8 + af2d845 commit 61b6583
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 38 deletions.
15 changes: 10 additions & 5 deletions packages/muelu/src/Utils/MueLu_Utilities_kokkos_decl.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -160,11 +160,11 @@ namespace MueLu {

/*! @brief Extract Matrix Diagonal
Returns Matrix diagonal in ArrayRCP.
Returns Matrix diagonal in RCP<Vector>.
NOTE -- it's assumed that A has been fillComplete'd.
*/
static Teuchos::ArrayRCP<SC> GetMatrixDiagonal(const Matrix& A); // FIXME
static RCP<Vector> GetMatrixDiagonal(const Matrix& A); // FIXME

/*! @brief Extract Matrix Diagonal
Expand Down Expand Up @@ -269,7 +269,7 @@ namespace MueLu {
static void ZeroDirichletRows(RCP<MultiVector>& X, const Kokkos::View<const bool*, typename NO::device_type>& dirichletRows, SC replaceWith=Teuchos::ScalarTraits<SC>::zero());

static void ZeroDirichletCols(RCP<Matrix>& A, const Kokkos::View<const bool*, typename NO::device_type>& dirichletCols, SC replaceWith=Teuchos::ScalarTraits<SC>::zero());

static void ApplyRowSumCriterion(const Matrix& A,
const typename Teuchos::ScalarTraits<Scalar>::magnitudeType rowSumTol,
Kokkos::View<bool*, typename NO::device_type> & dirichletRows);
Expand Down Expand Up @@ -380,8 +380,13 @@ namespace MueLu {
#endif
static RCP<Xpetra::Matrix<SC,LO,GO,NO> > Crs2Op(RCP<CrsMatrix> Op) { return Utilities::Crs2Op(Op); }

static ArrayRCP<SC> GetMatrixDiagonal(const Matrix& A) {
return UtilitiesBase::GetMatrixDiagonal(A);
static RCP<Vector> GetMatrixDiagonal(const Matrix& A) {
const auto rowMap = A.getRowMap();
auto diag = Xpetra::VectorFactory<Scalar,LocalOrdinal,GlobalOrdinal,Node>::Build(rowMap,true);

A.getLocalDiagCopy(*diag);

return diag;
}
static RCP<Vector> GetMatrixDiagonalInverse(const Matrix& A, Magnitude tol = Teuchos::ScalarTraits<SC>::eps()*100, const bool doLumped=false) {
return UtilitiesBase::GetMatrixDiagonalInverse(A, tol, doLumped);
Expand Down
47 changes: 14 additions & 33 deletions packages/muelu/src/Utils/MueLu_Utilities_kokkos_def.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -112,30 +112,13 @@ namespace MueLu {


template <class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
Teuchos::ArrayRCP<Scalar> Utilities_kokkos<Scalar, LocalOrdinal, GlobalOrdinal, Node>::
Teuchos::RCP<Xpetra::Vector<Scalar,LocalOrdinal,GlobalOrdinal,Node> >
Utilities_kokkos<Scalar, LocalOrdinal, GlobalOrdinal, Node>::
GetMatrixDiagonal(const Matrix& A) {
// FIXME_KOKKOS
const auto rowMap = A.getRowMap();
auto diag = Xpetra::VectorFactory<Scalar, LocalOrdinal, GlobalOrdinal, Node>::Build(rowMap, true);

size_t numRows = A.getRowMap()->getNodeNumElements();
Teuchos::ArrayRCP<SC> diag(numRows);

Teuchos::ArrayView<const LO> cols;
Teuchos::ArrayView<const SC> vals;
for (size_t i = 0; i < numRows; ++i) {
A.getLocalRowView(i, cols, vals);

LO j = 0;
for (; j < cols.size(); ++j) {
if (Teuchos::as<size_t>(cols[j]) == i) {
diag[i] = vals[j];
break;
}
}
if (j == cols.size()) {
// Diagonal entry is absent
diag[i] = Teuchos::ScalarTraits<SC>::zero();
}
}
A.getLocalDiagCopy(*diag);

return diag;
} //GetMatrixDiagonal
Expand Down Expand Up @@ -226,11 +209,9 @@ namespace MueLu {
crsOp->getLocalDiagCopy(*localDiag,offsets());
}
else {
ArrayRCP<SC> localDiagVals = localDiag->getDataNonConst(0);
Teuchos::ArrayRCP<SC> diagVals = GetMatrixDiagonal(A);
for (LO i = 0; i < localDiagVals.size(); i++)
localDiagVals[i] = diagVals[i];
localDiagVals = diagVals = null;
auto localDiagVals = localDiag->getDeviceLocalView(Xpetra::Access::ReadWrite);
const auto diagVals = GetMatrixDiagonal(A)->getDeviceLocalView(Xpetra::Access::ReadOnly);
Kokkos::deep_copy(localDiagVals, diagVals);
}

RCP<Vector> diagonal = VectorFactory::Build(colMap);
Expand Down Expand Up @@ -612,7 +593,7 @@ namespace MueLu {
return MueLu::ZeroDirichletCols<double,int,int,Node>(A, dirichletCols, replaceWith);
}

// Applies rowsum criterion
// Applies rowsum criterion
template<class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
void ApplyRowSumCriterion(const Xpetra::Matrix<Scalar,LocalOrdinal,GlobalOrdinal,Node>& A,
const typename Teuchos::ScalarTraits<Scalar>::magnitudeType rowSumTol,
Expand Down Expand Up @@ -711,7 +692,7 @@ namespace MueLu {
<device, typename local_graph_type::row_map_type, typename local_graph_type::entries_type, lno_nnz_view_t>
(localGraph.row_map, localGraph.entries);

RCP<Xpetra::Vector<LocalOrdinal,LocalOrdinal,GlobalOrdinal,Node> > retval =
RCP<Xpetra::Vector<LocalOrdinal,LocalOrdinal,GlobalOrdinal,Node> > retval =
Xpetra::VectorFactory<LocalOrdinal,LocalOrdinal,GlobalOrdinal,Node>::Build(Op.getRowMap());

// Copy out and reorder data
Expand All @@ -723,7 +704,7 @@ namespace MueLu {
});
return retval;
}

template <class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
Teuchos::RCP<Xpetra::Vector<LocalOrdinal,LocalOrdinal,GlobalOrdinal,Node> > CuthillMcKee(const Xpetra::Matrix<Scalar,LocalOrdinal,GlobalOrdinal,Node> &Op) {
using local_matrix_type = typename Xpetra::Matrix<Scalar,LocalOrdinal,GlobalOrdinal,Node>::local_matrix_type;
Expand All @@ -740,7 +721,7 @@ namespace MueLu {
<device, typename local_graph_type::row_map_type, typename local_graph_type::entries_type, lno_nnz_view_t>
(localGraph.row_map, localGraph.entries);

RCP<Xpetra::Vector<LocalOrdinal,LocalOrdinal,GlobalOrdinal,Node> > retval =
RCP<Xpetra::Vector<LocalOrdinal,LocalOrdinal,GlobalOrdinal,Node> > retval =
Xpetra::VectorFactory<LocalOrdinal,LocalOrdinal,GlobalOrdinal,Node>::Build(Op.getRowMap());

// Copy out data
Expand All @@ -761,7 +742,7 @@ namespace MueLu {
}

template <class Node>
Teuchos::RCP<Xpetra::Vector<int,int,int,Node> >
Teuchos::RCP<Xpetra::Vector<int,int,int,Node> >
Utilities_kokkos<double,int,int,Node>::ReverseCuthillMcKee(const Matrix &Op) {
return MueLu::ReverseCuthillMcKee<double,int,int,Node>(Op);
}
Expand All @@ -773,7 +754,7 @@ namespace MueLu {
}

template <class Node>
Teuchos::RCP<Xpetra::Vector<int,int,int,Node> >
Teuchos::RCP<Xpetra::Vector<int,int,int,Node> >
Utilities_kokkos<double,int,int,Node>::CuthillMcKee(const Matrix &Op) {
return MueLu::CuthillMcKee<double,int,int,Node>(Op);
}
Expand Down

0 comments on commit 61b6583

Please sign in to comment.