1 #include "kokkos_type.h"
2
3 namespace SPARTA_NS {
4
5 template <typename Device>
6 struct ExclScan {
7 using value_type = long;
8 using view_type = Kokkos::View<int*, Device>;
9 using total_type = Kokkos::View<int, Device>;
initSPARTA_NS::ExclScan10 KOKKOS_INLINE_FUNCTION void init(value_type& update) const { update = 0; }
joinSPARTA_NS::ExclScan11 KOKKOS_INLINE_FUNCTION void join(
12 volatile value_type& update, const volatile value_type& input) const {
13 update = update + input;
14 }
15 view_type in_;
16 view_type out_;
17 total_type total_;
ExclScanSPARTA_NS::ExclScan18 ExclScan(view_type &in, view_type &out, total_type &total):
19 in_(in), out_(out), total_(total) {}
operator ()SPARTA_NS::ExclScan20 KOKKOS_INLINE_FUNCTION void operator()(int i, value_type& update, bool final_pass) const {
21 update += in_[i];
22 if (final_pass) {
23 out_[0] = 0;
24 out_[i + 1] = static_cast<int>(update);
25 if (i + 1 == int(in_.extent(0))) total_() = static_cast<int>(update);
26 }
27 }
28 using execution_space = Device;
29 };
30
31 template <typename Device>
offset_scan(Kokkos::View<int *,Device> in,int & total)32 Kokkos::View<int*, Device> offset_scan(Kokkos::View<int*, Device> in, int& total) {
33
34 Kokkos::View<int*, Device> out;
35
36 if (in.size() == 0) {
37 total = 0;
38 out = Kokkos::View<int*, Device>(in.label() + "_scan",1);
39 } else {
40 out = Kokkos::View<int*, Device>(Kokkos::view_alloc(in.label() + "_scan",Kokkos::WithoutInitializing), in.size() + 1);
41 Kokkos::View<int, Device> total_dev(Kokkos::view_alloc("scan_total",Kokkos::WithoutInitializing));
42 typename Kokkos::View<int, Device>::HostMirror total_host(Kokkos::view_alloc("scan_total_mirror",Kokkos::WithoutInitializing));
43 Kokkos::parallel_scan(in.size(), ExclScan<Device>(in, out, total_dev));
44 Kokkos::deep_copy(total_host, total_dev);
45 total = total_host();
46 }
47 return out;
48 }
49
50 template Kokkos::View<int*, SPAHostType> offset_scan(
51 Kokkos::View<int*, SPAHostType> in, int& total);
52 #ifdef SPARTA_KOKKOS_GPU
53 template Kokkos::View<int*, DeviceType> offset_scan(
54 Kokkos::View<int*, DeviceType> in, int& total);
55 #endif
56
57 }
58