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
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
// SPDX-FileCopyrightText: Copyright Contributors to the Kokkos project
#include <Kokkos_Macros.hpp>
#ifdef KOKKOS_ENABLE_EXPERIMENTAL_CXX20_MODULES
import kokkos.core;
#else
#include <Kokkos_Core.hpp>
#endif
namespace Test {
// FIXME_SYCL This doesn't work yet for SYCL+CUDA
#if !defined(KOKKOS_ENABLE_SYCL) || defined(KOKKOS_ARCH_INTEL_GPU)
template <typename ExecutionSpace>
struct TestSharedAtomicsFunctor {
Kokkos::View<int, typename ExecutionSpace::memory_space> m_view;
TestSharedAtomicsFunctor(
Kokkos::View<int, typename ExecutionSpace::memory_space>& view)
: m_view(view) {}
KOKKOS_INLINE_FUNCTION void operator()(
const typename Kokkos::TeamPolicy<ExecutionSpace>::member_type t) const {
int* x = (int*)t.team_shmem().get_shmem(sizeof(int));
Kokkos::single(Kokkos::PerTeam(t), [&]() { *x = 0; });
t.team_barrier();
Kokkos::atomic_add(x, 1);
t.team_barrier();
Kokkos::single(Kokkos::PerTeam(t), [&]() { m_view() = *x; });
}
};
TEST(TEST_CATEGORY, atomic_shared) {
// FIXME_OPENMPTARGET - causes runtime failure with CrayClang compiler
#if defined(KOKKOS_COMPILER_CRAY_LLVM) && defined(KOKKOS_ENABLE_OPENMPTARGET)
GTEST_SKIP() << "known to fail with OpenMPTarget+Cray LLVM";
#endif
TEST_EXECSPACE exec;
Kokkos::View<int, typename TEST_EXECSPACE::memory_space> view("ref_value");
auto team_size =
Kokkos::TeamPolicy<TEST_EXECSPACE>(exec, 1, Kokkos::AUTO)
.team_size_recommended(TestSharedAtomicsFunctor<TEST_EXECSPACE>(view),
Kokkos::ParallelForTag{});
Kokkos::parallel_for(Kokkos::TeamPolicy<TEST_EXECSPACE>(exec, 1, team_size)
.set_scratch_size(0, Kokkos::PerTeam(8)),
TestSharedAtomicsFunctor<TEST_EXECSPACE>(view));
exec.fence("Fence after test kernel");
int i = 0;
Kokkos::deep_copy(i, view);
ASSERT_EQ(i, team_size);
}
#endif
} // namespace Test
|
