mdspan/concepts.hpp

Go to the documentation of this file.

#pragma once
 
#include <uni20/core/compiler_attributes.hpp>
#include <uni20/mdspan/mdspan.hpp>
 
#include <concepts>
#include <type_traits>
#include <utility>
 
namespace uni20
{
 
template <typename AP, typename = void> struct span_offset_type
{
    using type = std::size_t;
};
 
template <typename AP> struct span_offset_type<AP, std::void_t<typename AP::offset_type>>
{
    using type = typename AP::offset_type;
};
 
template <typename AP> using span_offset_t = typename span_offset_type<AP>::type;
 
template <class AP>
concept AccessorPolicy = requires {
  typename AP::element_type;
  typename AP::data_handle_type;
  typename AP::offset_policy;
  typename AP::reference;
} && requires(AP a, typename AP::data_handle_type dh, span_offset_t<AP> off) {
  { a.offset(dh, off) } -> std::convertible_to<typename AP::offset_policy::data_handle_type>;
  { a.access(dh, off) } -> std::same_as<typename AP::reference>;
};
 
struct host_access_domain
{};
 
struct cuda_access_domain
{};
 
template <class Accessor, class Domain> inline constexpr bool enable_accessor_in_domain = false;
 
template <class ElementType>
inline constexpr bool enable_accessor_in_domain<stdex::default_accessor<ElementType>, host_access_domain> = true;
 
template <class Accessor, class Domain>
concept AccessorInDomain = AccessorPolicy<std::remove_cvref_t<Accessor>> &&
                           enable_accessor_in_domain<std::remove_cvref_t<Accessor>, std::remove_cvref_t<Domain>>;
 
template <class Accessor>
concept HostAccessibleAccessor = AccessorInDomain<Accessor, host_access_domain>;
 
template <class Accessor>
concept CudaAccessibleAccessor = AccessorInDomain<Accessor, cuda_access_domain>;
 
template <AccessorPolicy Accessor>
  requires std::is_same_v<typename Accessor::reference, typename Accessor::element_type&>
class const_accessor_adaptor {
  public:
    using element_type = typename Accessor::element_type const;
    using reference = element_type&;
    using data_handle_type = typename Accessor::data_handle_type;
    using offset_policy = const_accessor_adaptor;
    using offset_type = span_offset_t<Accessor>;
 
    UNI20_HOST_DEVICE constexpr const_accessor_adaptor(Accessor const& to_be_wrapped) : wrapped_{to_be_wrapped} {}
 
    UNI20_HOST_DEVICE constexpr reference access(data_handle_type p, offset_type i) const
    {
      return wrapped_.access(p, i);
    }
 
    UNI20_HOST_DEVICE constexpr data_handle_type offset(data_handle_type p, offset_type i) const
    {
      return wrapped_.offset(p, i);
    }
 
