ViewUtils.h

Go to the documentation of this file.

//
// View Utilities
//   Utility functions relating to Kokkos views
//
 
#ifndef IPPL_VIEW_UTILS_H
#define IPPL_VIEW_UTILS_H
 
#include <Kokkos_Core.hpp>
 
#include "Types/ViewTypes.h"
 
namespace ippl {
    namespace detail {
        template <unsigned Dim, unsigned Current = 0, class BeginFunctor, class EndFunctor,
                  class Functor, typename Check  = std::nullptr_t>
        constexpr void nestedLoop(BeginFunctor&& begin, EndFunctor&& end, Functor&& body,
                                  Check&& check = nullptr) {
            for (size_t i = begin(Current); i < end(Current); ++i) {
                if constexpr (Dim - 1 == Current) {
                    body(i);
                } else {
                    auto inner = [i, &body](auto... args) {
                        body(i, args...);
                    };
                    nestedLoop<Dim, Current + 1>(begin, end, inner, check);
                }
            }
            if constexpr (!std::is_null_pointer_v<std::decay_t<Check>>) {
                check(Current);
            }
        }

        template <typename View, class Functor, typename Check = std::nullptr_t>
        constexpr void nestedViewLoop(View& view, int shift, Functor&& body,
                                      Check&& check = nullptr) {
            nestedLoop<View::rank>(
                [&](unsigned) {
                    return shift;
                },
                [&](unsigned d) {
                    return view.extent(d) - shift;
                },
                body, check);
        }

        template <typename T, unsigned Dim, class... Properties>
        void write(const typename ViewType<T, Dim, Properties...>::view_type& view,
                   std::ostream& out = std::cout) {
            using view_type = typename ViewType<T, Dim, Properties...>::view_type;
            typename view_type::HostMirror hview = Kokkos::create_mirror_view(view);
            Kokkos::deep_copy(hview, view);
 
            nestedViewLoop(
                view, 0,
                [&]<typename... Args>(Args&&... args) {
                    out << hview(args...) << " ";
                },
                [&](unsigned axis) {
                    if (axis + 1 >= 2 || axis == 0) {
                        out << std::endl;
                    }
                });
        }

        template <typename View, size_t... Idx>
        decltype(auto) shrinkView_impl(std::string label, const View& view, int nghost,
                                       const std::index_sequence<Idx...>&) {
            using view_type = typename Kokkos::View<typename View::data_type, Kokkos::LayoutLeft,
                                                    typename View::memory_space>::uniform_type;
            return view_type(label, (view.extent(Idx) - 2 * nghost)...);
        }

        template <typename View>
        decltype(auto) shrinkView(std::string label, const View& view, int nghost) {
            return shrinkView_impl(label, view, nghost, std::make_index_sequence<View::rank>{});
        }
    }  // namespace detail
}  // namespace ippl
 
#endif