Skip to content

Commit

Permalink
std_algos: improving min, max, minmax (kokkos#6421)
Browse files Browse the repository at this point in the history
  • Loading branch information
@fnrizzi
fnrizzi committed Sep 19, 2023
1 parent 6f9e50c commit d8fa856
Show file tree
Hide file tree
Showing 3 changed files with 111 additions and 44 deletions.
44 changes: 33 additions & 11 deletions algorithms/unit_tests/TestStdAlgorithmsTeamMaxElement.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,33 +23,40 @@ namespace TeamMaxElement {

namespace KE = Kokkos::Experimental;

template <class ViewType, class DistancesViewType>
template <class ViewType, class DistancesViewType, class IntraTeamSentinelView>
struct TestFunctorA {
ViewType m_view;
DistancesViewType m_distancesView;
IntraTeamSentinelView m_intraTeamSentinelView;
int m_apiPick;

TestFunctorA(const ViewType view, const DistancesViewType distancesView,
int apiPick)
: m_view(view), m_distancesView(distancesView), m_apiPick(apiPick) {}
IntraTeamSentinelView intraTeamSentinelView, int apiPick)
: m_view(view),
m_distancesView(distancesView),
m_intraTeamSentinelView(intraTeamSentinelView),
m_apiPick(apiPick) {}

template <class MemberType>
KOKKOS_INLINE_FUNCTION void operator()(const MemberType& member) const {
const auto myRowIndex = member.league_rank();
auto myRowView = Kokkos::subview(m_view, myRowIndex, Kokkos::ALL());
ptrdiff_t resultDist = 0;

if (m_apiPick == 0) {
auto it =
KE::max_element(member, KE::cbegin(myRowView), KE::cend(myRowView));
resultDist = KE::distance(KE::cbegin(myRowView), it);
Kokkos::single(Kokkos::PerTeam(member), [=, *this]() {
m_distancesView(myRowIndex) = KE::distance(KE::cbegin(myRowView), it);
m_distancesView(myRowIndex) = resultDist;
});
}

else if (m_apiPick == 1) {
auto it = KE::max_element(member, myRowView);
auto it = KE::max_element(member, myRowView);
resultDist = KE::distance(KE::begin(myRowView), it);
Kokkos::single(Kokkos::PerTeam(member), [=, *this]() {
m_distancesView(myRowIndex) = KE::distance(KE::begin(myRowView), it);
m_distancesView(myRowIndex) = resultDist;
});
}
#if not defined KOKKOS_ENABLE_OPENMPTARGET
Expand All @@ -58,20 +65,31 @@ struct TestFunctorA {
auto it =
KE::max_element(member, KE::cbegin(myRowView), KE::cend(myRowView),
CustomLessThanComparator<value_type>{});
resultDist = KE::distance(KE::cbegin(myRowView), it);
Kokkos::single(Kokkos::PerTeam(member), [=, *this]() {
m_distancesView(myRowIndex) = KE::distance(KE::cbegin(myRowView), it);
m_distancesView(myRowIndex) = resultDist;
});
}

else if (m_apiPick == 3) {
using value_type = typename ViewType::value_type;
auto it = KE::max_element(member, myRowView,
CustomLessThanComparator<value_type>{});
resultDist = KE::distance(KE::begin(myRowView), it);
Kokkos::single(Kokkos::PerTeam(member), [=, *this]() {
m_distancesView(myRowIndex) = KE::distance(KE::begin(myRowView), it);
m_distancesView(myRowIndex) = resultDist;
});
}
#endif

// store result of checking if all members have their local
// values matching the one stored in m_distancesView
member.team_barrier();
const bool intraTeamCheck = team_members_have_matching_result(
member, resultDist, m_distancesView(myRowIndex));
Kokkos::single(Kokkos::PerTeam(member), [=, *this]() {
m_intraTeamSentinelView(myRowIndex) = intraTeamCheck;
});
}
};

Expand Down Expand Up @@ -102,18 +120,21 @@ void test_A(std::size_t numTeams, std::size_t numCols, int apiId) {
// beginning of the interval that team operates on and then we check
// that these distances match the expectation
Kokkos::View<std::size_t*> distancesView("distancesView", numTeams);
// sentinel to check if all members of the team compute the same result
Kokkos::View<bool*> intraTeamSentinelView("intraTeamSameResult", numTeams);

// use CTAD for functor
TestFunctorA fnc(dataView, distancesView, apiId);
TestFunctorA fnc(dataView, distancesView, intraTeamSentinelView, apiId);
Kokkos::parallel_for(policy, fnc);

// -----------------------------------------------
// run std algo and check
// -----------------------------------------------
// here I can use cloneOfDataViewBeforeOp_h to run std algo on
// since that contains a valid copy of the data
auto distancesView_h = create_host_space_copy(distancesView);
auto dataViewAfterOp_h = create_host_space_copy(dataView);
auto distancesView_h = create_host_space_copy(distancesView);
auto intraTeamSentinelView_h = create_host_space_copy(intraTeamSentinelView);
auto dataViewAfterOp_h = create_host_space_copy(dataView);
for (std::size_t i = 0; i < cloneOfDataViewBeforeOp_h.extent(0); ++i) {
auto myRow = Kokkos::subview(cloneOfDataViewBeforeOp_h, i, Kokkos::ALL());

Expand All @@ -128,6 +149,7 @@ void test_A(std::size_t numTeams, std::size_t numCols, int apiId) {
}

ASSERT_EQ(stdDistance, distancesView_h(i));
ASSERT_TRUE(intraTeamSentinelView_h(i));
}

// dataView should remain unchanged
Expand Down
44 changes: 33 additions & 11 deletions algorithms/unit_tests/TestStdAlgorithmsTeamMinElement.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,33 +23,40 @@ namespace TeamMinElement {

namespace KE = Kokkos::Experimental;

template <class ViewType, class DistancesViewType>
template <class ViewType, class DistancesViewType, class IntraTeamSentinelView>
struct TestFunctorA {
ViewType m_view;
DistancesViewType m_distancesView;
IntraTeamSentinelView m_intraTeamSentinelView;
int m_apiPick;

TestFunctorA(const ViewType view, const DistancesViewType distancesView,
int apiPick)
: m_view(view), m_distancesView(distancesView), m_apiPick(apiPick) {}
IntraTeamSentinelView intraTeamSentinelView, int apiPick)
: m_view(view),
m_distancesView(distancesView),
m_intraTeamSentinelView(intraTeamSentinelView),
m_apiPick(apiPick) {}

template <class MemberType>
KOKKOS_INLINE_FUNCTION void operator()(const MemberType& member) const {
const auto myRowIndex = member.league_rank();
auto myRowView = Kokkos::subview(m_view, myRowIndex, Kokkos::ALL());
ptrdiff_t resultDist = 0;

if (m_apiPick == 0) {
auto it =
KE::min_element(member, KE::cbegin(myRowView), KE::cend(myRowView));
resultDist = KE::distance(KE::cbegin(myRowView), it);
Kokkos::single(Kokkos::PerTeam(member), [=, *this]() {
m_distancesView(myRowIndex) = KE::distance(KE::cbegin(myRowView), it);
m_distancesView(myRowIndex) = resultDist;
});
}

else if (m_apiPick == 1) {
auto it = KE::min_element(member, myRowView);
auto it = KE::min_element(member, myRowView);
resultDist = KE::distance(KE::begin(myRowView), it);
Kokkos::single(Kokkos::PerTeam(member), [=, *this]() {
m_distancesView(myRowIndex) = KE::distance(KE::begin(myRowView), it);
m_distancesView(myRowIndex) = resultDist;
});
}
#if not defined KOKKOS_ENABLE_OPENMPTARGET
Expand All @@ -58,20 +65,31 @@ struct TestFunctorA {
auto it =
KE::min_element(member, KE::cbegin(myRowView), KE::cend(myRowView),
CustomLessThanComparator<value_type>{});
resultDist = KE::distance(KE::cbegin(myRowView), it);
Kokkos::single(Kokkos::PerTeam(member), [=, *this]() {
m_distancesView(myRowIndex) = KE::distance(KE::cbegin(myRowView), it);
m_distancesView(myRowIndex) = resultDist;
});
}

else if (m_apiPick == 3) {
using value_type = typename ViewType::value_type;
auto it = KE::min_element(member, myRowView,
CustomLessThanComparator<value_type>{});
resultDist = KE::distance(KE::begin(myRowView), it);
Kokkos::single(Kokkos::PerTeam(member), [=, *this]() {
m_distancesView(myRowIndex) = KE::distance(KE::begin(myRowView), it);
m_distancesView(myRowIndex) = resultDist;
});
}
#endif

// store result of checking if all members have their local
// values matching the one stored in m_distancesView
member.team_barrier();
const bool intraTeamCheck = team_members_have_matching_result(
member, resultDist, m_distancesView(myRowIndex));
Kokkos::single(Kokkos::PerTeam(member), [=, *this]() {
m_intraTeamSentinelView(myRowIndex) = intraTeamCheck;
});
}
};

Expand Down Expand Up @@ -102,18 +120,21 @@ void test_A(std::size_t numTeams, std::size_t numCols, int apiId) {
// beginning of the interval that team operates on and then we check
// that these distances match the expectation
Kokkos::View<std::size_t*> distancesView("distancesView", numTeams);
// sentinel to check if all members of the team compute the same result
Kokkos::View<bool*> intraTeamSentinelView("intraTeamSameResult", numTeams);

// use CTAD for functor
TestFunctorA fnc(dataView, distancesView, apiId);
TestFunctorA fnc(dataView, distancesView, intraTeamSentinelView, apiId);
Kokkos::parallel_for(policy, fnc);

// -----------------------------------------------
// run std algo and check
// -----------------------------------------------
// here I can use cloneOfDataViewBeforeOp_h to run std algo on
// since that contains a valid copy of the data
auto distancesView_h = create_host_space_copy(distancesView);
auto dataViewAfterOp_h = create_host_space_copy(dataView);
auto distancesView_h = create_host_space_copy(distancesView);
auto intraTeamSentinelView_h = create_host_space_copy(intraTeamSentinelView);
auto dataViewAfterOp_h = create_host_space_copy(dataView);
for (std::size_t i = 0; i < cloneOfDataViewBeforeOp_h.extent(0); ++i) {
auto myRow = Kokkos::subview(cloneOfDataViewBeforeOp_h, i, Kokkos::ALL());

Expand All @@ -127,6 +148,7 @@ void test_A(std::size_t numTeams, std::size_t numCols, int apiId) {
stdDistance = KE::distance(KE::cbegin(myRow), it);
}
ASSERT_EQ(stdDistance, distancesView_h(i));
ASSERT_TRUE(intraTeamSentinelView_h(i));
}

// dataView should remain unchanged
Expand Down
67 changes: 45 additions & 22 deletions algorithms/unit_tests/TestStdAlgorithmsTeamMinMaxElement.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,41 +23,47 @@ namespace TeamMinMaxElement {

namespace KE = Kokkos::Experimental;

template <class ViewType, class DistancesViewType>
template <class ViewType, class DistancesViewType, class IntraTeamSentinelView>
struct TestFunctorA {
ViewType m_view;
DistancesViewType m_distancesView;
IntraTeamSentinelView m_intraTeamSentinelView;
int m_apiPick;

TestFunctorA(const ViewType view, const DistancesViewType distancesView,
int apiPick)
: m_view(view), m_distancesView(distancesView), m_apiPick(apiPick) {}
IntraTeamSentinelView intraTeamSentinelView, int apiPick)
: m_view(view),
m_distancesView(distancesView),
m_intraTeamSentinelView(intraTeamSentinelView),
m_apiPick(apiPick) {}

template <class MemberType>
KOKKOS_INLINE_FUNCTION void operator()(const MemberType& member) const {
const auto myRowIndex = member.league_rank();
auto myRowView = Kokkos::subview(m_view, myRowIndex, Kokkos::ALL());
ptrdiff_t resultDist1 = 0;
ptrdiff_t resultDist2 = 0;

if (m_apiPick == 0) {
auto itPair = KE::minmax_element(member, KE::cbegin(myRowView),
KE::cend(myRowView));
resultDist1 = KE::distance(KE::cbegin(myRowView), itPair.first);
resultDist2 = KE::distance(KE::cbegin(myRowView), itPair.second);

Kokkos::single(Kokkos::PerTeam(member), [=, *this]() {
m_distancesView(myRowIndex, 0) =
KE::distance(KE::cbegin(myRowView), itPair.first);
m_distancesView(myRowIndex, 1) =
KE::distance(KE::cbegin(myRowView), itPair.second);
m_distancesView(myRowIndex, 0) = resultDist1;
m_distancesView(myRowIndex, 1) = resultDist2;
});
}

else if (m_apiPick == 1) {
auto itPair = KE::minmax_element(member, myRowView);
resultDist1 = KE::distance(KE::begin(myRowView), itPair.first);
resultDist2 = KE::distance(KE::begin(myRowView), itPair.second);

Kokkos::single(Kokkos::PerTeam(member), [=, *this]() {
m_distancesView(myRowIndex, 0) =
KE::distance(KE::begin(myRowView), itPair.first);
m_distancesView(myRowIndex, 1) =
KE::distance(KE::begin(myRowView), itPair.second);
m_distancesView(myRowIndex, 0) = resultDist1;
m_distancesView(myRowIndex, 1) = resultDist2;
});
}
#if not defined KOKKOS_ENABLE_OPENMPTARGET
Expand All @@ -66,26 +72,39 @@ struct TestFunctorA {
auto itPair =
KE::minmax_element(member, KE::cbegin(myRowView), KE::cend(myRowView),
CustomLessThanComparator<value_type>{});
resultDist1 = KE::distance(KE::cbegin(myRowView), itPair.first);
resultDist2 = KE::distance(KE::cbegin(myRowView), itPair.second);

Kokkos::single(Kokkos::PerTeam(member), [=, *this]() {
m_distancesView(myRowIndex, 0) =
KE::distance(KE::cbegin(myRowView), itPair.first);
m_distancesView(myRowIndex, 1) =
KE::distance(KE::cbegin(myRowView), itPair.second);
m_distancesView(myRowIndex, 0) = resultDist1;
m_distancesView(myRowIndex, 1) = resultDist2;
});
}

else if (m_apiPick == 3) {
using value_type = typename ViewType::value_type;
auto itPair = KE::minmax_element(member, myRowView,
CustomLessThanComparator<value_type>{});
resultDist1 = KE::distance(KE::begin(myRowView), itPair.first);
resultDist2 = KE::distance(KE::begin(myRowView), itPair.second);

Kokkos::single(Kokkos::PerTeam(member), [=, *this]() {
m_distancesView(myRowIndex, 0) =
KE::distance(KE::begin(myRowView), itPair.first);
m_distancesView(myRowIndex, 1) =
KE::distance(KE::begin(myRowView), itPair.second);
m_distancesView(myRowIndex, 0) = resultDist1;
m_distancesView(myRowIndex, 1) = resultDist2;
});
}
#endif

// store result of checking if all members have their local
// values matching the one stored in m_distancesView
member.team_barrier();
const bool intraTeamCheck1 = team_members_have_matching_result(
member, resultDist1, m_distancesView(myRowIndex, 0));
const bool intraTeamCheck2 = team_members_have_matching_result(
member, resultDist2, m_distancesView(myRowIndex, 1));
Kokkos::single(Kokkos::PerTeam(member), [=, *this]() {
m_intraTeamSentinelView(myRowIndex) = intraTeamCheck1 && intraTeamCheck2;
});
}
};

Expand Down Expand Up @@ -116,18 +135,21 @@ void test_A(std::size_t numTeams, std::size_t numCols, int apiId) {
// beginning of the interval that team operates on and then we check
// that these distances match the expectation
Kokkos::View<std::size_t**> distancesView("distancesView", numTeams, 2);
// sentinel to check if all members of the team compute the same result
Kokkos::View<bool*> intraTeamSentinelView("intraTeamSameResult", numTeams);

// use CTAD for functor
TestFunctorA fnc(dataView, distancesView, apiId);
TestFunctorA fnc(dataView, distancesView, intraTeamSentinelView, apiId);
Kokkos::parallel_for(policy, fnc);

// -----------------------------------------------
// run std algo and check
// -----------------------------------------------
// here I can use cloneOfDataViewBeforeOp_h to run std algo on
// since that contains a valid copy of the data
auto distancesView_h = create_host_space_copy(distancesView);
auto dataViewAfterOp_h = create_host_space_copy(dataView);
auto distancesView_h = create_host_space_copy(distancesView);
auto intraTeamSentinelView_h = create_host_space_copy(intraTeamSentinelView);
auto dataViewAfterOp_h = create_host_space_copy(dataView);
for (std::size_t i = 0; i < cloneOfDataViewBeforeOp_h.extent(0); ++i) {
auto myRow = Kokkos::subview(cloneOfDataViewBeforeOp_h, i, Kokkos::ALL());

Expand All @@ -145,6 +167,7 @@ void test_A(std::size_t numTeams, std::size_t numCols, int apiId) {

ASSERT_EQ(stdDistance[0], distancesView_h(i, 0));
ASSERT_EQ(stdDistance[1], distancesView_h(i, 1));
ASSERT_TRUE(intraTeamSentinelView_h(i));
}

// dataView should remain unchanged
Expand Down

0 comments on commit d8fa856

Please sign in to comment.