__cxx03/__mdspan/layout_stride.h

0001 // -*- C++ -*-
0002 //===----------------------------------------------------------------------===//
0003 //
0004 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
0005 // See https://llvm.org/LICENSE.txt for license information.
0006 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
0007 //
0008 //                        Kokkos v. 4.0
0009 //       Copyright (2022) National Technology & Engineering
0010 //               Solutions of Sandia, LLC (NTESS).
0011 //
0012 // Under the terms of Contract DE-NA0003525 with NTESS,
0013 // the U.S. Government retains certain rights in this software.
0014 //
0015 //===---------------------------------------------------------------------===//
0016
0017 #ifndef _LIBCPP___CXX03___MDSPAN_LAYOUT_STRIDE_H
0018 #define _LIBCPP___CXX03___MDSPAN_LAYOUT_STRIDE_H
0019
0020 #include <__cxx03/__assert>
0021 #include <__cxx03/__config>
0022 #include <__cxx03/__fwd/mdspan.h>
0023 #include <__cxx03/__mdspan/extents.h>
0024 #include <__cxx03/__type_traits/is_constructible.h>
0025 #include <__cxx03/__type_traits/is_convertible.h>
0026 #include <__cxx03/__type_traits/is_nothrow_constructible.h>
0027 #include <__cxx03/__utility/as_const.h>
0028 #include <__cxx03/__utility/integer_sequence.h>
0029 #include <__cxx03/__utility/swap.h>
0030 #include <__cxx03/array>
0031 #include <__cxx03/cinttypes>
0032 #include <__cxx03/cstddef>
0033 #include <__cxx03/limits>
0034
0035 #if !defined(_LIBCPP_HAS_NO_PRAGMA_SYSTEM_HEADER)
0036 #  pragma GCC system_header
0037 #endif
0038
0039 _LIBCPP_PUSH_MACROS
0040 #include <__cxx03/__undef_macros>
0041
0042 _LIBCPP_BEGIN_NAMESPACE_STD
0043
0044 #if _LIBCPP_STD_VER >= 23
0045
0046 namespace __mdspan_detail {
0047 template <class _Layout, class _Mapping>
0048 constexpr bool __is_mapping_of =
0049     is_same_v<typename _Layout::template mapping<typename _Mapping::extents_type>, _Mapping>;
0050
0051 template <class _Mapping>
0052 concept __layout_mapping_alike = requires {
0053   requires __is_mapping_of<typename _Mapping::layout_type, _Mapping>;
0054   requires __is_extents_v<typename _Mapping::extents_type>;
0055   { _Mapping::is_always_strided() } -> same_as<bool>;
0056   { _Mapping::is_always_exhaustive() } -> same_as<bool>;
0057   { _Mapping::is_always_unique() } -> same_as<bool>;
0058   bool_constant<_Mapping::is_always_strided()>::value;
0059   bool_constant<_Mapping::is_always_exhaustive()>::value;
0060   bool_constant<_Mapping::is_always_unique()>::value;
0061 };
0062 } // namespace __mdspan_detail
0063
0064 template <class _Extents>
0065 class layout_stride::mapping {
0066 public:
0067   static_assert(__mdspan_detail::__is_extents<_Extents>::value,
0068                 "layout_stride::mapping template argument must be a specialization of extents.");
0069
0070   using extents_type = _Extents;
0071   using index_type   = typename extents_type::index_type;
0072   using size_type    = typename extents_type::size_type;
0073   using rank_type    = typename extents_type::rank_type;
0074   using layout_type  = layout_stride;
0075
0076 private:
0077   static constexpr rank_type __rank_ = extents_type::rank();
0078
0079   // Used for default construction check and mandates
0080   _LIBCPP_HIDE_FROM_ABI static constexpr bool __required_span_size_is_representable(const extents_type& __ext) {
0081     if constexpr (__rank_ == 0)
0082       return true;
0083
0084     index_type __prod = __ext.extent(0);
0085     for (rank_type __r = 1; __r < __rank_; __r++) {
0086       bool __overflowed = __builtin_mul_overflow(__prod, __ext.extent(__r), &__prod);
0087       if (__overflowed)
0088         return false;
0089     }
0090     return true;
0091   }
0092
0093   template <class _OtherIndexType>
0094   _LIBCPP_HIDE_FROM_ABI static constexpr bool
0095   __required_span_size_is_representable(const extents_type& __ext, span<_OtherIndexType, __rank_> __strides) {
0096     if constexpr (__rank_ == 0)
0097       return true;
0098
0099     index_type __size = 1;
0100     for (rank_type __r = 0; __r < __rank_; __r++) {
0101       // We can only check correct conversion of _OtherIndexType if it is an integral
0102       if constexpr (is_integral_v<_OtherIndexType>) {
0103         using _CommonType = common_type_t<index_type, _OtherIndexType>;
0104         if (static_cast<_CommonType>(__strides[__r]) > static_cast<_CommonType>(numeric_limits<index_type>::max()))
0105           return false;
0106       }
0107       if (__ext.extent(__r) == static_cast<index_type>(0))
0108         return true;
0109       index_type __prod     = (__ext.extent(__r) - 1);
0110       bool __overflowed_mul = __builtin_mul_overflow(__prod, static_cast<index_type>(__strides[__r]), &__prod);
0111       if (__overflowed_mul)
0112         return false;
0113       bool __overflowed_add = __builtin_add_overflow(__size, __prod, &__size);
0114       if (__overflowed_add)
0115         return false;
0116     }
0117     return true;
0118   }
0119
0120   // compute offset of a strided layout mapping
0121   template <class _StridedMapping>
0122   _LIBCPP_HIDE_FROM_ABI static constexpr index_type __offset(const _StridedMapping& __mapping) {
0123     if constexpr (_StridedMapping::extents_type::rank() == 0) {
0124       return static_cast<index_type>(__mapping());
0125     } else if (__mapping.required_span_size() == static_cast<typename _StridedMapping::index_type>(0)) {
0126       return static_cast<index_type>(0);
0127     } else {
0128       return [&]<size_t... _Pos>(index_sequence<_Pos...>) {
0129         return static_cast<index_type>(__mapping((_Pos ? 0 : 0)...));
0130       }(make_index_sequence<__rank_>());
0131     }
0132   }
0133
0134   // compute the permutation for sorting the stride array
0135   // we never actually sort the stride array
0136   _LIBCPP_HIDE_FROM_ABI constexpr void __bubble_sort_by_strides(array<rank_type, __rank_>& __permute) const {
0137     for (rank_type __i = __rank_ - 1; __i > 0; __i--) {
0138       for (rank_type __r = 0; __r < __i; __r++) {
0139         if (__strides_[__permute[__r]] > __strides_[__permute[__r + 1]]) {
0140           swap(__permute[__r], __permute[__r + 1]);
0141         } else {
0142           // if two strides are the same then one of the associated extents must be 1 or 0
0143           // both could be, but you can't have one larger than 1 come first
0144           if ((__strides_[__permute[__r]] == __strides_[__permute[__r + 1]]) &&
0145               (__extents_.extent(__permute[__r]) > static_cast<index_type>(1)))
0146             swap(__permute[__r], __permute[__r + 1]);
0147         }
0148       }
0149     }
0150   }
0151
0152   static_assert(extents_type::rank_dynamic() > 0 || __required_span_size_is_representable(extents_type()),
0153                 "layout_stride::mapping product of static extents must be representable as index_type.");
0154
0155 public:
0156   // [mdspan.layout.stride.cons], constructors
0157   _LIBCPP_HIDE_FROM_ABI constexpr mapping() noexcept : __extents_(extents_type()) {
0158     // Note the nominal precondition is covered by above static assert since
0159     // if rank_dynamic is != 0 required_span_size is zero for default construction
0160     if constexpr (__rank_ > 0) {
0161       index_type __stride = 1;
0162       for (rank_type __r = __rank_ - 1; __r > static_cast<rank_type>(0); __r--) {
0163         __strides_[__r] = __stride;
0164         __stride *= __extents_.extent(__r);
0165       }
0166       __strides_[0] = __stride;
0167     }
0168   }
0169
0170   _LIBCPP_HIDE_FROM_ABI constexpr mapping(const mapping&) noexcept = default;
0171
0172   template <class _OtherIndexType>
0173     requires(is_convertible_v<const _OtherIndexType&, index_type> &&
0174              is_nothrow_constructible_v<index_type, const _OtherIndexType&>)
0175   _LIBCPP_HIDE_FROM_ABI constexpr mapping(const extents_type& __ext, span<_OtherIndexType, __rank_> __strides) noexcept
0176       : __extents_(__ext), __strides_([&]<size_t... _Pos>(index_sequence<_Pos...>) {
0177           return __mdspan_detail::__possibly_empty_array<index_type, __rank_>{
0178               static_cast<index_type>(std::as_const(__strides[_Pos]))...};
0179         }(make_index_sequence<__rank_>())) {
0180     _LIBCPP_ASSERT_VALID_ELEMENT_ACCESS(
0181         ([&]<size_t... _Pos>(index_sequence<_Pos...>) {
0182           // For integrals we can do a pre-conversion check, for other types not
0183           if constexpr (is_integral_v<_OtherIndexType>) {
0184             return ((__strides[_Pos] > static_cast<_OtherIndexType>(0)) && ... && true);
0185           } else {
0186             return ((static_cast<index_type>(__strides[_Pos]) > static_cast<index_type>(0)) && ... && true);
0187           }
0188         }(make_index_sequence<__rank_>())),
0189         "layout_stride::mapping ctor: all strides must be greater than 0");
0190     _LIBCPP_ASSERT_VALID_ELEMENT_ACCESS(
0191         __required_span_size_is_representable(__ext, __strides),
0192         "layout_stride::mapping ctor: required span size is not representable as index_type.");
0193     if constexpr (__rank_ > 1) {
0194       _LIBCPP_ASSERT_UNCATEGORIZED(
0195           ([&]<size_t... _Pos>(index_sequence<_Pos...>) {
0196             // basically sort the dimensions based on strides and extents, sorting is represented in permute array
0197             array<rank_type, __rank_> __permute{_Pos...};
0198             __bubble_sort_by_strides(__permute);
0199
0200             // check that this permutations represents a growing set
0201             for (rank_type __i = 1; __i < __rank_; __i++)
0202               if (static_cast<index_type>(__strides[__permute[__i]]) <
0203                   static_cast<index_type>(__strides[__permute[__i - 1]]) * __extents_.extent(__permute[__i - 1]))
0204                 return false;
0205             return true;
0206           }(make_index_sequence<__rank_>())),
0207           "layout_stride::mapping ctor: the provided extents and strides lead to a non-unique mapping");
0208     }
0209   }
0210
0211   template <class _OtherIndexType>
0212     requires(is_convertible_v<const _OtherIndexType&, index_type> &&
0213              is_nothrow_constructible_v<index_type, const _OtherIndexType&>)
0214   _LIBCPP_HIDE_FROM_ABI constexpr mapping(const extents_type& __ext,
0215                                           const array<_OtherIndexType, __rank_>& __strides) noexcept
0216       : mapping(__ext, span(__strides)) {}
0217
0218   template <class _StridedLayoutMapping>
0219     requires(__mdspan_detail::__layout_mapping_alike<_StridedLayoutMapping> &&
0220              is_constructible_v<extents_type, typename _StridedLayoutMapping::extents_type> &&
0221              _StridedLayoutMapping::is_always_unique() && _StridedLayoutMapping::is_always_strided())
0222   _LIBCPP_HIDE_FROM_ABI constexpr explicit(
0223       !(is_convertible_v<typename _StridedLayoutMapping::extents_type, extents_type> &&
0224         (__mdspan_detail::__is_mapping_of<layout_left, _StridedLayoutMapping> ||
0225          __mdspan_detail::__is_mapping_of<layout_right, _StridedLayoutMapping> ||
0226          __mdspan_detail::__is_mapping_of<layout_stride, _StridedLayoutMapping>)))
0227       mapping(const _StridedLayoutMapping& __other) noexcept
0228       : __extents_(__other.extents()), __strides_([&]<size_t... _Pos>(index_sequence<_Pos...>) {
0229           // stride() only compiles for rank > 0
0230           if constexpr (__rank_ > 0) {
0231             return __mdspan_detail::__possibly_empty_array<index_type, __rank_>{
0232                 static_cast<index_type>(__other.stride(_Pos))...};
0233           } else {
0234             return __mdspan_detail::__possibly_empty_array<index_type, 0>{};
0235           }
0236         }(make_index_sequence<__rank_>())) {
0237     // stride() only compiles for rank > 0
0238     if constexpr (__rank_ > 0) {
0239       _LIBCPP_ASSERT_VALID_ELEMENT_ACCESS(
0240           ([&]<size_t... _Pos>(index_sequence<_Pos...>) {
0241             return ((static_cast<index_type>(__other.stride(_Pos)) > static_cast<index_type>(0)) && ... && true);
0242           }(make_index_sequence<__rank_>())),
0243           "layout_stride::mapping converting ctor: all strides must be greater than 0");
0244     }
0245     _LIBCPP_ASSERT_VALID_ELEMENT_ACCESS(
0246         __mdspan_detail::__is_representable_as<index_type>(__other.required_span_size()),
0247         "layout_stride::mapping converting ctor: other.required_span_size() must be representable as index_type.");
0248     _LIBCPP_ASSERT_VALID_ELEMENT_ACCESS(static_cast<index_type>(0) == __offset(__other),
0249                                         "layout_stride::mapping converting ctor: base offset of mapping must be zero.");
0250   }
0251
0252   _LIBCPP_HIDE_FROM_ABI constexpr mapping& operator=(const mapping&) noexcept = default;
0253
0254   // [mdspan.layout.stride.obs], observers
0255   _LIBCPP_HIDE_FROM_ABI constexpr const extents_type& extents() const noexcept { return __extents_; }
0256
0257   _LIBCPP_HIDE_FROM_ABI constexpr array<index_type, __rank_> strides() const noexcept {
0258     return [&]<size_t... _Pos>(index_sequence<_Pos...>) {
0259       return array<index_type, __rank_>{__strides_[_Pos]...};
0260     }(make_index_sequence<__rank_>());
0261   }
0262
0263   _LIBCPP_HIDE_FROM_ABI constexpr index_type required_span_size() const noexcept {
0264     if constexpr (__rank_ == 0) {
0265       return static_cast<index_type>(1);
0266     } else {
0267       return [&]<size_t... _Pos>(index_sequence<_Pos...>) {
0268         if ((__extents_.extent(_Pos) * ... * 1) == 0)
0269           return static_cast<index_type>(0);
0270         else
0271           return static_cast<index_type>(
0272               static_cast<index_type>(1) +
0273               (((__extents_.extent(_Pos) - static_cast<index_type>(1)) * __strides_[_Pos]) + ... +
0274                static_cast<index_type>(0)));
0275       }(make_index_sequence<__rank_>());
0276     }
0277   }
0278
0279   template <class... _Indices>
0280     requires((sizeof...(_Indices) == __rank_) && (is_convertible_v<_Indices, index_type> && ...) &&
0281              (is_nothrow_constructible_v<index_type, _Indices> && ...))
0282   _LIBCPP_HIDE_FROM_ABI constexpr index_type operator()(_Indices... __idx) const noexcept {
0283     // Mappings are generally meant to be used for accessing allocations and are meant to guarantee to never
0284     // return a value exceeding required_span_size(), which is used to know how large an allocation one needs
0285     // Thus, this is a canonical point in multi-dimensional data structures to make invalid element access checks
0286     // However, mdspan does check this on its own, so for now we avoid double checking in hardened mode
0287     _LIBCPP_ASSERT_UNCATEGORIZED(__mdspan_detail::__is_multidimensional_index_in(__extents_, __idx...),
0288                                  "layout_stride::mapping: out of bounds indexing");
0289     return [&]<size_t... _Pos>(index_sequence<_Pos...>) {
0290       return ((static_cast<index_type>(__idx) * __strides_[_Pos]) + ... + index_type(0));
0291     }(make_index_sequence<sizeof...(_Indices)>());
0292   }
0293
0294   _LIBCPP_HIDE_FROM_ABI static constexpr bool is_always_unique() noexcept { return true; }
0295   _LIBCPP_HIDE_FROM_ABI static constexpr bool is_always_exhaustive() noexcept { return false; }
0296   _LIBCPP_HIDE_FROM_ABI static constexpr bool is_always_strided() noexcept { return true; }
0297
0298   _LIBCPP_HIDE_FROM_ABI static constexpr bool is_unique() noexcept { return true; }
0299   // The answer of this function is fairly complex in the case where one or more
0300   // extents are zero.
0301   // Technically it is meaningless to query is_exhaustive() in that case, but unfortunately
0302   // the way the standard defines this function, we can't give a simple true or false then.
0303   _LIBCPP_HIDE_FROM_ABI constexpr bool is_exhaustive() const noexcept {
0304     if constexpr (__rank_ == 0)
0305       return true;
0306     else {
0307       index_type __span_size = required_span_size();
0308       if (__span_size == static_cast<index_type>(0)) {
0309         if constexpr (__rank_ == 1)
0310           return __strides_[0] == 1;
0311         else {
0312           rank_type __r_largest = 0;
0313           for (rank_type __r = 1; __r < __rank_; __r++)
0314             if (__strides_[__r] > __strides_[__r_largest])
0315               __r_largest = __r;
0316           for (rank_type __r = 0; __r < __rank_; __r++)
0317             if (__extents_.extent(__r) == 0 && __r != __r_largest)
0318               return false;
0319           return true;
0320         }
0321       } else {
0322         return required_span_size() == [&]<size_t... _Pos>(index_sequence<_Pos...>) {
0323           return (__extents_.extent(_Pos) * ... * static_cast<index_type>(1));
0324         }(make_index_sequence<__rank_>());
0325       }
0326     }
0327   }
0328   _LIBCPP_HIDE_FROM_ABI static constexpr bool is_strided() noexcept { return true; }
0329
0330   // according to the standard layout_stride does not have a constraint on stride(r) for rank>0
0331   // it still has the precondition though
0332   _LIBCPP_HIDE_FROM_ABI constexpr index_type stride(rank_type __r) const noexcept {
0333     _LIBCPP_ASSERT_VALID_ELEMENT_ACCESS(__r < __rank_, "layout_stride::mapping::stride(): invalid rank index");
0334     return __strides_[__r];
0335   }
0336
0337   template <class _OtherMapping>
0338     requires(__mdspan_detail::__layout_mapping_alike<_OtherMapping> &&
0339              (_OtherMapping::extents_type::rank() == __rank_) && _OtherMapping::is_always_strided())
0340   _LIBCPP_HIDE_FROM_ABI friend constexpr bool operator==(const mapping& __lhs, const _OtherMapping& __rhs) noexcept {
0341     if (__offset(__rhs))
0342       return false;
0343     if constexpr (__rank_ == 0)
0344       return true;
0345     else {
0346       return __lhs.extents() == __rhs.extents() && [&]<size_t... _Pos>(index_sequence<_Pos...>) {
0347         // avoid warning when comparing signed and unsigner integers and pick the wider of two types
0348         using _CommonType = common_type_t<index_type, typename _OtherMapping::index_type>;
0349         return ((static_cast<_CommonType>(__lhs.stride(_Pos)) == static_cast<_CommonType>(__rhs.stride(_Pos))) && ... &&
0350                 true);
0351       }(make_index_sequence<__rank_>());
0352     }
0353   }
0354
0355 private:
0356   _LIBCPP_NO_UNIQUE_ADDRESS extents_type __extents_{};
0357   _LIBCPP_NO_UNIQUE_ADDRESS __mdspan_detail::__possibly_empty_array<index_type, __rank_> __strides_{};
0358 };
0359
0360 #endif // _LIBCPP_STD_VER >= 23
0361
0362 _LIBCPP_END_NAMESPACE_STD
0363
0364 _LIBCPP_POP_MACROS
0365
0366 #endif // _LIBCPP___CXX03___MDSPAN_LAYOUT_STRIDE_H