Skip to content

Commit 56b6402

Browse files
WeiqunZhangWeiqunZhang
authored
ParallelFor with compile time optimization of kernels with run time parameters (AMReX-Codes#2954)
Branches inside ParallelFor can be very expensive. If a branch uses a lot of resources (e.g., registers), it can significantly affect the performance even if at run time the branch is never executed because it affects the GPU occupancy. For CPUs, it can affect vectorization of the kernel. The new ParallelFor functions use C++17 fold expression to generate kernel launches for all run time variants. Only one will be executed. Which one is chosen at run time depends the run time parameters. The kernel function can use constexpr if to discard unused code blocks for better run time performance. Here are two examples of how to use them. int runtime_option = ...; enum All_options : int { A0, A1, A2, A3}; // Four ParallelFors will be generated. ParallelFor(TypeList<CompileTimeOptions<A0,A1,A2,A3>>{}, {runtime_option}, box, [=] AMREX_GPU_DEVICE (int i, int j, int k, auto control) { ... if constexpr (control.value == A0) { ... } else if constexpr (control.value == A1) { ... } else if constexpr (control.value == A2) { ... else { ... } ... }); and int A_runtime_option = ...; int B_runtime_option = ...; enum A_options : int { A0, A1, A2, A3}; enum B_options : int { B0, B1 }; // 4*2=8 ParallelFors will be generated. ParallelFor(TypeList<CompileTimeOptions<A0,A1,A2,A3>, CompileTimeOptions<B0,B1> > {}, {A_runtime_option, B_runtime_option}, N, [=] AMREX_GPU_DEVICE (int i, auto A_control, auto B_control) { ... if constexpr (A_control.value == A0) { ... } else if constexpr (A_control.value == A1) { ... } else if constexpr (A_control.value == A2) { ... else { ... } if constexpr (A_control.value != A3 && B_control.value == B1) { ... } ... }); Note that that due to a limitation of CUDA's extended device lambda, the constexpr if block cannot be the one that captures a variable first. If nvcc complains about it, you will have to manually capture it outside constexpr if. The data type for the parameters is int. Thank Maikel Nadolski and Alex Sinn for showing us the meta-programming techniques used here.
1 parent bcbf17f commit 56b6402

File tree

9 files changed

+431
-1
lines changed

9 files changed

+431
-1
lines changed

‎Src/Base/AMReX_CTOParallelForImpl.H‎

Lines changed: 331 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,331 @@
1+
#ifndef AMREX_CTO_PARALLEL_FOR_H_
2+
#define AMREX_CTO_PARALLEL_FOR_H_
3+
4+
#include <AMReX_BLassert.H>
5+
#include <AMReX_Box.H>
6+
#include <AMReX_Tuple.H>
7+
8+
#include <array>
9+
#include <type_traits>
10+
11+
/* This header is not for the users to include directly. It's meant to be
12+
* included in AMReX_GpuLaunch.H, which has included the headers needed
13+
* here. */
14+
15+
/* Thank Maikel Nadolski and Alex Sinn for the techniques used here! */
16+
17+
namespace amrex {
18+
19+
template <int... ctr>
20+
struct CompileTimeOptions {
21+
// TypeList is defined in AMReX_Tuple.H
22+
using list_type = TypeList<std::integral_constant<int, ctr>...>;
23+
};
24+
25+
#if (__cplusplus >= 201703L)
26+
27+
namespace meta
28+
{
29+
template <typename... As, typename... Bs>
30+
constexpr auto operator+ (TypeList<As...>, TypeList<Bs...>) {
31+
return TypeList<As..., Bs...>{};
32+
}
33+
34+
template <typename... Ls, typename A>
35+
constexpr auto single_product (TypeList<Ls...>, A) {
36+
return TypeList<decltype(Ls{} + TypeList<A>{})...>{};
37+
}
38+
39+
template <typename LLs, typename... As>
40+
constexpr auto operator* (LLs, TypeList<As...>) {
41+
return (TypeList<>{} + ... + single_product(LLs{}, As{}));
42+
}
43+
44+
template <typename... Ls>
45+
constexpr auto cartesian_product_n (TypeList<Ls...>) {
46+
return (TypeList<TypeList<>>{} * ... * Ls{});
47+
}
48+
}
49+
50+
namespace detail
51+
{
52+
template <int MT, typename T, class F, typename... As>
53+
std::enable_if_t<std::is_integral<T>::value || std::is_same<T,Box>::value, bool>
54+
ParallelFor_helper2 (T const& N, F&& f, TypeList<As...>,
55+
std::array<int,sizeof...(As)> const& runtime_options)
56+
{
57+
if (runtime_options == std::array<int,sizeof...(As)>{As::value...}) {
58+
if constexpr (std::is_integral<T>::value) {
59+
ParallelFor<MT>(N, [f] AMREX_GPU_DEVICE (T i) noexcept
60+
{
61+
f(i, As{}...);
62+
});
63+
} else {
64+
ParallelFor<MT>(N, [f] AMREX_GPU_DEVICE (int i, int j, int k) noexcept
65+
{
66+
f(i, j, k, As{}...);
67+
});
68+
}
69+
return true;
70+
} else {
71+
return false;
72+
}
73+
}
74+
75+
template <int MT, typename T, class F, typename... As>
76+
std::enable_if_t<std::is_integral<T>::value, bool>
77+
ParallelFor_helper2 (Box const& box, T ncomp, F&& f, TypeList<As...>,
78+
std::array<int,sizeof...(As)> const& runtime_options)
79+
{
80+
if (runtime_options == std::array<int,sizeof...(As)>{As::value...}) {
81+
ParallelFor<MT>(box, ncomp, [f] AMREX_GPU_DEVICE (int i, int j, int k, T n) noexcept
82+
{
83+
f(i, j, k, n, As{}...);
84+
});
85+
return true;
86+
} else {
87+
return false;
88+
}
89+
}
90+
91+
template <int MT, typename T, class F, typename... PPs, typename RO>
92+
std::enable_if_t<std::is_integral<T>::value || std::is_same<T,Box>::value>
93+
ParallelFor_helper1 (T const& N, F&& f, TypeList<PPs...>,
94+
RO const& runtime_options)
95+
{
96+
bool found_option = (false || ... ||
97+
ParallelFor_helper2<MT>(N, std::forward<F>(f),
98+
PPs{}, runtime_options));
99+
amrex::ignore_unused(found_option);
100+
AMREX_ASSERT(found_option);
101+
}
102+
103+
template <int MT, typename T, class F, typename... PPs, typename RO>
104+
std::enable_if_t<std::is_integral<T>::value>
105+
ParallelFor_helper1 (Box const& box, T ncomp, F&& f, TypeList<PPs...>,
106+
RO const& runtime_options)
107+
{
108+
bool found_option = (false || ... ||
109+
ParallelFor_helper2<MT>(box, ncomp, std::forward<F>(f),
110+
PPs{}, runtime_options));
111+
amrex::ignore_unused(found_option);
112+
AMREX_ASSERT(found_option);
113+
}
114+
}
115+
116+
#endif
117+
118+
template <int MT, typename T, class F, typename... CTOs>
119+
std::enable_if_t<std::is_integral<T>::value>
120+
ParallelFor (TypeList<CTOs...> /*list_of_compile_time_options*/,
121+
std::array<int,sizeof...(CTOs)> const& runtime_options,
122+
T N, F&& f)
123+
{
124+
#if (__cplusplus >= 201703L)
125+
using OptionsListList = TypeList<typename CTOs::list_type...>;
126+
detail::ParallelFor_helper1<MT>(N, std::forward<F>(f),
127+
meta::cartesian_product_n(OptionsListList{}),
128+
runtime_options);
129+
#else
130+
amrex::ignore_unused(N, f, runtime_options);
131+
static_assert(std::is_integral<F>::value, "This requires C++17");
132+
#endif
133+
}
134+
135+
template <int MT, class F, typename... CTOs>
136+
void ParallelFor (TypeList<CTOs...> /*list_of_compile_time_options*/,
137+
std::array<int,sizeof...(CTOs)> const& runtime_options,
138+
Box const& box, F&& f)
139+
{
140+
#if (__cplusplus >= 201703L)
141+
using OptionsListList = TypeList<typename CTOs::list_type...>;
142+
detail::ParallelFor_helper1<MT>(box, std::forward<F>(f),
143+
meta::cartesian_product_n(OptionsListList{}),
144+
runtime_options);
145+
#else
146+
amrex::ignore_unused(box, f, runtime_options);
147+
static_assert(std::is_integral<F>::value, "This requires C++17");
148+
#endif
149+
}
150+
151+
template <int MT, typename T, class F, typename... CTOs>
152+
std::enable_if_t<std::is_integral<T>::value>
153+
ParallelFor (TypeList<CTOs...> /*list_of_compile_time_options*/,
154+
std::array<int,sizeof...(CTOs)> const& runtime_options,
155+
Box const& box, T ncomp, F&& f)
156+
{
157+
#if (__cplusplus >= 201703L)
158+
using OptionsListList = TypeList<typename CTOs::list_type...>;
159+
detail::ParallelFor_helper1<MT>(box, ncomp, std::forward<F>(f),
160+
meta::cartesian_product_n(OptionsListList{}),
161+
runtime_options);
162+
#else
163+
amrex::ignore_unused(box, ncomp, f, runtime_options);
164+
static_assert(std::is_integral<F>::value, "This requires C++17");
165+
#endif
166+
}
167+
168+
/**
169+
* \brief ParallelFor with compile time optimization of kernels with run time options.
170+
*
171+
* It uses fold expression to generate kernel launches for all combinations
172+
* of the run time options. The kernel function can use constexpr if to
173+
* discard unused code blocks for better run time performance. In the
174+
* example below, the code will be expanded into 4*2=8 normal ParallelFors
175+
* for all combinations of the run time parameters.
176+
\verbatim
177+
int A_runtime_option = ...;
178+
int B_runtime_option = ...;
179+
enum A_options : int { A0, A1, A2, A3};
180+
enum B_options : int { B0, B1 };
181+
ParallelFor(TypeList<CompileTimeOptions<A0,A1,A2,A3>,
182+
CompileTimeOptions<B0,B1>>{},
183+
{A_runtime_option, B_runtime_option},
184+
N, [=] AMREX_GPU_DEVICE (int i, auto A_control, auto B_control)
185+
{
186+
...
187+
if constexpr (A_control.value == A0) {
188+
...
189+
} else if constexpr (A_control.value == A1) {
190+
...
191+
} else if constexpr (A_control.value == A2) {
192+
...
193+
else {
194+
...
195+
}
196+
if constexpr (A_control.value != A3 && B_control.value == B1) {
197+
...
198+
}
199+
...
200+
});
201+
\endverbatim
202+
* Note that due to a limitation of CUDA's extended device lambda, the
203+
* constexpr if block cannot be the one that captures a variable first.
204+
* If nvcc complains about it, you will have to manually capture it outside
205+
* constexpr if. The data type for the parameters is int.
206+
*
207+
* \param ctos list of all possible values of the parameters.
208+
* \param option the run time parameters.
209+
* \param N an interger specifying the 1D for loop's range.
210+
* \param f a callable object taking an integer and working on that iteration.
211+
*/
212+
template <typename T, class F, typename... CTOs>
213+
std::enable_if_t<std::is_integral<T>::value>
214+
ParallelFor (TypeList<CTOs...> ctos,
215+
std::array<int,sizeof...(CTOs)> const& option,
216+
T N, F&& f)
217+
{
218+
ParallelFor<AMREX_GPU_MAX_THREADS>(ctos, option, N, std::forward<F>(f));
219+
}
220+
221+
/**
222+
* \brief ParallelFor with compile time optimization of kernels with run time options.
223+
*
224+
* It uses fold expression to generate kernel launches for all combinations
225+
* of the run time options. The kernel function can use constexpr if to
226+
* discard unused code blocks for better run time performance. In the
227+
* example below, the code will be expanded into 4*2=8 normal ParallelFors
228+
* for all combinations of the run time parameters.
229+
\verbatim
230+
int A_runtime_option = ...;
231+
int B_runtime_option = ...;
232+
enum A_options : int { A0, A1, A2, A3};
233+
enum B_options : int { B0, B1 };
234+
ParallelFor(TypeList<CompileTimeOptions<A0,A1,A2,A3>,
235+
CompileTimeOptions<B0,B1>>{},
236+
{A_runtime_option, B_runtime_option},
237+
box, [=] AMREX_GPU_DEVICE (int i, int j, int k,
238+
auto A_control, auto B_control)
239+
{
240+
...
241+
if constexpr (A_control.value == A0) {
242+
...
243+
} else if constexpr (A_control.value == A1) {
244+
...
245+
} else if constexpr (A_control.value == A2) {
246+
...
247+
else {
248+
...
249+
}
250+
if constexpr (A_control.value != A3 && B_control.value == B1) {
251+
...
252+
}
253+
...
254+
});
255+
\endverbatim
256+
* Note that due to a limitation of CUDA's extended device lambda, the
257+
* constexpr if block cannot be the one that captures a variable first.
258+
* If nvcc complains about it, you will have to manually capture it outside
259+
* constexpr if. The data type for the parameters is int.
260+
*
261+
* \param ctos list of all possible values of the parameters.
262+
* \param option the run time parameters.
263+
* \param box a Box specifying the 3D for loop's range.
264+
* \param f a callable object taking three integers and working on the given cell.
265+
*/
266+
template <class F, typename... CTOs>
267+
void ParallelFor (TypeList<CTOs...> ctos,
268+
std::array<int,sizeof...(CTOs)> const& option,
269+
Box const& box, F&& f)
270+
{
271+
ParallelFor<AMREX_GPU_MAX_THREADS>(ctos, option, box, std::forward<F>(f));
272+
}
273+
274+
/**
275+
* \brief ParallelFor with compile time optimization of kernels with run time options.
276+
*
277+
* It uses fold expression to generate kernel launches for all combinations
278+
* of the run time options. The kernel function can use constexpr if to
279+
* discard unused code blocks for better run time performance. In the
280+
* example below, the code will be expanded into 4*2=8 normal ParallelFors
281+
* for all combinations of the run time parameters.
282+
\verbatim
283+
int A_runtime_option = ...;
284+
int B_runtime_option = ...;
285+
enum A_options : int { A0, A1, A2, A3};
286+
enum B_options : int { B0, B1 };
287+
ParallelFor(TypeList<CompileTimeOptions<A0,A1,A2,A3>,
288+
CompileTimeOptions<B0,B1>>{},
289+
{A_runtime_option, B_runtime_option},
290+
box, ncomp, [=] AMREX_GPU_DEVICE (int i, int j, int k, int n,
291+
auto A_control, auto B_control)
292+
{
293+
...
294+
if constexpr (A_control.value == A0) {
295+
...
296+
} else if constexpr (A_control.value == A1) {
297+
...
298+
} else if constexpr (A_control.value == A2) {
299+
...
300+
else {
301+
...
302+
}
303+
if constexpr (A_control.value != A3 && B_control.value == B1) {
304+
...
305+
}
306+
...
307+
});
308+
\endverbatim
309+
* Note that due to a limitation of CUDA's extended device lambda, the
310+
* constexpr if block cannot be the one that captures a variable first.
311+
* If nvcc complains about it, you will have to manually capture it outside
312+
* constexpr if. The data type for the parameters is int.
313+
*
314+
* \param ctos list of all possible values of the parameters.
315+
* \param option the run time parameters.
316+
* \param box a Box specifying the iteration in 3D space.
317+
* \param ncomp an integer specifying the range for iteration over components.
318+
* \param f a callable object taking three integers and working on the given cell.
319+
*/
320+
template <typename T, class F, typename... CTOs>
321+
std::enable_if_t<std::is_integral<T>::value>
322+
ParallelFor (TypeList<CTOs...> ctos,
323+
std::array<int,sizeof...(CTOs)> const& option,
324+
Box const& box, T ncomp, F&& f)
325+
{
326+
ParallelFor<AMREX_GPU_MAX_THREADS>(ctos, option, box, ncomp, std::forward<F>(f));
327+
}
328+
329+
}
330+
331+
#endif

‎Src/Base/AMReX_GpuLaunch.H‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -550,4 +550,6 @@ namespace Gpu {
550550

551551
#endif
552552

553+
#include <AMReX_CTOParallelForImpl.H>
554+
553555
#endif

‎Src/Base/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -224,6 +224,7 @@ target_sources( amrex
224224
AMReX_MFParallelForC.H
225225
AMReX_MFParallelForG.H
226226
AMReX_TagParallelFor.H
227+
AMReX_CTOParallelForImpl.H
227228
AMReX_ParReduce.H
228229
# CUDA --------------------------------------------------------------------
229230
AMReX_CudaGraph.H

‎Src/Base/Make.package‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,7 @@ C$(AMREX_BASE)_headers += AMReX_MFParallelForC.H
100100
C$(AMREX_BASE)_headers += AMReX_MFParallelForG.H
101101

102102
C$(AMREX_BASE)_headers += AMReX_TagParallelFor.H
103+
C$(AMREX_BASE)_headers += AMReX_CTOParallelForImpl.H
103104

104105
C$(AMREX_BASE)_headers += AMReX_ParReduce.H
105106

‎Tests/CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
#
22
# List of subdirectories to search for CMakeLists.
33
#
4-
set( AMREX_TESTS_SUBDIRS AsyncOut MultiBlock Amr CLZ Parser)
4+
set( AMREX_TESTS_SUBDIRS AsyncOut MultiBlock Amr CLZ Parser CTOParFor)
55

66
if (AMReX_PARTICLES)
77
list(APPEND AMREX_TESTS_SUBDIRS Particles)

‎Tests/CTOParFor/CMakeLists.txt‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
set(_sources main.cpp)
2+
set(_input_files)
3+
4+
setup_test(_sources _input_files)
5+
6+
unset(_sources)
7+
unset(_input_files)

‎Tests/CTOParFor/GNUmakefile‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,20 @@
1+
AMREX_HOME = ../../
2+
3+
DEBUG = FALSE
4+
DIM = 3
5+
COMP = gcc
6+
7+
USE_MPI = FALSE
8+
USE_OMP = FALSE
9+
USE_CUDA = FALSE
10+
11+
TINY_PROFILE = FALSE
12+
13+
CXXSTD = c++17
14+
15+
include $(AMREX_HOME)/Tools/GNUMake/Make.defs
16+
17+
include ./Make.package
18+
include $(AMREX_HOME)/Src/Base/Make.package
19+
20+
include $(AMREX_HOME)/Tools/GNUMake/Make.rules

‎Tests/CTOParFor/Make.package‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
CEXE_sources += main.cpp
2+
3+
4+

0 commit comments

Comments
 (0)