17#ifndef KOKKOS_IMPL_PUBLIC_INCLUDE
18#include <Kokkos_Macros.hpp>
20 "Including non-public Kokkos header files is not allowed.");
22#ifndef KOKKOS_CORE_EXP_MD_RANGE_POLICY_HPP
23#define KOKKOS_CORE_EXP_MD_RANGE_POLICY_HPP
25#include <initializer_list>
28#include <Kokkos_Rank.hpp>
29#include <Kokkos_Array.hpp>
30#include <impl/KokkosExp_Host_IterateTile.hpp>
31#include <Kokkos_ExecPolicy.hpp>
47template <
typename ExecSpace>
48struct default_outer_direction {
50 static constexpr Iterate value = Iterate::Right;
53template <
typename ExecSpace>
54struct default_inner_direction {
56 static constexpr Iterate value = Iterate::Right;
63constexpr std::enable_if_t<!std::is_signed<T>::value,
bool>
64is_less_than_value_initialized_variable(T) {
69constexpr std::enable_if_t<std::is_signed<T>::value,
bool>
70is_less_than_value_initialized_variable(T arg) {
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))) {
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;
98 constexpr std::size_t N = a.
size();
99 static_assert(M <= N,
"");
100 auto* ptr = a.data();
105 for (
auto x : init) {
106 *ptr++ = checked_narrow_cast<T>(x);
107 (void)checked_narrow_cast<IndexType>(x);
116template <
class IndexType,
class NVCC_WONT_LET_ME_CALL_YOU_Array,
class U,
118constexpr NVCC_WONT_LET_ME_CALL_YOU_Array to_array_potentially_narrowing(
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]);
131struct TileSizeProperties {
133 int default_largest_tile_size;
134 int default_tile_size;
135 int max_total_tile_size;
138template <
typename ExecutionSpace>
139TileSizeProperties get_tile_size_properties(
const ExecutionSpace&) {
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();
152template <
typename... Properties>
153struct MDRangePolicy :
public Kokkos::Impl::PolicyTraits<Properties...> {
154 using traits = Kokkos::Impl::PolicyTraits<Properties...>;
155 using range_policy = RangePolicy<Properties...>;
157 typename traits::execution_space m_space;
159 using impl_range_policy =
160 RangePolicy<
typename traits::execution_space,
161 typename traits::schedule_type,
typename traits::index_type>;
163 using execution_policy =
164 MDRangePolicy<Properties...>;
167 template <
class... OtherProperties>
168 friend struct MDRangePolicy;
170 static_assert(!std::is_void<typename traits::iteration_pattern>::value,
171 "Kokkos Error: MD iteration pattern not defined");
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;
178 static constexpr int rank = iteration_pattern::rank;
179 static_assert(rank < 7,
"Kokkos MDRangePolicy Error: Unsupported rank...");
181 using index_type =
typename traits::index_type;
182 using array_index_type = std::int64_t;
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;
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;
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;
214 static constexpr auto Right = Iterate::Right;
215 static constexpr auto Left = Iterate::Left;
217 KOKKOS_INLINE_FUNCTION
const typename traits::execution_space& space()
const {
221 MDRangePolicy() =
default;
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] = {})
231 Impl::to_array_potentially_narrowing<index_type, decltype(m_lower)>(
233 Impl::to_array_potentially_narrowing<index_type, decltype(m_upper)>(
235 Impl::to_array_potentially_narrowing<index_type, decltype(m_tile)>(
238 LN == rank && UN == rank && TN <= rank,
239 "MDRangePolicy: Constructor initializer lists have wrong size");
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] = {})
252 Impl::to_array_potentially_narrowing<index_type, decltype(m_lower)>(
254 Impl::to_array_potentially_narrowing<index_type, decltype(m_upper)>(
256 Impl::to_array_potentially_narrowing<index_type, decltype(m_tile)>(
259 LN == rank && UN == rank && TN <= rank,
260 "MDRangePolicy: Constructor initializer lists have wrong size");
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) {}
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));
277 template <
typename T, std::size_t NT = rank,
278 typename = std::enable_if_t<std::is_integral<T>::value>>
282 : MDRangePolicy(typename traits::execution_space(), lower, upper, tile) {}
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,
292 Impl::to_array_potentially_narrowing<index_type, decltype(m_lower)>(
294 Impl::to_array_potentially_narrowing<index_type, decltype(m_upper)>(
296 Impl::to_array_potentially_narrowing<index_type, decltype(m_tile)>(
299 template <
class... OtherProperties>
300 MDRangePolicy(
const MDRangePolicy<OtherProperties...> p)
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) {}
311 void impl_change_tile_size(
const point_type& tile) {
313 init_helper(Impl::get_tile_size_properties(m_space));
315 bool impl_tune_tile_size()
const {
return m_tune_tile_size; }
318 void init_helper(Impl::TileSizeProperties properties) {
319 m_prod_tile_dims = 1;
323 if (inner_direction == Iterate::Right) {
325 rank_start = rank - 1;
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;
341 m_tile[i] = properties.default_largest_tile_size == 0
342 ? std::max<int>(length, 1)
343 : properties.default_largest_tile_size;
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];
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");
Declaration of various MemoryLayout options.
A thread safe view to a bitset.
KOKKOS_FORCEINLINE_FUNCTION unsigned size() const