TypeUtils.h

//
// Type Utilities
//   Metaprogramming utility functions for type manipulation
//
 
#ifndef IPPL_TYPE_UTILS_H
#define IPPL_TYPE_UTILS_H
 
#include <Kokkos_Core.hpp>
 
#include "Types/Variant.h"
 
#include "Utility/IpplException.h"
 
namespace ippl {
    namespace detail {
 
        template <typename Check, typename... Collection>
        struct IsUnique {
            constexpr static bool enable = !std::disjunction_v<std::is_same<Check, Collection>...>;
            typedef Check type;
        };
 
        template <template <typename> class Wrapper>
        struct WrapUnique {
            template <typename Check, typename... Collection>
            struct Verifier {
                typedef Wrapper<Check> type;
                constexpr static bool enable =
                    !std::disjunction_v<std::is_same<type, Collection>...>;
            };
        };
 
        template <bool B, typename T>
        using ConditionalType = std::conditional_t<B, T, void>;
 
        template <typename Type, typename...>
        struct IsEnabled {
            constexpr static bool enable = !std::is_void_v<Type>;
            typedef Type type;
        };
 
        template <typename, typename, template <typename...> class Verifier = IsUnique>
        struct ConstructVariant;
 
        template <template <typename...> class Verifier>
        struct ConstructVariant<std::variant<>, std::variant<>, Verifier> {
            typedef std::variant<> type;
        };
 
        template <typename... T, template <typename...> class Verifier>
        struct ConstructVariant<std::variant<>, std::variant<T...>, Verifier> {
            typedef std::variant<T...> type;
        };
 
        template <typename Next, typename... ToAdd, typename... Added,
                  template <typename...> class Verifier>
        struct ConstructVariant<std::variant<Next, ToAdd...>, std::variant<Added...>, Verifier> {
            // Convenience aliases
            template <bool B, class T, class F>
            using cond = std::conditional_t<B, T, F>;
            template <typename... T>
            using variant = std::variant<T...>;
 
            using Check = Verifier<Next, Added...>;
 
            typedef cond<
                Check::enable,
                // The verifier has indicated that this type should be added
                typename ConstructVariant<variant<ToAdd...>,
                                          variant<typename Check::type, Added...>, Verifier>::type,
                // The verifier has indicated that the type should not be added
                typename ConstructVariant<variant<ToAdd...>, variant<Added...>, Verifier>::type>
                type;
        };
 
        template <typename... Types>
        using VariantFromConditionalTypes =
            typename ConstructVariant<std::variant<Types...>, std::variant<>, IsEnabled>::type;
 
        template <typename... Types>
        using VariantFromUniqueTypes =
            typename ConstructVariant<std::variant<Types...>, std::variant<>, IsUnique>::type;
 
        template <template <typename...> class Verifier, typename... Types>
        using VariantWithVerifier =
            typename ConstructVariant<std::variant<Types...>, std::variant<>, Verifier>::type;
 
        template <template <typename...> class Type, typename Pack>
        struct Forward;
 
        template <template <typename...> class Type, typename... Spaces>
        struct Forward<Type, std::variant<Spaces...>> {
            using type = Type<Spaces...>;
        };
 
        template <template <typename...> class Type, typename T, typename... Properties>
        struct Forward<Type, Kokkos::View<T, Properties...>> {
            using type = Type<Properties...>;
        };
 
        template <template <typename...> class Type, typename View>
        struct CreateUniformType {
            using view_type = typename View::uniform_type;
            using type      = typename Forward<Type, view_type>::type;
        };
 
        template <template <typename...> class Type>
        struct TypeForAllSpaces {
            using unique_memory_spaces = VariantFromUniqueTypes<
                Kokkos::HostSpace, Kokkos::SharedSpace, Kokkos::SharedHostPinnedSpace
#ifdef KOKKOS_ENABLE_CUDA
                ,
                Kokkos::CudaSpace, Kokkos::CudaHostPinnedSpace, Kokkos::CudaUVMSpace
#endif
#ifdef KOKKOS_ENABLE_HIP
                ,
                Kokkos::HIPSpace, Kokkos::HIPHostPinnedSpace, Kokkos::HIPManagedSpace
#endif
#ifdef KOKKOS_ENABLE_SYCL
                ,
                Kokkos::Experimental::SYCLDeviceUSMSpace, Kokkos::Experimental::SYCLHostUSMSpace,
                Kokkos::Experimental::SYCLSharedUSMSpace
#endif
                >;
 
