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 8095574 : constexpr void nestedLoop(BeginFunctor&& begin, EndFunctor&& end, Functor&& body,
34 : Check&& check = nullptr) {
35 25194238 : for (size_t i = begin(Current); i < end(Current); ++i) {
36 : if constexpr (Dim - 1 == Current) {
37 9228022 : body(i);
38 : } else {
39 91677022 : auto inner = [i, &body](auto... args) {
40 41899536 : body(i, args...);
41 : };
42 7870642 : 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 8095574 : }
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 388 : constexpr void nestedViewLoop(View& view, int shift, Functor&& body,
63 : Check&& check = nullptr) {
64 388 : nestedLoop<View::rank>(
65 2597118 : [&](unsigned) {
66 2596792 : return shift;
67 : },
68 9982842 : [&](unsigned d) {
69 9979896 : return view.extent(d) - shift;
70 : },
71 : body, check);
72 388 : }
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 24 : 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 24 : 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 24 : decltype(auto) shrinkView(std::string label, const View& view, int nghost) {
123 24 : return shrinkView_impl(label, view, nghost, std::make_index_sequence<View::rank>{});
124 : }
125 : } // namespace detail
126 : } // namespace ippl
127 :
128 : #endif
|
