vector_field.hpp

// SPDX-License-Identifier: MIT
 
#pragma once
#include "ddc_alias_inline_functions.hpp"
#include "ddc_aliases.hpp"
#include "vector_field_mem.hpp"
 
 
template <
        class ElementType,
        class IdxRangeType,
        class NDTag,
        class MemorySpace = Kokkos::DefaultExecutionSpace::memory_space,
        class LayoutStridedPolicy = std::experimental::layout_right>
class VectorField;
 
template <
        class ElementType,
        class IdxRangeType,
        class NDTag,
        class MemorySpace,
        class LayoutStridedPolicy>
inline constexpr bool enable_vector_field<
        VectorField<ElementType, IdxRangeType, NDTag, MemorySpace, LayoutStridedPolicy>> = true;
 
template <
        class ElementType,
        class IdxRangeType,
        class NDTag,
        class MemorySpace,
        class LayoutStridedPolicy>
inline constexpr bool enable_data_access_methods<
        VectorField<ElementType, IdxRangeType, NDTag, MemorySpace, LayoutStridedPolicy>> = true;
 
template <
        class ElementType,
        class IdxRangeType,
        class NDTag,
        class MemorySpace,
        class LayoutStridedPolicy>
inline constexpr bool enable_borrowed_vector_field<
        VectorField<ElementType, IdxRangeType, NDTag, MemorySpace, LayoutStridedPolicy>> = true;
 
 
template <
        class ElementType,
        class IdxRangeType,
        class NDTag,
        class MemorySpace,
        class LayoutStridedPolicy>
