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 179 180 181 182 183
|
//@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>
namespace Test {
namespace stdalgos {
namespace TeamReplaceCopyIf {
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;
DistancesViewType m_distancesView;
IntraTeamSentinelView m_intraTeamSentinelView;
ValueType m_threshold;
ValueType m_newValue;
int m_apiPick;
TestFunctorA(const SourceViewType sourceView, const DestViewType destView,
const DistancesViewType distancesView,
const IntraTeamSentinelView intraTeamSentinelView,
ValueType threshold, ValueType newVal, int apiPick)
: m_sourceView(sourceView),
m_destView(destView),
m_distancesView(distancesView),
m_intraTeamSentinelView(intraTeamSentinelView),
m_threshold(threshold),
m_newValue(newVal),
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::replace_copy_if(
member, KE::begin(myRowViewFrom), KE::end(myRowViewFrom),
KE::begin(myRowViewDest), predicate, m_newValue);
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::replace_copy_if(member, myRowViewFrom, myRowViewDest,
predicate, m_newValue);
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 with values,
and run a team-level replace_copy_if where the values strictly
greater than a threshold are replaced/copied into a new view
with a new value, while those <= threshold are copied but unchanged
*/
const auto threshold = static_cast<ValueType>(151);
const auto newVal = static_cast<ValueType>(1);
// -----------------------------------------------
// prepare data
// -----------------------------------------------
// create a view in the memory space associated with default exespace
// with as many rows as the number of teams and fill it with random
// values from an arbitrary range
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);
// replace_copy_if returns an iterator so to verify that it is correct
// 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, distancesView, intraTeamSentinelView,
threshold, newVal, apiId);
Kokkos::parallel_for(policy, fnc);
// -----------------------------------------------
// run cpp-std kernel and check
// -----------------------------------------------
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 < sourceView.extent(0); ++i) {
auto rowFrom =
Kokkos::subview(cloneOfSourceViewBeforeOp_h, i, Kokkos::ALL());
auto rowDest = Kokkos::subview(stdDestView, i, Kokkos::ALL());
auto it = std::replace_copy_if(KE::cbegin(rowFrom), KE::cend(rowFrom),
KE::begin(rowDest), predicate, newVal);
const std::size_t stdDistance = KE::distance(KE::begin(rowDest), it);
ASSERT_EQ(stdDistance, distancesView_h(i));
ASSERT_TRUE(intraTeamSentinelView_h(i));
}
auto dataViewAfterOp_h = create_host_space_copy(destView);
expect_equal_host_views(stdDestView, dataViewAfterOp_h);
}
template <class LayoutTag, class ValueType>
void run_all_scenarios() {
for (int numTeams : teamSizesToTest) {
for (const auto& numCols : {0, 1, 2, 13, 101, 1444, 8153}) {
for (int apiId : {0, 1}) {
test_A<LayoutTag, ValueType>(numTeams, numCols, apiId);
}
}
}
}
TEST(std_algorithms_replace_copy_if_team_test, test) {
run_all_scenarios<DynamicTag, double>();
run_all_scenarios<StridedTwoRowsTag, int>();
run_all_scenarios<StridedThreeRowsTag, unsigned>();
}
} // namespace TeamReplaceCopyIf
} // namespace stdalgos
} // namespace Test
|