1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178
|
//@HEADER
// ************************************************************************
//
// Kokkos v. 4.0
// Copyright (2022) National Technology & Engineering
// Solutions of Sandia, LLC (NTESS).
//
// Under the terms of Contract DE-NA0003525 with NTESS,
// the U.S. Government retains certain rights in this software.
//
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
// See https://kokkos.org/LICENSE for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//@HEADER
#include <TestStdAlgorithmsCommon.hpp>
#include <algorithm>
namespace Test {
namespace stdalgos {
namespace TeamRemoveCopyIf {
namespace KE = Kokkos::Experimental;
template <class ValueType>
struct GreaterThanValueFunctor {
ValueType m_val;
KOKKOS_INLINE_FUNCTION
GreaterThanValueFunctor(ValueType val) : m_val(val) {}
KOKKOS_INLINE_FUNCTION
bool operator()(ValueType val) const { return (val > m_val); }
};
template <class SourceViewType, class DestViewType, class DistancesViewType,
class IntraTeamSentinelView, class ValueType>
struct TestFunctorA {
SourceViewType m_sourceView;
DestViewType m_destView;
ValueType m_threshold;
DistancesViewType m_distancesView;
IntraTeamSentinelView m_intraTeamSentinelView;
int m_apiPick;
TestFunctorA(const SourceViewType sourceView, const DestViewType destView,
ValueType threshold, const DistancesViewType distancesView,
const IntraTeamSentinelView intraTeamSentinelView, int apiPick)
: m_sourceView(sourceView),
m_destView(destView),
m_threshold(threshold),
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 myRowViewFrom =
Kokkos::subview(m_sourceView, myRowIndex, Kokkos::ALL());
auto myRowViewDest = Kokkos::subview(m_destView, myRowIndex, Kokkos::ALL());
ptrdiff_t resultDist = 0;
GreaterThanValueFunctor predicate(m_threshold);
if (m_apiPick == 0) {
auto it = KE::remove_copy_if(member, KE::cbegin(myRowViewFrom),
KE::cend(myRowViewFrom),
KE::begin(myRowViewDest), predicate);
resultDist = KE::distance(KE::begin(myRowViewDest), it);
Kokkos::single(Kokkos::PerTeam(member), [=, *this]() {
m_distancesView(myRowIndex) = resultDist;
});
} else if (m_apiPick == 1) {
auto it =
KE::remove_copy_if(member, myRowViewFrom, myRowViewDest, predicate);
resultDist = KE::distance(KE::begin(myRowViewDest), it);
Kokkos::single(Kokkos::PerTeam(member), [=, *this]() {
m_distancesView(myRowIndex) = resultDist;
});
}
// 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;
});
}
};
template <class LayoutTag, class ValueType>
void test_A(std::size_t numTeams, std::size_t numCols, int apiId) {
/* description:
use a rank-2 view randomly filled,
and run a team-level remove_copy_if where the values copied
are those strictly greater than a threshold.
*/
const auto threshold = static_cast<ValueType>(531);
// -----------------------------------------------
// prepare data
// -----------------------------------------------
auto [sourceView, cloneOfSourceViewBeforeOp_h] =
create_random_view_and_host_clone(
LayoutTag{}, numTeams, numCols,
Kokkos::pair<ValueType, ValueType>{5, 523}, "sourceView");
// -----------------------------------------------
// launch kokkos kernel
// -----------------------------------------------
using space_t = Kokkos::DefaultExecutionSpace;
Kokkos::TeamPolicy<space_t> policy(numTeams, Kokkos::AUTO());
// create the destination view
Kokkos::View<ValueType**> destView("destView", numTeams, numCols);
// each team stores the distance of the returned iterator from the
// beginning of the interval that team operates on and then we check
// that these distances match the std result
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(sourceView, destView, threshold, distancesView,
intraTeamSentinelView, apiId);
Kokkos::parallel_for(policy, fnc);
// -----------------------------------------------
// check against std
// -----------------------------------------------
auto destViewAfterOp_h = create_host_space_copy(destView);
auto distancesView_h = create_host_space_copy(distancesView);
auto intraTeamSentinelView_h = create_host_space_copy(intraTeamSentinelView);
Kokkos::View<ValueType**, Kokkos::HostSpace> stdDestView("stdDestView",
numTeams, numCols);
GreaterThanValueFunctor<ValueType> predicate(threshold);
for (std::size_t i = 0; i < destViewAfterOp_h.extent(0); ++i) {
auto rowFrom =
Kokkos::subview(cloneOfSourceViewBeforeOp_h, i, Kokkos::ALL());
auto rowDest = Kokkos::subview(stdDestView, i, Kokkos::ALL());
auto stdIt = std::remove_copy_if(KE::cbegin(rowFrom), KE::cend(rowFrom),
KE::begin(rowDest), predicate);
const std::size_t stdDistance = KE::distance(KE::begin(rowDest), stdIt);
EXPECT_TRUE(distancesView_h(i) == stdDistance);
for (std::size_t j = 0; j < distancesView_h(i); ++j) {
EXPECT_TRUE(destViewAfterOp_h(i, j) == stdDestView(i, j));
}
ASSERT_TRUE(intraTeamSentinelView_h(i));
}
}
template <class LayoutTag, class ValueType>
void run_all_scenarios() {
for (int numTeams : teamSizesToTest) {
for (const auto& numCols : {0, 1, 2, 13, 101, 1444, 8113}) {
for (int apiId : {0, 1}) {
test_A<LayoutTag, ValueType>(numTeams, numCols, apiId);
}
}
}
}
TEST(std_algorithms_remove_copy_if_team_test, test) {
run_all_scenarios<DynamicTag, double>();
run_all_scenarios<StridedTwoRowsTag, int>();
run_all_scenarios<StridedThreeRowsTag, unsigned>();
}
} // namespace TeamRemoveCopyIf
} // namespace stdalgos
} // namespace Test
|