  private:
    Accessor wrapped_;
};
 
template <class Accessor, class Domain>
inline constexpr bool enable_accessor_in_domain<const_accessor_adaptor<Accessor>, Domain> =
    enable_accessor_in_domain<Accessor, Domain>;
 
//
// const_accessor overloads: build a read-only accessor from a mutable one.
//
 
template <typename T>
constexpr stdex::default_accessor<T const> const_accessor(stdex::default_accessor<T> const& accessor_policy)
{
  (void)accessor_policy;
  return stdex::default_accessor<T const>();
}
 
template <class Accessor> struct is_default_accessor : std::false_type
{};
 
template <class ElementType> struct is_default_accessor<stdex::default_accessor<ElementType>> : std::true_type
{};
 
template <class Accessor>
inline constexpr bool is_default_accessor_v = is_default_accessor<std::remove_cvref_t<Accessor>>::value;
 
// /// \brief const_accessor on const_default_accessor yields itself.
// template <typename T> constexpr default_accessor<T const> const_accessor(default_accessor<T const> const&)
// {
//   return default_accessor<T const>();
// }
 
template <AccessorPolicy Acc>
  requires(!std::is_const_v<typename Acc::element_type> &&
           std::is_same_v<typename Acc::reference, typename Acc::element_type&>)
constexpr auto const_accessor(Acc const& acc)
{
  return const_accessor_adaptor<Acc>{acc};
}
 
template <AccessorPolicy Acc>
  requires(std::is_const_v<typename Acc::element_type> && !is_default_accessor_v<Acc>)
constexpr Acc const_accessor(Acc const& acc)
{
  return acc;
}
 
template <AccessorPolicy Acc> using const_accessor_t = decltype(const_accessor(std::declval<Acc>()));
 
namespace detail
{
 
template <class S> using span_type_t = std::remove_cvref_t<S>;
 
template <class S, std::size_t... Axis> consteval bool span_has_ranked_subscript(std::index_sequence<Axis...>)
{
  using span_type = span_type_t<S>;
  return requires(span_type const& span, typename span_type::index_type index) {
    { span.operator[](((void)Axis, index)...) } -> std::same_as<typename span_type::reference>;
  };
}
 
template <class S, std::size_t... Axis> consteval bool span_has_ranked_assignment(std::index_sequence<Axis...>)
{
  using span_type = span_type_t<S>;
  return requires(span_type& span, typename span_type::index_type index, typename span_type::value_type value) {
    span.operator[](((void)Axis, index)...) = value;
  };
}
 
template <class S>
concept SpanMetadata =
    requires(span_type_t<S> const& span, std::size_t axis) {
      typename span_type_t<S>::element_type;
      typename span_type_t<S>::value_type;
      typename span_type_t<S>::index_type;
      typename span_type_t<S>::extents_type;
      typename span_type_t<S>::layout_type;
      typename span_type_t<S>::mapping_type;
      typename span_type_t<S>::accessor_type;
      typename span_type_t<S>::data_handle_type;
      typename span_type_t<S>::reference;
      typename std::integral_constant<std::size_t, span_type_t<S>::rank()>;
 
      { span.extents() } -> std::convertible_to<typename span_type_t<S>::extents_type>;
      { span.extent(axis) } -> std::convertible_to<typename span_type_t<S>::index_type>;
      { span.mapping() } -> std::convertible_to<typename span_type_t<S>::mapping_type>;
      { span.accessor() } -> std::convertible_to<typename span_type_t<S>::accessor_type>;
    } && std::same_as<typename span_type_t<S>::value_type, std::remove_cv_t<typename span_type_t<S>::element_type>> &&
    std::same_as<typename span_type_t<S>::index_type, typename span_type_t<S>::extents_type::index_type> &&
    std::same_as<typename span_type_t<S>::mapping_type,
                 typename span_type_t<S>::layout_type::template mapping<typename span_type_t<S>::extents_type>> &&
    std::same_as<typename span_type_t<S>::data_handle_type, typename span_type_t<S>::accessor_type::data_handle_type> &&
    std::same_as<typename span_type_t<S>::reference, typename span_type_t<S>::accessor_type::reference> &&
    (span_type_t<S>::rank() == span_type_t<S>::extents_type::rank());
 
template <class S>
concept SpanDescriptor = SpanMetadata<S> && requires(span_type_t<S> const& span) {
  { span.data_handle() } -> std::convertible_to<typename span_type_t<S>::data_handle_type>;
};
 
template <class S>
concept MdspecDescriptor = SpanMetadata<S> && requires(span_type_t<S> const& span) {
  typename span_type_t<S>::data_descriptor_type;
  { span.data_descriptor() } -> std::convertible_to<typename span_type_t<S>::data_descriptor_type const&>;
};
 
} // namespace detail
 
template <class S>
concept MdspanLike = detail::SpanDescriptor<S> && AccessorPolicy<typename detail::span_type_t<S>::accessor_type> &&
                     detail::span_has_ranked_subscript<S>(std::make_index_sequence<detail::span_type_t<S>::rank()>{});
 
template <class S>
concept MdspecLike =
    MdspanLike<S> || (detail::MdspecDescriptor<S> && AccessorPolicy<typename detail::span_type_t<S>::accessor_type>);
 
template <class S>
concept HostAccessibleMdspan = MdspanLike<S> && HostAccessibleAccessor<typename detail::span_type_t<S>::accessor_type>;
 
template <class S>
concept CudaAccessibleMdspan = MdspanLike<S> && CudaAccessibleAccessor<typename detail::span_type_t<S>::accessor_type>;
 
template <class S>
concept HostAccessibleMdspec = MdspecLike<S> && HostAccessibleAccessor<typename detail::span_type_t<S>::accessor_type>;
 
template <class S>
concept CudaAccessibleMdspec = MdspecLike<S> && CudaAccessibleAccessor<typename detail::span_type_t<S>::accessor_type>;
 
namespace detail
{
 
template <class Span, class... Index>
concept MdspanIndexPack = MdspanLike<Span> && (sizeof...(Index) == span_type_t<Span>::rank()) &&
                          (std::is_convertible_v<Index, typename span_type_t<Span>::index_type> && ...) &&
                          (std::is_nothrow_constructible_v<typename span_type_t<Span>::index_type, Index> && ...);
 
template <class Accessor>
concept AccessorReferenceAssignable =
    requires(Accessor const& accessor, typename Accessor::data_handle_type handle, span_offset_t<Accessor> offset,
             std::remove_const_t<typename Accessor::element_type> value) { accessor.access(handle, offset) = value; };
 
} // namespace detail
 
template <MdspanLike Span, class... Index>
  requires detail::MdspanIndexPack<Span, Index...>
constexpr decltype(auto) const_access(Span const& span, Index... indices)
{
  using index_type = typename detail::span_type_t<Span>::index_type;
  auto accessor = const_accessor(span.accessor());
  return accessor.access(span.data_handle(), span.mapping()(static_cast<index_type>(std::move(indices))...));
}
 
template <MdspanLike Span> [[nodiscard]] constexpr auto make_const_mdspan(Span const& span)
{
  auto accessor = const_accessor(span.accessor());
  using accessor_type = decltype(accessor);
  using span_type = std::remove_cvref_t<Span>;
  using const_mdspan_type = stdex::mdspan<typename accessor_type::element_type, typename span_type::extents_type,
                                          typename span_type::layout_type, accessor_type>;
  return const_mdspan_type{span.data_handle(), span.mapping(), std::move(accessor)};
}
 
template <class S>
concept MutableMdspanLike =
    MdspanLike<S> && (!std::is_const_v<typename detail::span_type_t<S>::element_type>) &&
    detail::span_has_ranked_assignment<S>(std::make_index_sequence<detail::span_type_t<S>::rank()>{});
 
template <class S>
concept MutableMdspecLike =
    MdspecLike<S> && (!std::is_const_v<typename detail::span_type_t<S>::element_type>) &&
    (MutableMdspanLike<S> || detail::AccessorReferenceAssignable<typename detail::span_type_t<S>::accessor_type>);
 
template <class MDS>
concept StridedMdspanLike = MdspanLike<MDS> && detail::span_type_t<MDS>::is_always_strided() &&
                            requires(detail::span_type_t<MDS> const& span, std::size_t axis) {
                              {
                                span.stride(axis)
                              } -> std::convertible_to<typename detail::span_type_t<MDS>::index_type>;
                              {
                                span.mapping().stride(axis)
                              } -> std::convertible_to<typename detail::span_type_t<MDS>::index_type>;
                            };
 
template <class MDS>
concept MutableStridedMdspanLike = MutableMdspanLike<MDS> && StridedMdspanLike<MDS>;
 
template <class S>
concept StridedMdspecLike = MdspecLike<S> && detail::span_type_t<S>::is_always_strided() &&
                            requires(detail::span_type_t<S> const& span, std::size_t axis) {
                              { span.stride(axis) } -> std::convertible_to<typename detail::span_type_t<S>::index_type>;
                              {
                                span.mapping().stride(axis)
                              } -> std::convertible_to<typename detail::span_type_t<S>::index_type>;
                            };
 
template <class S>
concept MutableStridedMdspecLike = MutableMdspecLike<S> && StridedMdspecLike<S>;
 
template <class MDS, std::size_t Rank>
concept RankedMdspanLike = MdspanLike<std::remove_cvref_t<MDS>> && (std::remove_cvref_t<MDS>::rank() == Rank);
 
template <class S, std::size_t Rank>
concept RankedMdspecLike = MdspecLike<std::remove_cvref_t<S>> && (std::remove_cvref_t<S>::rank() == Rank);
 
template <class S>
concept DiagonalMdspecLike = MdspecLike<std::remove_cvref_t<S>> && requires(std::remove_cvref_t<S> const& span) {
  { diagonal_components(span) };
  requires RankedMdspecLike<decltype(diagonal_components(span)), 1>;
  requires std::same_as<typename std::remove_cvref_t<decltype(diagonal_components(span))>::element_type,
                        typename std::remove_cvref_t<S>::element_type>;
};
 
template <class S>
concept MutableDiagonalMdspecLike =
    MutableMdspecLike<std::remove_cvref_t<S>> && DiagonalMdspecLike<std::remove_cvref_t<S>> &&
    MutableMdspecLike<decltype(diagonal_components(std::declval<std::remove_cvref_t<S> const&>()))>;
 
template <class S, std::size_t Rank>
concept MutableRankedMdspecLike = MutableMdspecLike<S> && RankedMdspecLike<S, Rank>;
 
template <class MDS, std::size_t Rank>
concept MutableRankedMdspanLike =
    MutableMdspanLike<std::remove_cvref_t<MDS>> && (std::remove_cvref_t<MDS>::rank() == Rank);
 
template <class MDS>
concept DefaultAccessorMdspanLike =
    MdspanLike<std::remove_cvref_t<MDS>> && is_default_accessor_v<typename std::remove_cvref_t<MDS>::accessor_type>;
 
template <class MDS, std::size_t Rank>
concept RankedStridedMdspanLike = RankedMdspanLike<MDS, Rank> && StridedMdspanLike<std::remove_cvref_t<MDS>>;
 
template <class S, std::size_t Rank>
concept RankedStridedMdspecLike = RankedMdspecLike<S, Rank> && StridedMdspecLike<std::remove_cvref_t<S>>;
 
template <class S, std::size_t Rank>
concept MutableRankedStridedMdspecLike =
    MutableRankedMdspecLike<S, Rank> && MutableStridedMdspecLike<std::remove_cvref_t<S>>;
 
template <class MDS, std::size_t Rank>
concept MutableRankedStridedMdspanLike =
    MutableRankedMdspanLike<MDS, Rank> && MutableStridedMdspanLike<std::remove_cvref_t<MDS>>;
 
namespace detail
{
 
template <StridedMdspecLike S, size_t... I> constexpr auto strides_impl(S const& s, std::index_sequence<I...>)
{
  using index_type = typename S::index_type;
  // fold the pack I... into an array by calling s.mapping().stride(I) for each I
  return std::array<index_type, sizeof...(I)>{s.mapping().stride(I)...};
}
 
} // namespace detail
 
template <StridedMdspecLike S> auto strides(S const& s)
{
  return detail::strides_impl(s, std::make_index_sequence<S::rank()>{});
}
 
} // namespace uni20