forked from arborx/ArborX
-
Notifications
You must be signed in to change notification settings - Fork 0
/
tstDetailsKokkosExtKernelStdAlgorithms.cpp
67 lines (55 loc) · 2.53 KB
/
tstDetailsKokkosExtKernelStdAlgorithms.cpp
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
/****************************************************************************
* Copyright (c) 2017-2022 by the ArborX authors *
* All rights reserved. *
* *
* This file is part of the ArborX library. ArborX is *
* distributed under a BSD 3-clause license. For the licensing terms see *
* the LICENSE file in the top-level directory. *
* *
* SPDX-License-Identifier: BSD-3-Clause *
****************************************************************************/
#include "ArborX_EnableDeviceTypes.hpp" // ARBORX_DEVICE_TYPES
#include "ArborX_EnableViewComparison.hpp"
#include <ArborX_DetailsKokkosExtKernelStdAlgorithms.hpp>
#include <ArborX_DetailsKokkosExtViewHelpers.hpp>
#include <ArborX_Exception.hpp>
#include <Kokkos_Core.hpp>
#include <boost/test/unit_test.hpp>
#include <algorithm>
#include <numeric>
#include <vector>
#define BOOST_TEST_MODULE KokkosExtKernelStdAlgorithms
namespace tt = boost::test_tools;
BOOST_AUTO_TEST_CASE_TEMPLATE(nth_element, DeviceType, ARBORX_DEVICE_TYPES)
{
using ExecutionSpace = typename DeviceType::execution_space;
ExecutionSpace space;
using ArborX::Details::KokkosExt::nth_element;
for (auto v_ref : {std::vector<float>{}, std::vector<float>{0.5f},
std::vector<float>{0.1f, 0.1f, 0.1f},
std::vector<float>{0.1f, 0.2f, 0.3f},
std::vector<float>{0.1f, 0.3f, -0.5f, 1.0f, -0.9f, -1.2f}})
{
int const n = v_ref.size();
Kokkos::View<float *, DeviceType> v("v", n);
Kokkos::deep_copy(
space, v,
Kokkos::View<float *, Kokkos::HostSpace,
Kokkos::MemoryTraits<Kokkos::Unmanaged>>(v_ref.data(), n));
Kokkos::View<float *, DeviceType> nth("nth", n);
for (int i = 0; i < n; ++i)
{
auto v_copy = ArborX::Details::KokkosExt::clone(space, v);
Kokkos::parallel_for(
Kokkos::RangePolicy<ExecutionSpace>(space, 0, 1), KOKKOS_LAMBDA(int) {
nth_element(v_copy.data(), v_copy.data() + i, v_copy.data() + n);
nth(i) = v_copy(i);
});
}
space.fence();
auto nth_host =
Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace{}, nth);
std::sort(v_ref.begin(), v_ref.end());
BOOST_TEST(nth_host == v_ref, tt::per_element());
}
}