Skip to content

Commit

Permalink
MueLu: std::complex Replaced With Kokkos::complex
Browse files Browse the repository at this point in the history 
Signed-off-by: Ian Halim <[email protected]>
  • Loading branch information
Ian Halim committed Jul 26, 2024
1 parent cdee728 commit e2fcfd1
Showing 1 changed file with 26 additions and 23 deletions.
49 changes: 26 additions & 23 deletions packages/muelu/src/Graph/MatrixTransformation/MueLu_CoalesceDropFactory_def.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -582,10 +582,13 @@ void CoalesceDropFactory<Scalar, LocalOrdinal, GlobalOrdinal, Node>::Build(Level
using ExecSpace = typename Node::execution_space;
using TeamPol = Kokkos::TeamPolicy<ExecSpace>;
using TeamMem = typename TeamPol::member_type;

using ATS = Kokkos::ArithTraits<Scalar>;
using impl_scalar_type = typename ATS::val_type;
using implATS = Kokkos::ArithTraits<impl_scalar_type>;

//move from host to device
ArrayView<const SC> ghostedDiagValsArrayView = ghostedDiagVals.view(ghostedDiagVals.lowerOffset(), ghostedDiagVals.size());
Kokkos::View<const SC*, ExecSpace> ghostedDiagValsView = Kokkos::Compat::getKokkosViewDeepCopy<ExecSpace>(ghostedDiagValsArrayView);
ArrayView<const impl_scalar_type> ghostedDiagValsArrayView = ghostedDiagVals.view(ghostedDiagVals.lowerOffset(), ghostedDiagVals.size());
Kokkos::View<const impl_scalar_type*, ExecSpace> ghostedDiagValsView = Kokkos::Compat::getKokkosViewDeepCopy<ExecSpace>(ghostedDiagValsArrayView);
auto boundaryNodesDevice = Kokkos::create_mirror_view(ExecSpace(), boundaryNodes);

auto At = Utilities::Op2TpetraCrs(A);
Expand Down Expand Up @@ -627,28 +630,28 @@ void CoalesceDropFactory<Scalar, LocalOrdinal, GlobalOrdinal, Node>::Build(Level
size_t dropStart = n;
if (algorithm == unscaled_cut) {
//push diagonals and boundaries to the right, sort everything else by aij on the left
Kokkos::Experimental::sort_team(teamMember, index_view, [=](size_t& x, size_t& y) {
Kokkos::Experimental::sort_team(teamMember, index_view, [=](size_t& x, size_t& y) -> bool {
if(drop_view(x) || drop_view(y)) {
return drop_view(x) < drop_view(y);
}
else {
typename STS::magnitudeType x_aij = static_cast<SC>(std::fabs(static_cast<double>(rowView.value(x) * rowView.value(x))));
typename STS::magnitudeType y_aij = static_cast<SC>(std::fabs(static_cast<double>(rowView.value(y) * rowView.value(y))));
return x_aij > y_aij;
auto x_aij = implATS::magnitude(rowView.value(x) * rowView.value(x));
auto y_aij = implATS::magnitude(rowView.value(y) * rowView.value(y));
return x_aij > y_aij;
}
});

//find index where dropping starts
Kokkos::parallel_reduce(Kokkos::TeamThreadRange(teamMember, 1, n), [=](size_t i, size_t& min) {
auto const& x = index_view(i - 1);
auto const& y = index_view(i);
typename STS::magnitudeType x_aij = 0;
typename STS::magnitudeType y_aij = 0;
typename implATS::magnitudeType x_aij = 0;
typename implATS::magnitudeType y_aij = 0;
if(!drop_view(x)) {
x_aij = static_cast<SC>(std::fabs(static_cast<double>(rowView.value(x) * rowView.value(x))));
x_aij = implATS::magnitude(rowView.value(x) * rowView.value(x));
}
if(!drop_view(y)) {
y_aij = static_cast<SC>(std::fabs(static_cast<double>(rowView.value(y) * rowView.value(y))));
y_aij = implATS::magnitude(rowView.value(y) * rowView.value(y));
}

if(x_aij > realThreshold * y_aij) {
Expand All @@ -659,33 +662,33 @@ void CoalesceDropFactory<Scalar, LocalOrdinal, GlobalOrdinal, Node>::Build(Level
}, Kokkos::Min<size_t>(dropStart));
} else if (algorithm == scaled_cut) {
//push diagonals and boundaries to the right, sort everything else by aij/aiiajj on the left
Kokkos::Experimental::sort_team(teamMember, index_view, [=](size_t& x, size_t& y) {
Kokkos::Experimental::sort_team(teamMember, index_view, [=](size_t& x, size_t& y) -> bool {
if(drop_view(x) || drop_view(y)) {
return drop_view(x) < drop_view(y);
}
else {
typename STS::magnitudeType x_aij = static_cast<SC>(std::fabs(static_cast<double>(rowView.value(x) * rowView.value(x))));
typename STS::magnitudeType y_aij = static_cast<SC>(std::fabs(static_cast<double>(rowView.value(y) * rowView.value(y))));
typename STS::magnitudeType x_aiiajj = static_cast<SC>(std::fabs(static_cast<double>(threshold * threshold * ghostedDiagValsView(rowView.colidx(x)) * ghostedDiagValsView(row))));
typename STS::magnitudeType y_aiiajj = static_cast<SC>(std::fabs(static_cast<double>(threshold * threshold * ghostedDiagValsView(rowView.colidx(y)) * ghostedDiagValsView(row))));
return x_aij / x_aiiajj > y_aij / y_aiiajj;
auto x_aij = implATS::magnitude(rowView.value(x) * rowView.value(x));
auto y_aij = implATS::magnitude(rowView.value(y) * rowView.value(y));
auto x_aiiajj = implATS::magnitude(threshold * threshold * ghostedDiagValsView(rowView.colidx(x)) * ghostedDiagValsView(row));
auto y_aiiajj = implATS::magnitude(threshold * threshold * ghostedDiagValsView(rowView.colidx(y)) * ghostedDiagValsView(row));
return (x_aij / x_aiiajj) > (y_aij / y_aiiajj);
}
});

//find index where dropping starts
Kokkos::parallel_reduce(Kokkos::TeamThreadRange(teamMember, 1, n), [=](size_t i, size_t& min) {
auto const& x = index_view(i - 1);
auto const& y = index_view(i);
typename STS::magnitudeType x_val = 0;
typename STS::magnitudeType y_val = 0;
typename implATS::magnitudeType x_val = 0;
typename implATS::magnitudeType y_val = 0;
if(!drop_view(x)) {
typename STS::magnitudeType x_aij = static_cast<SC>(std::fabs(static_cast<double>(rowView.value(x) * rowView.value(x))));
typename STS::magnitudeType x_aiiajj = static_cast<SC>(std::fabs(static_cast<double>(threshold * threshold * ghostedDiagValsView(rowView.colidx(x)) * ghostedDiagValsView(row))));
typename implATS::magnitudeType x_aij = implATS::magnitude(rowView.value(x) * rowView.value(x));
typename implATS::magnitudeType x_aiiajj = implATS::magnitude(threshold * threshold * ghostedDiagValsView(rowView.colidx(x)) * ghostedDiagValsView(row));
x_val = x_aij / x_aiiajj;
}
if(!drop_view(y)) {
typename STS::magnitudeType y_aij = static_cast<SC>(std::fabs(static_cast<double>(rowView.value(y) * rowView.value(y))));
typename STS::magnitudeType y_aiiajj = static_cast<SC>(std::fabs(static_cast<double>(threshold * threshold * ghostedDiagValsView(rowView.colidx(y)) * ghostedDiagValsView(row))));
typename implATS::magnitudeType y_aij = implATS::magnitude(rowView.value(y) * rowView.value(y));
typename implATS::magnitudeType y_aiiajj = implATS::magnitude(threshold * threshold * ghostedDiagValsView(rowView.colidx(y)) * ghostedDiagValsView(row));
y_val = y_aij / y_aiiajj;
}

Expand Down

0 comments on commit e2fcfd1

Please sign in to comment.