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
|
//@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 TeamAnyOf {
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 DataViewType, class AnyOfResultsViewType,
class IntraTeamSentinelView, class UnaryPredType>
struct TestFunctorA {
DataViewType m_dataView;
AnyOfResultsViewType m_anyOfResultsView;
IntraTeamSentinelView m_intraTeamSentinelView;
int m_apiPick;
UnaryPredType m_unaryPred;
TestFunctorA(const DataViewType dataView,
const AnyOfResultsViewType anyOfResultsView,
const IntraTeamSentinelView intraTeamSentinelView, int apiPick,
UnaryPredType unaryPred)
: m_dataView(dataView),
m_anyOfResultsView(anyOfResultsView),
m_intraTeamSentinelView(intraTeamSentinelView),
m_apiPick(apiPick),
m_unaryPred(unaryPred) {}
template <class MemberType>
KOKKOS_INLINE_FUNCTION void operator()(const MemberType& member) const {
const auto myRowIndex = member.league_rank();
auto myRowViewFrom = Kokkos::subview(m_dataView, myRowIndex, Kokkos::ALL());
bool result = false;
switch (m_apiPick) {
case 0: {
result = KE::any_of(member, KE::cbegin(myRowViewFrom),
KE::cend(myRowViewFrom), m_unaryPred);
Kokkos::single(Kokkos::PerTeam(member), [=, *this]() {
m_anyOfResultsView(myRowIndex) = result;
});
break;
}
case 1: {
result = KE::any_of(member, myRowViewFrom, m_unaryPred);
Kokkos::single(Kokkos::PerTeam(member), [=, *this]() {
m_anyOfResultsView(myRowIndex) = result;
});
break;
}
}
// 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, result, m_anyOfResultsView(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 any_of
*/
// -----------------------------------------------
// 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.
constexpr ValueType lowerBound = 5;
constexpr ValueType upperBound = 523;
const auto bounds = make_bounds(lowerBound, upperBound);
auto [dataView, dataViewBeforeOp_h] = create_random_view_and_host_clone(
LayoutTag{}, numTeams, numCols, bounds, "dataView");
// -----------------------------------------------
// launch kokkos kernel
// -----------------------------------------------
using space_t = Kokkos::DefaultExecutionSpace;
Kokkos::TeamPolicy<space_t> policy(numTeams, Kokkos::AUTO());
// to verify that things work, each team stores the result of its any_of call,
// and then we check that these match what we expect
Kokkos::View<bool*> anyOfResultsView("anyOfResultsView", numTeams);
// sentinel to check if all members of the team compute the same result
Kokkos::View<bool*> intraTeamSentinelView("intraTeamSameResult", numTeams);
GreaterThanValueFunctor unaryPred{lowerBound};
// use CTAD for functor
TestFunctorA fnc(dataView, anyOfResultsView, intraTeamSentinelView, apiId,
unaryPred);
Kokkos::parallel_for(policy, fnc);
// -----------------------------------------------
// run cpp-std kernel and check
// -----------------------------------------------
auto anyOfResultsView_h = create_host_space_copy(anyOfResultsView);
auto intraTeamSentinelView_h = create_host_space_copy(intraTeamSentinelView);
for (std::size_t i = 0; i < dataView.extent(0); ++i) {
auto rowFrom = Kokkos::subview(dataViewBeforeOp_h, i, Kokkos::ALL());
const bool result =
std::any_of(KE::cbegin(rowFrom), KE::cend(rowFrom), unaryPred);
ASSERT_EQ(result, anyOfResultsView_h(i));
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, 8153}) {
for (int apiId : {0, 1}) {
test_A<LayoutTag, ValueType>(numTeams, numCols, apiId);
}
}
}
}
TEST(std_algorithms_any_of_team_test, test) {
run_all_scenarios<DynamicTag, double>();
run_all_scenarios<StridedTwoRowsTag, int>();
run_all_scenarios<StridedThreeRowsTag, unsigned>();
}
} // namespace TeamAnyOf
} // namespace stdalgos
} // namespace Test
|