ViewUtils.h

//
// View Utilities
//   Utility functions relating to Kokkos views
//
 
#ifndef IPPL_VIEW_UTILS_H
#define IPPL_VIEW_UTILS_H
 
#include <Kokkos_Core.hpp>
 
#include <iostream>
 
#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::host_mirror_type 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 <unsigned Dim, typename View>
        void write_as_list_impl(const View& view, std::ostream& out = std::cout) {
            auto N = view.extent(0);
            out << "[";
            for (std::size_t i = 0; i < N; ++i) {
                if constexpr (Dim == 1) {
                    out << view(i);
                } else {
                    auto make_subview = [&]<std::size_t... Is>(std::index_sequence<Is...>) {
                        return Kokkos::subview(view, i, (static_cast<void>(Is), Kokkos::ALL)...);
                    };
 
                    auto subview = make_subview(std::make_index_sequence<Dim - 1>());
                    write_as_list_impl<Dim - 1>(subview, out);
                }
                if (i != N - 1)
                    out << ", ";
            }
            out << "]";
        };
 
        template <typename T, unsigned Dim, class... Properties>
        void write_as_list(const typename ViewType<T, Dim, Properties...>::view_type& view,
                           std::ostream& out = std::cout) {
            auto hview =
                Kokkos::create_mirror_view_and_copy(Kokkos::DefaultHostExecutionSpace(), view);
 
            write_as_list_impl<Dim>(hview, out);
        }
 
        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