[INFRA] Import NVIDIA/CCCL upstream as optimization reference library
CCCL (CUDA C++ Core Libraries) provides: - CUB: device/block/warp-level GPU primitives (reduce, scan, sort, topk) - Thrust: high-level parallel algorithms (transform_reduce, sort, scan) - libcudacxx: CUDA C++ standard library (atomics, barriers, memory) - cudax: experimental features (memory resources, allocators) - Tuning policies: per-SM hardware-specific algorithm parameters Competition optimization vectors mapped to CCCL: - Output TPS (83% weight): warp_reduce, block_reduce, device_topk - Input TPS (14% weight): device_scan, block_load, prefetch - Cache TPS (3% weight): prefix caching strategy patterns - Memory (0.9 util): pooled/cached/buddy allocators Source: https://github.com/NVIDIA/cccl (shallow clone, HEAD only) License: Apache-2.0
This commit is contained in:
@@ -0,0 +1,76 @@
|
||||
//@HEADER
|
||||
// ************************************************************************
|
||||
//
|
||||
// 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.
|
||||
//
|
||||
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://kokkos.org/LICENSE for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
// ************************************************************************
|
||||
//@HEADER
|
||||
|
||||
#ifndef _CUDA_STD___LINALG_CONJUGATE_IF_NEEDED_H
|
||||
#define _CUDA_STD___LINALG_CONJUGATE_IF_NEEDED_H
|
||||
|
||||
#include <cuda/std/detail/__config>
|
||||
|
||||
#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 <cuda/std/__concepts/concept_macros.h>
|
||||
#include <cuda/std/__type_traits/is_arithmetic.h>
|
||||
#include <cuda/std/complex>
|
||||
|
||||
#include <cuda/std/__cccl/prologue.h>
|
||||
|
||||
_CCCL_BEGIN_NAMESPACE_CUDA_STD
|
||||
|
||||
namespace linalg
|
||||
{
|
||||
_CCCL_BEGIN_NAMESPACE_CPO(__conj_if_needed)
|
||||
|
||||
// Need newline for concept emulation :_(
|
||||
template <class _Type>
|
||||
_CCCL_CONCEPT _HasConj = _CCCL_REQUIRES_EXPR((_Type), _Type __a)(static_cast<void>(::cuda::std::conj(__a)));
|
||||
|
||||
struct __conj_if_needed
|
||||
{
|
||||
template <class _Type>
|
||||
_CCCL_API constexpr auto operator()(const _Type& __t) const
|
||||
{
|
||||
if constexpr (is_arithmetic_v<_Type> || !_HasConj<_Type>)
|
||||
{
|
||||
return __t;
|
||||
}
|
||||
else
|
||||
{
|
||||
return ::cuda::std::conj(__t);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
_CCCL_END_NAMESPACE_CPO
|
||||
|
||||
inline namespace __cpo
|
||||
{
|
||||
_CCCL_GLOBAL_CONSTANT auto conj_if_needed = __conj_if_needed::__conj_if_needed{};
|
||||
} // namespace __cpo
|
||||
} // end namespace linalg
|
||||
|
||||
_CCCL_END_NAMESPACE_CUDA_STD
|
||||
|
||||
#include <cuda/std/__cccl/epilogue.h>
|
||||
|
||||
#endif // _CUDA_STD___LINALG_CONJUGATED_HPP
|
||||
@@ -0,0 +1,52 @@
|
||||
//@HEADER
|
||||
// ************************************************************************
|
||||
//
|
||||
// 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.
|
||||
//
|
||||
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://kokkos.org/LICENSE for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
// ************************************************************************
|
||||
//@HEADER
|
||||
|
||||
#ifndef _CUDA_STD___LINALG_CONJUGATE_TRANSPOSED_H
|
||||
#define _CUDA_STD___LINALG_CONJUGATE_TRANSPOSED_H
|
||||
|
||||
#include <cuda/std/detail/__config>
|
||||
|
||||
#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 <cuda/std/__linalg/conjugated.h>
|
||||
#include <cuda/std/__linalg/transposed.h>
|
||||
|
||||
#include <cuda/std/__cccl/prologue.h>
|
||||
|
||||
_CCCL_BEGIN_NAMESPACE_CUDA_STD
|
||||
|
||||
namespace linalg
|
||||
{
|
||||
template <class _ElementType, class _Extents, class _Layout, class _Accessor>
|
||||
[[nodiscard]] _CCCL_API constexpr auto conjugate_transposed(mdspan<_ElementType, _Extents, _Layout, _Accessor> __a)
|
||||
{
|
||||
return conjugated(transposed(__a));
|
||||
}
|
||||
} // end namespace linalg
|
||||
|
||||
_CCCL_END_NAMESPACE_CUDA_STD
|
||||
|
||||
#include <cuda/std/__cccl/epilogue.h>
|
||||
|
||||
#endif // _CUDA_STD___LINALG_CONJUGATE_TRANSPOSED_HPP
|
||||
136
cccl_upstream/libcudacxx/include/cuda/std/__linalg/conjugated.h
Normal file
136
cccl_upstream/libcudacxx/include/cuda/std/__linalg/conjugated.h
Normal file
@@ -0,0 +1,136 @@
|
||||
//@HEADER
|
||||
// ************************************************************************
|
||||
//
|
||||
// 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.
|
||||
//
|
||||
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://kokkos.org/LICENSE for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
// ************************************************************************
|
||||
//@HEADER
|
||||
|
||||
#ifndef _CUDA_STD___LINALG_CONJUGATED_H
|
||||
#define _CUDA_STD___LINALG_CONJUGATED_H
|
||||
|
||||
#include <cuda/std/detail/__config>
|
||||
|
||||
#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 <cuda/std/__linalg/conj_if_needed.h>
|
||||
#include <cuda/std/__type_traits/add_const.h>
|
||||
#include <cuda/std/__type_traits/is_arithmetic.h>
|
||||
#include <cuda/std/__type_traits/remove_const.h>
|
||||
#include <cuda/std/__utility/declval.h>
|
||||
#include <cuda/std/mdspan>
|
||||
|
||||
#include <cuda/std/__cccl/prologue.h>
|
||||
|
||||
_CCCL_BEGIN_NAMESPACE_CUDA_STD
|
||||
|
||||
namespace linalg
|
||||
{
|
||||
template <class _NestedAccessor>
|
||||
class conjugated_accessor
|
||||
{
|
||||
private:
|
||||
using __nested_element_type = typename _NestedAccessor::element_type;
|
||||
using __nc_result_type = decltype(conj_if_needed(::cuda::std::declval<__nested_element_type>()));
|
||||
|
||||
public:
|
||||
using element_type = add_const_t<__nc_result_type>;
|
||||
using reference = remove_const_t<element_type>;
|
||||
using data_handle_type = typename _NestedAccessor::data_handle_type;
|
||||
using offset_policy = conjugated_accessor<typename _NestedAccessor::offset_policy>;
|
||||
|
||||
_CCCL_HIDE_FROM_ABI constexpr conjugated_accessor() = default;
|
||||
|
||||
_CCCL_API constexpr conjugated_accessor(const _NestedAccessor& __acc)
|
||||
: __nested_accessor_(__acc)
|
||||
{}
|
||||
|
||||
_CCCL_TEMPLATE(class _OtherNestedAccessor)
|
||||
_CCCL_REQUIRES(is_constructible_v<_NestedAccessor, const _OtherNestedAccessor&> _CCCL_AND
|
||||
is_convertible_v<_OtherNestedAccessor, _NestedAccessor>)
|
||||
_CCCL_API constexpr conjugated_accessor(const conjugated_accessor<_OtherNestedAccessor>& __other)
|
||||
: __nested_accessor_(__other.nested_accessor())
|
||||
{}
|
||||
|
||||
_CCCL_TEMPLATE(class _OtherNestedAccessor)
|
||||
_CCCL_REQUIRES(is_constructible_v<_NestedAccessor, const _OtherNestedAccessor&> _CCCL_AND(
|
||||
!is_convertible_v<_OtherNestedAccessor, _NestedAccessor>))
|
||||
_CCCL_API explicit constexpr conjugated_accessor(const conjugated_accessor<_OtherNestedAccessor>& __other)
|
||||
: __nested_accessor_(__other.nested_accessor())
|
||||
{}
|
||||
|
||||
_CCCL_API constexpr reference access(data_handle_type __p, size_t __i) const noexcept
|
||||
{
|
||||
return conj_if_needed(__nested_element_type(__nested_accessor_.access(__p, __i)));
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_API constexpr typename offset_policy::data_handle_type
|
||||
offset(data_handle_type __p, size_t __i) const noexcept
|
||||
{
|
||||
return __nested_accessor_.offset(__p, __i);
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_API constexpr const _NestedAccessor& nested_accessor() const noexcept
|
||||
{
|
||||
return __nested_accessor_;
|
||||
}
|
||||
|
||||
private:
|
||||
_NestedAccessor __nested_accessor_;
|
||||
};
|
||||
|
||||
template <class _ElementType, class _Extents, class _Layout, class _Accessor>
|
||||
[[nodiscard]] _CCCL_API constexpr auto conjugated(mdspan<_ElementType, _Extents, _Layout, _Accessor> __a)
|
||||
{
|
||||
using __value_type = typename decltype(__a)::value_type;
|
||||
// Current status of [linalg] only optimizes if _Accessor is conjugated_accessor<_Accessor> for some _Accessor.
|
||||
// There's a separate specialization for that case below.
|
||||
|
||||
// P3050 optimizes conjugated's accessor type for when we know that it can't be complex: arithmetic types,
|
||||
// and types for which `conj` is not ADL-findable.
|
||||
if constexpr (is_arithmetic_v<__value_type> || !__conj_if_needed::_HasConj<__value_type>)
|
||||
{
|
||||
return mdspan<_ElementType, _Extents, _Layout, _Accessor>(__a.data_handle(), __a.mapping(), __a.accessor());
|
||||
}
|
||||
else
|
||||
{
|
||||
using __return_element_type = typename conjugated_accessor<_Accessor>::element_type;
|
||||
using __return_accessor_type = conjugated_accessor<_Accessor>;
|
||||
return mdspan<__return_element_type, _Extents, _Layout, __return_accessor_type>{
|
||||
__a.data_handle(), __a.mapping(), __return_accessor_type(__a.accessor())};
|
||||
}
|
||||
}
|
||||
|
||||
// Conjugation is self-annihilating
|
||||
template <class _ElementType, class _Extents, class _Layout, class _NestedAccessor>
|
||||
[[nodiscard]] _CCCL_API constexpr auto
|
||||
conjugated(mdspan<_ElementType, _Extents, _Layout, conjugated_accessor<_NestedAccessor>> __a)
|
||||
{
|
||||
using __return_element_type = typename _NestedAccessor::element_type;
|
||||
using __return_accessor_type = _NestedAccessor;
|
||||
return mdspan<__return_element_type, _Extents, _Layout, __return_accessor_type>(
|
||||
__a.data_handle(), __a.mapping(), __a.accessor().nested_accessor());
|
||||
}
|
||||
} // end namespace linalg
|
||||
|
||||
_CCCL_END_NAMESPACE_CUDA_STD
|
||||
|
||||
#include <cuda/std/__cccl/epilogue.h>
|
||||
|
||||
#endif // _CUDA_STD___LINALG_CONJUGATED_HPP
|
||||
128
cccl_upstream/libcudacxx/include/cuda/std/__linalg/scaled.h
Normal file
128
cccl_upstream/libcudacxx/include/cuda/std/__linalg/scaled.h
Normal file
@@ -0,0 +1,128 @@
|
||||
//@HEADER
|
||||
// ************************************************************************
|
||||
//
|
||||
// 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.
|
||||
//
|
||||
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://kokkos.org/LICENSE for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
// ************************************************************************
|
||||
//@HEADER
|
||||
|
||||
#ifndef _CUDA_STD___LINALG_SCALED_H
|
||||
#define _CUDA_STD___LINALG_SCALED_H
|
||||
|
||||
#include <cuda/std/detail/__config>
|
||||
|
||||
#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 <cuda/std/__concepts/concept_macros.h>
|
||||
#include <cuda/std/__type_traits/add_const.h>
|
||||
#include <cuda/std/__type_traits/remove_const.h>
|
||||
#include <cuda/std/__utility/declval.h>
|
||||
#include <cuda/std/mdspan>
|
||||
|
||||
#include <cuda/std/__cccl/prologue.h>
|
||||
|
||||
_CCCL_BEGIN_NAMESPACE_CUDA_STD
|
||||
|
||||
namespace linalg
|
||||
{
|
||||
template <class _ScalingFactor, class _NestedAccessor>
|
||||
class scaled_accessor
|
||||
{
|
||||
public:
|
||||
using element_type = add_const_t<
|
||||
decltype(::cuda::std::declval<_ScalingFactor>() * ::cuda::std::declval<typename _NestedAccessor::element_type>())>;
|
||||
using reference = remove_const_t<element_type>;
|
||||
using data_handle_type = typename _NestedAccessor::data_handle_type;
|
||||
using offset_policy = scaled_accessor<_ScalingFactor, typename _NestedAccessor::offset_policy>;
|
||||
|
||||
_CCCL_HIDE_FROM_ABI constexpr scaled_accessor() = default;
|
||||
|
||||
_CCCL_TEMPLATE(class _OtherScalingFactor, class _OtherNestedAccessor)
|
||||
_CCCL_REQUIRES(is_constructible_v<_NestedAccessor, const _OtherNestedAccessor&> _CCCL_AND
|
||||
is_constructible_v<_ScalingFactor, _OtherScalingFactor> _CCCL_AND(
|
||||
!is_convertible_v<_OtherNestedAccessor, _NestedAccessor>))
|
||||
_CCCL_API explicit constexpr scaled_accessor(const scaled_accessor<_OtherScalingFactor, _OtherNestedAccessor>& __other)
|
||||
: __scaling_factor_(__other.scaling_factor())
|
||||
, __nested_accessor_(__other.nested_accessor())
|
||||
{}
|
||||
|
||||
_CCCL_TEMPLATE(class _OtherScalingFactor, class _OtherNestedAccessor)
|
||||
_CCCL_REQUIRES(is_constructible_v<_NestedAccessor, const _OtherNestedAccessor&> _CCCL_AND
|
||||
is_constructible_v<_ScalingFactor, _OtherScalingFactor> _CCCL_AND
|
||||
is_convertible_v<_OtherNestedAccessor, _NestedAccessor>)
|
||||
_CCCL_API constexpr scaled_accessor(const scaled_accessor<_OtherScalingFactor, _OtherNestedAccessor>& __other)
|
||||
: __scaling_factor_(__other.scaling_factor())
|
||||
, __nested_accessor_(__other.nested_accessor())
|
||||
{}
|
||||
|
||||
_CCCL_API constexpr scaled_accessor(const _ScalingFactor& __s, const _NestedAccessor& __a)
|
||||
: __scaling_factor_(__s)
|
||||
, __nested_accessor_(__a)
|
||||
{}
|
||||
|
||||
_CCCL_API constexpr reference access(data_handle_type __p, size_t __i) const
|
||||
{
|
||||
return __scaling_factor_ * typename _NestedAccessor::element_type(__nested_accessor_.access(__p, __i));
|
||||
}
|
||||
|
||||
[[nodiscard]]
|
||||
_CCCL_API inline typename offset_policy::data_handle_type constexpr offset(data_handle_type __p, size_t __i) const
|
||||
{
|
||||
return __nested_accessor_.offset(__p, __i);
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_API constexpr _NestedAccessor nested_accessor() const noexcept
|
||||
{
|
||||
return __nested_accessor_;
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_API constexpr _ScalingFactor scaling_factor() const noexcept
|
||||
{
|
||||
return __scaling_factor_;
|
||||
}
|
||||
|
||||
private:
|
||||
_ScalingFactor __scaling_factor_;
|
||||
_NestedAccessor __nested_accessor_;
|
||||
};
|
||||
|
||||
namespace __detail
|
||||
{
|
||||
template <class _ScalingFactor, class _NestedAccessor>
|
||||
using __scaled_element_type = add_const_t<typename scaled_accessor<_ScalingFactor, _NestedAccessor>::element_type>;
|
||||
} // namespace __detail
|
||||
|
||||
template <class _ScalingFactor, class _ElementType, class _Extents, class _Layout, class _Accessor>
|
||||
[[nodiscard]]
|
||||
_CCCL_API constexpr mdspan<__detail::__scaled_element_type<_ScalingFactor, _Accessor>,
|
||||
_Extents,
|
||||
_Layout,
|
||||
scaled_accessor<_ScalingFactor, _Accessor>>
|
||||
scaled(_ScalingFactor __scaling_factor, mdspan<_ElementType, _Extents, _Layout, _Accessor> __x)
|
||||
{
|
||||
using __acc_type = scaled_accessor<_ScalingFactor, _Accessor>;
|
||||
return {__x.data_handle(), __x.mapping(), __acc_type{__scaling_factor, __x.accessor()}};
|
||||
}
|
||||
} // end namespace linalg
|
||||
|
||||
_CCCL_END_NAMESPACE_CUDA_STD
|
||||
|
||||
#include <cuda/std/__cccl/epilogue.h>
|
||||
|
||||
#endif // _CUDA_STD___LINALG_SCALED_HPP
|
||||
315
cccl_upstream/libcudacxx/include/cuda/std/__linalg/transposed.h
Normal file
315
cccl_upstream/libcudacxx/include/cuda/std/__linalg/transposed.h
Normal file
@@ -0,0 +1,315 @@
|
||||
//@HEADER
|
||||
// ************************************************************************
|
||||
//
|
||||
// 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.
|
||||
//
|
||||
// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
|
||||
// See https://kokkos.org/LICENSE for license information.
|
||||
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
||||
// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES.
|
||||
//
|
||||
// ************************************************************************
|
||||
//@HEADER
|
||||
|
||||
#ifndef _CUDA_STD___LINALG_TRANSPOSED_H
|
||||
#define _CUDA_STD___LINALG_TRANSPOSED_H
|
||||
|
||||
#include <cuda/std/detail/__config>
|
||||
|
||||
#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 <cuda/std/__concepts/concept_macros.h>
|
||||
#include <cuda/std/__type_traits/is_convertible.h>
|
||||
#include <cuda/std/__type_traits/is_same.h>
|
||||
#include <cuda/std/array>
|
||||
#include <cuda/std/mdspan>
|
||||
|
||||
#include <cuda/std/__cccl/prologue.h>
|
||||
|
||||
_CCCL_BEGIN_NAMESPACE_CUDA_STD
|
||||
|
||||
namespace linalg
|
||||
{
|
||||
namespace __detail
|
||||
{
|
||||
// This struct helps us impose the rank constraint on the __type alias itself.
|
||||
_CCCL_TEMPLATE(class _Extents)
|
||||
_CCCL_REQUIRES((_Extents::rank() == 2))
|
||||
struct __transpose_extents_t_impl
|
||||
{
|
||||
using __type = extents<typename _Extents::index_type, _Extents::static_extent(1), _Extents::static_extent(0)>;
|
||||
};
|
||||
|
||||
template <class _Extents>
|
||||
using __transpose_extents_t = typename __transpose_extents_t_impl<_Extents>::__type;
|
||||
|
||||
_CCCL_TEMPLATE(class _Extents)
|
||||
_CCCL_REQUIRES((_Extents::rank() == 2))
|
||||
_CCCL_API constexpr __transpose_extents_t<_Extents> __transpose_extents(const _Extents& __e)
|
||||
{
|
||||
static_assert(is_same_v<typename __transpose_extents_t<_Extents>::index_type, typename _Extents::index_type>,
|
||||
"Please fix __transpose_extents_t to account for P2553, which adds a template parameter SizeType to "
|
||||
"extents.");
|
||||
constexpr size_t __ext0 = _Extents::static_extent(0);
|
||||
constexpr size_t __ext1 = _Extents::static_extent(1);
|
||||
if constexpr (__ext0 == dynamic_extent)
|
||||
{
|
||||
if constexpr (__ext1 == dynamic_extent)
|
||||
{
|
||||
return __transpose_extents_t<_Extents>{__e.extent(1), __e.extent(0)};
|
||||
}
|
||||
else
|
||||
{
|
||||
return __transpose_extents_t<_Extents>{/* __e.extent(1), */ __e.extent(0)};
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
if constexpr (__ext1 == dynamic_extent)
|
||||
{
|
||||
return __transpose_extents_t<_Extents>{__e.extent(1) /* , __e.extent(0) */};
|
||||
}
|
||||
else
|
||||
{
|
||||
return __transpose_extents_t<_Extents>{}; // all extents are static
|
||||
}
|
||||
}
|
||||
}
|
||||
} // namespace __detail
|
||||
|
||||
template <class _Layout>
|
||||
class layout_transpose
|
||||
{
|
||||
public:
|
||||
using nested_layout_type = _Layout;
|
||||
|
||||
template <class _Extents>
|
||||
struct mapping
|
||||
{
|
||||
private:
|
||||
using __nested_mapping_type = typename _Layout::template mapping<__detail::__transpose_extents_t<_Extents>>;
|
||||
|
||||
static constexpr bool __required_span_size_noexcept = noexcept(__nested_mapping_type{}.required_span_size());
|
||||
|
||||
static constexpr bool __is_nested_unique_noexcept = noexcept(__nested_mapping_type{}.is_unique());
|
||||
|
||||
static constexpr bool __is_exhaustive_noexcept = noexcept(__nested_mapping_type{}.is_exhaustive());
|
||||
|
||||
static constexpr bool __is_strided_noexcept = noexcept(__nested_mapping_type{}.is_strided());
|
||||
|
||||
public:
|
||||
using extents_type = _Extents;
|
||||
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 layout_type = layout_transpose;
|
||||
|
||||
_CCCL_API constexpr explicit mapping(const __nested_mapping_type& __map)
|
||||
: __nested_mapping_(__map)
|
||||
, __extents_(__detail::__transpose_extents(__map.extents()))
|
||||
{}
|
||||
|
||||
[[nodiscard]] _CCCL_API constexpr const extents_type& extents() const noexcept
|
||||
{
|
||||
return __extents_;
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_API constexpr index_type required_span_size() const noexcept(__required_span_size_noexcept)
|
||||
{
|
||||
return __nested_mapping_.required_span_size();
|
||||
}
|
||||
|
||||
_CCCL_TEMPLATE(class _IndexType0, class _IndexType1)
|
||||
_CCCL_REQUIRES(is_convertible_v<_IndexType0, index_type> _CCCL_AND is_convertible_v<_IndexType1, index_type>)
|
||||
_CCCL_API constexpr index_type operator()(_IndexType0 __i, _IndexType1 __j) const
|
||||
{
|
||||
return __nested_mapping_(__j, __i);
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_API constexpr const __nested_mapping_type& nested_mapping() const noexcept
|
||||
{
|
||||
return __nested_mapping_;
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_API static constexpr bool is_always_unique() noexcept
|
||||
{
|
||||
return __nested_mapping_type::is_always_unique();
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_API static constexpr bool is_always_exhaustive() noexcept
|
||||
{
|
||||
return __nested_mapping_type::is_always_exhaustive();
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_API static constexpr bool is_always_strided() noexcept
|
||||
{
|
||||
return __nested_mapping_type::is_always_strided();
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_API constexpr bool is_unique() const noexcept(__is_nested_unique_noexcept)
|
||||
{
|
||||
return __nested_mapping_.is_unique();
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_API constexpr bool is_exhaustive() const noexcept(__is_exhaustive_noexcept)
|
||||
{
|
||||
return __nested_mapping_.is_exhaustive();
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_API constexpr bool is_strided() const noexcept(__is_strided_noexcept)
|
||||
{
|
||||
return __nested_mapping_.is_strided();
|
||||
}
|
||||
|
||||
[[nodiscard]] _CCCL_API constexpr index_type stride(size_t __r) const
|
||||
{
|
||||
_CCCL_ASSERT(this->is_strided(), "layout must be strided");
|
||||
_CCCL_ASSERT(__r < extents_type::rank(), "rank must be less than extents rank");
|
||||
return __nested_mapping_.stride(__r == 0 ? 1 : 0);
|
||||
}
|
||||
|
||||
template <class _OtherExtents>
|
||||
_CCCL_API friend constexpr bool operator==(const mapping& __lhs, const mapping<_OtherExtents>& __rhs) noexcept
|
||||
{
|
||||
return __lhs.__nested_mapping_ == __rhs.__nested_mapping_;
|
||||
}
|
||||
|
||||
template <class _OtherExtents>
|
||||
_CCCL_API friend constexpr bool operator!=(const mapping& __lhs, const mapping<_OtherExtents>& __rhs) noexcept
|
||||
{
|
||||
return __lhs.__nested_mapping_ != __rhs.__nested_mapping_;
|
||||
}
|
||||
|
||||
private:
|
||||
__nested_mapping_type __nested_mapping_;
|
||||
extents_type __extents_;
|
||||
};
|
||||
};
|
||||
|
||||
namespace __detail
|
||||
{
|
||||
template <class _ElementType, class _Accessor>
|
||||
struct __transposed_element_accessor
|
||||
{
|
||||
using __element_type = _ElementType;
|
||||
using __accessor_type = _Accessor;
|
||||
|
||||
_CCCL_API static constexpr __accessor_type __accessor(const _Accessor& __a)
|
||||
{
|
||||
return __accessor_type(__a);
|
||||
}
|
||||
};
|
||||
|
||||
template <class _ElementType>
|
||||
struct __transposed_element_accessor<_ElementType, default_accessor<_ElementType>>
|
||||
{
|
||||
using __element_type = _ElementType;
|
||||
using __accessor_type = default_accessor<__element_type>;
|
||||
|
||||
_CCCL_API static constexpr __accessor_type __accessor(const default_accessor<_ElementType>& __a)
|
||||
{
|
||||
return __accessor_type(__a);
|
||||
}
|
||||
};
|
||||
|
||||
template <class _Layout>
|
||||
struct __transposed_layout
|
||||
{
|
||||
using __layout_type = layout_transpose<_Layout>;
|
||||
|
||||
template <class __OriginalMapping>
|
||||
_CCCL_API static constexpr auto __mapping(const __OriginalMapping& __orig_map)
|
||||
{
|
||||
using __extents_type = __transpose_extents_t<typename __OriginalMapping::__extents_type>;
|
||||
using __return_mapping_type = typename __layout_type::template __mapping<__extents_type>;
|
||||
return __return_mapping_type{__orig_map};
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct __transposed_layout<layout_left>
|
||||
{
|
||||
using __layout_type = layout_right;
|
||||
|
||||
template <class _OriginalExtents>
|
||||
_CCCL_API static constexpr auto __mapping(const typename layout_left::template mapping<_OriginalExtents>& __orig_map)
|
||||
{
|
||||
using __original_mapping_type = typename layout_left::template mapping<_OriginalExtents>;
|
||||
using __extents_type = __transpose_extents_t<typename __original_mapping_type::extents_type>;
|
||||
using __return_mapping_type = typename __layout_type::template mapping<__extents_type>;
|
||||
return __return_mapping_type{__transpose_extents(__orig_map.extents())};
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct __transposed_layout<layout_right>
|
||||
{
|
||||
using __layout_type = layout_left;
|
||||
|
||||
template <class _OriginalExtents>
|
||||
_CCCL_API static constexpr auto __mapping(const typename layout_right::template mapping<_OriginalExtents>& __orig_map)
|
||||
{
|
||||
using __original_mapping_type = typename layout_right::template mapping<_OriginalExtents>;
|
||||
using __extents_type = __transpose_extents_t<typename __original_mapping_type::extents_type>;
|
||||
using __return_mapping_type = typename __layout_type::template mapping<__extents_type>;
|
||||
return __return_mapping_type{__transpose_extents(__orig_map.extents())};
|
||||
}
|
||||
};
|
||||
|
||||
template <>
|
||||
struct __transposed_layout<layout_stride>
|
||||
{
|
||||
using __layout_type = layout_stride;
|
||||
|
||||
template <class _OriginalExtents>
|
||||
_CCCL_API static constexpr auto __mapping(const typename layout_stride::template mapping<_OriginalExtents>& __orig_map)
|
||||
{
|
||||
using __original_mapping_type = typename layout_stride::template mapping<_OriginalExtents>;
|
||||
using __original_extents_type = typename __original_mapping_type::extents_type;
|
||||
using __extents_type = __transpose_extents_t<__original_extents_type>;
|
||||
using __return_mapping_type = typename __layout_type::template mapping<__extents_type>;
|
||||
return __return_mapping_type{
|
||||
__transpose_extents(__orig_map.extents()),
|
||||
array<typename __extents_type::index_type, _OriginalExtents::rank() /* __orig_map.rank() */>{
|
||||
__orig_map.stride(1), __orig_map.stride(0)}};
|
||||
}
|
||||
};
|
||||
|
||||
// TODO add support for padded layouts
|
||||
|
||||
template <class _NestedLayout>
|
||||
struct __transposed_layout<layout_transpose<_NestedLayout>>
|
||||
{
|
||||
using __layout_type = _NestedLayout;
|
||||
};
|
||||
} // namespace __detail
|
||||
|
||||
template <class _ElementType, class _Extents, class _Layout, class _Accessor>
|
||||
[[nodiscard]] _CCCL_API constexpr auto transposed(mdspan<_ElementType, _Extents, _Layout, _Accessor> __a)
|
||||
{
|
||||
using __element_type = typename __detail::__transposed_element_accessor<_ElementType, _Accessor>::__element_type;
|
||||
using __layout_type = typename __detail::__transposed_layout<_Layout>::__layout_type;
|
||||
using __accessor_type = typename __detail::__transposed_element_accessor<_ElementType, _Accessor>::__accessor_type;
|
||||
auto __mapping = __detail::__transposed_layout<_Layout>::__mapping(__a.mapping());
|
||||
auto __accessor = __detail::__transposed_element_accessor<_ElementType, _Accessor>::__accessor(__a.accessor());
|
||||
return mdspan<__element_type, typename decltype(__mapping)::extents_type, __layout_type, __accessor_type>{
|
||||
__a.data_handle(), __mapping, __accessor};
|
||||
}
|
||||
} // end namespace linalg
|
||||
|
||||
_CCCL_END_NAMESPACE_CUDA_STD
|
||||
|
||||
#include <cuda/std/__cccl/epilogue.h>
|
||||
|
||||
#endif // _CUDA_STD___LINALG_TRANSPOSED_HPP
|
||||
Reference in New Issue
Block a user