Kokkos Core Kernels Package Version of the Day
Loading...
Searching...
No Matches
Kokkos_Layout.hpp
Go to the documentation of this file.
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
19
20#ifndef KOKKOS_IMPL_PUBLIC_INCLUDE
21#include <Kokkos_Macros.hpp>
22static_assert(false,
23 "Including non-public Kokkos header files is not allowed.");
24#endif
25#ifndef KOKKOS_LAYOUT_HPP
26#define KOKKOS_LAYOUT_HPP
27
28#include <cstddef>
29#include <impl/Kokkos_Traits.hpp>
30
31namespace Kokkos {
32
33enum { ARRAY_LAYOUT_MAX_RANK = 8 };
34
35//----------------------------------------------------------------------------
50struct LayoutLeft {
53
54 size_t dimension[ARRAY_LAYOUT_MAX_RANK];
55
56 enum : bool { is_extent_constructible = true };
57
58 LayoutLeft(LayoutLeft const&) = default;
59 LayoutLeft(LayoutLeft&&) = default;
60 LayoutLeft& operator=(LayoutLeft const&) = default;
61 LayoutLeft& operator=(LayoutLeft&&) = default;
62
63 KOKKOS_INLINE_FUNCTION
64 explicit constexpr LayoutLeft(size_t N0 = KOKKOS_IMPL_CTOR_DEFAULT_ARG,
65 size_t N1 = KOKKOS_IMPL_CTOR_DEFAULT_ARG,
66 size_t N2 = KOKKOS_IMPL_CTOR_DEFAULT_ARG,
67 size_t N3 = KOKKOS_IMPL_CTOR_DEFAULT_ARG,
68 size_t N4 = KOKKOS_IMPL_CTOR_DEFAULT_ARG,
69 size_t N5 = KOKKOS_IMPL_CTOR_DEFAULT_ARG,
70 size_t N6 = KOKKOS_IMPL_CTOR_DEFAULT_ARG,
71 size_t N7 = KOKKOS_IMPL_CTOR_DEFAULT_ARG)
72 : dimension{N0, N1, N2, N3, N4, N5, N6, N7} {}
73
74 friend bool operator==(const LayoutLeft& left, const LayoutLeft& right) {
75 for (unsigned int rank = 0; rank < ARRAY_LAYOUT_MAX_RANK; ++rank)
76 if (left.dimension[rank] != right.dimension[rank]) return false;
77 return true;
78 }
79
80 friend bool operator!=(const LayoutLeft& left, const LayoutLeft& right) {
81 return !(left == right);
82 }
83};
84
85//----------------------------------------------------------------------------
102
103 size_t dimension[ARRAY_LAYOUT_MAX_RANK];
104
105 enum : bool { is_extent_constructible = true };
106
107 LayoutRight(LayoutRight const&) = default;
108 LayoutRight(LayoutRight&&) = default;
109 LayoutRight& operator=(LayoutRight const&) = default;
110 LayoutRight& operator=(LayoutRight&&) = default;
111
112 KOKKOS_INLINE_FUNCTION
113 explicit constexpr LayoutRight(size_t N0 = KOKKOS_IMPL_CTOR_DEFAULT_ARG,
114 size_t N1 = KOKKOS_IMPL_CTOR_DEFAULT_ARG,
115 size_t N2 = KOKKOS_IMPL_CTOR_DEFAULT_ARG,
116 size_t N3 = KOKKOS_IMPL_CTOR_DEFAULT_ARG,
117 size_t N4 = KOKKOS_IMPL_CTOR_DEFAULT_ARG,
118 size_t N5 = KOKKOS_IMPL_CTOR_DEFAULT_ARG,
119 size_t N6 = KOKKOS_IMPL_CTOR_DEFAULT_ARG,
120 size_t N7 = KOKKOS_IMPL_CTOR_DEFAULT_ARG)
121 : dimension{N0, N1, N2, N3, N4, N5, N6, N7} {}
122
123 friend bool operator==(const LayoutRight& left, const LayoutRight& right) {
124 for (unsigned int rank = 0; rank < ARRAY_LAYOUT_MAX_RANK; ++rank)
125 if (left.dimension[rank] != right.dimension[rank]) return false;
126 return true;
127 }
128
129 friend bool operator!=(const LayoutRight& left, const LayoutRight& right) {
130 return !(left == right);
131 }
132};
133
134//----------------------------------------------------------------------------
141
142 size_t dimension[ARRAY_LAYOUT_MAX_RANK];
143 size_t stride[ARRAY_LAYOUT_MAX_RANK];
144
145 enum : bool { is_extent_constructible = false };
146
147 LayoutStride(LayoutStride const&) = default;
148 LayoutStride(LayoutStride&&) = default;
149 LayoutStride& operator=(LayoutStride const&) = default;
150 LayoutStride& operator=(LayoutStride&&) = default;
151
159 template <typename iTypeOrder, typename iTypeDimen>
160 KOKKOS_INLINE_FUNCTION static LayoutStride order_dimensions(
161 int const rank, iTypeOrder const* const order,
162 iTypeDimen const* const dimen) {
163 LayoutStride tmp;
164 // Verify valid rank order:
165 int check_input = ARRAY_LAYOUT_MAX_RANK < rank ? 0 : int(1 << rank) - 1;
166 for (int r = 0; r < ARRAY_LAYOUT_MAX_RANK; ++r) {
167 tmp.dimension[r] = KOKKOS_IMPL_CTOR_DEFAULT_ARG;
168 tmp.stride[r] = 0;
169 }
170 for (int r = 0; r < rank; ++r) {
171 check_input &= ~int(1 << order[r]);
172 }
173 if (0 == check_input) {
174 size_t n = 1;
175 for (int r = 0; r < rank; ++r) {
176 tmp.stride[order[r]] = n;
177 n *= (dimen[order[r]]);
178 tmp.dimension[r] = dimen[r];
179 }
180 }
181 return tmp;
182 }
183
184 KOKKOS_INLINE_FUNCTION
185 explicit constexpr LayoutStride(
186 size_t N0 = KOKKOS_IMPL_CTOR_DEFAULT_ARG, size_t S0 = 0,
187 size_t N1 = KOKKOS_IMPL_CTOR_DEFAULT_ARG, size_t S1 = 0,
188 size_t N2 = KOKKOS_IMPL_CTOR_DEFAULT_ARG, size_t S2 = 0,
189 size_t N3 = KOKKOS_IMPL_CTOR_DEFAULT_ARG, size_t S3 = 0,
190 size_t N4 = KOKKOS_IMPL_CTOR_DEFAULT_ARG, size_t S4 = 0,
191 size_t N5 = KOKKOS_IMPL_CTOR_DEFAULT_ARG, size_t S5 = 0,
192 size_t N6 = KOKKOS_IMPL_CTOR_DEFAULT_ARG, size_t S6 = 0,
193 size_t N7 = KOKKOS_IMPL_CTOR_DEFAULT_ARG, size_t S7 = 0)
194 : dimension{N0, N1, N2, N3, N4, N5, N6, N7}, stride{S0, S1, S2, S3,
195 S4, S5, S6, S7} {}
196
197 friend bool operator==(const LayoutStride& left, const LayoutStride& right) {
198 for (unsigned int rank = 0; rank < ARRAY_LAYOUT_MAX_RANK; ++rank)
199 if (left.dimension[rank] != right.dimension[rank] ||
200 left.stride[rank] != right.stride[rank])
201 return false;
202 return true;
203 }
204
205 friend bool operator!=(const LayoutStride& left, const LayoutStride& right) {
206 return !(left == right);
207 }
208};
209
210// ===================================================================================
211
213
214enum class Iterate {
215 Default,
216 Left, // Left indices stride fastest
217 Right // Right indices stride fastest
218};
219
220// To check for LayoutTiled
221// This is to hide extra compile-time 'identifier' info within the LayoutTiled
222// class by not relying on template specialization to include the ArgN*'s
223template <typename LayoutTiledCheck, class Enable = void>
224struct is_layouttiled : std::false_type {};
225
226template <typename LayoutTiledCheck>
227struct is_layouttiled<LayoutTiledCheck,
228 std::enable_if_t<LayoutTiledCheck::is_array_layout_tiled>>
229 : std::true_type {};
230
231namespace Experimental {
232
234// Must have Rank >= 2
235template <
236 Kokkos::Iterate OuterP, Kokkos::Iterate InnerP, unsigned ArgN0,
237 unsigned ArgN1, unsigned ArgN2 = 0, unsigned ArgN3 = 0, unsigned ArgN4 = 0,
238 unsigned ArgN5 = 0, unsigned ArgN6 = 0, unsigned ArgN7 = 0,
239 bool IsPowerOfTwo =
240 (Kokkos::Impl::is_integral_power_of_two(ArgN0) &&
241 Kokkos::Impl::is_integral_power_of_two(ArgN1) &&
242 (Kokkos::Impl::is_integral_power_of_two(ArgN2) || (ArgN2 == 0)) &&
243 (Kokkos::Impl::is_integral_power_of_two(ArgN3) || (ArgN3 == 0)) &&
244 (Kokkos::Impl::is_integral_power_of_two(ArgN4) || (ArgN4 == 0)) &&
245 (Kokkos::Impl::is_integral_power_of_two(ArgN5) || (ArgN5 == 0)) &&
246 (Kokkos::Impl::is_integral_power_of_two(ArgN6) || (ArgN6 == 0)) &&
247 (Kokkos::Impl::is_integral_power_of_two(ArgN7) || (ArgN7 == 0)))>
249 static_assert(IsPowerOfTwo,
250 "LayoutTiled must be given power-of-two tile dimensions");
251
252 using array_layout = LayoutTiled<OuterP, InnerP, ArgN0, ArgN1, ArgN2, ArgN3,
253 ArgN4, ArgN5, ArgN6, ArgN7, IsPowerOfTwo>;
254 static constexpr Iterate outer_pattern = OuterP;
255 static constexpr Iterate inner_pattern = InnerP;
256
257 enum { N0 = ArgN0 };
258 enum { N1 = ArgN1 };
259 enum { N2 = ArgN2 };
260 enum { N3 = ArgN3 };
261 enum { N4 = ArgN4 };
262 enum { N5 = ArgN5 };
263 enum { N6 = ArgN6 };
264 enum { N7 = ArgN7 };
265
266 size_t dimension[ARRAY_LAYOUT_MAX_RANK];
267
268 enum : bool { is_extent_constructible = true };
269
270 LayoutTiled(LayoutTiled const&) = default;
271 LayoutTiled(LayoutTiled&&) = default;
272 LayoutTiled& operator=(LayoutTiled const&) = default;
273 LayoutTiled& operator=(LayoutTiled&&) = default;
274
275 KOKKOS_INLINE_FUNCTION
276 explicit constexpr LayoutTiled(size_t argN0 = 0, size_t argN1 = 0,
277 size_t argN2 = 0, size_t argN3 = 0,
278 size_t argN4 = 0, size_t argN5 = 0,
279 size_t argN6 = 0, size_t argN7 = 0)
280 : dimension{argN0, argN1, argN2, argN3, argN4, argN5, argN6, argN7} {}
281
282 friend bool operator==(const LayoutTiled& left, const LayoutTiled& right) {
283 for (unsigned int rank = 0; rank < ARRAY_LAYOUT_MAX_RANK; ++rank)
284 if (left.dimension[rank] != right.dimension[rank]) return false;
285 return true;
286 }
287
288 friend bool operator!=(const LayoutTiled& left, const LayoutTiled& right) {
289 return !(left == right);
290 }
291};
292
293} // namespace Experimental
294
295// For use with view_copy
296template <typename... Layout>
297struct layout_iterate_type_selector {
298 static const Kokkos::Iterate outer_iteration_pattern =
299 Kokkos::Iterate::Default;
300 static const Kokkos::Iterate inner_iteration_pattern =
301 Kokkos::Iterate::Default;
302};
303
304template <>
305struct layout_iterate_type_selector<Kokkos::LayoutRight> {
306 static const Kokkos::Iterate outer_iteration_pattern = Kokkos::Iterate::Right;
307 static const Kokkos::Iterate inner_iteration_pattern = Kokkos::Iterate::Right;
308};
309
310template <>
311struct layout_iterate_type_selector<Kokkos::LayoutLeft> {
312 static const Kokkos::Iterate outer_iteration_pattern = Kokkos::Iterate::Left;
313 static const Kokkos::Iterate inner_iteration_pattern = Kokkos::Iterate::Left;
314};
315
316template <>
317struct layout_iterate_type_selector<Kokkos::LayoutStride> {
318 static const Kokkos::Iterate outer_iteration_pattern =
319 Kokkos::Iterate::Default;
320 static const Kokkos::Iterate inner_iteration_pattern =
321 Kokkos::Iterate::Default;
322};
323
324template <unsigned ArgN0, unsigned ArgN1, unsigned ArgN2, unsigned ArgN3,
325 unsigned ArgN4, unsigned ArgN5, unsigned ArgN6, unsigned ArgN7>
326struct layout_iterate_type_selector<Kokkos::Experimental::LayoutTiled<
327 Kokkos::Iterate::Left, Kokkos::Iterate::Left, ArgN0, ArgN1, ArgN2, ArgN3,
328 ArgN4, ArgN5, ArgN6, ArgN7, true>> {
329 static const Kokkos::Iterate outer_iteration_pattern = Kokkos::Iterate::Left;
330 static const Kokkos::Iterate inner_iteration_pattern = Kokkos::Iterate::Left;
331};
332
333template <unsigned ArgN0, unsigned ArgN1, unsigned ArgN2, unsigned ArgN3,
334 unsigned ArgN4, unsigned ArgN5, unsigned ArgN6, unsigned ArgN7>
335struct layout_iterate_type_selector<Kokkos::Experimental::LayoutTiled<
336 Kokkos::Iterate::Right, Kokkos::Iterate::Left, ArgN0, ArgN1, ArgN2, ArgN3,
337 ArgN4, ArgN5, ArgN6, ArgN7, true>> {
338 static const Kokkos::Iterate outer_iteration_pattern = Kokkos::Iterate::Right;
339 static const Kokkos::Iterate inner_iteration_pattern = Kokkos::Iterate::Left;
340};
341
342template <unsigned ArgN0, unsigned ArgN1, unsigned ArgN2, unsigned ArgN3,
343 unsigned ArgN4, unsigned ArgN5, unsigned ArgN6, unsigned ArgN7>
344struct layout_iterate_type_selector<Kokkos::Experimental::LayoutTiled<
345 Kokkos::Iterate::Left, Kokkos::Iterate::Right, ArgN0, ArgN1, ArgN2, ArgN3,
346 ArgN4, ArgN5, ArgN6, ArgN7, true>> {
347 static const Kokkos::Iterate outer_iteration_pattern = Kokkos::Iterate::Left;
348 static const Kokkos::Iterate inner_iteration_pattern = Kokkos::Iterate::Right;
349};
350
351template <unsigned ArgN0, unsigned ArgN1, unsigned ArgN2, unsigned ArgN3,
352 unsigned ArgN4, unsigned ArgN5, unsigned ArgN6, unsigned ArgN7>
353struct layout_iterate_type_selector<Kokkos::Experimental::LayoutTiled<
354 Kokkos::Iterate::Right, Kokkos::Iterate::Right, ArgN0, ArgN1, ArgN2, ArgN3,
355 ArgN4, ArgN5, ArgN6, ArgN7, true>> {
356 static const Kokkos::Iterate outer_iteration_pattern = Kokkos::Iterate::Right;
357 static const Kokkos::Iterate inner_iteration_pattern = Kokkos::Iterate::Right;
358};
359
360} // namespace Kokkos
361
362#endif // #ifndef KOKKOS_LAYOUT_HPP
Memory layout tag indicating left-to-right (Fortran scheme) striding of multi-indices.
Memory layout tag indicating right-to-left (C or lexigraphical scheme) striding of multi-indices.
Memory layout tag indicated arbitrarily strided multi-index mapping into contiguous memory.
static KOKKOS_INLINE_FUNCTION LayoutStride order_dimensions(int const rank, iTypeOrder const *const order, iTypeDimen const *const dimen)
Compute strides from ordered dimensions.