Kokkos Core Kernels Package Version of the Day
Loading...
Searching...
No Matches
KokkosExp_MDRangePolicy.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_CORE_EXP_MD_RANGE_POLICY_HPP
23#define KOKKOS_CORE_EXP_MD_RANGE_POLICY_HPP
24
25#include <initializer_list>
26
27#include <Kokkos_Layout.hpp>
28#include <Kokkos_Rank.hpp>
29#include <Kokkos_Array.hpp>
30#include <impl/KokkosExp_Host_IterateTile.hpp>
31#include <Kokkos_ExecPolicy.hpp>
32#include <type_traits>
33
34namespace Kokkos {
35
36// ------------------------------------------------------------------ //
37// Moved to Kokkos_Layout.hpp for more general accessibility
38/*
39enum class Iterate
40{
41 Default, // Default for the device
42 Left, // Left indices stride fastest
43 Right, // Right indices stride fastest
44};
45*/
46
47template <typename ExecSpace>
48struct default_outer_direction {
49 using type = Iterate;
50 static constexpr Iterate value = Iterate::Right;
51};
52
53template <typename ExecSpace>
54struct default_inner_direction {
55 using type = Iterate;
56 static constexpr Iterate value = Iterate::Right;
57};
58
59namespace Impl {
60// NOTE the comparison below is encapsulated to silent warnings about pointless
61// comparison of unsigned integer with zero
62template <class T>
63constexpr std::enable_if_t<!std::is_signed<T>::value, bool>
64is_less_than_value_initialized_variable(T) {
65 return false;
66}
67
68template <class T>
69constexpr std::enable_if_t<std::is_signed<T>::value, bool>
70is_less_than_value_initialized_variable(T arg) {
71 return arg < T{};
72}
73
74// Checked narrowing conversion that calls abort if the cast changes the value
75template <class To, class From>
76constexpr To checked_narrow_cast(From arg) {
77 constexpr const bool is_different_signedness =
78 (std::is_signed<To>::value != std::is_signed<From>::value);
79 auto const ret = static_cast<To>(arg);
80 if (static_cast<From>(ret) != arg ||
81 (is_different_signedness &&
82 is_less_than_value_initialized_variable(arg) !=
83 is_less_than_value_initialized_variable(ret))) {
84 Kokkos::abort("unsafe narrowing conversion");
85 }
86 return ret;
87}
88// NOTE prefer C array U[M] to std::initalizer_list<U> so that the number of
89// elements can be deduced (https://stackoverflow.com/q/40241370)
90// NOTE for some unfortunate reason the policy bounds are stored as signed
91// integer arrays (point_type which is Kokkos::Array<std::int64_t>) so we
92// specify the index type (actual policy index_type from the traits) and check
93// ahead of time that narrowing conversions will be safe.
94template <class IndexType, class Array, class U, std::size_t M>
95constexpr Array to_array_potentially_narrowing(const U (&init)[M]) {
96 using T = typename Array::value_type;
97 Array a{};
98 constexpr std::size_t N = a.size();
99 static_assert(M <= N, "");
100 auto* ptr = a.data();
101 // NOTE equivalent to
102 // std::transform(std::begin(init), std::end(init), a.data(),
103 // [](U x) { return static_cast<T>(x); });
104 // except that std::transform is not constexpr.
105 for (auto x : init) {
106 *ptr++ = checked_narrow_cast<T>(x);
107 (void)checked_narrow_cast<IndexType>(x); // see note above
108 }
109 return a;
110}
111
112// NOTE Making a copy even when std::is_same<Array, Kokkos::Array<U, M>>::value
113// is true to reduce code complexity. You may change this if you have a good
114// reason to. Intentionally not enabling std::array at this time but this may
115// change too.
116template <class IndexType, class NVCC_WONT_LET_ME_CALL_YOU_Array, class U,
117 std::size_t M>
118constexpr NVCC_WONT_LET_ME_CALL_YOU_Array to_array_potentially_narrowing(
119 Kokkos::Array<U, M> const& other) {
120 using T = typename NVCC_WONT_LET_ME_CALL_YOU_Array::value_type;
121 NVCC_WONT_LET_ME_CALL_YOU_Array a{};
122 constexpr std::size_t N = a.size();
123 static_assert(M <= N, "");
124 for (std::size_t i = 0; i < M; ++i) {
125 a[i] = checked_narrow_cast<T>(other[i]);
126 (void)checked_narrow_cast<IndexType>(other[i]); // see note above
127 }
128 return a;
129}
130
131struct TileSizeProperties {
132 int max_threads;
133 int default_largest_tile_size;
134 int default_tile_size;
135 int max_total_tile_size;
136};
137
138template <typename ExecutionSpace>
139TileSizeProperties get_tile_size_properties(const ExecutionSpace&) {
140 // Host settings
141 TileSizeProperties properties;
142 properties.max_threads = std::numeric_limits<int>::max();
143 properties.default_largest_tile_size = 0;
144 properties.default_tile_size = 2;
145 properties.max_total_tile_size = std::numeric_limits<int>::max();
146 return properties;
147}
148
149} // namespace Impl
150
151// multi-dimensional iteration pattern
152template <typename... Properties>
153struct MDRangePolicy : public Kokkos::Impl::PolicyTraits<Properties...> {
154 using traits = Kokkos::Impl::PolicyTraits<Properties...>;
155 using range_policy = RangePolicy<Properties...>;
156
157 typename traits::execution_space m_space;
158
159 using impl_range_policy =
160 RangePolicy<typename traits::execution_space,
161 typename traits::schedule_type, typename traits::index_type>;
162
163 using execution_policy =
164 MDRangePolicy<Properties...>; // needed for is_execution_space
165 // interrogation
166
167 template <class... OtherProperties>
168 friend struct MDRangePolicy;
169
170 static_assert(!std::is_void<typename traits::iteration_pattern>::value,
171 "Kokkos Error: MD iteration pattern not defined");
172
173 using iteration_pattern = typename traits::iteration_pattern;
174 using work_tag = typename traits::work_tag;
175 using launch_bounds = typename traits::launch_bounds;
176 using member_type = typename range_policy::member_type;
177
178 static constexpr int rank = iteration_pattern::rank;
179 static_assert(rank < 7, "Kokkos MDRangePolicy Error: Unsupported rank...");
180
181 using index_type = typename traits::index_type;
182 using array_index_type = std::int64_t;
183 using point_type = Kokkos::Array<array_index_type, rank>; // was index_type
185 // If point_type or tile_type is not templated on a signed integral type (if
186 // it is unsigned), then if user passes in intializer_list of
187 // runtime-determined values of signed integral type that are not const will
188 // receive a compiler error due to an invalid case for implicit conversion -
189 // "conversion from integer or unscoped enumeration type to integer type that
190 // cannot represent all values of the original, except where source is a
191 // constant expression whose value can be stored exactly in the target type"
192 // This would require the user to either pass a matching index_type parameter
193 // as template parameter to the MDRangePolicy or static_cast the individual
194 // values
195
196 point_type m_lower = {};
197 point_type m_upper = {};
198 tile_type m_tile = {};
199 point_type m_tile_end = {};
200 index_type m_num_tiles = 1;
201 index_type m_prod_tile_dims = 1;
202 bool m_tune_tile_size = false;
203
204 static constexpr auto outer_direction =
205 (iteration_pattern::outer_direction != Iterate::Default)
206 ? iteration_pattern::outer_direction
207 : default_outer_direction<typename traits::execution_space>::value;
208
209 static constexpr auto inner_direction =
210 iteration_pattern::inner_direction != Iterate::Default
211 ? iteration_pattern::inner_direction
212 : default_inner_direction<typename traits::execution_space>::value;
213
214 static constexpr auto Right = Iterate::Right;
215 static constexpr auto Left = Iterate::Left;
216
217 KOKKOS_INLINE_FUNCTION const typename traits::execution_space& space() const {
218 return m_space;
219 }
220
221 MDRangePolicy() = default;
222
223 template <typename LT, std::size_t LN, typename UT, std::size_t UN,
224 typename TT = array_index_type, std::size_t TN = rank,
225 typename = std::enable_if_t<std::is_integral<LT>::value &&
226 std::is_integral<UT>::value &&
227 std::is_integral<TT>::value>>
228 MDRangePolicy(const LT (&lower)[LN], const UT (&upper)[UN],
229 const TT (&tile)[TN] = {})
230 : MDRangePolicy(
231 Impl::to_array_potentially_narrowing<index_type, decltype(m_lower)>(
232 lower),
233 Impl::to_array_potentially_narrowing<index_type, decltype(m_upper)>(
234 upper),
235 Impl::to_array_potentially_narrowing<index_type, decltype(m_tile)>(
236 tile)) {
237 static_assert(
238 LN == rank && UN == rank && TN <= rank,
239 "MDRangePolicy: Constructor initializer lists have wrong size");
240 }
241
242 template <typename LT, std::size_t LN, typename UT, std::size_t UN,
243 typename TT = array_index_type, std::size_t TN = rank,
244 typename = std::enable_if_t<std::is_integral<LT>::value &&
245 std::is_integral<UT>::value &&
246 std::is_integral<TT>::value>>
247 MDRangePolicy(const typename traits::execution_space& work_space,
248 const LT (&lower)[LN], const UT (&upper)[UN],
249 const TT (&tile)[TN] = {})
250 : MDRangePolicy(
251 work_space,
252 Impl::to_array_potentially_narrowing<index_type, decltype(m_lower)>(
253 lower),
254 Impl::to_array_potentially_narrowing<index_type, decltype(m_upper)>(
255 upper),
256 Impl::to_array_potentially_narrowing<index_type, decltype(m_tile)>(
257 tile)) {
258 static_assert(
259 LN == rank && UN == rank && TN <= rank,
260 "MDRangePolicy: Constructor initializer lists have wrong size");
261 }
262
263 // NOTE: Keeping these two constructor despite the templated constructors
264 // from Kokkos arrays for backwards compability to allow construction from
265 // double-braced initializer lists.
266 MDRangePolicy(point_type const& lower, point_type const& upper,
267 tile_type const& tile = tile_type{})
268 : MDRangePolicy(typename traits::execution_space(), lower, upper, tile) {}
269
270 MDRangePolicy(const typename traits::execution_space& work_space,
271 point_type const& lower, point_type const& upper,
272 tile_type const& tile = tile_type{})
273 : m_space(work_space), m_lower(lower), m_upper(upper), m_tile(tile) {
274 init_helper(Impl::get_tile_size_properties(work_space));
275 }
276
277 template <typename T, std::size_t NT = rank,
278 typename = std::enable_if_t<std::is_integral<T>::value>>
279 MDRangePolicy(Kokkos::Array<T, rank> const& lower,
280 Kokkos::Array<T, rank> const& upper,
282 : MDRangePolicy(typename traits::execution_space(), lower, upper, tile) {}
283
284 template <typename T, std::size_t NT = rank,
285 typename = std::enable_if_t<std::is_integral<T>::value>>
286 MDRangePolicy(const typename traits::execution_space& work_space,
287 Kokkos::Array<T, rank> const& lower,
288 Kokkos::Array<T, rank> const& upper,
290 : MDRangePolicy(
291 work_space,
292 Impl::to_array_potentially_narrowing<index_type, decltype(m_lower)>(
293 lower),
294 Impl::to_array_potentially_narrowing<index_type, decltype(m_upper)>(
295 upper),
296 Impl::to_array_potentially_narrowing<index_type, decltype(m_tile)>(
297 tile)) {}
298
299 template <class... OtherProperties>
300 MDRangePolicy(const MDRangePolicy<OtherProperties...> p)
301 : traits(p), // base class may contain data such as desired occupancy
302 m_space(p.m_space),
303 m_lower(p.m_lower),
304 m_upper(p.m_upper),
305 m_tile(p.m_tile),
306 m_tile_end(p.m_tile_end),
307 m_num_tiles(p.m_num_tiles),
308 m_prod_tile_dims(p.m_prod_tile_dims),
309 m_tune_tile_size(p.m_tune_tile_size) {}
310
311 void impl_change_tile_size(const point_type& tile) {
312 m_tile = tile;
313 init_helper(Impl::get_tile_size_properties(m_space));
314 }
315 bool impl_tune_tile_size() const { return m_tune_tile_size; }
316
317 private:
318 void init_helper(Impl::TileSizeProperties properties) {
319 m_prod_tile_dims = 1;
320 int increment = 1;
321 int rank_start = 0;
322 int rank_end = rank;
323 if (inner_direction == Iterate::Right) {
324 increment = -1;
325 rank_start = rank - 1;
326 rank_end = -1;
327 }
328 for (int i = rank_start; i != rank_end; i += increment) {
329 const index_type length = m_upper[i] - m_lower[i];
330 if (m_tile[i] <= 0) {
331 m_tune_tile_size = true;
332 if ((inner_direction == Iterate::Right && (i < rank - 1)) ||
333 (inner_direction == Iterate::Left && (i > 0))) {
334 if (m_prod_tile_dims * properties.default_tile_size <
335 static_cast<index_type>(properties.max_total_tile_size)) {
336 m_tile[i] = properties.default_tile_size;
337 } else {
338 m_tile[i] = 1;
339 }
340 } else {
341 m_tile[i] = properties.default_largest_tile_size == 0
342 ? std::max<int>(length, 1)
343 : properties.default_largest_tile_size;
344 }
345 }
346 m_tile_end[i] =
347 static_cast<index_type>((length + m_tile[i] - 1) / m_tile[i]);
348 m_num_tiles *= m_tile_end[i];
349 m_prod_tile_dims *= m_tile[i];
350 }
351 if (m_prod_tile_dims > static_cast<index_type>(properties.max_threads)) {
352 printf(" Product of tile dimensions exceed maximum limit: %d\n",
353 static_cast<int>(properties.max_threads));
355 "ExecSpace Error: MDRange tile dims exceed maximum number "
356 "of threads per block - choose smaller tile dims");
357 }
358 }
359};
360
361} // namespace Kokkos
362
363#endif // KOKKOS_CORE_EXP_MD_RANGE_POLICY_HPP
Declaration of various MemoryLayout options.
A thread safe view to a bitset.
KOKKOS_FORCEINLINE_FUNCTION unsigned size() const