Kokkos Core Kernels Package  Version of the Day
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Groups Pages
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>
19 static_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 
34 namespace Kokkos {
35 
36 // ------------------------------------------------------------------ //
37 // Moved to Kokkos_Layout.hpp for more general accessibility
38 /*
39 enum class Iterate
40 {
41  Default, // Default for the device
42  Left, // Left indices stride fastest
43  Right, // Right indices stride fastest
44 };
45 */
46 
47 template <typename ExecSpace>
48 struct default_outer_direction {
49  using type = Iterate;
50  static constexpr Iterate value = Iterate::Right;
51 };
52 
53 template <typename ExecSpace>
54 struct default_inner_direction {
55  using type = Iterate;
56  static constexpr Iterate value = Iterate::Right;
57 };
58 
59 namespace Impl {
60 // NOTE the comparison below is encapsulated to silent warnings about pointless
61 // comparison of unsigned integer with zero
62 template <class T>
63 constexpr std::enable_if_t<!std::is_signed<T>::value, bool>
64 is_less_than_value_initialized_variable(T) {
65  return false;
66 }
67 
68 template <class T>
69 constexpr std::enable_if_t<std::is_signed<T>::value, bool>
70 is_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
75 template <class To, class From>
76 constexpr To checked_narrow_cast(From arg, std::size_t idx) {
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  auto msg =
85  "Kokkos::MDRangePolicy bound type error: an unsafe implicit conversion "
86  "is performed on a bound (" +
87  std::to_string(arg) + ") in dimension (" + std::to_string(idx) +
88  "), which may not preserve its original value.\n";
89  Kokkos::abort(msg.c_str());
90  }
91  return ret;
92 }
93 // NOTE prefer C array U[M] to std::initalizer_list<U> so that the number of
94 // elements can be deduced (https://stackoverflow.com/q/40241370)
95 // NOTE for some unfortunate reason the policy bounds are stored as signed
96 // integer arrays (point_type which is Kokkos::Array<std::int64_t>) so we
97 // specify the index type (actual policy index_type from the traits) and check
98 // ahead of time that narrowing conversions will be safe.
99 template <class IndexType, class Array, class U, std::size_t M>
100 constexpr Array to_array_potentially_narrowing(const U (&init)[M]) {
101  using T = typename Array::value_type;
102  Array a{};
103  constexpr std::size_t N = a.size();
104  static_assert(M <= N);
105  auto* ptr = a.data();
106  // NOTE equivalent to
107  // std::transform(std::begin(init), std::end(init), a.data(),
108  // [](U x) { return static_cast<T>(x); });
109  // except that std::transform is not constexpr.
110  for (std::size_t i = 0; i < M; ++i) {
111  *ptr++ = checked_narrow_cast<T>(init[i], i);
112  (void)checked_narrow_cast<IndexType>(init[i], i); // see note above
113  }
114  return a;
115 }
116 
117 // NOTE Making a copy even when std::is_same<Array, Kokkos::Array<U, M>>::value
118 // is true to reduce code complexity. You may change this if you have a good
119 // reason to. Intentionally not enabling std::array at this time but this may
120 // change too.
121 template <class IndexType, class NVCC_WONT_LET_ME_CALL_YOU_Array, class U,
122  std::size_t M>
123 constexpr NVCC_WONT_LET_ME_CALL_YOU_Array to_array_potentially_narrowing(
124  Kokkos::Array<U, M> const& other) {
125  using T = typename NVCC_WONT_LET_ME_CALL_YOU_Array::value_type;
126  NVCC_WONT_LET_ME_CALL_YOU_Array a{};
127  constexpr std::size_t N = a.size();
128  static_assert(M <= N);
129  for (std::size_t i = 0; i < M; ++i) {
130  a[i] = checked_narrow_cast<T>(other[i], i);
131  (void)checked_narrow_cast<IndexType>(other[i], i); // see note above
132  }
133  return a;
134 }
135 
136 struct TileSizeProperties {
137  int max_threads;
138  int default_largest_tile_size;
139  int default_tile_size;
140  int max_total_tile_size;
141 };
142 
143 template <typename ExecutionSpace>
144 TileSizeProperties get_tile_size_properties(const ExecutionSpace&) {
145  // Host settings
146  TileSizeProperties properties;
147  properties.max_threads = std::numeric_limits<int>::max();
148  properties.default_largest_tile_size = 0;
149  properties.default_tile_size = 2;
150  properties.max_total_tile_size = std::numeric_limits<int>::max();
151  return properties;
152 }
153 
154 } // namespace Impl
155 
156 // multi-dimensional iteration pattern
157 template <typename... Properties>
158 struct MDRangePolicy;
159 
160 // Note: If MDRangePolicy has a primary template, implicit CTAD (deduction
161 // guides) are generated -> MDRangePolicy<> by some compilers, which is
162 // incorrect. By making it a template specialization instead, no implicit CTAD
163 // is generated. This works because there has to be at least one property
164 // specified (which is Rank<...>); otherwise, we'd get the static_assert
165 // "Kokkos::Error: MD iteration pattern not defined". This template
166 // specialization uses <P, Properties...> in all places for correctness.
167 template <typename P, typename... Properties>
168 struct MDRangePolicy<P, Properties...>
169  : public Kokkos::Impl::PolicyTraits<P, Properties...> {
170  using traits = Kokkos::Impl::PolicyTraits<P, Properties...>;
171  using range_policy = RangePolicy<P, Properties...>;
172 
173  typename traits::execution_space m_space;
174 
175  using impl_range_policy =
176  RangePolicy<typename traits::execution_space,
177  typename traits::schedule_type, typename traits::index_type>;
178 
179  using execution_policy =
180  MDRangePolicy<P, Properties...>; // needed for is_execution_policy
181  // interrogation
182 
183  template <class... OtherProperties>
184  friend struct MDRangePolicy;
185 
186  static_assert(!std::is_void<typename traits::iteration_pattern>::value,
187  "Kokkos Error: MD iteration pattern not defined");
188 
189  using iteration_pattern = typename traits::iteration_pattern;
190  using work_tag = typename traits::work_tag;
191  using launch_bounds = typename traits::launch_bounds;
192  using member_type = typename range_policy::member_type;
193 
194  static constexpr int rank = iteration_pattern::rank;
195  static_assert(rank < 7, "Kokkos MDRangePolicy Error: Unsupported rank...");
196 
197  using index_type = typename traits::index_type;
198  using array_index_type = std::int64_t;
199  using point_type = Kokkos::Array<array_index_type, rank>; // was index_type
200  using tile_type = Kokkos::Array<array_index_type, rank>;
201  // If point_type or tile_type is not templated on a signed integral type (if
202  // it is unsigned), then if user passes in intializer_list of
203  // runtime-determined values of signed integral type that are not const will
204  // receive a compiler error due to an invalid case for implicit conversion -
205  // "conversion from integer or unscoped enumeration type to integer type that
206  // cannot represent all values of the original, except where source is a
207  // constant expression whose value can be stored exactly in the target type"
208  // This would require the user to either pass a matching index_type parameter
209  // as template parameter to the MDRangePolicy or static_cast the individual
210  // values
211 
212  point_type m_lower = {};
213  point_type m_upper = {};
214  tile_type m_tile = {};
215  point_type m_tile_end = {};
216  index_type m_num_tiles = 1;
217  index_type m_prod_tile_dims = 1;
218  bool m_tune_tile_size = false;
219 
220  static constexpr auto outer_direction =
221  (iteration_pattern::outer_direction != Iterate::Default)
222  ? iteration_pattern::outer_direction
223  : default_outer_direction<typename traits::execution_space>::value;
224 
225  static constexpr auto inner_direction =
226  iteration_pattern::inner_direction != Iterate::Default
227  ? iteration_pattern::inner_direction
228  : default_inner_direction<typename traits::execution_space>::value;
229 
230  static constexpr auto Right = Iterate::Right;
231  static constexpr auto Left = Iterate::Left;
232 
233  KOKKOS_INLINE_FUNCTION const typename traits::execution_space& space() const {
234  return m_space;
235  }
236 
237  MDRangePolicy() = default;
238 
239  template <typename LT, std::size_t LN, typename UT, std::size_t UN,
240  typename TT = array_index_type, std::size_t TN = rank,
241  typename = std::enable_if_t<std::is_integral<LT>::value &&
242  std::is_integral<UT>::value &&
243  std::is_integral<TT>::value>>
244  MDRangePolicy(const LT (&lower)[LN], const UT (&upper)[UN],
245  const TT (&tile)[TN] = {})
246  : MDRangePolicy(
247  Impl::to_array_potentially_narrowing<index_type, decltype(m_lower)>(
248  lower),
249  Impl::to_array_potentially_narrowing<index_type, decltype(m_upper)>(
250  upper),
251  Impl::to_array_potentially_narrowing<index_type, decltype(m_tile)>(
252  tile)) {
253  static_assert(
254  LN == rank && UN == rank && TN <= rank,
255  "MDRangePolicy: Constructor initializer lists have wrong size");
256  }
257 
258  template <typename LT, std::size_t LN, typename UT, std::size_t UN,
259  typename TT = array_index_type, std::size_t TN = rank,
260  typename = std::enable_if_t<std::is_integral<LT>::value &&
261  std::is_integral<UT>::value &&
262  std::is_integral<TT>::value>>
263  MDRangePolicy(const typename traits::execution_space& work_space,
264  const LT (&lower)[LN], const UT (&upper)[UN],
265  const TT (&tile)[TN] = {})
266  : MDRangePolicy(
267  work_space,
268  Impl::to_array_potentially_narrowing<index_type, decltype(m_lower)>(
269  lower),
270  Impl::to_array_potentially_narrowing<index_type, decltype(m_upper)>(
271  upper),
272  Impl::to_array_potentially_narrowing<index_type, decltype(m_tile)>(
273  tile)) {
274  static_assert(
275  LN == rank && UN == rank && TN <= rank,
276  "MDRangePolicy: Constructor initializer lists have wrong size");
277  }
278 
279  // NOTE: Keeping these two constructor despite the templated constructors
280  // from Kokkos arrays for backwards compability to allow construction from
281  // double-braced initializer lists.
282  MDRangePolicy(point_type const& lower, point_type const& upper,
283  tile_type const& tile = tile_type{})
284  : MDRangePolicy(typename traits::execution_space(), lower, upper, tile) {}
285 
286  MDRangePolicy(const typename traits::execution_space& work_space,
287  point_type const& lower, point_type const& upper,
288  tile_type const& tile = tile_type{})
289  : m_space(work_space), m_lower(lower), m_upper(upper), m_tile(tile) {
290  init_helper(Impl::get_tile_size_properties(work_space));
291  }
292 
293  template <typename T, std::size_t NT = rank,
294  typename = std::enable_if_t<std::is_integral<T>::value>>
295  MDRangePolicy(Kokkos::Array<T, rank> const& lower,
296  Kokkos::Array<T, rank> const& upper,
298  : MDRangePolicy(typename traits::execution_space(), lower, upper, tile) {}
299 
300  template <typename T, std::size_t NT = rank,
301  typename = std::enable_if_t<std::is_integral<T>::value>>
302  MDRangePolicy(const typename traits::execution_space& work_space,
303  Kokkos::Array<T, rank> const& lower,
304  Kokkos::Array<T, rank> const& upper,
306  : MDRangePolicy(
307  work_space,
308  Impl::to_array_potentially_narrowing<index_type, decltype(m_lower)>(
309  lower),
310  Impl::to_array_potentially_narrowing<index_type, decltype(m_upper)>(
311  upper),
312  Impl::to_array_potentially_narrowing<index_type, decltype(m_tile)>(
313  tile)) {}
314 
315  template <class... OtherProperties>
316  MDRangePolicy(const MDRangePolicy<OtherProperties...> p)
317  : traits(p), // base class may contain data such as desired occupancy
318  m_space(p.m_space),
319  m_lower(p.m_lower),
320  m_upper(p.m_upper),
321  m_tile(p.m_tile),
322  m_tile_end(p.m_tile_end),
323  m_num_tiles(p.m_num_tiles),
324  m_prod_tile_dims(p.m_prod_tile_dims),
325  m_tune_tile_size(p.m_tune_tile_size) {}
326 
327  void impl_change_tile_size(const point_type& tile) {
328  m_tile = tile;
329  init_helper(Impl::get_tile_size_properties(m_space));
330  }
331  bool impl_tune_tile_size() const { return m_tune_tile_size; }
332 
333  private:
334  void init_helper(Impl::TileSizeProperties properties) {
335  m_prod_tile_dims = 1;
336  int increment = 1;
337  int rank_start = 0;
338  int rank_end = rank;
339  if (inner_direction == Iterate::Right) {
340  increment = -1;
341  rank_start = rank - 1;
342  rank_end = -1;
343  }
344  for (int i = rank_start; i != rank_end; i += increment) {
345  const index_type length = m_upper[i] - m_lower[i];
346 
347  if (m_upper[i] < m_lower[i]) {
348  std::string msg =
349  "Kokkos::MDRangePolicy bounds error: The lower bound (" +
350  std::to_string(m_lower[i]) + ") is greater than its upper bound (" +
351  std::to_string(m_upper[i]) + ") in dimension " + std::to_string(i) +
352  ".\n";
353 #if !defined(KOKKOS_ENABLE_DEPRECATED_CODE_4)
354  Kokkos::abort(msg.c_str());
355 #elif defined(KOKKOS_ENABLE_DEPRECATION_WARNINGS)
356  Kokkos::Impl::log_warning(msg);
357 #endif
358  }
359 
360  if (m_tile[i] <= 0) {
361  m_tune_tile_size = true;
362  if ((inner_direction == Iterate::Right && (i < rank - 1)) ||
363  (inner_direction == Iterate::Left && (i > 0))) {
364  if (m_prod_tile_dims * properties.default_tile_size <
365  static_cast<index_type>(properties.max_total_tile_size)) {
366  m_tile[i] = properties.default_tile_size;
367  } else {
368  m_tile[i] = 1;
369  }
370  } else {
371  m_tile[i] = properties.default_largest_tile_size == 0
372  ? std::max<int>(length, 1)
373  : properties.default_largest_tile_size;
374  }
375  }
376  m_tile_end[i] =
377  static_cast<index_type>((length + m_tile[i] - 1) / m_tile[i]);
378  m_num_tiles *= m_tile_end[i];
379  m_prod_tile_dims *= m_tile[i];
380  }
381  if (m_prod_tile_dims > static_cast<index_type>(properties.max_threads)) {
382  printf(" Product of tile dimensions exceed maximum limit: %d\n",
383  static_cast<int>(properties.max_threads));
384  Kokkos::abort(
385  "ExecSpace Error: MDRange tile dims exceed maximum number "
386  "of threads per block - choose smaller tile dims");
387  }
388  }
389 };
390 
391 template <typename LT, size_t N, typename UT>
392 MDRangePolicy(const LT (&)[N], const UT (&)[N])->MDRangePolicy<Rank<N>>;
393 
394 template <typename LT, size_t N, typename UT, typename TT, size_t TN>
395 MDRangePolicy(const LT (&)[N], const UT (&)[N], const TT (&)[TN])
396  ->MDRangePolicy<Rank<N>>;
397 
398 template <typename LT, size_t N, typename UT>
399 MDRangePolicy(DefaultExecutionSpace const&, const LT (&)[N], const UT (&)[N])
400  ->MDRangePolicy<Rank<N>>;
401 
402 template <typename LT, size_t N, typename UT, typename TT, size_t TN>
403 MDRangePolicy(DefaultExecutionSpace const&, const LT (&)[N], const UT (&)[N],
404  const TT (&)[TN])
405  ->MDRangePolicy<Rank<N>>;
406 
407 template <typename ES, typename LT, size_t N, typename UT,
408  typename = std::enable_if_t<is_execution_space_v<ES>>>
409 MDRangePolicy(ES const&, const LT (&)[N], const UT (&)[N])
410  ->MDRangePolicy<ES, Rank<N>>;
411 
412 template <typename ES, typename LT, size_t N, typename UT, typename TT,
413  size_t TN, typename = std::enable_if_t<is_execution_space_v<ES>>>
414 MDRangePolicy(ES const&, const LT (&)[N], const UT (&)[N], const TT (&)[TN])
415  ->MDRangePolicy<ES, Rank<N>>;
416 
417 template <typename T, size_t N>
418 MDRangePolicy(Array<T, N> const&, Array<T, N> const&)->MDRangePolicy<Rank<N>>;
419 
420 template <typename T, size_t N, size_t NT>
421 MDRangePolicy(Array<T, N> const&, Array<T, N> const&, Array<T, NT> const&)
422  ->MDRangePolicy<Rank<N>>;
423 
424 template <typename T, size_t N>
425 MDRangePolicy(DefaultExecutionSpace const&, Array<T, N> const&,
426  Array<T, N> const&)
427  ->MDRangePolicy<Rank<N>>;
428 
429 template <typename T, size_t N, size_t NT>
430 MDRangePolicy(DefaultExecutionSpace const&, Array<T, N> const&,
431  Array<T, N> const&, Array<T, NT> const&)
432  ->MDRangePolicy<Rank<N>>;
433 
434 template <typename ES, typename T, size_t N,
435  typename = std::enable_if_t<is_execution_space_v<ES>>>
436 MDRangePolicy(ES const&, Array<T, N> const&, Array<T, N> const&)
437  ->MDRangePolicy<ES, Rank<N>>;
438 
439 template <typename ES, typename T, size_t N, size_t NT,
440  typename = std::enable_if_t<is_execution_space_v<ES>>>
441 MDRangePolicy(ES const&, Array<T, N> const&, Array<T, N> const&,
442  Array<T, NT> const&)
443  ->MDRangePolicy<ES, Rank<N>>;
444 
445 } // namespace Kokkos
446 
447 #endif // KOKKOS_CORE_EXP_MD_RANGE_POLICY_HPP
Derived from the C++17 &#39;std::array&#39;. Dropping the iterator interface.
Declaration of various MemoryLayout options.