class VectorField
    : public VectorFieldCommon<
              Field<ElementType, IdxRangeType, MemorySpace, LayoutStridedPolicy>,
              NDTag>
{
public:
    using field_type = Field<ElementType, IdxRangeType, MemorySpace, LayoutStridedPolicy>;
 
private:
    using base_type = VectorFieldCommon<field_type, NDTag>;
 
public:
    using element_type = typename base_type::element_type;
 
    using typename base_type::NDTypeTag;
 
    using typename base_type::chunk_span_type;
    using typename base_type::chunk_view_type;
 
public:
    using span_type
            = VectorField<ElementType, IdxRangeType, NDTag, MemorySpace, LayoutStridedPolicy>;
    using view_type
            = VectorField<const ElementType, IdxRangeType, NDTag, MemorySpace, LayoutStridedPolicy>;
 
    using layout_type = LayoutStridedPolicy;
 
    using discrete_domain_type = typename base_type::discrete_domain_type;
    using index_range_type = discrete_domain_type;
 
    using memory_space = typename field_type::memory_space;
 
private:
    template <class, class, class, class, class>
    friend class VectorField;
 
    template <class OElementType, class Allocator, std::size_t... Is>
    KOKKOS_FUNCTION constexpr VectorField(
            VectorFieldMem<OElementType, IdxRangeType, NDTag, Allocator>& other,
            std::index_sequence<Is...> const&) noexcept
        : base_type((field_type(ddcHelper::get<ddc::type_seq_element_t<Is, NDTag>>(other)))...)
    {
    }
 
    // Disabled by SFINAE in the case of `ElementType` is not `const` to avoid write access
    template <
            class OElementType,
            std::size_t... Is,
            class SFINAEElementType = ElementType,
            class = std::enable_if_t<std::is_const_v<SFINAEElementType>>,
            class Allocator>
    KOKKOS_FUNCTION constexpr VectorField(
            VectorFieldMem<OElementType, IdxRangeType, NDTag, Allocator> const& other,
            std::index_sequence<Is...> const&) noexcept
        : base_type((field_type(ddcHelper::get<ddc::type_seq_element_t<Is, NDTag>>(other)))...)
    {
    }
 
    template <class OElementType, std::size_t... Is>
    KOKKOS_FUNCTION constexpr VectorField(
            VectorField<
                    OElementType,
                    index_range_type,
                    NDTag,
                    MemorySpace,
                    LayoutStridedPolicy> const& other,
            std::index_sequence<Is...> const&) noexcept
        : base_type((field_type(ddcHelper::get<ddc::type_seq_element_t<Is, NDTag>>(other)))...)
    {
    }
 
    template <class SliceType, std::size_t... Is>
    constexpr auto get_slice(SliceType const& slice_spec, std::index_sequence<Is...> const&)
    {
        auto chunk_slices = std::make_tuple(
                this->template get<ddc::type_seq_element_t<Is, NDTag>>()[slice_spec]...);
        using FieldType = std::tuple_element_t<0, decltype(chunk_slices)>;
        return VectorField<
                ElementType,
                typename FieldType::discrete_domain_type,
                NDTag,
                typename FieldType::memory_space,
                typename FieldType::layout_type>(std::move(std::get<Is>(chunk_slices))...);
    }
 
    template <class... ODDims, typename T, T... ints>
    KOKKOS_FUNCTION element_type
    operator()(Idx<ODDims...> const& delems, std::integer_sequence<T, ints...>) const noexcept
    {
        return element_type((base_type::m_values[ints](delems))...);
    }
 
public:
    KOKKOS_DEFAULTED_FUNCTION constexpr VectorField() = default;
 
    KOKKOS_DEFAULTED_FUNCTION ~VectorField() = default;
 
    KOKKOS_DEFAULTED_FUNCTION constexpr VectorField(VectorField const& other) = default;
 
    KOKKOS_DEFAULTED_FUNCTION constexpr VectorField(VectorField&& other) = default;
 
    template <class OElementType, class Allocator>
    KOKKOS_FUNCTION explicit constexpr VectorField(
            VectorFieldMem<OElementType, IdxRangeType, NDTag, Allocator>& other) noexcept
        : VectorField(other, std::make_index_sequence<base_type::NDims> {})
    {
    }
 
    // Disabled by SFINAE in the case of `ElementType` is not `const` to avoid write access
    template <
            class OElementType,
            class SFINAEElementType = ElementType,
            class = std::enable_if_t<std::is_const_v<SFINAEElementType>>,
            class Allocator>
    KOKKOS_FUNCTION explicit constexpr VectorField(
            VectorFieldMem<OElementType, IdxRangeType, NDTag, Allocator> const& other) noexcept
        : VectorField(other, std::make_index_sequence<base_type::NDims> {})
    {
    }
 
    template <class OElementType, class Allocator>
    VectorField(VectorFieldMem<OElementType, IdxRangeType, NDTag, Allocator>&& other) = delete;
 
    template <class OElementType>
    KOKKOS_FUNCTION constexpr VectorField(VectorField<
                                          OElementType,
                                          index_range_type,
                                          NDTag,
                                          MemorySpace,
                                          LayoutStridedPolicy> const& other) noexcept
        : VectorField(other, std::make_index_sequence<base_type::NDims> {})
    {
    }
 
    template <
            class... OElementType,
            class = std::enable_if_t<
                    std::conjunction_v<std::is_same<OElementType, ElementType>...>>,
            class = std::enable_if_t<sizeof...(OElementType) == base_type::NDims>>
    KOKKOS_FUNCTION VectorField(index_range_type const& idx_range, OElementType*... ptr)
        : base_type((field_type(ptr, idx_range))...)
    {
    }
 
    template <
            class... FieldType,
            class = std::enable_if_t<std::conjunction_v<std::is_same<FieldType, field_type>...>>>
    KOKKOS_FUNCTION constexpr VectorField(FieldType... fields) : base_type(std::move(fields)...)
    {
    }
 
    KOKKOS_DEFAULTED_FUNCTION constexpr VectorField& operator=(VectorField const& other) = default;
 
    KOKKOS_DEFAULTED_FUNCTION constexpr VectorField& operator=(VectorField&& other) = default;
 
    constexpr view_type span_cview() const
    {
        return view_type(*this);
    }
 
    constexpr span_type span_view() const
    {
        return *this;
    }
 
    template <class... ODDims>
    KOKKOS_FUNCTION element_type operator()(Idx<ODDims> const&... delems) const noexcept
    {
        Idx<ODDims...> delem_idx(delems...);
        return this->
        operator()(delem_idx, std::make_integer_sequence<int, element_type::size()> {});
    }
 
    template <class... ODDims, class = std::enable_if_t<sizeof...(ODDims) != 1>>
    KOKKOS_FUNCTION element_type operator()(Idx<ODDims...> const& delems) const noexcept
    {
        return this->operator()(delems, std::make_integer_sequence<int, element_type::size()> {});
    }
 
 
    template <class... QueryDDims>
    constexpr auto operator[](Idx<QueryDDims...> const& slice_spec)
    {
        return get_slice(slice_spec, std::make_index_sequence<base_type::NDims> {});
    }
 
    template <class... QueryDDims>
    constexpr auto operator[](IdxRange<QueryDDims...> const& oidx_range)
    {
        return get_slice(oidx_range, std::make_index_sequence<base_type::NDims> {});
    }
 
    template <class QueryTag>
    inline constexpr chunk_span_type get() const noexcept
    {
        static_assert(
                ddc::in_tags_v<QueryTag, NDTypeTag>,
                "requested Tag absent from TaggedVector");
        return base_type::m_values[ddc::type_seq_rank_v<QueryTag, NDTypeTag>].span_view();
    }
};
 
template <
        class ElementType,
        class IdxRangeType,
        class NDTag,
        class MemorySpace = Kokkos::DefaultExecutionSpace::memory_space,
        class LayoutStridedPolicy = std::experimental::layout_right>
using VectorConstField
        = VectorField<const ElementType, IdxRangeType, NDTag, MemorySpace, LayoutStridedPolicy>;