diff --git a/packages/muelu/src/Graph/MatrixTransformation/MueLu_ClassicalDropping.hpp b/packages/muelu/src/Graph/MatrixTransformation/MueLu_ClassicalDropping.hpp index 29676c46355d..8ecf72c2aa9c 100644 --- a/packages/muelu/src/Graph/MatrixTransformation/MueLu_ClassicalDropping.hpp +++ b/packages/muelu/src/Graph/MatrixTransformation/MueLu_ClassicalDropping.hpp @@ -13,32 +13,43 @@ #include "MueLu_DroppingCommon.hpp" #include "Kokkos_Core.hpp" #include "Kokkos_ArithTraits.hpp" +#include "Xpetra_Matrix.hpp" +#include "MueLu_Utilities.hpp" namespace MueLu::ClassicalDropping { -template +template class AbsDropFunctor { private: + using matrix_type = Xpetra::Matrix; + using diag_vec_type = Xpetra::MultiVector; + using local_matrix_type = typename matrix_type::local_matrix_type; using scalar_type = typename local_matrix_type::value_type; using local_ordinal_type = typename local_matrix_type::ordinal_type; using memory_space = typename local_matrix_type::memory_space; - using results_view = Kokkos::View; + using diag_view_type = typename Kokkos::DualView::t_dev; + + using results_view = Kokkos::View; using ATS = Kokkos::ArithTraits; using magnitudeType = typename ATS::magnitudeType; using boundary_nodes_view = Kokkos::View; local_matrix_type A; + Teuchos::RCP diagVec; diag_view_type diag; // corresponds to overlapped diagonal magnitudeType eps; results_view results; public: - AbsDropFunctor(local_matrix_type& A_, magnitudeType threshold, diag_view_type diag_, results_view& results_) - : A(A_) - , diag(diag_) + AbsDropFunctor(matrix_type& A_, magnitudeType threshold, results_view& results_) + : A(A_.getLocalMatrixDevice()) , eps(threshold) - , results(results_) {} + , results(results_) { + diagVec = Utilities::GetMatrixOverlappedDiagonal(A_); + auto lclDiag2d = diagVec->getDeviceLocalView(Xpetra::Access::ReadOnly); + diag = Kokkos::subview(lclDiag2d, Kokkos::ALL(), 0); + } KOKKOS_FORCEINLINE_FUNCTION bool operator()(const local_ordinal_type rlid) const { @@ -58,29 +69,39 @@ class AbsDropFunctor { } }; -template +template class SignedClassicalRSDropFunctor { private: + using matrix_type = Xpetra::Matrix; + using local_matrix_type = typename matrix_type::local_matrix_type; using scalar_type = typename local_matrix_type::value_type; using local_ordinal_type = typename local_matrix_type::ordinal_type; using memory_space = typename local_matrix_type::memory_space; - using results_view = Kokkos::View; + + using results_view = Kokkos::View; using ATS = Kokkos::ArithTraits; using magnitudeType = typename ATS::magnitudeType; using boundary_nodes_view = Kokkos::View; + using diag_vec_type = Xpetra::MultiVector; + using diag_view_type = typename Kokkos::DualView::t_dev; + local_matrix_type A; + Teuchos::RCP diagVec; diag_view_type diag; // corresponds to overlapped diagonal magnitudeType eps; results_view results; public: - SignedClassicalRSDropFunctor(local_matrix_type& A_, magnitudeType threshold, diag_view_type diag_, results_view& results_) - : A(A_) - , diag(diag_) + SignedClassicalRSDropFunctor(matrix_type& A_, magnitudeType threshold, results_view& results_) + : A(A_.getLocalMatrixDevice()) , eps(threshold) - , results(results_) {} + , results(results_) { + diagVec = Utilities::GetMatrixMaxMinusOffDiagonal(A_); + auto lclDiag2d = diagVec->getDeviceLocalView(Xpetra::Access::ReadOnly); + diag = Kokkos::subview(lclDiag2d, Kokkos::ALL(), 0); + } KOKKOS_FORCEINLINE_FUNCTION bool operator()(const local_ordinal_type rlid) const { @@ -97,13 +118,18 @@ class SignedClassicalRSDropFunctor { } }; -template +template class SignedClassicalSADropFunctor { private: + using matrix_type = Xpetra::Matrix; + using diag_vec_type = Xpetra::MultiVector; + using local_matrix_type = typename matrix_type::local_matrix_type; using scalar_type = typename local_matrix_type::value_type; using local_ordinal_type = typename local_matrix_type::ordinal_type; using memory_space = typename local_matrix_type::memory_space; - using results_view = Kokkos::View; + using diag_view_type = typename Kokkos::DualView::t_dev; + + using results_view = Kokkos::View; using ATS = Kokkos::ArithTraits; using magnitudeType = typename ATS::magnitudeType; @@ -111,16 +137,21 @@ class SignedClassicalSADropFunctor { using boundary_nodes_view = Kokkos::View; local_matrix_type A; + Teuchos::RCP diagVec; diag_view_type diag; // corresponds to overlapped diagonal magnitudeType eps; results_view results; public: - SignedClassicalSADropFunctor(local_matrix_type& A_, magnitudeType threshold, diag_view_type diag_, results_view& results_) - : A(A_) - , diag(diag_) + SignedClassicalSADropFunctor(matrix_type& A_, magnitudeType threshold, results_view& results_) + : A(A_.getLocalMatrixDevice()) , eps(threshold) - , results(results_) {} + , results(results_) { + // Construct ghosted matrix diagonal + diagVec = Utilities::GetMatrixOverlappedDiagonal(A_); + auto lclDiag2d = diagVec->getDeviceLocalView(Xpetra::Access::ReadOnly); + diag = Kokkos::subview(lclDiag2d, Kokkos::ALL(), 0); + } KOKKOS_FORCEINLINE_FUNCTION bool operator()(const local_ordinal_type rlid) const { diff --git a/packages/muelu/src/Graph/MatrixTransformation/MueLu_CoalesceDropFactory_kokkos_def.hpp b/packages/muelu/src/Graph/MatrixTransformation/MueLu_CoalesceDropFactory_kokkos_def.hpp index e6d3868672ff..6ee03b4d50df 100644 --- a/packages/muelu/src/Graph/MatrixTransformation/MueLu_CoalesceDropFactory_kokkos_def.hpp +++ b/packages/muelu/src/Graph/MatrixTransformation/MueLu_CoalesceDropFactory_kokkos_def.hpp @@ -14,7 +14,6 @@ #include #include -#include "MueLu_DroppingCommon.hpp" #include "Xpetra_Matrix.hpp" #include "MueLu_CoalesceDropFactory_kokkos_decl.hpp" @@ -25,7 +24,6 @@ #include "MueLu_LWGraph_kokkos.hpp" #include "MueLu_MasterList.hpp" #include "MueLu_Monitor.hpp" -#include "MueLu_Utilities.hpp" // #define MUELU_COALESCE_DROP_DEBUG 1 @@ -36,10 +34,6 @@ #include "MueLu_DistanceLaplacianDropping.hpp" #include "MueLu_MatrixConstruction.hpp" -#ifdef HAVE_XPETRA_TPETRA -#include "Tpetra_CrsGraphTransposer.hpp" -#endif - namespace MueLu { template @@ -338,18 +332,12 @@ std::tuplegetDeviceLocalView(Xpetra::Access::ReadOnly); - auto lclDiag1d = Kokkos::subview(lclDiag2d, Kokkos::ALL(), 0); - if (algo == "block diagonal classical") { auto BlockNumbers = GetBlockNumberMVs(currentLevel); - auto block_diagonalize = Misc::BlockDiagonalizeFunctor(lclA, - std::get<0>(BlockNumbers)->getDeviceLocalView(Xpetra::Access::ReadOnly), - std::get<1>(BlockNumbers)->getDeviceLocalView(Xpetra::Access::ReadOnly), results); + auto block_diagonalize = Misc::BlockDiagonalizeFunctor(*A, *std::get<0>(BlockNumbers), *std::get<1>(BlockNumbers), results); + if (classicalAlgoStr == "default") { - auto classical_dropping = ClassicalDropping::AbsDropFunctor(lclA, threshold, lclDiag1d, results); + auto classical_dropping = ClassicalDropping::AbsDropFunctor(*A, threshold, results); if (aggregationMayCreateDirichlet) { runCountingFunctor(block_diagonalize, @@ -364,24 +352,24 @@ std::tuplegetDeviceLocalView(Xpetra::Access::ReadOnly); - auto lclDiag1d = Kokkos::subview(lclDiag2d, Kokkos::ALL(), 0); - - auto signed_classical_rs_dropping = ClassicalDropping::SignedClassicalRSDropFunctor(lclA, threshold, lclDiag1d, results); + auto signed_classical_rs_dropping = ClassicalDropping::SignedClassicalRSDropFunctor(*A, threshold, results); if (algo == "block diagonal signed classical" || algo == "block diagonal colored signed classical") { auto BlockNumbers = GetBlockNumberMVs(currentLevel); - auto block_diagonalize = Misc::BlockDiagonalizeFunctor(lclA, - std::get<0>(BlockNumbers)->getDeviceLocalView(Xpetra::Access::ReadOnly), - std::get<1>(BlockNumbers)->getDeviceLocalView(Xpetra::Access::ReadOnly), results); + auto block_diagonalize = Misc::BlockDiagonalizeFunctor(*A, *std::get<0>(BlockNumbers), *std::get<1>(BlockNumbers), results); if (aggregationMayCreateDirichlet) { runCountingFunctor(block_diagonalize, @@ -475,12 +455,7 @@ std::tuplegetDeviceLocalView(Xpetra::Access::ReadOnly); - auto lclDiag1d = Kokkos::subview(lclDiag2d, Kokkos::ALL(), 0); - - auto signed_classical_sa_dropping = ClassicalDropping::SignedClassicalSADropFunctor(lclA, threshold, lclDiag1d, results); + auto signed_classical_sa_dropping = ClassicalDropping::SignedClassicalSADropFunctor(*A, threshold, results); if (aggregationMayCreateDirichlet) { runCountingFunctor(signed_classical_sa_dropping, @@ -496,29 +471,15 @@ std::tuple::magnitudeType, LO, GO, NO>; auto coords = Get>(currentLevel, "Coordinates"); - RCP ghostedCoords; - auto importer = A->getCrsGraph()->getImporter(); - if (!importer.is_null()) { - ghostedCoords = Xpetra::MultiVectorFactory::magnitudeType, LO, GO, NO>::Build(importer->getTargetMap(), coords->getNumVectors()); - ghostedCoords->doImport(*coords, *importer, Xpetra::INSERT); - } else - ghostedCoords = coords; - auto lclCoords = coords->getDeviceLocalView(Xpetra::Access::ReadOnly); - auto lclGhostedCoordinates = ghostedCoords->getDeviceLocalView(Xpetra::Access::ReadOnly); - auto dist2 = DistanceLaplacian::DistanceFunctor(lclCoords, lclGhostedCoordinates); - // Construct ghosted distance Laplacian diagonal - auto diag = DistanceLaplacian::getDiagonal(A, dist2); - auto lclDiag2d = diag->getDeviceLocalView(Xpetra::Access::ReadOnly); - auto lclDiag1d = Kokkos::subview(lclDiag2d, Kokkos::ALL(), 0); + + auto dist2 = DistanceLaplacian::DistanceFunctor(*A, *coords); if (algo == "block diagonal distance laplacian") { auto BlockNumbers = GetBlockNumberMVs(currentLevel); - auto block_diagonalize = Misc::BlockDiagonalizeFunctor(lclA, - std::get<0>(BlockNumbers)->getDeviceLocalView(Xpetra::Access::ReadOnly), - std::get<1>(BlockNumbers)->getDeviceLocalView(Xpetra::Access::ReadOnly), results); + auto block_diagonalize = Misc::BlockDiagonalizeFunctor(*A, *std::get<0>(BlockNumbers), *std::get<1>(BlockNumbers), results); if (distanceLaplacianAlgoStr == "default") { - auto dist_laplacian_dropping = DistanceLaplacian::DropFunctor(lclA, threshold, lclDiag1d, dist2, results); + auto dist_laplacian_dropping = DistanceLaplacian::DropFunctor(*A, threshold, dist2, results); if (aggregationMayCreateDirichlet) { runCountingFunctor(block_diagonalize, @@ -535,16 +496,16 @@ std::tuple(BlockNumbers)->getDeviceLocalView(Xpetra::Access::ReadOnly), - std::get<1>(BlockNumbers)->getDeviceLocalView(Xpetra::Access::ReadOnly), results); + auto block_diagonalize = Misc::BlockDiagonalizeFunctor(*A, *std::get<0>(BlockNumbers), *std::get<1>(BlockNumbers), results); runCountingFunctor(block_diagonalize); } else { @@ -919,13 +878,8 @@ std::tuplegetDeviceLocalView(Xpetra::Access::ReadOnly); - auto lclDiag1d = Kokkos::subview(lclDiag2d, Kokkos::ALL(), 0); - if (classicalAlgoStr == "default") { - auto classical_dropping = ClassicalDropping::AbsDropFunctor(lclA, threshold, lclDiag1d, results); + auto classical_dropping = ClassicalDropping::AbsDropFunctor(*A, threshold, results); if (aggregationMayCreateDirichlet) { runCountingFunctor(classical_dropping, @@ -947,13 +901,7 @@ std::tuplegetDeviceLocalView(Xpetra::Access::ReadOnly); - auto lclDiag1d = Kokkos::subview(lclDiag2d, Kokkos::ALL(), 0); - - auto signed_classical_rs_dropping = ClassicalDropping::SignedClassicalRSDropFunctor(lclA, threshold, lclDiag1d, results); + auto signed_classical_rs_dropping = ClassicalDropping::SignedClassicalRSDropFunctor(*A, threshold, results); if (aggregationMayCreateDirichlet) { runCountingFunctor(signed_classical_rs_dropping, @@ -967,12 +915,7 @@ std::tuplegetDeviceLocalView(Xpetra::Access::ReadOnly); - auto lclDiag1d = Kokkos::subview(lclDiag2d, Kokkos::ALL(), 0); - - auto signed_classical_sa_dropping = ClassicalDropping::SignedClassicalSADropFunctor(lclA, threshold, lclDiag1d, results); + auto signed_classical_sa_dropping = ClassicalDropping::SignedClassicalSADropFunctor(*A, threshold, results); if (aggregationMayCreateDirichlet) { runCountingFunctor(signed_classical_sa_dropping, @@ -988,23 +931,11 @@ std::tuple::magnitudeType, LO, GO, NO>; auto coords = Get>(currentLevel, "Coordinates"); - RCP ghostedCoords; - auto importer = A->getCrsGraph()->getImporter(); - if (!importer.is_null()) { - ghostedCoords = Xpetra::MultiVectorFactory::magnitudeType, LO, GO, NO>::Build(importer->getTargetMap(), coords->getNumVectors()); - ghostedCoords->doImport(*coords, *importer, Xpetra::INSERT); - } else - ghostedCoords = coords; - auto lclCoords = coords->getDeviceLocalView(Xpetra::Access::ReadOnly); - auto lclGhostedCoordinates = ghostedCoords->getDeviceLocalView(Xpetra::Access::ReadOnly); - auto dist2 = DistanceLaplacian::DistanceFunctor(lclCoords, lclGhostedCoordinates); - // Construct ghosted distance Laplacian diagonal - auto diag = DistanceLaplacian::getDiagonal(A, dist2); - auto lclDiag2d = diag->getDeviceLocalView(Xpetra::Access::ReadOnly); - auto lclDiag1d = Kokkos::subview(lclDiag2d, Kokkos::ALL(), 0); + + auto dist2 = DistanceLaplacian::DistanceFunctor(*A, *coords); if (distanceLaplacianAlgoStr == "default") { - auto dist_laplacian_dropping = DistanceLaplacian::DropFunctor(lclA, threshold, lclDiag1d, dist2, results); + auto dist_laplacian_dropping = DistanceLaplacian::DropFunctor(*A, threshold, dist2, results); if (aggregationMayCreateDirichlet) { runCountingFunctor(dist_laplacian_dropping, diff --git a/packages/muelu/src/Graph/MatrixTransformation/MueLu_CutDrop.hpp b/packages/muelu/src/Graph/MatrixTransformation/MueLu_CutDrop.hpp index 758a23213721..aa2c93d97e35 100644 --- a/packages/muelu/src/Graph/MatrixTransformation/MueLu_CutDrop.hpp +++ b/packages/muelu/src/Graph/MatrixTransformation/MueLu_CutDrop.hpp @@ -14,6 +14,10 @@ #include "Kokkos_ArithTraits.hpp" #include "MueLu_DroppingCommon.hpp" #include "Kokkos_Sort.hpp" +#include "MueLu_Utilities.hpp" +#include "Xpetra_Matrix.hpp" +#include "Xpetra_MultiVector.hpp" +#include "MueLu_DistanceLaplacianDropping.hpp" namespace MueLu::CutDrop { @@ -22,23 +26,27 @@ enum decisionAlgoType { defaultAlgo, scaled_cut, scaled_cut_symmetric }; -template +template class UnscaledComparison { - private: + public: + using matrix_type = Xpetra::Matrix; + + using local_matrix_type = typename matrix_type::local_matrix_type; using scalar_type = typename local_matrix_type::value_type; using local_ordinal_type = typename local_matrix_type::ordinal_type; using memory_space = typename local_matrix_type::memory_space; using results_view = Kokkos::View; - using ATS = Kokkos::ArithTraits; - using magnitudeType = typename ATS::magnitudeType; - local_matrix_type A; results_view results; + private: + using ATS = Kokkos::ArithTraits; + using magnitudeType = typename ATS::magnitudeType; + public: - UnscaledComparison(local_matrix_type& A_, results_view& results_) - : A(A_) + UnscaledComparison(matrix_type& A_, results_view& results_) + : A(A_.getLocalMatrixDevice()) , results(results_) {} template @@ -95,27 +103,37 @@ class UnscaledComparison { } }; -template +template class ScaledComparison { - private: + public: + using matrix_type = Xpetra::Matrix; + using local_matrix_type = typename matrix_type::local_matrix_type; using scalar_type = typename local_matrix_type::value_type; using local_ordinal_type = typename local_matrix_type::ordinal_type; using memory_space = typename local_matrix_type::memory_space; + using diag_vec_type = Xpetra::MultiVector; + using diag_view_type = typename Kokkos::DualView::t_dev; using results_view = Kokkos::View; + local_matrix_type A; + results_view results; + + private: using ATS = Kokkos::ArithTraits; using magnitudeType = typename ATS::magnitudeType; - local_matrix_type A; magnitudeType eps; + Teuchos::RCP diagVec; diag_view_type diag; - results_view results; public: - ScaledComparison(local_matrix_type& A_, diag_view_type& diag_, results_view& results_) - : A(A_) - , diag(diag_) - , results(results_) {} + ScaledComparison(matrix_type& A_, results_view& results_) + : A(A_.getLocalMatrixDevice()) + , results(results_) { + diagVec = Utilities::GetMatrixOverlappedDiagonal(A_); + auto lclDiag2d = diagVec->getDeviceLocalView(Xpetra::Access::ReadOnly); + diag = Kokkos::subview(lclDiag2d, Kokkos::ALL(), 0); + } template struct Comparator { @@ -177,29 +195,40 @@ class ScaledComparison { } }; -template +template class ScaledDistanceLaplacianComparison { - private: + public: + using matrix_type = Xpetra::Matrix; + using local_matrix_type = typename matrix_type::local_matrix_type; using scalar_type = typename local_matrix_type::value_type; using local_ordinal_type = typename local_matrix_type::ordinal_type; using memory_space = typename local_matrix_type::memory_space; + using diag_vec_type = Xpetra::MultiVector; + using diag_view_type = typename Kokkos::DualView::t_dev; using results_view = Kokkos::View; + local_matrix_type A; + results_view results; + + private: using ATS = Kokkos::ArithTraits; using magnitudeType = typename ATS::magnitudeType; - local_matrix_type A; magnitudeType eps; + Teuchos::RCP diagVec; diag_view_type diag; DistanceFunctorType dist2; - results_view results; public: - ScaledDistanceLaplacianComparison(local_matrix_type& A_, diag_view_type& diag_, DistanceFunctorType& dist2_, results_view& results_) - : A(A_) - , diag(diag_) - , dist2(dist2_) - , results(results_) {} + ScaledDistanceLaplacianComparison(matrix_type& A_, DistanceFunctorType& dist2_, results_view& results_) + : A(A_.getLocalMatrixDevice()) + , results(results_) + , dist2(dist2_) { + // Construct ghosted distance Laplacian diagonal + diagVec = DistanceLaplacian::getDiagonal(A_, dist2); + auto lclDiag2d = diagVec->getDeviceLocalView(Xpetra::Access::ReadOnly); + diag = Kokkos::subview(lclDiag2d, Kokkos::ALL(), 0); + } template struct Comparator { @@ -213,8 +242,8 @@ class ScaledDistanceLaplacianComparison { using magnitudeType = typename ATS::magnitudeType; const local_matrix_type2 A; - const DistanceFunctorType2 dist2; const diag_view_type2 diag; + const DistanceFunctorType2 dist2; const local_ordinal_type rlid; const local_ordinal_type offset; const results_view results; @@ -310,9 +339,10 @@ class PermutationFunctor { } }; -template +template class CutDropFunctor { private: + using local_matrix_type = typename comparison_type::local_matrix_type; using scalar_type = typename local_matrix_type::value_type; using local_ordinal_type = typename local_matrix_type::ordinal_type; using memory_space = typename local_matrix_type::memory_space; @@ -329,11 +359,11 @@ class CutDropFunctor { Kokkos::View index; public: - CutDropFunctor(local_matrix_type& A_, comparison_type& comparison_, magnitudeType threshold, results_view& results_) - : A(A_) + CutDropFunctor(comparison_type& comparison_, magnitudeType threshold) + : A(comparison_.A) , comparison(comparison_) , eps(threshold) - , results(results_) { + , results(comparison_.results) { index = Kokkos::View("indices", A.nnz()); using TeamPol = Kokkos::TeamPolicy; diff --git a/packages/muelu/src/Graph/MatrixTransformation/MueLu_DistanceLaplacianDropping.hpp b/packages/muelu/src/Graph/MatrixTransformation/MueLu_DistanceLaplacianDropping.hpp index bf164bf5defb..49b6740d1d7e 100644 --- a/packages/muelu/src/Graph/MatrixTransformation/MueLu_DistanceLaplacianDropping.hpp +++ b/packages/muelu/src/Graph/MatrixTransformation/MueLu_DistanceLaplacianDropping.hpp @@ -20,20 +20,33 @@ namespace MueLu::DistanceLaplacian { -template +template class DistanceFunctor { private: - using scalar_type = typename CoordsType::non_const_value_type; - using ATS = Kokkos::ArithTraits; - using magnitudeType = typename ATS::magnitudeType; + using matrix_type = Xpetra::Matrix; + using local_matrix_type = typename matrix_type::local_matrix_type; + using scalar_type = typename local_matrix_type::value_type; + using local_ordinal_type = LocalOrdinal; + using ATS = Kokkos::ArithTraits; + using magnitudeType = typename ATS::magnitudeType; + using coords_type = Xpetra::MultiVector; + using local_coords_type = typename coords_type::dual_view_type_const::t_dev; - CoordsType coords; - CoordsType ghostedCoords; + local_coords_type coords; + local_coords_type ghostedCoords; public: - DistanceFunctor(CoordsType coords_, CoordsType ghostedCoords_) - : coords(coords_) - , ghostedCoords(ghostedCoords_) {} + DistanceFunctor(matrix_type& A, coords_type& coords_) { + Teuchos::RCP ghosted_coords; + auto importer = A.getCrsGraph()->getImporter(); + if (!importer.is_null()) { + ghosted_coords = Xpetra::MultiVectorFactory::Build(importer->getTargetMap(), coords_.getNumVectors()); + ghosted_coords->doImport(coords_, *importer, Xpetra::INSERT); + } else + ghosted_coords = Teuchos::rcpFromRef(coords_); + coords = coords_.getDeviceLocalView(Xpetra::Access::ReadOnly); + ghostedCoords = ghosted_coords->getDeviceLocalView(Xpetra::Access::ReadOnly); + } KOKKOS_INLINE_FUNCTION magnitudeType distance2(const local_ordinal_type row, const local_ordinal_type col) const { @@ -49,7 +62,7 @@ class DistanceFunctor { template Teuchos::RCP > -getDiagonal(const Teuchos::RCP >& A, +getDiagonal(Xpetra::Matrix& A, DistanceFunctorType& distFunctor) { using scalar_type = Scalar; using local_ordinal_type = LocalOrdinal; @@ -59,9 +72,9 @@ getDiagonal(const Teuchos::RCP; - auto diag = Xpetra::MultiVectorFactory::Build(A->getRowMap(), 1); + auto diag = Xpetra::MultiVectorFactory::Build(A.getRowMap(), 1); { - auto lclA = A->getLocalMatrixDevice(); + auto lclA = A.getLocalMatrixDevice(); auto lclDiag = diag->getDeviceLocalView(Xpetra::Access::OverwriteAll); Kokkos::parallel_for( @@ -80,9 +93,9 @@ getDiagonal(const Teuchos::RCPgetCrsGraph()->getImporter(); + auto importer = A.getCrsGraph()->getImporter(); if (!importer.is_null()) { - auto ghostedDiag = Xpetra::MultiVectorFactory::Build(A->getColMap(), 1); + auto ghostedDiag = Xpetra::MultiVectorFactory::Build(A.getColMap(), 1); ghostedDiag->doImport(*diag, *importer, Xpetra::INSERT); return ghostedDiag; } else { @@ -90,13 +103,18 @@ getDiagonal(const Teuchos::RCP +template class DropFunctor { private: + using matrix_type = Xpetra::Matrix; + using local_matrix_type = typename matrix_type::local_matrix_type; using scalar_type = typename local_matrix_type::value_type; using local_ordinal_type = typename local_matrix_type::ordinal_type; using memory_space = typename local_matrix_type::memory_space; - using results_view = Kokkos::View; + using diag_vec_type = Xpetra::MultiVector; + using diag_view_type = typename Kokkos::DualView::t_dev; + + using results_view = Kokkos::View; using ATS = Kokkos::ArithTraits; using magnitudeType = typename ATS::magnitudeType; @@ -104,18 +122,22 @@ class DropFunctor { local_matrix_type A; magnitudeType eps; + Teuchos::RCP diagVec; diag_view_type diag; // corresponds to overlapped diagonal DistanceFunctorType dist2; results_view results; const scalar_type one = ATS::one(); public: - DropFunctor(local_matrix_type& A_, magnitudeType threshold, diag_view_type diag_, DistanceFunctorType& dist2_, results_view& results_) - : A(A_) + DropFunctor(matrix_type& A_, magnitudeType threshold, DistanceFunctorType& dist2_, results_view& results_) + : A(A_.getLocalMatrixDevice()) , eps(threshold) - , diag(diag_) , dist2(dist2_) - , results(results_) {} + , results(results_) { + diagVec = getDiagonal(A_, dist2); + auto lclDiag2d = diagVec->getDeviceLocalView(Xpetra::Access::ReadOnly); + diag = Kokkos::subview(lclDiag2d, Kokkos::ALL(), 0); + } KOKKOS_FORCEINLINE_FUNCTION bool operator()(local_ordinal_type rlid) const { diff --git a/packages/muelu/src/Graph/MatrixTransformation/MueLu_DroppingCommon.hpp b/packages/muelu/src/Graph/MatrixTransformation/MueLu_DroppingCommon.hpp index d4dd221958e2..431b434f4bca 100644 --- a/packages/muelu/src/Graph/MatrixTransformation/MueLu_DroppingCommon.hpp +++ b/packages/muelu/src/Graph/MatrixTransformation/MueLu_DroppingCommon.hpp @@ -12,6 +12,8 @@ #include "Kokkos_Core.hpp" #include "Kokkos_ArithTraits.hpp" +#include "Xpetra_Access.hpp" +#include "Xpetra_Matrix.hpp" namespace MueLu { @@ -255,24 +257,30 @@ class MarkSingletonVectorFunctor { } }; -template +template class BlockDiagonalizeFunctor { private: + using matrix_type = Xpetra::Matrix; + using local_matrix_type = typename matrix_type::local_matrix_type; + using scalar_type = typename local_matrix_type::value_type; using local_ordinal_type = typename local_matrix_type::ordinal_type; using memory_space = typename local_matrix_type::memory_space; using results_view = Kokkos::View; + using block_indices_type = Xpetra::MultiVector; + using local_block_indices_view_type = typename block_indices_type::dual_view_type_const::t_dev; + local_matrix_type A; - block_indices_view_type point_to_block; - block_indices_view_type ghosted_point_to_block; + local_block_indices_view_type point_to_block; + local_block_indices_view_type ghosted_point_to_block; results_view results; public: - BlockDiagonalizeFunctor(local_matrix_type& A_, block_indices_view_type point_to_block_, block_indices_view_type ghosted_point_to_block_, results_view& results_) - : A(A_) - , point_to_block(point_to_block_) - , ghosted_point_to_block(ghosted_point_to_block_) + BlockDiagonalizeFunctor(matrix_type& A_, block_indices_type& point_to_block_, block_indices_type& ghosted_point_to_block_, results_view& results_) + : A(A_.getLocalMatrixDevice()) + , point_to_block(point_to_block_.getDeviceLocalView(Xpetra::Access::ReadOnly)) + , ghosted_point_to_block(ghosted_point_to_block_.getDeviceLocalView(Xpetra::Access::ReadOnly)) , results(results_) {} KOKKOS_FORCEINLINE_FUNCTION