Branch data Line data Source code
1 : : //
2 : : // View Utilities
3 : : // Utility functions relating to Kokkos views
4 : : //
5 : :
6 : : #ifndef IPPL_VIEW_UTILS_H
7 : : #define IPPL_VIEW_UTILS_H
8 : :
9 : : #include <Kokkos_Core.hpp>
10 : :
11 : : #include "Types/ViewTypes.h"
12 : :
13 : : namespace ippl {
14 : : namespace detail {
15 : : /*!
16 : : * Expands into a nested loop via templating
17 : : * Source:
18 : : * https://stackoverflow.com/questions/34535795/n-dimensionally-nested-metaloops-with-templates
19 : : * @tparam Dim the number of nested levels
20 : : * @tparam BeginFunctor functor type for determining the start index of each loop
21 : : * @tparam EndFunctor functor type for determining the end index of each loop
22 : : * @tparam Functor functor type for the loop body
23 : : * @tparam Check functor type for loop check
24 : : * @param begin a functor that returns the starting index for each level of the loop
25 : : * @param end a functor that returns the ending index (exclusive) for each level of the loop
26 : : * @param body a functor to be called in each iteration of the loop with the indices as
27 : : * arguments
28 : : * @param check a check function to be run after each loop; takes the current axis as an
29 : : * argument (default none)
30 : : */
31 : : template <unsigned Dim, unsigned Current = 0, class BeginFunctor, class EndFunctor,
32 : : class Functor, typename Check = std::nullptr_t>
33 : 5650266 : constexpr void nestedLoop(BeginFunctor&& begin, EndFunctor&& end, Functor&& body,
34 : : Check&& check = nullptr) {
35 [ + + ]: 18488960 : for (size_t i = begin(Current); i < end(Current); ++i) {
36 : : if constexpr (Dim - 1 == Current) {
37 : 7189022 : body(i);
38 : : } else {
39 : 70504088 : auto inner = [i, &body](auto... args) {
40 : 24097728 : body(i, args...);
41 : : };
42 [ + - ]: 5649672 : nestedLoop<Dim, Current + 1>(begin, end, inner, check);
43 : : }
44 : : }
45 : : if constexpr (!std::is_null_pointer_v<std::decay_t<Check>>) {
46 : 0 : check(Current);
47 : : }
48 : 5650266 : }
49 : :
50 : : /*!
51 : : * Convenience function for nested looping through a view
52 : : * @tparam View the view type
53 : : * @tparam Functor the loop body functor type
54 : : * @tparam Check functor type for loop check
55 : : * @param view the view
56 : : * @param shift the number of ghost cells
57 : : * @param body the functor to be called in each iteration
58 : : * @param check a check function to be run after each loop; takes the current axis as an
59 : : * argument (default none)
60 : : */
61 : : template <typename View, class Functor, typename Check = std::nullptr_t>
62 : 174 : constexpr void nestedViewLoop(View& view, int shift, Functor&& body,
63 : : Check&& check = nullptr) {
64 : 174 : nestedLoop<View::rank>(
65 : 2408146 : [&](unsigned) {
66 : 1765460 : return shift;
67 : : },
68 [ + - ]: 9272814 : [&](unsigned d) {
69 : 6192852 : return view.extent(d) - shift;
70 : : },
71 : : body, check);
72 : 174 : }
73 : :
74 : : /*!
75 : : * Writes a view to an output stream
76 : : * @tparam T view data type
77 : : * @tparam Dim view dimension
78 : : * @tparam Properties further template parameters of Kokkos
79 : : *
80 : : * @param view to write
81 : : * @param out stream
82 : : */
83 : : template <typename T, unsigned Dim, class... Properties>
84 : 0 : void write(const typename ViewType<T, Dim, Properties...>::view_type& view,
85 : : std::ostream& out = std::cout) {
86 : : using view_type = typename ViewType<T, Dim, Properties...>::view_type;
87 [ # # ]: 0 : typename view_type::HostMirror hview = Kokkos::create_mirror_view(view);
88 [ # # ]: 0 : Kokkos::deep_copy(hview, view);
89 : :
90 : 0 : nestedViewLoop(
91 : : view, 0,
92 : 0 : [&]<typename... Args>(Args&&... args) {
93 : 0 : out << hview(args...) << " ";
94 : : },
95 [ # # ]: 0 : [&](unsigned axis) {
96 [ # # # # : 0 : if (axis + 1 >= 2 || axis == 0) {
# # # # ]
[ # # # # ]
97 : 0 : out << std::endl;
98 : : }
99 : : });
100 : 0 : }
101 : :
102 : : /*!
103 : : * Utility function for shrinkView
104 : : */
105 : : template <typename View, size_t... Idx>
106 : 12 : decltype(auto) shrinkView_impl(std::string label, const View& view, int nghost,
107 : : const std::index_sequence<Idx...>&) {
108 : : using view_type = typename Kokkos::View<typename View::data_type, Kokkos::LayoutLeft,
109 : : typename View::memory_space>::uniform_type;
110 [ + - ]: 12 : return view_type(label, (view.extent(Idx) - 2 * nghost)...);
111 : : }
112 : :
113 : : /*!
114 : : * Constructs a new view with size equal to that of the given view, minus the ghost cells
115 : : * (used for heFFTe, which expects the data to have a certain layout and no ghost cells)
116 : : * @param label the new view's name
117 : : * @param view the view to shrink
118 : : * @param nghost the number of ghost cells on the view's boundary
119 : : * @return The shrunken view
120 : : */
121 : : template <typename View>
122 : 12 : decltype(auto) shrinkView(std::string label, const View& view, int nghost) {
123 [ + - + - ]: 12 : return shrinkView_impl(label, view, nghost, std::make_index_sequence<View::rank>{});
124 : : }
125 : : } // namespace detail
126 : : } // namespace ippl
127 : :
128 : : #endif
|
