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
|
//@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 <Kokkos_Core.hpp>
namespace Test {
namespace {
struct TestIsAsynchFunctor {
Kokkos::View<double, TEST_EXECSPACE> atomic_test;
TestIsAsynchFunctor(Kokkos::View<double, TEST_EXECSPACE> atomic_test_)
: atomic_test(atomic_test_) {}
KOKKOS_INLINE_FUNCTION
void operator()(const int) const { Kokkos::atomic_add(&atomic_test(), 1.0); }
};
template <class PolicyType, class ReduceFunctor>
void test_reduce_device_view(int64_t N, PolicyType policy,
ReduceFunctor functor) {
using ExecSpace = TEST_EXECSPACE;
Kokkos::View<int64_t, TEST_EXECSPACE> result("Result");
Kokkos::View<double, TEST_EXECSPACE> atomic_test("Atomic");
int64_t reducer_result, view_result, scalar_result;
Kokkos::Timer timer;
// Establish whether execspace is asynchronous
Kokkos::parallel_for("Test::ReduceDeviceView::TestIsAsynch",
Kokkos::RangePolicy<TEST_EXECSPACE>(0, 1000000),
TestIsAsynchFunctor(atomic_test));
double time0 = timer.seconds();
timer.reset();
typename ExecSpace::execution_space().fence();
double time_fence0 = timer.seconds();
Kokkos::deep_copy(result, 0);
// We need a warm-up to get reasonable results
Kokkos::parallel_reduce("Test::ReduceDeviceView::TestReducer", policy,
functor,
Kokkos::Sum<int64_t, TEST_EXECSPACE>(result));
Kokkos::fence();
timer.reset();
bool is_async = time0 < time_fence0;
// Test Reducer
Kokkos::parallel_reduce("Test::ReduceDeviceView::TestReducer", policy,
functor,
Kokkos::Sum<int64_t, TEST_EXECSPACE>(result));
double time1 = timer.seconds();
// Check whether it was asyncronous
timer.reset();
typename ExecSpace::execution_space().fence();
double time_fence1 = timer.seconds();
Kokkos::deep_copy(reducer_result, result);
Kokkos::deep_copy(result, 0);
ASSERT_EQ(N, reducer_result);
// We need a warm-up to get reasonable results
Kokkos::parallel_reduce("Test::ReduceDeviceView::TestView", policy, functor,
result);
Kokkos::fence();
timer.reset();
// Test View
Kokkos::parallel_reduce("Test::ReduceDeviceView::TestView", policy, functor,
result);
double time2 = timer.seconds();
// Check whether it was asyncronous
timer.reset();
typename ExecSpace::execution_space().fence();
double time_fence2 = timer.seconds();
Kokkos::deep_copy(view_result, result);
Kokkos::deep_copy(result, 0);
ASSERT_EQ(N, view_result);
timer.reset();
// Test Scalar
Kokkos::parallel_reduce("Test::ReduceDeviceView::TestScalar", policy, functor,
scalar_result);
double time3 = timer.seconds();
// Check whether it was asyncronous
timer.reset();
typename ExecSpace::execution_space().fence();
double time_fence3 = timer.seconds();
ASSERT_EQ(N, scalar_result);
if (is_async) {
ASSERT_LT(time1, time_fence1);
}
if (is_async) {
ASSERT_LT(time2, time_fence2);
ASSERT_GT(time3, time_fence3);
}
}
struct RangePolicyFunctor {
KOKKOS_INLINE_FUNCTION
void operator()(const int, int64_t& lsum) const { lsum += 1; }
};
struct MDRangePolicyFunctor {
KOKKOS_INLINE_FUNCTION
void operator()(const int, const int, const int, int64_t& lsum) const {
lsum += 1;
}
};
struct TeamPolicyFunctor {
int M;
TeamPolicyFunctor(int M_) : M(M_) {}
KOKKOS_INLINE_FUNCTION
void operator()(const Kokkos::TeamPolicy<TEST_EXECSPACE>::member_type& team,
int64_t& lsum) const {
for (int i = team.team_rank(); i < M; i += team.team_size()) lsum += 1;
}
};
} // namespace
TEST(TEST_CATEGORY, reduce_device_view_range_policy) {
// Avoid running out of memory
#ifdef KOKKOS_ENABLE_SYCL
int N = 100 * 1024 * 1024;
#else
int N = 1000 * 1024 * 1024;
#endif
test_reduce_device_view(N, Kokkos::RangePolicy<TEST_EXECSPACE>(0, N),
RangePolicyFunctor());
}
TEST(TEST_CATEGORY, reduce_device_view_mdrange_policy) {
int N = 1000 * 1024 * 1024;
test_reduce_device_view(
N,
Kokkos::MDRangePolicy<TEST_EXECSPACE, Kokkos::Rank<3>>(
{0, 0, 0}, {1000, 1024, 1024}),
MDRangePolicyFunctor());
}
TEST(TEST_CATEGORY, reduce_device_view_team_policy) {
// FIXME_SYCL The number of workgroups on CUDA devices can not be larger than
// 65535
#ifdef KOKKOS_ENABLE_SYCL
int N = 63 * 1024 * 1024;
test_reduce_device_view(
N, Kokkos::TeamPolicy<TEST_EXECSPACE>(63 * 1024, Kokkos::AUTO),
TeamPolicyFunctor(1024));
#else
int N = 1000 * 1024 * 1024;
test_reduce_device_view(
N, Kokkos::TeamPolicy<TEST_EXECSPACE>(1000 * 1024, Kokkos::AUTO),
TeamPolicyFunctor(1024));
#endif
}
} // namespace Test
|