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 25, 2024
1 parent cdee728 commit bdc6d3d
Showing 1 changed file with 22 additions and 19 deletions.
41 changes: 22 additions & 19 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 @@ -632,8 +635,8 @@ void CoalesceDropFactory<Scalar, LocalOrdinal, GlobalOrdinal, Node>::Build(Level
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 implATS::magnitudeType x_aij = static_cast<impl_scalar_type>(implATS::magnitude(static_cast<double>(rowView.value(x) * rowView.value(x))));
typename implATS::magnitudeType y_aij = static_cast<impl_scalar_type>(implATS::magnitude(static_cast<double>(rowView.value(y) * rowView.value(y))));
return x_aij > y_aij;
}
});
Expand All @@ -642,13 +645,13 @@ void CoalesceDropFactory<Scalar, LocalOrdinal, GlobalOrdinal, Node>::Build(Level
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 = static_cast<impl_scalar_type>(implATS::magnitude(static_cast<double>(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 = static_cast<impl_scalar_type>(implATS::magnitude(static_cast<double>(rowView.value(y) * rowView.value(y))));
}

if(x_aij > realThreshold * y_aij) {
Expand All @@ -664,10 +667,10 @@ void CoalesceDropFactory<Scalar, LocalOrdinal, GlobalOrdinal, Node>::Build(Level
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))));
typename implATS::magnitudeType x_aij = static_cast<impl_scalar_type>(implATS::magnitude(static_cast<double>(rowView.value(x) * rowView.value(x))));
typename implATS::magnitudeType y_aij = static_cast<impl_scalar_type>(implATS::magnitude(static_cast<double>(rowView.value(y) * rowView.value(y))));
typename implATS::magnitudeType x_aiiajj = static_cast<impl_scalar_type>(implATS::magnitude(static_cast<double>(threshold * threshold * ghostedDiagValsView(rowView.colidx(x)) * ghostedDiagValsView(row))));
typename implATS::magnitudeType y_aiiajj = static_cast<impl_scalar_type>(implATS::magnitude(static_cast<double>(threshold * threshold * ghostedDiagValsView(rowView.colidx(y)) * ghostedDiagValsView(row))));
return x_aij / x_aiiajj > y_aij / y_aiiajj;
}
});
Expand All @@ -676,16 +679,16 @@ void CoalesceDropFactory<Scalar, LocalOrdinal, GlobalOrdinal, Node>::Build(Level
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 = static_cast<impl_scalar_type>(implATS::magnitude(static_cast<double>(rowView.value(x) * rowView.value(x))));
typename implATS::magnitudeType x_aiiajj = static_cast<impl_scalar_type>(implATS::magnitude(static_cast<double>(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 = static_cast<impl_scalar_type>(implATS::magnitude(static_cast<double>(rowView.value(y) * rowView.value(y))));
typename implATS::magnitudeType y_aiiajj = static_cast<impl_scalar_type>(implATS::magnitude(static_cast<double>(threshold * threshold * ghostedDiagValsView(rowView.colidx(y)) * ghostedDiagValsView(row))));
y_val = y_aij / y_aiiajj;
}

Expand Down

0 comments on commit bdc6d3d

Please sign in to comment.