Kokkos Core Kernels Package Version of the Day
Loading...
Searching...
No Matches
Kokkos_Extents.hpp
1//@HEADER
2// ************************************************************************
3//
4// Kokkos v. 4.0
5// Copyright (2022) National Technology & Engineering
6// Solutions of Sandia, LLC (NTESS).
7//
8// Under the terms of Contract DE-NA0003525 with NTESS,
9// the U.S. Government retains certain rights in this software.
10//
11// Part of Kokkos, under the Apache License v2.0 with LLVM Exceptions.
12// See https://kokkos.org/LICENSE for license information.
13// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
14//
15//@HEADER
16
17#ifndef KOKKOS_IMPL_PUBLIC_INCLUDE
18#include <Kokkos_Macros.hpp>
19static_assert(false,
20 "Including non-public Kokkos header files is not allowed.");
21#endif
22#ifndef KOKKOS_KOKKOS_EXTENTS_HPP
23#define KOKKOS_KOKKOS_EXTENTS_HPP
24
25#include <cstddef>
26#include <type_traits>
27#include <Kokkos_Macros.hpp>
28
29namespace Kokkos {
30namespace Experimental {
31
32constexpr ptrdiff_t dynamic_extent = -1;
33
34template <ptrdiff_t... ExtentSpecs>
35struct Extents {
36 /* TODO @enhancement flesh this out more */
37};
38
39template <class Exts, ptrdiff_t NewExtent>
40struct PrependExtent;
41
42template <ptrdiff_t... Exts, ptrdiff_t NewExtent>
43struct PrependExtent<Extents<Exts...>, NewExtent> {
44 using type = Extents<NewExtent, Exts...>;
45};
46
47template <class Exts, ptrdiff_t NewExtent>
48struct AppendExtent;
49
50template <ptrdiff_t... Exts, ptrdiff_t NewExtent>
51struct AppendExtent<Extents<Exts...>, NewExtent> {
52 using type = Extents<Exts..., NewExtent>;
53};
54
55} // end namespace Experimental
56
57namespace Impl {
58
59namespace _parse_view_extents_impl {
60
61template <class T>
62struct _all_remaining_extents_dynamic : std::true_type {};
63
64template <class T>
65struct _all_remaining_extents_dynamic<T*> : _all_remaining_extents_dynamic<T> {
66};
67
68template <class T, unsigned N>
69struct _all_remaining_extents_dynamic<T[N]> : std::false_type {};
70
71template <class T, class Result, class = void>
72struct _parse_impl {
73 using type = Result;
74};
75
76// We have to treat the case of int**[x] specially, since it *doesn't* go
77// backwards
78template <class T, ptrdiff_t... ExtentSpec>
79struct _parse_impl<T*, Kokkos::Experimental::Extents<ExtentSpec...>,
80 std::enable_if_t<_all_remaining_extents_dynamic<T>::value>>
81 : _parse_impl<T, Kokkos::Experimental::Extents<
82 Kokkos::Experimental::dynamic_extent, ExtentSpec...>> {
83};
84
85// int*(*[x])[y] should still work also (meaning int[][x][][y])
86template <class T, ptrdiff_t... ExtentSpec>
87struct _parse_impl<
88 T*, Kokkos::Experimental::Extents<ExtentSpec...>,
89 std::enable_if_t<!_all_remaining_extents_dynamic<T>::value>> {
90 using _next = Kokkos::Experimental::AppendExtent<
91 typename _parse_impl<T, Kokkos::Experimental::Extents<ExtentSpec...>,
92 void>::type,
93 Kokkos::Experimental::dynamic_extent>;
94 using type = typename _next::type;
95};
96
97template <class T, ptrdiff_t... ExtentSpec, unsigned N>
98struct _parse_impl<T[N], Kokkos::Experimental::Extents<ExtentSpec...>, void>
99 : _parse_impl<
100 T, Kokkos::Experimental::Extents<ExtentSpec...,
101 ptrdiff_t(N)> // TODO @pedantic this
102 // could be a
103 // narrowing cast
104 > {};
105
106} // end namespace _parse_view_extents_impl
107
108template <class DataType>
109struct ParseViewExtents {
110 using type = typename _parse_view_extents_impl ::_parse_impl<
111 DataType, Kokkos::Experimental::Extents<>>::type;
112};
113
114template <class ValueType, ptrdiff_t Ext>
115struct ApplyExtent {
116 using type = ValueType[Ext];
117};
118
119template <class ValueType>
120struct ApplyExtent<ValueType, Kokkos::Experimental::dynamic_extent> {
121 using type = ValueType*;
122};
123
124template <class ValueType, unsigned N, ptrdiff_t Ext>
125struct ApplyExtent<ValueType[N], Ext> {
126 using type = typename ApplyExtent<ValueType, Ext>::type[N];
127};
128
129template <class ValueType, ptrdiff_t Ext>
130struct ApplyExtent<ValueType*, Ext> {
131 using type = ValueType * [Ext];
132};
133
134template <class ValueType>
135struct ApplyExtent<ValueType*, Kokkos::Experimental::dynamic_extent> {
136 using type =
137 typename ApplyExtent<ValueType,
138 Kokkos::Experimental::dynamic_extent>::type*;
139};
140
141template <class ValueType, unsigned N>
142struct ApplyExtent<ValueType[N], Kokkos::Experimental::dynamic_extent> {
143 using type =
144 typename ApplyExtent<ValueType,
145 Kokkos::Experimental::dynamic_extent>::type[N];
146};
147
148} // end namespace Impl
149
150} // end namespace Kokkos
151
152#endif // KOKKOS_KOKKOS_EXTENTS_HPP