-
Notifications
You must be signed in to change notification settings - Fork 8
/
user.cpp
59 lines (48 loc) · 1.45 KB
/
user.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
#include "user.hpp"
#include "Kokkos_Core.hpp"
#include <cstdint>
#include <iostream>
struct InitView {
explicit InitView(view_type _v) : m_view(_v) {}
KOKKOS_INLINE_FUNCTION
void operator()(const int i) const {
m_view(i, 0) = -(i + 1);
m_view(i, 1) = (i + 1);
}
private:
view_type m_view;
};
struct ModifyView {
explicit ModifyView(view_type _v) : m_view(_v) {}
KOKKOS_INLINE_FUNCTION
void operator()(const int i) const {
m_view(i, 0) *= 2;
m_view(i, 1) *= 2;
}
private:
view_type m_view;
};
using exec_space = typename view_type::traits::execution_space;
///
/// \fn generate_view
/// \brief This is meant to emulate some function that exists in a user library
/// which returns a Kokkos::View and will have a python binding
///
view_type generate_view(size_t n) {
if (!Kokkos::is_initialized()) {
std::cerr << "[user-bindings]> Initializing Kokkos..." << std::endl;
Kokkos::initialize();
}
std::cerr << "[user-bindings]> Generating View..." << std::flush;
view_type _v("user_view", n, 2);
Kokkos::RangePolicy<exec_space, int> range(0, n);
Kokkos::parallel_for("generate_view", range, InitView{_v});
std::cerr << " Done." << std::endl;
return _v;
}
void modify_view(view_type _v) {
std::cerr << "[user-bindings]> Modifying View..." << std::flush;
Kokkos::RangePolicy<exec_space, int> range(0, _v.extent(0));
Kokkos::parallel_for("modify_view", range, ModifyView{_v});
std::cerr << " Done." << std::endl;
}