// -*- C++ -*- //===----------------------------------------------------------------------===// // // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. // See https://llvm.org/LICENSE.txt for license information. // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception // SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. // // Kokkos v. 4.0 // Copyright (2022) National Technology & Engineering // Solutions of Sandia, LLC (NTESS). // // Under the terms of Contract DE-NA0003525 with NTESS, // the U.S. Government retains certain rights in this software. // //===---------------------------------------------------------------------===// #ifndef _LIBCUDACXX___MDSPAN_MDSPAN_HPP #define _LIBCUDACXX___MDSPAN_MDSPAN_HPP #include #if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC) # pragma GCC system_header #elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG) # pragma clang system_header #elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC) # pragma system_header #endif // no system header #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include #include _LIBCUDACXX_BEGIN_NAMESPACE_STD template > class mdspan : private __mdspan_ebco, _AccessorPolicy> { private: static_assert(__mdspan_detail::__is_extents_v<_Extents>, "mdspan: Extents template parameter must be a specialization of extents."); static_assert(!_CCCL_TRAIT(is_array, _ElementType), "mdspan: ElementType template parameter may not be an array type"); static_assert(!_CCCL_TRAIT(is_abstract, _ElementType), "mdspan: ElementType template parameter may not be an abstract class"); static_assert(_CCCL_TRAIT(is_same, _ElementType, typename _AccessorPolicy::element_type), "mdspan: ElementType template parameter must match AccessorPolicy::element_type"); static_assert(__mdspan_detail::__is_valid_layout_mapping<_LayoutPolicy, _Extents>, "mdspan: LayoutPolicy template parameter is invalid. A common mistake is to pass a layout mapping " "instead of a layout policy"); template friend class mdspan; public: using extents_type = _Extents; using layout_type = _LayoutPolicy; using accessor_type = _AccessorPolicy; using mapping_type = typename layout_type::template mapping; using element_type = _ElementType; using value_type = remove_cv_t; using index_type = typename extents_type::index_type; using size_type = typename extents_type::size_type; using rank_type = typename extents_type::rank_type; using data_handle_type = typename accessor_type::data_handle_type; using reference = typename accessor_type::reference; using __base = __mdspan_ebco, _AccessorPolicy>; _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI static constexpr rank_type rank() noexcept { return extents_type::rank(); } _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI static constexpr rank_type rank_dynamic() noexcept { return extents_type::rank_dynamic(); } _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI static constexpr size_t static_extent(rank_type __r) noexcept { return extents_type::static_extent(__r); } _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr index_type extent(rank_type __r) const noexcept { return mapping().extents().extent(__r); }; public: //-------------------------------------------------------------------------------- // [mdspan.mdspan.cons], mdspan constructors, assignment, and destructor template static constexpr bool __can_default_construct = (_Extents2::rank_dynamic() > 0) && _CCCL_TRAIT(is_default_constructible, data_handle_type) && _CCCL_TRAIT(is_default_constructible, mapping_type) && _CCCL_TRAIT(is_default_constructible, accessor_type); _CCCL_TEMPLATE(class _Extents2 = _Extents) _CCCL_REQUIRES(__can_default_construct<_Extents2>) _LIBCUDACXX_HIDE_FROM_ABI constexpr mdspan() noexcept( _CCCL_TRAIT(is_nothrow_default_constructible, data_handle_type) && _CCCL_TRAIT(is_nothrow_default_constructible, mapping_type) && _CCCL_TRAIT(is_nothrow_default_constructible, accessor_type)) : __base() {} _CCCL_HIDE_FROM_ABI constexpr mdspan(const mdspan&) = default; _CCCL_HIDE_FROM_ABI constexpr mdspan(mdspan&&) = default; template static constexpr bool __matches_dynamic_rank = (_Size == extents_type::rank_dynamic()); template static constexpr bool __matches_static_rank = (_Size == extents_type::rank()) && (_Size != extents_type::rank_dynamic()); template static constexpr bool __can_construct_from_handle_and_variadic = (__matches_dynamic_rank || __matches_static_rank) && __mdspan_detail::__all_convertible_to_index_type && _CCCL_TRAIT(is_constructible, mapping_type, extents_type) && _CCCL_TRAIT(is_default_constructible, accessor_type); _CCCL_TEMPLATE(class... _OtherIndexTypes) _CCCL_REQUIRES(__can_construct_from_handle_and_variadic<_OtherIndexTypes...>) _LIBCUDACXX_HIDE_FROM_ABI explicit constexpr mdspan(data_handle_type __p, _OtherIndexTypes... __exts) : __base(_CUDA_VSTD::move(__p), extents_type(static_cast(_CUDA_VSTD::move(__exts))...)) {} template static constexpr bool __is_constructible_from_index_type = _CCCL_TRAIT(is_convertible, const _OtherIndexType&, index_type) && _CCCL_TRAIT(is_nothrow_constructible, index_type, const _OtherIndexType&) && _CCCL_TRAIT(is_constructible, mapping_type, extents_type) && _CCCL_TRAIT(is_default_constructible, accessor_type); _CCCL_TEMPLATE(class _OtherIndexType, size_t _Size) _CCCL_REQUIRES(__matches_dynamic_rank<_Size> _CCCL_AND __is_constructible_from_index_type<_OtherIndexType>) _LIBCUDACXX_HIDE_FROM_ABI constexpr mdspan(data_handle_type __p, const array<_OtherIndexType, _Size>& __exts) : __base(_CUDA_VSTD::move(__p), extents_type{__exts}) {} _CCCL_TEMPLATE(class _OtherIndexType, size_t _Size) _CCCL_REQUIRES(__matches_static_rank<_Size> _CCCL_AND __is_constructible_from_index_type<_OtherIndexType>) _LIBCUDACXX_HIDE_FROM_ABI explicit constexpr mdspan(data_handle_type __p, const array<_OtherIndexType, _Size>& __exts) : __base(_CUDA_VSTD::move(__p), extents_type{__exts}) {} _CCCL_TEMPLATE(class _OtherIndexType, size_t _Size) _CCCL_REQUIRES(__matches_dynamic_rank<_Size> _CCCL_AND __is_constructible_from_index_type<_OtherIndexType>) _LIBCUDACXX_HIDE_FROM_ABI constexpr mdspan(data_handle_type __p, span<_OtherIndexType, _Size> __exts) : __base(_CUDA_VSTD::move(__p), extents_type{__exts}) {} _CCCL_TEMPLATE(class _OtherIndexType, size_t _Size) _CCCL_REQUIRES(__matches_static_rank<_Size> _CCCL_AND __is_constructible_from_index_type<_OtherIndexType>) _LIBCUDACXX_HIDE_FROM_ABI explicit constexpr mdspan(data_handle_type __p, span<_OtherIndexType, _Size> __exts) : __base(_CUDA_VSTD::move(__p), extents_type{__exts}) {} _CCCL_TEMPLATE(class _AccessorPolicy2 = _AccessorPolicy, class _Mapping2 = mapping_type) _CCCL_REQUIRES(_CCCL_TRAIT(is_default_constructible, _AccessorPolicy2) _CCCL_AND _CCCL_TRAIT(is_constructible, _Mapping2, const extents_type&)) _LIBCUDACXX_HIDE_FROM_ABI constexpr mdspan(data_handle_type __p, const extents_type& __exts) : __base(_CUDA_VSTD::move(__p), __exts) {} _CCCL_TEMPLATE(class _AccessorPolicy2 = _AccessorPolicy) _CCCL_REQUIRES(_CCCL_TRAIT(is_default_constructible, _AccessorPolicy2)) _LIBCUDACXX_HIDE_FROM_ABI constexpr mdspan(data_handle_type __p, const mapping_type& __m) : __base(_CUDA_VSTD::move(__p), __m) {} _LIBCUDACXX_HIDE_FROM_ABI constexpr mdspan(data_handle_type __p, const mapping_type& __m, const accessor_type& __a) : __base(_CUDA_VSTD::move(__p), __m, __a) {} template static constexpr bool __is_convertible_from = _CCCL_TRAIT(is_constructible, mapping_type, const typename _OtherLayoutPolicy::template mapping<_OtherExtents>&) && _CCCL_TRAIT(is_constructible, accessor_type, const _OtherAccessor&); template static constexpr bool __is_implicit_convertible_from = _CCCL_TRAIT(is_convertible, const typename _OtherLayoutPolicy::template mapping<_OtherExtents>&, mapping_type) && _CCCL_TRAIT(is_convertible, const _OtherAccessor&, accessor_type); _CCCL_TEMPLATE(class _OtherElementType, class _OtherExtents, class _OtherLayoutPolicy, class _OtherAccessor) _CCCL_REQUIRES(__is_convertible_from<_OtherExtents, _OtherLayoutPolicy, _OtherAccessor> // _CCCL_AND __is_implicit_convertible_from<_OtherExtents, _OtherLayoutPolicy, _OtherAccessor>) _LIBCUDACXX_HIDE_FROM_ABI constexpr mdspan( const mdspan<_OtherElementType, _OtherExtents, _OtherLayoutPolicy, _OtherAccessor>& __other) : __base(__other.data_handle(), __other.mapping(), __other.accessor()) { static_assert(_CCCL_TRAIT(is_constructible, data_handle_type, const typename _OtherAccessor::data_handle_type&), "mdspan: incompatible data_handle_type for mdspan construction"); static_assert(_CCCL_TRAIT(is_constructible, extents_type, _OtherExtents), "mdspan: incompatible extents for mdspan construction"); if constexpr (extents_type::rank() != 0) { // The following precondition is part of the standard, but is unlikely to be triggered. // The extents constructor checks this and the mapping must be storing the extents, since // its extents() function returns a const reference to extents_type. // The only way this can be triggered is if the mapping conversion constructor would for example // always construct its extents() only from the dynamic extents, instead of from the other extents. for (size_t __r = 0; __r != extents_type::rank(); __r++) { // Not catching this could lead to out of bounds errors later // e.g. mdspan, non_checking_layout> m = // mdspan, non_checking_layout>(ptr, 200); leads to an extent of -56 on m _CCCL_ASSERT((static_extent(__r) == dynamic_extent) || (static_cast(__other.extent(__r)) == static_cast(static_extent(__r))), "mdspan: conversion mismatch of source dynamic extents with static extents"); } } } _CCCL_TEMPLATE(class _OtherElementType, class _OtherExtents, class _OtherLayoutPolicy, class _OtherAccessor) _CCCL_REQUIRES(__is_convertible_from<_OtherExtents, _OtherLayoutPolicy, _OtherAccessor> _CCCL_AND( !__is_implicit_convertible_from<_OtherExtents, _OtherLayoutPolicy, _OtherAccessor>)) _LIBCUDACXX_HIDE_FROM_ABI explicit constexpr mdspan( const mdspan<_OtherElementType, _OtherExtents, _OtherLayoutPolicy, _OtherAccessor>& __other) : __base(__other.data_handle(), __other.mapping(), __other.accessor()) { static_assert(_CCCL_TRAIT(is_constructible, data_handle_type, const typename _OtherAccessor::data_handle_type&), "mdspan: incompatible data_handle_type for mdspan construction"); static_assert(_CCCL_TRAIT(is_constructible, extents_type, _OtherExtents), "mdspan: incompatible extents for mdspan construction"); if constexpr (extents_type::rank() != 0) { // The following precondition is part of the standard, but is unlikely to be triggered. // The extents constructor checks this and the mapping must be storing the extents, since // its extents() function returns a const reference to extents_type. // The only way this can be triggered is if the mapping conversion constructor would for example // always construct its extents() only from the dynamic extents, instead of from the other extents. for (size_t __r = 0; __r < extents_type::rank(); __r++) { // Not catching this could lead to out of bounds errors later // e.g. mdspan, non_checking_layout> m = // mdspan, non_checking_layout>(ptr, 200); leads to an extent of -56 on m _CCCL_ASSERT((static_extent(__r) == dynamic_extent) || (static_cast(__other.extent(__r)) == static_cast(static_extent(__r))), "mdspan: conversion mismatch of source dynamic extents with static extents"); } } } _CCCL_HIDE_FROM_ABI constexpr mdspan& operator=(const mdspan&) = default; _CCCL_HIDE_FROM_ABI constexpr mdspan& operator=(mdspan&&) = default; //-------------------------------------------------------------------------------- // [mdspan.mdspan.members], members #if defined(_LIBCUDACXX_HAS_MULTIARG_OPERATOR_BRACKETS) _CCCL_TEMPLATE(class... _OtherIndexTypes) _CCCL_REQUIRES((sizeof...(_OtherIndexTypes) == extents_type::rank()) _CCCL_AND __mdspan_detail::__all_convertible_to_index_type) _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr reference operator[](_OtherIndexTypes... __indices) const { // Note the standard layouts would also check this, but user provided ones may not, so we // check the precondition here _CCCL_ASSERT(__mdspan_detail::__is_multidimensional_index_in(extents(), __indices...), "mdspan: operator[] out of bounds access"); return accessor().access(data_handle(), mapping()(static_cast(_CUDA_VSTD::move(__indices))...)); } #else _CCCL_TEMPLATE(class _OtherIndexType) _CCCL_REQUIRES((extents_type::rank() == 1) _CCCL_AND _CCCL_TRAIT(is_convertible, _OtherIndexType, index_type) _CCCL_AND _CCCL_TRAIT(is_nothrow_constructible, index_type, _OtherIndexType)) _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr reference operator[](_OtherIndexType __index) const { return accessor().access(data_handle(), mapping()(static_cast(_CUDA_VSTD::move(__index)))); } #endif // _LIBCUDACXX_HAS_MULTIARG_OPERATOR_BRACKETS template _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr decltype(auto) __op_bracket(const array<_OtherIndexType, _Extents::rank()>& __indices, index_sequence<_Idxs...>) const noexcept { // Note the standard layouts would also check this, but user provided ones may not, so we // check the precondition here _CCCL_ASSERT(__mdspan_detail::__is_multidimensional_index_in(extents(), __indices[_Idxs]...), "mdspan: operator[array] out of bounds access"); return mapping()(__indices[_Idxs]...); } template _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr decltype(auto) __op_bracket(span<_OtherIndexType, _Extents::rank()> __indices, index_sequence<_Idxs...>) const noexcept { // Note the standard layouts would also check this, but user provided ones may not, so we // check the precondition here _CCCL_ASSERT(__mdspan_detail::__is_multidimensional_index_in(extents(), __indices[_Idxs]...), "mdspan: operator[span] out of bounds access"); return mapping()(__indices[_Idxs]...); } _CCCL_TEMPLATE(class _OtherIndexType) _CCCL_REQUIRES(_CCCL_TRAIT(is_convertible, const _OtherIndexType&, index_type) _CCCL_AND _CCCL_TRAIT(is_nothrow_constructible, index_type, const _OtherIndexType&)) _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr reference operator[](const array<_OtherIndexType, extents_type::rank()>& __indices) const { return accessor().access(data_handle(), __op_bracket(__indices, make_index_sequence())); } _CCCL_TEMPLATE(class _OtherIndexType) _CCCL_REQUIRES(_CCCL_TRAIT(is_convertible, const _OtherIndexType&, index_type) _CCCL_AND _CCCL_TRAIT(is_nothrow_constructible, index_type, const _OtherIndexType&)) _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr reference operator[](span<_OtherIndexType, extents_type::rank()> __indices) const { return accessor().access(data_handle(), __op_bracket(__indices, make_index_sequence())); } //! Nonstandard extension to no break our users too hard _CCCL_TEMPLATE(class... _Indices) _CCCL_REQUIRES(__mdspan_detail::__all_convertible_to_index_type) _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr reference operator()(_Indices... __indices) const { // Note the standard layouts would also check this, but user provided ones may not, so we // check the precondition here _CCCL_ASSERT(__mdspan_detail::__is_multidimensional_index_in(extents(), __indices...), "mdspan: operator() out of bounds access"); return accessor().access(data_handle(), mapping()(__indices...)); } _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI static constexpr bool __mul_overflow(size_t x, size_t y, size_t* res) noexcept { *res = x * y; return x && ((*res / x) != y); } template _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr bool __check_size() const noexcept { size_t __prod = 1; for (size_t __r = 0; __r != extents_type::rank(); ++__r) { if (__mul_overflow(__prod, mapping().extents().extent(__r), &__prod)) { return false; } } return true; } template _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr size_type __op_size(index_sequence<_Idxs...>) const noexcept { return (size_type{1} * ... * static_cast(mapping().extents().extent(_Idxs))); } _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr size_type size() const noexcept { // Could leave this as only checked in debug mode: semantically size() is never // guaranteed to be related to any accessible range _CCCL_ASSERT(__check_size(), "mdspan: size() is not representable as size_type"); return __op_size(make_index_sequence()); } template _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr bool __op_empty(index_sequence<_Idxs...>) const noexcept { return (((mapping().extents().extent(_Idxs) == index_type{0})) || ...); } _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr bool empty() const noexcept { return __op_empty(make_index_sequence()); } _LIBCUDACXX_HIDE_FROM_ABI friend constexpr void swap(mdspan& __x, mdspan& __y) noexcept { swap(static_cast<__base&>(__x), static_cast<__base&>(__y)); } _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr const extents_type& extents() const noexcept { return mapping().extents(); }; _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr const data_handle_type& data_handle() const noexcept { return this->template __get<0>(); }; _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr const mapping_type& mapping() const noexcept { return this->template __get<1>(); }; _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr const accessor_type& accessor() const noexcept { return this->template __get<2>(); }; _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI static constexpr bool is_always_unique() noexcept(noexcept(mapping_type::is_always_unique())) { return mapping_type::is_always_unique(); }; _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI static constexpr bool is_always_exhaustive() noexcept(noexcept(mapping_type::is_always_exhaustive())) { return mapping_type::is_always_exhaustive(); }; _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI static constexpr bool is_always_strided() noexcept(noexcept(mapping_type::is_always_strided())) { return mapping_type::is_always_strided(); }; _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr bool is_unique() const noexcept(noexcept(_CUDA_VSTD::declval().is_unique())) { return mapping().is_unique(); }; _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr bool is_exhaustive() const noexcept(noexcept(_CUDA_VSTD::declval().is_exhaustive())) { return mapping().is_exhaustive(); }; _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr bool is_strided() const noexcept(noexcept(_CUDA_VSTD::declval().is_strided())) { return mapping().is_strided(); }; _CCCL_NODISCARD _LIBCUDACXX_HIDE_FROM_ABI constexpr index_type stride(rank_type __r) const { return mapping().stride(__r); }; }; _CCCL_TEMPLATE(class _ElementType, class... _OtherIndexTypes) _CCCL_REQUIRES((sizeof...(_OtherIndexTypes) > 0) _CCCL_AND(_CCCL_TRAIT(is_convertible, _OtherIndexTypes, size_t) && ... && true)) _CCCL_HOST_DEVICE explicit mdspan(_ElementType*, _OtherIndexTypes...) -> mdspan<_ElementType, extents...>>; _CCCL_TEMPLATE(class _Pointer) _CCCL_REQUIRES(_CCCL_TRAIT(is_pointer, remove_reference_t<_Pointer>)) _CCCL_HOST_DEVICE mdspan(_Pointer&&) -> mdspan>, extents>; _CCCL_TEMPLATE(class _CArray) _CCCL_REQUIRES(_CCCL_TRAIT(is_array, _CArray) _CCCL_AND(rank_v<_CArray> == 1)) _CCCL_HOST_DEVICE mdspan(_CArray&) -> mdspan, extents>>; template _CCCL_HOST_DEVICE mdspan(_ElementType*, const array<_OtherIndexType, _Size>&) -> mdspan<_ElementType, dextents>; template _CCCL_HOST_DEVICE mdspan(_ElementType*, span<_OtherIndexType, _Size>) -> mdspan<_ElementType, dextents>; // This one is necessary because all the constructors take `data_handle_type`s, not // `_ElementType*`s, and `data_handle_type` is taken from `accessor_type::data_handle_type`, which // seems to throw off automatic deduction guides. template _CCCL_HOST_DEVICE mdspan(_ElementType*, const extents<_OtherIndexType, _ExtentsPack...>&) -> mdspan<_ElementType, extents<_OtherIndexType, _ExtentsPack...>>; template _CCCL_HOST_DEVICE mdspan(_ElementType*, const _MappingType&) -> mdspan<_ElementType, typename _MappingType::extents_type, typename _MappingType::layout_type>; template _CCCL_HOST_DEVICE mdspan(const typename _AccessorType::data_handle_type, const _MappingType&, const _AccessorType&) -> mdspan; _LIBCUDACXX_END_NAMESPACE_STD #endif // _LIBCUDACXX___MDSPAN_MDSPAN_H