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
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project
#include <TestSYCL_Category.hpp>
#include <Kokkos_Macros.hpp>
#ifdef KOKKOS_ENABLE_EXPERIMENTAL_CXX20_MODULES
import kokkos.core;
#else
#include <Kokkos_Core.hpp>
#endif
namespace Test {
namespace Impl {
struct SYCLQueueScratchTestFunctor {
using team_t = Kokkos::TeamPolicy<Kokkos::SYCL>::member_type;
using scratch_t = Kokkos::View<int64_t*, Kokkos::SYCL::scratch_memory_space>;
Kokkos::View<int64_t, Kokkos::SYCLDeviceUSMSpace,
Kokkos::MemoryTraits<Kokkos::Atomic>>
counter;
int N, M;
SYCLQueueScratchTestFunctor(
Kokkos::View<int64_t, Kokkos::SYCLDeviceUSMSpace> counter_, int N_,
int M_)
: counter(counter_), N(N_), M(M_) {}
KOKKOS_FUNCTION
void operator()(const team_t& team) const {
scratch_t scr(team.team_scratch(1), M);
Kokkos::parallel_for(Kokkos::TeamThreadRange(team, 0, M),
[&](int i) { scr[i] = 0; });
team.team_barrier();
for (int i = 0; i < N; i++) {
Kokkos::parallel_for(Kokkos::TeamThreadRange(team, 0, M),
[&](int j) { scr[j] += 1; });
}
team.team_barrier();
Kokkos::parallel_for(Kokkos::TeamThreadRange(team, 0, M), [&](int i) {
if (scr[i] != N) counter()++;
});
}
};
void sycl_queue_scratch_test_one(
int N, int T, int M_base,
Kokkos::View<int64_t, Kokkos::SYCLDeviceUSMSpace> counter,
Kokkos::SYCL sycl, int tid) {
int M = M_base + tid * 5;
Kokkos::TeamPolicy<Kokkos::SYCL> p(sycl, T, 64);
using scratch_t = Kokkos::View<int64_t*, Kokkos::SYCL::scratch_memory_space>;
int bytes = scratch_t::shmem_size(M);
for (int r = 0; r < 15; r++) {
Kokkos::parallel_for("Run", p.set_scratch_size(1, Kokkos::PerTeam(bytes)),
SYCLQueueScratchTestFunctor(counter, N, M));
}
}
void sycl_queue_scratch_test(
int N, int T, int M_base,
Kokkos::View<int64_t, Kokkos::SYCLDeviceUSMSpace> counter) {
constexpr int K = 4;
Kokkos::SYCL default_space;
sycl::context default_context = default_space.sycl_queue().get_context();
sycl::queue queue(default_context, sycl::default_selector_v,
sycl::property::queue::in_order());
std::array<Kokkos::SYCL, K> sycl;
for (int i = 0; i < K; i++) {
sycl[i] =
Kokkos::SYCL(sycl::queue(default_context, sycl::default_selector_v,
sycl::property::queue::in_order()));
}
// Test that growing scratch size in subsequent calls doesn't crash things
#if defined(KOKKOS_ENABLE_OPENMP)
#pragma omp parallel
{
int tid = omp_get_thread_num();
// Limit how many threads submit
if (tid < 4) {
sycl_queue_scratch_test_one(N, T, M_base, counter, sycl[tid], tid);
}
}
#else
for (int tid = 0; tid < K; tid++) {
sycl_queue_scratch_test_one(N, T, M_base, counter, sycl[tid], tid);
}
#endif
// Test that if everything is large enough, multiple launches with different
// scratch sizes don't step on each other
for (int tid = K - 1; tid >= 0; tid--) {
sycl_queue_scratch_test_one(N, T, M_base, counter, sycl[tid], tid);
}
Kokkos::fence();
}
} // namespace Impl
TEST(sycl, team_scratch_1_queues) {
int N = 1000000;
int T = 10;
int M_base = 150;
Kokkos::View<int64_t, Kokkos::SYCLDeviceUSMSpace> counter("C");
Impl::sycl_queue_scratch_test(N, T, M_base, counter);
int64_t result;
Kokkos::deep_copy(result, counter);
ASSERT_EQ(0, result);
}
} // namespace Test
|
