From 9db41cde025579225e9c9d2f730d6f9d3f672217 Mon Sep 17 00:00:00 2001 From: Sebastian Kinnewig Date: Fri, 9 Feb 2024 11:40:23 +0100 Subject: [PATCH] Template TpetraWrappers::SparseMatrix::vmult on the vector type. --- .../lac/trilinos_tpetra_sparse_matrix.h | 26 ++- .../trilinos_tpetra_sparse_matrix.templates.h | 183 +++++++++++++----- source/lac/trilinos_tpetra_sparse_matrix.cc | 31 +++ 3 files changed, 182 insertions(+), 58 deletions(-) diff --git a/include/deal.II/lac/trilinos_tpetra_sparse_matrix.h b/include/deal.II/lac/trilinos_tpetra_sparse_matrix.h index f29536400e..e466f89851 100644 --- a/include/deal.II/lac/trilinos_tpetra_sparse_matrix.h +++ b/include/deal.II/lac/trilinos_tpetra_sparse_matrix.h @@ -27,11 +27,14 @@ # include # include # include +# include // Tpetra includes # include # include +# include + DEAL_II_NAMESPACE_OPEN @@ -162,6 +165,12 @@ namespace LinearAlgebra using GraphType = Tpetra::CrsGraph; + /** + * Typedef for Tpetra::Vector + */ + using VectorType = Tpetra:: + Vector; + /** * @name Constructors and initialization. */ @@ -799,9 +808,9 @@ namespace LinearAlgebra * initialized with the same IndexSet that was used for the column indices * of the matrix. */ + template void - vmult(Vector &dst, - const Vector &src) const; + vmult(InputVectorType &dst, const InputVectorType &src) const; /* * Matrix-vector multiplication: let dst = MT*src with @@ -810,9 +819,9 @@ namespace LinearAlgebra * * Source and destination must not be the same vector. */ + template void - Tvmult(Vector &dst, - const Vector &src) const; + Tvmult(InputVectorType &dst, const InputVectorType &src) const; /** * Adding matrix-vector multiplication. Add M*src on dst @@ -820,10 +829,9 @@ namespace LinearAlgebra * * Source and destination must not be the same vector. */ + template void - vmult_add(Vector &dst, - const Vector &src) const; - + vmult_add(InputVectorType &dst, const InputVectorType &src) const; /** * Adding matrix-vector multiplication. Add MT*src to @@ -832,9 +840,9 @@ namespace LinearAlgebra * * Source and destination must not be the same vector. */ + template void - Tvmult_add(Vector &dst, - const Vector &src) const; + Tvmult_add(InputVectorType &dst, const InputVectorType &src) const; /** * Return the square of the norm of the vector $v$ with respect to the diff --git a/include/deal.II/lac/trilinos_tpetra_sparse_matrix.templates.h b/include/deal.II/lac/trilinos_tpetra_sparse_matrix.templates.h index e0660c7729..4a49d36982 100644 --- a/include/deal.II/lac/trilinos_tpetra_sparse_matrix.templates.h +++ b/include/deal.II/lac/trilinos_tpetra_sparse_matrix.templates.h @@ -31,6 +31,114 @@ namespace LinearAlgebra namespace TpetraWrappers { + namespace internal + { + template + void + apply(const SparseMatrix &M, + const Vector &src, + Vector &dst, + Teuchos::ETransp mode = Teuchos::NO_TRANS, + Number alpha = Teuchos::ScalarTraits::one(), + Number beta = Teuchos::ScalarTraits::zero()) + { + Assert(&src != &dst, + SparseMatrix::ExcSourceEqualsDestination()); + Assert(M.trilinos_matrix().isFillComplete(), + SparseMatrix::ExcMatrixNotCompressed()); + + if (mode == Teuchos::NO_TRANS) + { + Assert(src.trilinos_vector().getMap()->isSameAs( + *M.trilinos_matrix().getDomainMap()), + SparseMatrix::ExcColMapMissmatch()); + Assert(dst.trilinos_vector().getMap()->isSameAs( + *M.trilinos_matrix().getRangeMap()), + SparseMatrix::ExcDomainMapMissmatch()); + } + else + { + Assert(dst.trilinos_vector().getMap()->isSameAs( + *M.trilinos_matrix().getDomainMap()), + SparseMatrix::ExcColMapMissmatch()); + Assert(src.trilinos_vector().getMap()->isSameAs( + *M.trilinos_matrix().getRangeMap()), + SparseMatrix::ExcDomainMapMissmatch()); + } + + M.trilinos_matrix().apply( + src.trilinos_vector(), dst.trilinos_vector(), mode, alpha, beta); + } + + + + template + void + apply(const SparseMatrix &M, + const dealii::Vector &src, + dealii::Vector &dst, + Teuchos::ETransp mode = Teuchos::NO_TRANS, + Number alpha = Teuchos::ScalarTraits::one(), + Number beta = Teuchos::ScalarTraits::zero()) + { + Assert(&src != &dst, + SparseMatrix::ExcSourceEqualsDestination()); + Assert(M.trilinos_matrix().isFillComplete(), + SparseMatrix::ExcMatrixNotCompressed()); + + // get the size of the input vectors: + const size_type dst_local_size = dst.end() - dst.begin(); + const size_type src_local_size = src.end() - src.begin(); + + // For the dst vector: + Kokkos::View + kokkos_view_dst(dst.begin(), dst_local_size, 1); + + // get a Kokkos::DualView + auto mirror_view_dst = Kokkos::create_mirror_view_and_copy( + typename MemorySpace::kokkos_space{}, kokkos_view_dst); + typename SparseMatrix::VectorType::dual_view_type + kokkos_dual_view_dst(mirror_view_dst, kokkos_view_dst); + + // create the Tpetra::Vector + typename SparseMatrix::VectorType tpetra_dst( + M.trilinos_matrix().getRangeMap(), kokkos_dual_view_dst); + + // For the src vector: + // create a Kokkos::View from the src vector + Kokkos::View + kokkos_view_src(const_cast(src.begin()), src_local_size, 1); + + // get a Kokkos::DualView + auto mirror_view_src = Kokkos::create_mirror_view_and_copy( + typename MemorySpace::kokkos_space{}, kokkos_view_src); + typename SparseMatrix::VectorType::dual_view_type + kokkos_dual_view_src(mirror_view_src, kokkos_view_src); + + // create the Tpetra::Vector + typename SparseMatrix::VectorType tpetra_src( + M.trilinos_matrix().getDomainMap(), kokkos_dual_view_src); + + M.trilinos_matrix().apply(tpetra_src, tpetra_dst, mode, alpha, beta); + } + + + + template + void + apply(SparseMatrix &, + const VectorType &, + VectorType &, + Teuchos::ETransp, + Number, + Number) + { + DEAL_II_NOT_IMPLEMENTED(); + } + } // namespace internal + + + // reinit_matrix(): namespace { @@ -1170,81 +1278,58 @@ namespace LinearAlgebra // Multiplications - template + template void - SparseMatrix::vmult( - Vector &dst, - const Vector &src) const + SparseMatrix::vmult(InputVectorType &dst, + const InputVectorType &src) const { - Assert(&src != &dst, ExcSourceEqualsDestination()); - Assert(matrix->isFillComplete(), ExcMatrixNotCompressed()); - Assert(src.trilinos_vector().getMap()->isSameAs(*matrix->getDomainMap()), - ExcColMapMissmatch()); - Assert(dst.trilinos_vector().getMap()->isSameAs(*matrix->getRangeMap()), - ExcDomainMapMissmatch()); - matrix->apply(src.trilinos_vector(), dst.trilinos_vector()); + internal::apply(*this, src, dst); } template + template void - SparseMatrix::Tvmult( - Vector &dst, - const Vector &src) const + SparseMatrix::Tvmult(InputVectorType &dst, + const InputVectorType &src) const { - Assert(&src != &dst, ExcSourceEqualsDestination()); - Assert(matrix->isFillComplete(), ExcMatrixNotCompressed()); - Assert(dst.trilinos_vector().getMap()->isSameAs(*matrix->getDomainMap()), - ExcColMapMissmatch()); - Assert(src.trilinos_vector().getMap()->isSameAs(*matrix->getRangeMap()), - ExcDomainMapMissmatch()); - matrix->apply(src.trilinos_vector(), - dst.trilinos_vector(), - Teuchos::TRANS); + internal::apply(*this, src, dst, Teuchos::TRANS); } template + template void SparseMatrix::vmult_add( - Vector &dst, - const Vector &src) const + InputVectorType &dst, + const InputVectorType &src) const { - Assert(&src != &dst, ExcSourceEqualsDestination()); - Assert(matrix->isFillComplete(), ExcMatrixNotCompressed()); - Assert(src.trilinos_vector().getMap()->isSameAs(*matrix->getDomainMap()), - ExcColMapMissmatch()); - Assert(dst.trilinos_vector().getMap()->isSameAs(*matrix->getRangeMap()), - ExcDomainMapMissmatch()); - matrix->apply(src.trilinos_vector(), - dst.trilinos_vector(), - Teuchos::NO_TRANS, - Teuchos::ScalarTraits::one(), - Teuchos::ScalarTraits::one()); + internal::apply(*this, + src, + dst, + Teuchos::NO_TRANS, + Teuchos::ScalarTraits::one(), + Teuchos::ScalarTraits::one()); } template + template void SparseMatrix::Tvmult_add( - Vector &dst, - const Vector &src) const + InputVectorType &dst, + const InputVectorType &src) const { - Assert(&src != &dst, ExcSourceEqualsDestination()); - Assert(matrix->isFillComplete(), ExcMatrixNotCompressed()); - Assert(dst.trilinos_vector().getMap()->isSameAs(*matrix->getDomainMap()), - ExcColMapMissmatch()); - Assert(src.trilinos_vector().getMap()->isSameAs(*matrix->getRangeMap()), - ExcDomainMapMissmatch()); - matrix->apply(src.trilinos_vector(), - dst.trilinos_vector(), - Teuchos::TRANS, - Teuchos::ScalarTraits::one(), - Teuchos::ScalarTraits::one()); + internal::apply(*this, + src, + dst, + Teuchos::TRANS, + Teuchos::ScalarTraits::one(), + Teuchos::ScalarTraits::one()); } diff --git a/source/lac/trilinos_tpetra_sparse_matrix.cc b/source/lac/trilinos_tpetra_sparse_matrix.cc index 76c5335856..3cc43d6226 100644 --- a/source/lac/trilinos_tpetra_sparse_matrix.cc +++ b/source/lac/trilinos_tpetra_sparse_matrix.cc @@ -46,6 +46,37 @@ namespace LinearAlgebra template void SparseMatrix::reinit(const dealii::DynamicSparsityPattern &); + template void + SparseMatrix::vmult(Vector &dst, + const Vector &src) const; + + template void + SparseMatrix::Tvmult(Vector &dst, + const Vector &src) const; + + template void + SparseMatrix::vmult_add(Vector &dst, + const Vector &src) const; + + template void + SparseMatrix::Tvmult_add(Vector &dst, + const Vector &src) const; + + template void + SparseMatrix::vmult(::dealii::Vector &dst, + const ::dealii::Vector &src) const; + + template void + SparseMatrix::Tvmult(::dealii::Vector &dst, + const ::dealii::Vector &src) const; + + template void + SparseMatrix::vmult_add(::dealii::Vector &dst, + const ::dealii::Vector &src) const; + + template void + SparseMatrix::Tvmult_add(::dealii::Vector &dst, + const ::dealii::Vector &src) const; } // namespace TpetraWrappers } // namespace LinearAlgebra # endif // DOXYGEN -- 2.39.5