            using unique_exec_spaces = VariantFromUniqueTypes<Kokkos::DefaultExecutionSpace
#ifdef KOKKOS_ENABLE_OPENMP
                                                              ,
                                                              Kokkos::OpenMP
#endif
#ifdef KOKKOS_ENABLE_OPENMPTARGET
                                                              ,
                                                              Kokkos::OpenMPTarget
#endif
#ifdef KOKKOS_ENABLE_THREADS
                                                              ,
                                                              Kokkos::Threads
#endif
#ifdef KOKKOS_ENABLE_SERIAL
                                                              ,
                                                              Kokkos::Serial
#endif
#ifdef KOKKOS_ENABLE_CUDA
                                                              ,
                                                              Kokkos::Cuda
#endif
#ifdef KOKKOS_ENABLE_HIP
                                                              ,
                                                              Kokkos::HIP
#endif
#ifdef KOKKOS_ENABLE_SYCL
                                                              ,
                                                              Kokkos::Experimental::SYCL
#endif
#ifdef KOKKOS_ENABLE_HPX
                                                              ,
                                                              Kokkos::HPX
#endif
                                                              >;
 
            using memory_spaces_type = typename Forward<Type, unique_memory_spaces>::type;
            using exec_spaces_type   = typename Forward<Type, unique_exec_spaces>::type;
        };
 
        template <template <typename> class Type, typename... Spaces>
        class MultispaceContainer {
            template <typename T, typename... Ts>
            using Verifier = typename WrapUnique<Type>::template Verifier<T, Ts...>;
 
            using Types = VariantWithVerifier<Verifier, Spaces...>;
 
            std::array<Types, sizeof...(Spaces)> elements_m;
 
            template <typename Space, unsigned Idx = 0>
            constexpr static unsigned spaceToIndex() {
                static_assert(Idx < sizeof...(Spaces));
                if constexpr (std::is_same_v<Space,
                                             std::tuple_element_t<Idx, std::tuple<Spaces...>>>) {
                    return Idx;
                } else {
                    return spaceToIndex<Space, Idx + 1>();
                }
                // Silences incorrect nvcc warning: missing return statement at end of non-void
                // function
                throw IpplException("detail::MultispaceContainer::spaceToIndex",
                                    "Unreachable state");
            }
 
            template <typename Space>
            void initElements() {
                elements_m[spaceToIndex<Space>()] = Type<Space>{};
            }
 
            template <typename MemorySpace, typename Filter,
                      std::enable_if_t<std::is_null_pointer_v<std::decay_t<Filter>>, int> = 0>
            constexpr bool copyToSpace(Filter&&) {
                return true;
            }
 
            template <typename MemorySpace, typename Filter,
                      std::enable_if_t<!std::is_null_pointer_v<std::decay_t<Filter>>, int> = 0>
            bool copyToSpace(Filter&& predicate) {
                return predicate.template operator()<MemorySpace>();
            }
 
        public:
            MultispaceContainer() { (initElements<Spaces>(), ...); }
 
            template <typename DataType, typename Filter = std::nullptr_t>
            MultispaceContainer(const DataType& data, Filter&& predicate = nullptr)
                : MultispaceContainer() {
                using space = typename DataType::memory_space;
                static_assert(std::is_same_v<DataType, Type<space>>);
 
                elements_m[spaceToIndex<space>()] = data;
                copyToOtherSpaces<space>(predicate);
            }
 
            template <typename Space, typename Filter = std::nullptr_t>
            void copyToOtherSpaces(Filter&& predicate = nullptr) {
                forAll([&]<typename DataType>(DataType& dst) {
                    using memory_space = typename DataType::memory_space;
                    if constexpr (!std::is_same_v<Space, memory_space>) {
                        if (copyToSpace<memory_space>(predicate)) {
                            dst = Kokkos::create_mirror_view_and_copy(
                                Kokkos::view_alloc(memory_space{}, Kokkos::WithoutInitializing),
                                get<Space>());
                        }
                    }
                });
            }
 
            template <typename Space>
            const Type<Space>& get() const {
                return std::get<Type<Space>>(elements_m[spaceToIndex<Space>()]);
            }
 
            template <typename Space>
            Type<Space>& get() {
                return std::get<Type<Space>>(elements_m[spaceToIndex<Space>()]);
            }
 
            template <typename Functor>
            void forAll(Functor&& f) const {
                (f(get<Spaces>()), ...);
            }
 
            template <typename Functor>
            void forAll(Functor&& f) {
                (f(get<Spaces>()), ...);
            }
        };
 
        template <template <typename> class Type>
        struct ContainerForAllSpaces {
            template <typename... Spaces>
            using container_type = MultispaceContainer<Type, Spaces...>;
 
            using type = typename TypeForAllSpaces<container_type>::memory_spaces_type;
 
            // Static factory function that takes a lambda to initialize each memory space
            template <typename Functor>
            static type createContainer(Functor&& initFunc) {
                return type{std::forward<Functor>(initFunc)};
            }
        };
 
        template <typename Functor>
        void runForAllSpaces(Functor&& f) {
            using all_spaces = typename TypeForAllSpaces<std::variant>::memory_spaces_type;
            auto runner      = [&]<typename... Spaces>(const std::variant<Spaces...>&) {
                (f.template operator()<Spaces>(), ...);
            };
            runner(all_spaces{});
        }
    }  // namespace detail
}  // namespace ippl
 
#endif