zip_layout.hpp

Go to the documentation of this file.

#pragma once
 
#include <uni20/core/compiler_attributes.hpp>
#include <uni20/mdspan/concepts.hpp>
#include <uni20/mdspan/mdspan.hpp>
 
#include <algorithm>
#include <array>
#include <cassert>
#include <cstddef>
#include <initializer_list>
#include <tuple>
#include <type_traits>
#include <utility>
 
namespace uni20
{
 
namespace detail
{
 
// Helper function to merge extents. If the extents are different, then one of them must be dynamic,
// otherwise they are not compatible. If one is dynamic and one is static, then return the static extent.
 
template <std::size_t A, std::size_t B> static constexpr std::size_t merge_extent()
{
  if constexpr (A == B)
    return A;
  else if constexpr (A == stdex::dynamic_extent)
    return B;
  else if constexpr (B == stdex::dynamic_extent)
    return A;
  else
    static_assert(A == B || A == stdex::dynamic_extent || B == stdex::dynamic_extent, "Incompatible static extents");
  return 0;
}
 
// Primary template: merge_pack over a pack of N extents
template <std::size_t... Ns> struct merge_pack;
 
// Base case: single value
template <std::size_t A> struct merge_pack<A>
{
    static constexpr std::size_t value = A;
};
 
// Recursive case: merge first two, then fold in the rest
template <std::size_t A, std::size_t B, std::size_t... Rest> struct merge_pack<A, B, Rest...>
{
    static constexpr std::size_t value = merge_pack<merge_extent<A, B>(), Rest...>::value;
};
 
// Given a set of mdspans, determine the common extents using the merge_extent() function above
template <typename FirstSpan, typename... OtherSpans> struct common_extents
{
    static constexpr std::size_t NumSpans = 1 + sizeof...(OtherSpans);
 
    // 1) All spans must have the same rank R
    static constexpr std::size_t R = FirstSpan::rank();
    static_assert((... && (OtherSpans::rank() == R)), "common_extents: ranks must match");
 
    using index_type = FirstSpan::index_type;
 
  private:
    // Fold-merge the static extents at dimension I
    template <std::size_t I> static constexpr std::size_t merged_static_extent()
    {
      // build a pack of the static-extent<I> of each Span
      return merge_pack<FirstSpan::extents_type::static_extent(I),
                        OtherSpans::extents_type::static_extent(I)...>::value;
    }
 
    // Build the extents type by expanding merged_static_extent<Is>()...
    template <std::size_t... Is> static constexpr auto make_extents(std::index_sequence<Is...>)
    {
      return stdex::extents<index_type,
                            merged_static_extent<Is>()... // one per dimension
                            >{};
    }
 
  public:
    using type = decltype(make_extents(std::make_index_sequence<R>{}));
 
  private:
    template <std::size_t... Is>
 
    static type make_impl(std::index_sequence<Is...>, FirstSpan const& first, OtherSpans const&... otherspans)
    {
      // Construct the extents object. All of the extents of each span must agree.
      type ext{first.extent(Is)...};
 
      auto do_check = [&](auto const& sp) { assert(ext == sp.extents()); };
 
      (do_check(otherspans), ...);
 
      return ext;
    }
 
  public:
    static type make(FirstSpan const& firstspan, OtherSpans const&... otherspans)
    {
      return make_impl(std::make_index_sequence<R>{}, firstspan, otherspans...);
    }
};
 
} // namespace detail
 
template <typename... Spans> using common_extents_t = typename detail::common_extents<Spans...>::type;
 
template <typename... Spans> auto make_common_extents(Spans... spans)
{
  return detail::common_extents<Spans...>::make(spans...);
}
 
namespace detail
{
 
template <typename T, std::size_t N1, std::size_t N2, std::size_t... I, std::size_t... J>
UNI20_HOST_DEVICE constexpr std::array<T, N1 + N2> concat_impl(const std::array<T, N1>& a1, const std::array<T, N2>& a2,
                                                               std::index_sequence<I...>, std::index_sequence<J...>)
{
  // pack-expand a1[I]..., then a2[J]...
  return {{a1[I]..., a2[J]...}};
}
 
template <typename T, std::size_t N, std::size_t... I>
UNI20_HOST_DEVICE constexpr std::array<T, N + 1> concat_impl(const std::array<T, N>& a, const T& x,
                                                             std::index_sequence<I...>)
{
  return {{a[I]..., x}};
}
 
template <typename T, std::size_t N, std::size_t... I>
UNI20_HOST_DEVICE constexpr std::array<T, N + 1> concat_impl(const T& x, const std::array<T, N>& a,
                                                             std::index_sequence<I...>)
{
  return {{x, a[I]...}};
}
 
} // namespace detail
 
template <typename T, std::size_t N1, std::size_t N2>
UNI20_HOST_DEVICE constexpr std::array<T, N1 + N2> concat(const std::array<T, N1>& a1, const std::array<T, N2>& a2)
{
  return detail::concat_impl(a1, a2, std::make_index_sequence<N1>{}, std::make_index_sequence<N2>{});
}
 
template <typename T, std::size_t N>
UNI20_HOST_DEVICE constexpr std::array<T, N + 1> concat(const T& x, const std::array<T, N>& a)
{
  return detail::concat_impl(x, a, std::make_index_sequence<N>{});
}
 
template <typename T, std::size_t N>
UNI20_HOST_DEVICE constexpr std::array<T, N + 1> concat(const std::array<T, N>& a, const T& x)
{
  return detail::concat_impl(a, x, std::make_index_sequence<N>{});
}
 
namespace detail
{
 
template <typename IndexType, std::size_t Rank, typename Mapping, std::size_t... I>
UNI20_HOST_DEVICE constexpr std::array<IndexType, Rank> mapping_strides_impl(Mapping const& mapping,
                                                                             std::index_sequence<I...>)
{
  return {static_cast<IndexType>(mapping.stride(I))...};
}
 
template <typename IndexType, std::size_t Rank, typename Mapping>
UNI20_HOST_DEVICE constexpr std::array<IndexType, Rank> mapping_strides(Mapping const& mapping)
{
  return mapping_strides_impl<IndexType, Rank>(mapping, std::make_index_sequence<Rank>{});
}
 
// build a std::tuple of N default-constructed T’s
template <typename T, std::size_t... I> constexpr auto make_tuple_n_impl(std::index_sequence<I...>)
{
  // the comma-expression (static_cast<void>(I), T{}) ensures we produce T{} N times
  return std::tuple{(static_cast<void>(I), T{})...};
}
 
} // namespace detail
 
template <typename T, std::size_t N>
using tuple_n = decltype(detail::make_tuple_n_impl<T>(std::make_index_sequence<N>{}));
 
template <std::size_t NumSpans> struct StridedZipLayout
{
    static_assert(NumSpans >= 1, "StridedZipLayout requires at least one span");
 
    static constexpr std::size_t num_spans = NumSpans;
 
    template <typename Ext> struct mapping
    {
        using layout_type = StridedZipLayout; 
        using extents_type = Ext;             
        using index_type = typename Ext::index_type;
        using rank_type = typename Ext::rank_type;
 
        using offset_type = tuple_n<index_type, NumSpans>;
 
        using mapping_type = mapping<Ext>;
 
        static constexpr std::size_t num_spans = NumSpans;
 
        UNI20_HOST_DEVICE static constexpr bool is_always_unique() noexcept { return true; }
        UNI20_HOST_DEVICE static constexpr bool is_always_exhaustive() noexcept { return false; }
        UNI20_HOST_DEVICE static constexpr bool is_always_strided() noexcept { return false; }
 
        [[nodiscard]] UNI20_HOST_DEVICE constexpr bool is_unique() const noexcept { return true; }
        [[nodiscard]] UNI20_HOST_DEVICE constexpr bool is_exhaustive() const noexcept { return false; }
 
        // std::size_t required_span_size() const does not make sense for a StridedZipLayout
 
        [[nodiscard]] UNI20_HOST_DEVICE constexpr bool is_strided() const noexcept
        {
          for (std::size_t d = 0; d < Ext::rank(); ++d)
          {
            auto s0 = all_strides_[0][d];
            for (std::size_t i = 1; i < NumSpans; ++i)
            {
              if (all_strides_[i][d] != s0) return false;
            }
          }
          return true;
        }
 
        [[nodiscard]] UNI20_HOST_DEVICE constexpr std::array<index_type, Ext::rank()> strides() const noexcept
        {
          assert(this->is_strided());
          return all_strides_[0];
        }
 
        [[nodiscard]] UNI20_HOST_DEVICE constexpr index_type stride(rank_type r) const noexcept
        {
          assert(this->is_strided());
          return all_strides_[0][r];
        }
 
        [[nodiscard]] UNI20_HOST_DEVICE std::array<std::array<index_type, Ext::rank()>, num_spans> const&
        all_strides() const noexcept
        {
          return all_strides_;
        }
 
        template <typename... Mappings>
        UNI20_HOST_DEVICE constexpr mapping(extents_type const& exts, Mappings... maps) noexcept
            : extents_(exts), all_strides_{detail::mapping_strides<index_type, Ext::rank()>(maps)...}
        {}
 
        UNI20_HOST_DEVICE constexpr mapping(
            extents_type const& exts,
            std::array<std::array<index_type, Ext::rank()>, NumSpans> const& strides_pack) noexcept
            : extents_(exts), all_strides_(strides_pack)
        {}
 
        template <typename OtherExt>
        UNI20_HOST_DEVICE constexpr mapping(
            std::array<index_type, Ext::rank()> const& new_strides,
            typename StridedZipLayout<num_spans - 1>::template mapping<OtherExt> const& other) noexcept
            : extents_(other.extents()), all_strides_(concat(new_strides, other.all_strides()))
        {}
 
        template <typename OtherExt>
        UNI20_HOST_DEVICE constexpr mapping(
            typename StridedZipLayout<num_spans - 1>::template mapping<OtherExt> const& other,
            std::array<index_type, Ext::rank()> const& new_strides) noexcept
            : extents_(other.extents()), all_strides_(concat(other.all_strides(), new_strides))
        {}
 
        template <typename LeftMap, typename RightMap>
          requires(LeftMap::num_spans + RightMap::num_spans == num_spans)
        UNI20_HOST_DEVICE constexpr mapping(LeftMap const& left, RightMap const& right) noexcept
            : extents_(left.extents()), all_strides_(concat(left.all_strides(), right.all_strides()))
        {}
 
        [[nodiscard]] UNI20_HOST_DEVICE constexpr extents_type const& extents() const noexcept { return extents_; }
 
        template <typename... Idx> UNI20_HOST_DEVICE constexpr offset_type operator()(Idx... idxs) const noexcept
        {
          static_assert(sizeof...(Idx) == Ext::rank(), "Wrong number of indices");
          // pack the multi-index into an array
          std::array<index_type, Ext::rank()> ix{static_cast<index_type>(idxs)...};
          // dispatch over spans
          return offset_impl(ix, std::make_index_sequence<NumSpans>{});
        }
 
      private:
        // For a given span S, compute its offset by folding over all dims D:
        template <std::size_t S, std::size_t... D>
        [[nodiscard]] UNI20_HOST_DEVICE constexpr index_type span_offset(std::array<index_type, Ext::rank()> const& ix,
                                                                         std::index_sequence<D...>) const noexcept
        {
          // fold: 0 + (stride[S][0]*ix[0]) + (stride[S][1]*ix[1]) + …
          return ((all_strides_[S][D] * ix[D]) + ... + index_type(0));
        }
 
        // Build the tuple of all span-offsets in one go:
        template <std::size_t... S>
        [[nodiscard]] UNI20_HOST_DEVICE constexpr offset_type offset_impl(std::array<index_type, Ext::rank()> const& ix,
                                                                          std::index_sequence<S...>) const noexcept
        {
          // pack-expand: span_offset<0>(ix), span_offset<1>(ix), …
          return offset_type{span_offset<S>(ix, std::make_index_sequence<Ext::rank()>{})...};
        }
 
        extents_type extents_;
        std::array<std::array<index_type, Ext::rank()>, num_spans> all_strides_;
    };
};
 
template <typename... Layouts> struct GeneralZipLayout
{
    template <typename Extents> struct mapping
    {
        using layout_type = GeneralZipLayout;
        using extents_type = Extents;
        using index_type = typename Extents::index_type;
        using rank_type = typename Extents::rank_type;
        UNI20_HOST_DEVICE static constexpr rank_type rank() noexcept { return Extents::rank(); }
        UNI20_HOST_DEVICE static constexpr rank_type rank_dynamic() noexcept { return Extents::rank_dynamic(); }
        using offset_type = std::tuple<span_offset_t<typename Layouts::template mapping<Extents>>...>;
        using mapping_type = mapping<Extents>;
 
        UNI20_HOST_DEVICE static constexpr bool is_always_unique() noexcept { return true; }
        UNI20_HOST_DEVICE static constexpr bool is_always_exhaustive() noexcept { return false; }
        UNI20_HOST_DEVICE static constexpr bool is_always_strided() noexcept { return false; }
 
        UNI20_HOST_DEVICE explicit constexpr mapping(Extents const& ext,
                                                     typename Layouts::template mapping<Extents> const&... m) noexcept
            : extents_(ext), impls_{m...}
        {}
 
        [[nodiscard]] UNI20_HOST_DEVICE constexpr extents_type const& extents() const noexcept { return extents_; }
 
        [[nodiscard]] UNI20_HOST_DEVICE constexpr std::size_t required_span_size() const noexcept
        {
          std::size_t max_sz = 0;
          std::apply(
              [&](auto const&... m) {
                (void)std::initializer_list<int>{
                    (max_sz = std::max<std::size_t>(max_sz, m.required_span_size()), 0)...};
              },
              impls_);
          return max_sz;
        }
 
        template <typename... Idx> UNI20_HOST_DEVICE constexpr offset_type operator()(Idx... idxs) const noexcept
        {
          return std::apply([&](auto const&... m) { return offset_type{m(idxs...)...}; }, impls_);
        }
 
        UNI20_HOST_DEVICE constexpr bool operator==(mapping const& o) const noexcept
        {
          return extents_ == o.extents_ && impls_ == o.impls_;
        }
        UNI20_HOST_DEVICE constexpr bool operator!=(mapping const& o) const noexcept { return !(*this == o); }
 
      private:
        extents_type extents_;
        std::tuple<typename Layouts::template mapping<Extents>...> impls_;
    };
};
 
template <typename... Spans> struct zip_layout_selector
{
    using type = GeneralZipLayout<typename Spans::layout_type...>;
};
 
template <StridedMdspanLike... Spans> struct zip_layout_selector<Spans...>
{
    using type = StridedZipLayout<sizeof...(Spans)>;
};
 
template <typename... Spans> using zip_layout_t = typename zip_layout_selector<Spans...>::type;
 
} // namespace uni20