/
tut_nested-loop-reorder.cpp
141 lines (112 loc) · 4.78 KB
/
tut_nested-loop-reorder.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
// Copyright (c) 2016-19, Lawrence Livermore National Security, LLC
// and RAJA project contributors. See the RAJA/COPYRIGHT file for details.
//
// SPDX-License-Identifier: (BSD-3-Clause)
//~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~//
#include <cstdlib>
#include <iostream>
#include "RAJA/RAJA.hpp"
/*
* Nested Loop Reorder Example
*
* This example shows how to reorder RAJA nested loops by reordering
* nested policy arguments. It does no actual computation and just
* prints out the loop indices to show the different orderings.
*
* RAJA features shown:
* - Index range segment
* - 'RAJA::kernel' loop abstractions and execution policies
* - Nested loop reordering
* - Strongly-typed loop indices
*/
//
// Define three named loop index types used in the triply-nested loop examples.
// These will trigger compilation errors if lambda index argument ordering
// and types do not match the typed range index ordering.
//
// _nestedreorder_idxtypes_start
RAJA_INDEX_VALUE(KIDX, "KIDX");
RAJA_INDEX_VALUE(JIDX, "JIDX");
RAJA_INDEX_VALUE(IIDX, "IIDX");
// _nestedreorder_idxtypes_end
int main(int RAJA_UNUSED_ARG(argc), char **RAJA_UNUSED_ARG(argv[]))
{
std::cout << "\n\nRAJA nested loop reorder example...\n";
//
// Typed index ranges
//
// _nestedreorder_ranges_start
RAJA::TypedRangeSegment<KIDX> KRange(2, 4);
RAJA::TypedRangeSegment<JIDX> JRange(1, 3);
RAJA::TypedRangeSegment<IIDX> IRange(0, 2);
// _nestedreorder_ranges_end
//----------------------------------------------------------------------------//
std::cout << "\n Running loop reorder example (K-outer, J-middle, I-inner)"
<< "...\n\n" << " (I, J, K)\n" << " ---------\n";
// _nestedreorder_kji_start
using KJI_EXECPOL = RAJA::KernelPolicy<
RAJA::statement::For<2, RAJA::seq_exec, // k
RAJA::statement::For<1, RAJA::seq_exec, // j
RAJA::statement::For<0, RAJA::seq_exec,// i
RAJA::statement::Lambda<0>
>
>
>
>;
RAJA::kernel<KJI_EXECPOL>( RAJA::make_tuple(IRange, JRange, KRange),
[=] (IIDX i, JIDX j, KIDX k) {
printf( " (%d, %d, %d) \n", (int)(*i), (int)(*j), (int)(*k));
});
// _nestedreorder_kji_end
//----------------------------------------------------------------------------//
std::cout << "\n Running loop reorder example (J-outer, I-middle, K-inner)"
<< "...\n\n" << " (I, J, K)\n" << " ---------\n";
// _nestedreorder_jik_start
using JIK_EXECPOL = RAJA::KernelPolicy<
RAJA::statement::For<1, RAJA::seq_exec, // j
RAJA::statement::For<0, RAJA::seq_exec, // i
RAJA::statement::For<2, RAJA::seq_exec,// k
RAJA::statement::Lambda<0>
>
>
>
>;
// _nestedreorder_jik_end
RAJA::kernel<JIK_EXECPOL>( RAJA::make_tuple(IRange, JRange, KRange),
[=] (IIDX i, JIDX j, KIDX k) {
printf( " (%d, %d, %d) \n", (int)(*i), (int)(*j), (int)(*k));
});
//----------------------------------------------------------------------------//
std::cout << "\n Running loop reorder example (I-outer, K-middle, J-inner)"
<< "...\n\n" << " (I, J, K)\n" << " ---------\n";
// _nestedreorder_ikj_start
using IKJ_EXECPOL = RAJA::KernelPolicy<
RAJA::statement::For<0, RAJA::seq_exec, // i
RAJA::statement::For<2, RAJA::seq_exec, // k
RAJA::statement::For<1, RAJA::seq_exec,// j
RAJA::statement::Lambda<0>
>
>
>
>;
// _nestedreorder_ikj_end
RAJA::kernel<IKJ_EXECPOL>( RAJA::make_tuple(IRange, JRange, KRange),
[=] (IIDX i, JIDX j, KIDX k) {
printf( " (%d, %d, %d) \n", (int)(*i), (int)(*j), (int)(*k));
});
#if 0
//----------------------------------------------------------------------------//
// The following demonstrates that code will not compile if lambda argument
// types/order do not match the types/order of the For statements.
//----------------------------------------------------------------------------//
// _nestedreorder_typemismatch_start
RAJA::kernel<IKJ_EXECPOL>( RAJA::make_tuple(IRange, JRange, KRange),
[=] (JIDX i, IIDX j, KIDX k) {
printf( " (%d, %d, %d) \n", (int)(*i), (int)(*j), (int)(*k));
});
// _nestedreorder_typemismatch_end
#endif
std::cout << "\n DONE!...\n";
return 0;
}